diff --git a/src/auto-reply/reply/agent-runner.ts b/src/auto-reply/reply/agent-runner.ts index 227e6f17e..d7880c46e 100644 --- a/src/auto-reply/reply/agent-runner.ts +++ b/src/auto-reply/reply/agent-runner.ts @@ -418,41 +418,58 @@ export async function runReplyAgent(params: { await signalTypingIfNeeded(replyPayloads, typingSignals); - if (isDiagnosticsEnabled(cfg) && hasNonzeroUsage(usage)) { + // Compute usage stats and notify via callback if available + if (hasNonzeroUsage(usage)) { const input = usage.input ?? 0; const output = usage.output ?? 0; const cacheRead = usage.cacheRead ?? 0; const cacheWrite = usage.cacheWrite ?? 0; const promptTokens = input + cacheRead + cacheWrite; const totalTokens = usage.total ?? promptTokens + output; - const costConfig = resolveModelCostConfig({ - provider: providerUsed, + const durationMs = Date.now() - runStartedAt; + + // Notify usage callback (for presence updates, analytics, etc.) + opts?.onUsageAvailable?.({ + inputTokens: input, + outputTokens: output, + totalTokens, + cacheReadTokens: cacheRead, + cacheWriteTokens: cacheWrite, model: modelUsed, - config: cfg, - }); - const costUsd = estimateUsageCost({ usage, cost: costConfig }); - emitDiagnosticEvent({ - type: "model.usage", - sessionKey, - sessionId: followupRun.run.sessionId, - channel: replyToChannel, provider: providerUsed, - model: modelUsed, - usage: { - input, - output, - cacheRead, - cacheWrite, - promptTokens, - total: totalTokens, - }, - context: { - limit: contextTokensUsed, - used: totalTokens, - }, - costUsd, - durationMs: Date.now() - runStartedAt, + durationMs, }); + + if (isDiagnosticsEnabled(cfg)) { + const costConfig = resolveModelCostConfig({ + provider: providerUsed, + model: modelUsed, + config: cfg, + }); + const costUsd = estimateUsageCost({ usage, cost: costConfig }); + emitDiagnosticEvent({ + type: "model.usage", + sessionKey, + sessionId: followupRun.run.sessionId, + channel: replyToChannel, + provider: providerUsed, + model: modelUsed, + usage: { + input, + output, + cacheRead, + cacheWrite, + promptTokens, + total: totalTokens, + }, + context: { + limit: contextTokensUsed, + used: totalTokens, + }, + costUsd, + durationMs, + }); + } } const responseUsageRaw = diff --git a/src/auto-reply/reply/dispatch-from-config.ts b/src/auto-reply/reply/dispatch-from-config.ts index 428320a68..76e36f718 100644 --- a/src/auto-reply/reply/dispatch-from-config.ts +++ b/src/auto-reply/reply/dispatch-from-config.ts @@ -11,7 +11,7 @@ import { import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js"; import { getReplyFromConfig } from "../reply.js"; import type { FinalizedMsgContext } from "../templating.js"; -import type { GetReplyOptions, ReplyPayload } from "../types.js"; +import type { GetReplyOptions, ReplyPayload, ReplyUsage } from "../types.js"; import { formatAbortReplyText, tryFastAbortFromMessage } from "./abort.js"; import { shouldSkipDuplicateInbound } from "./inbound-dedupe.js"; import type { ReplyDispatcher, ReplyDispatchKind } from "./reply-dispatcher.js"; @@ -266,10 +266,18 @@ export async function dispatchReplyFromConfig(params: { return { queuedFinal, counts }; } + // Mutable ref to capture usage stats from the agent run + const usageRef: { value?: ReplyUsage } = {}; + const replyResult = await (params.replyResolver ?? getReplyFromConfig)( ctx, { ...params.replyOptions, + onUsageAvailable: (usage) => { + usageRef.value = usage; + // Also call the original callback if provided + params.replyOptions?.onUsageAvailable?.(usage); + }, onBlockReply: (payload: ReplyPayload, context) => { const run = async () => { const ttsPayload = await maybeApplyTtsToPayload({ @@ -332,13 +340,7 @@ export async function dispatchReplyFromConfig(params: { const counts = dispatcher.getQueuedCounts(); counts.final += routedFinalCount; - // Trigger message_sent hook for final replies - // NOTE: Usage stats (tokens, model, etc.) are not currently available here. - // They are computed in agent-runner.ts but not returned through the reply pipeline. - // To enable full usage tracking, getReplyFromConfig/replyResolver would need to - // return usage data alongside the reply text. For now, hooks receive the message - // content without usage stats - presence managers can track cumulative tokens - // via other means (e.g., diagnostic events or session metadata). + // Trigger message_sent hook for final replies with usage stats if (hookRunner?.hasHooks("message_sent") && replies.length > 0) { const targetSessionKey = ctx.CommandSource === "native" ? ctx.CommandTargetSessionKey?.trim() : undefined; @@ -349,6 +351,9 @@ export async function dispatchReplyFromConfig(params: { const channelId = (ctx.OriginatingChannel ?? ctx.Surface ?? ctx.Provider ?? "").toLowerCase(); const conversationId = ctx.OriginatingTo ?? ctx.To ?? ctx.From ?? undefined; + // Get captured usage stats from the agent run + const usage = usageRef.value; + for (const reply of replies) { void hookRunner .runMessageSent( @@ -356,6 +361,19 @@ export async function dispatchReplyFromConfig(params: { to: conversationId ?? "", content: reply.text ?? "", success: queuedFinal, + // Include usage stats if available + usage: usage + ? { + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + totalTokens: usage.totalTokens, + cacheReadTokens: usage.cacheReadTokens, + cacheCreationTokens: usage.cacheWriteTokens, + model: usage.model, + provider: usage.provider, + } + : undefined, + durationMs: usage?.durationMs, }, { channelId, diff --git a/src/auto-reply/types.ts b/src/auto-reply/types.ts index 1aa0fe067..887f4eccc 100644 --- a/src/auto-reply/types.ts +++ b/src/auto-reply/types.ts @@ -32,6 +32,8 @@ export type GetReplyOptions = { /** Called when the actual model is selected (including after fallback). * Use this to get model/provider/thinkLevel for responsePrefix template interpolation. */ onModelSelected?: (ctx: ModelSelectedContext) => void; + /** Called when usage statistics are available after the agent run completes. */ + onUsageAvailable?: (usage: ReplyUsage) => void; disableBlockStreaming?: boolean; /** Timeout for block reply delivery (ms). */ blockReplyTimeoutMs?: number; @@ -55,3 +57,21 @@ export type ReplyPayload = { /** Channel-specific payload data (per-channel envelope). */ channelData?: Record; }; + +/** Usage statistics from an agent run. */ +export type ReplyUsage = { + inputTokens?: number; + outputTokens?: number; + totalTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + model?: string; + provider?: string; + durationMs?: number; +}; + +/** Result from getReplyFromConfig including reply payload(s) and usage stats. */ +export type ReplyResult = { + replies: ReplyPayload[]; + usage?: ReplyUsage; +};