fix: persist session usage metadata on suppressed replies

This commit is contained in:
Peter Steinberger
2026-01-24 11:04:53 +00:00
parent c02204fd1e
commit 1bbbb10abf
5 changed files with 225 additions and 139 deletions

View File

@@ -4,12 +4,7 @@ import { lookupContextTokens } from "../../agents/context.js";
import { DEFAULT_CONTEXT_TOKENS } from "../../agents/defaults.js";
import { runWithModelFallback } from "../../agents/model-fallback.js";
import { runEmbeddedPiAgent } from "../../agents/pi-embedded.js";
import { hasNonzeroUsage } from "../../agents/usage.js";
import {
resolveAgentIdFromSessionKey,
type SessionEntry,
updateSessionStoreEntry,
} from "../../config/sessions.js";
import { resolveAgentIdFromSessionKey, type SessionEntry } from "../../config/sessions.js";
import type { TypingMode } from "../../config/types.js";
import { logVerbose } from "../../globals.js";
import { registerAgentRunContext } from "../../infra/agent-events.js";
@@ -26,6 +21,7 @@ import {
} from "./reply-payloads.js";
import { resolveReplyToMode } from "./reply-threading.js";
import { isRoutableChannel, routeReply } from "./route-reply.js";
import { persistSessionUsageUpdate } from "./session-usage.js";
import { incrementCompactionCount } from "./session-updates.js";
import type { TypingController } from "./typing.js";
import { createTypingSignaler } from "./typing-mode.js";
@@ -190,6 +186,26 @@ export function createFollowupRunner(params: {
return;
}
if (storePath && sessionKey) {
const usage = runResult.meta.agentMeta?.usage;
const modelUsed = runResult.meta.agentMeta?.model ?? fallbackModel ?? defaultModel;
const contextTokensUsed =
agentCfgContextTokens ??
lookupContextTokens(modelUsed) ??
sessionEntry?.contextTokens ??
DEFAULT_CONTEXT_TOKENS;
await persistSessionUsageUpdate({
storePath,
sessionKey,
usage,
modelUsed,
providerUsed: fallbackProvider,
contextTokensUsed,
logLabel: "followup",
});
}
const payloadArray = runResult.payloads ?? [];
if (payloadArray.length === 0) return;
const sanitizedPayloads = payloadArray.flatMap((payload) => {
@@ -245,56 +261,6 @@ export function createFollowupRunner(params: {
}
}
if (storePath && sessionKey) {
const usage = runResult.meta.agentMeta?.usage;
const modelUsed = runResult.meta.agentMeta?.model ?? fallbackModel ?? defaultModel;
const contextTokensUsed =
agentCfgContextTokens ??
lookupContextTokens(modelUsed) ??
sessionEntry?.contextTokens ??
DEFAULT_CONTEXT_TOKENS;
if (hasNonzeroUsage(usage)) {
try {
await updateSessionStoreEntry({
storePath,
sessionKey,
update: async (entry) => {
const input = usage.input ?? 0;
const output = usage.output ?? 0;
const promptTokens = input + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
return {
inputTokens: input,
outputTokens: output,
totalTokens: promptTokens > 0 ? promptTokens : (usage.total ?? input),
modelProvider: fallbackProvider ?? entry.modelProvider,
model: modelUsed,
contextTokens: contextTokensUsed ?? entry.contextTokens,
updatedAt: Date.now(),
};
},
});
} catch (err) {
logVerbose(`failed to persist followup usage update: ${String(err)}`);
}
} else if (modelUsed || contextTokensUsed) {
try {
await updateSessionStoreEntry({
storePath,
sessionKey,
update: async (entry) => ({
modelProvider: fallbackProvider ?? entry.modelProvider,
model: modelUsed ?? entry.model,
contextTokens: contextTokensUsed ?? entry.contextTokens,
updatedAt: Date.now(),
}),
});
} catch (err) {
logVerbose(`failed to persist followup model/context update: ${String(err)}`);
}
}
}
await sendFollowupPayloads(finalPayloads, queued);
} finally {
typing.markRunComplete();