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

@@ -1,6 +1,5 @@
import crypto from "node:crypto";
import fs from "node:fs";
import { setCliSessionId } from "../../agents/cli-session.js";
import { lookupContextTokens } from "../../agents/context.js";
import { DEFAULT_CONTEXT_TOKENS } from "../../agents/defaults.js";
import { resolveModelAuthMode } from "../../agents/model-auth.js";
@@ -38,6 +37,7 @@ import { resolveBlockStreamingCoalescing } from "./block-streaming.js";
import { createFollowupRunner } from "./followup-runner.js";
import { enqueueFollowupRun, type FollowupRun, type QueueSettings } from "./queue.js";
import { createReplyToModeFilterForChannel, resolveReplyToMode } from "./reply-threading.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";
@@ -365,6 +365,30 @@ export async function runReplyAgent(params: {
await Promise.allSettled(pendingToolTasks);
}
const usage = runResult.meta.agentMeta?.usage;
const modelUsed = runResult.meta.agentMeta?.model ?? fallbackModel ?? defaultModel;
const providerUsed =
runResult.meta.agentMeta?.provider ?? fallbackProvider ?? followupRun.run.provider;
const cliSessionId = isCliProvider(providerUsed, cfg)
? runResult.meta.agentMeta?.sessionId?.trim()
: undefined;
const contextTokensUsed =
agentCfgContextTokens ??
lookupContextTokens(modelUsed) ??
activeSessionEntry?.contextTokens ??
DEFAULT_CONTEXT_TOKENS;
await persistSessionUsageUpdate({
storePath,
sessionKey,
usage,
modelUsed,
providerUsed,
contextTokensUsed,
systemPromptReport: runResult.meta.systemPromptReport,
cliSessionId,
});
// Drain any late tool/block deliveries before deciding there's "nothing to send".
// Otherwise, a late typing trigger (e.g. from a tool callback) can outlive the run and
// keep the typing indicator stuck.
@@ -395,19 +419,6 @@ export async function runReplyAgent(params: {
await signalTypingIfNeeded(replyPayloads, typingSignals);
const usage = runResult.meta.agentMeta?.usage;
const modelUsed = runResult.meta.agentMeta?.model ?? fallbackModel ?? defaultModel;
const providerUsed =
runResult.meta.agentMeta?.provider ?? fallbackProvider ?? followupRun.run.provider;
const cliSessionId = isCliProvider(providerUsed, cfg)
? runResult.meta.agentMeta?.sessionId?.trim()
: undefined;
const contextTokensUsed =
agentCfgContextTokens ??
lookupContextTokens(modelUsed) ??
activeSessionEntry?.contextTokens ??
DEFAULT_CONTEXT_TOKENS;
if (isDiagnosticsEnabled(cfg) && hasNonzeroUsage(usage)) {
const input = usage.input ?? 0;
const output = usage.output ?? 0;
@@ -445,72 +456,6 @@ export async function runReplyAgent(params: {
});
}
if (storePath && sessionKey) {
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);
const patch: Partial<SessionEntry> = {
inputTokens: input,
outputTokens: output,
totalTokens: promptTokens > 0 ? promptTokens : (usage.total ?? input),
modelProvider: providerUsed,
model: modelUsed,
contextTokens: contextTokensUsed ?? entry.contextTokens,
systemPromptReport: runResult.meta.systemPromptReport ?? entry.systemPromptReport,
updatedAt: Date.now(),
};
if (cliSessionId) {
const nextEntry = { ...entry, ...patch };
setCliSessionId(nextEntry, providerUsed, cliSessionId);
return {
...patch,
cliSessionIds: nextEntry.cliSessionIds,
claudeCliSessionId: nextEntry.claudeCliSessionId,
};
}
return patch;
},
});
} catch (err) {
logVerbose(`failed to persist usage update: ${String(err)}`);
}
} else if (modelUsed || contextTokensUsed) {
try {
await updateSessionStoreEntry({
storePath,
sessionKey,
update: async (entry) => {
const patch: Partial<SessionEntry> = {
modelProvider: providerUsed ?? entry.modelProvider,
model: modelUsed ?? entry.model,
contextTokens: contextTokensUsed ?? entry.contextTokens,
systemPromptReport: runResult.meta.systemPromptReport ?? entry.systemPromptReport,
updatedAt: Date.now(),
};
if (cliSessionId) {
const nextEntry = { ...entry, ...patch };
setCliSessionId(nextEntry, providerUsed, cliSessionId);
return {
...patch,
cliSessionIds: nextEntry.cliSessionIds,
claudeCliSessionId: nextEntry.claudeCliSessionId,
};
}
return patch;
},
});
} catch (err) {
logVerbose(`failed to persist model/context update: ${String(err)}`);
}
}
}
const responseUsageRaw =
activeSessionEntry?.responseUsage ??
(sessionKey ? activeSessionStore?.[sessionKey]?.responseUsage : undefined);