diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index 3af435e..c993c5e 100644 --- a/packages/core/src/agent/agent.ts +++ b/packages/core/src/agent/agent.ts @@ -9,7 +9,7 @@ import { import { withFallback } from "../providers/fallback.js"; import { buildSystemPrompt } from "./system-prompt.js"; import { buildReadTools, buildWriteTools, formatTree } from "./tools.js"; -import { TraceRecorder, TraceStore } from "./trace.js"; +import { TraceRecorder, TraceStore, type TraceUsage } from "./trace.js"; const MAX_STEPS = 12; @@ -95,6 +95,23 @@ function errorMessage(err: unknown): string { return err instanceof Error ? err.message : String(err); } +/** Sum token usage across the run's steps (issue #15). Undefined when the provider reports none. */ +function sumStepsUsage( + steps: ReadonlyArray<{ usage?: { inputTokens?: number; outputTokens?: number } }> +): TraceUsage | undefined { + let inputTokens = 0; + let outputTokens = 0; + let reported = false; + for (const step of steps) { + const u = step.usage; + if (!u || (u.inputTokens == null && u.outputTokens == null)) continue; + reported = true; + inputTokens += u.inputTokens ?? 0; + outputTokens += u.outputTokens ?? 0; + } + return reported ? { inputTokens, outputTokens } : undefined; +} + /** Read-only Q&A over the bundle. */ export async function runQuery( kb: KnowledgeBase, @@ -114,7 +131,7 @@ export async function runQuery( tools: buildReadTools(kb, recorder), stopWhen: stepCountIs(MAX_STEPS), }); - const trace = recorder.finalize("query", question, result.text, "success", modelChain); + const trace = recorder.finalize("query", question, result.text, "success", modelChain, sumStepsUsage(result.steps)); await traceStore(kb).save(trace); return { answer: result.text, steps: result.steps.length, traceId: trace.id }; } catch (err) { @@ -145,7 +162,7 @@ export async function runMutation( stopWhen: stepCountIs(MAX_STEPS), temperature: 0.2, }); - const trace = recorder.finalize("mutation", instruction, result.text, "success", modelChain); + const trace = recorder.finalize("mutation", instruction, result.text, "success", modelChain, sumStepsUsage(result.steps)); await traceStore(kb).save(trace); return { ok: true, @@ -200,10 +217,14 @@ export async function streamChat( messages, tools: { ...buildReadTools(kb, recorder), ...buildWriteTools(kb, filesChanged, recorder) }, stopWhen: stepCountIs(MAX_STEPS), - onFinish: async ({ text }) => { + onFinish: async ({ text, totalUsage }) => { // Persist only turns that actually touched the bundle. if (recorder.steps.length > 0) { - await traceStore(kb).save(recorder.finalize("chat", input, text, "success", modelChain)); + const usage = + totalUsage && (totalUsage.inputTokens != null || totalUsage.outputTokens != null) + ? { inputTokens: totalUsage.inputTokens ?? 0, outputTokens: totalUsage.outputTokens ?? 0 } + : undefined; + await traceStore(kb).save(recorder.finalize("chat", input, text, "success", modelChain, usage)); } }, }); diff --git a/packages/core/src/agent/trace.ts b/packages/core/src/agent/trace.ts index 79b6bf2..5267a1a 100644 --- a/packages/core/src/agent/trace.ts +++ b/packages/core/src/agent/trace.ts @@ -15,6 +15,12 @@ export interface TraceStep { export type TraceOutcome = "success" | "partial" | "failed"; +/** Token accounting for the run (issue #15: context observability). */ +export interface TraceUsage { + inputTokens: number; + outputTokens: number; +} + export interface QueryTrace { id: string; kind: "query" | "mutation" | "chat"; @@ -29,6 +35,8 @@ export interface QueryTrace { notation: string; outcome: TraceOutcome; modelChain: string[]; + /** Total tokens consumed across all steps of the run, when the provider reports them. */ + usage?: TraceUsage; } /** Collects steps during one agent run. Thread one instance through the tools. */ @@ -45,7 +53,8 @@ export class TraceRecorder { input: string, answer: string, outcome: TraceOutcome = "success", - modelChain: string[] = [] + modelChain: string[] = [], + usage?: TraceUsage ): QueryTrace { return { id: `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 7)}`, @@ -58,6 +67,7 @@ export class TraceRecorder { notation: buildNotation(this.steps, outcome), outcome, modelChain, + usage, }; } } diff --git a/packages/core/test/trace-usage.test.ts b/packages/core/test/trace-usage.test.ts new file mode 100644 index 0000000..ccea3ee --- /dev/null +++ b/packages/core/test/trace-usage.test.ts @@ -0,0 +1,20 @@ +import { describe, it, expect } from "vitest"; +import { TraceRecorder } from "../src/agent/trace.js"; + +describe("trace usage (#15)", () => { + it("records token usage on the finalized trace", () => { + const r = new TraceRecorder(); + r.record("read_concept", "/a.md", ["/a.md"]); + const trace = r.finalize("query", "q", "a", "success", ["openai:m"], { + inputTokens: 3200, + outputTokens: 410, + }); + expect(trace.usage).toEqual({ inputTokens: 3200, outputTokens: 410 }); + }); + + it("leaves usage undefined when the provider reports none", () => { + const r = new TraceRecorder(); + const trace = r.finalize("query", "q", "a"); + expect(trace.usage).toBeUndefined(); + }); +}); diff --git a/packages/server/src/api/browse.ts b/packages/server/src/api/browse.ts index 072adcd..30f9237 100644 --- a/packages/server/src/api/browse.ts +++ b/packages/server/src/api/browse.ts @@ -53,7 +53,7 @@ export function browseRouter(kb: KnowledgeBase): Router { // List view: omit full steps/answers to keep the payload light. const all = await traces.list(); res.json( - all.map(({ id, kind, input, startedAt, durationMs, notation, steps }) => ({ + all.map(({ id, kind, input, startedAt, durationMs, notation, steps, usage }) => ({ id, kind, input, @@ -61,6 +61,7 @@ export function browseRouter(kb: KnowledgeBase): Router { durationMs, notation, stepCount: steps.length, + usage, })) ); }); diff --git a/packages/web/src/api.ts b/packages/web/src/api.ts index 95efbda..d594818 100644 --- a/packages/web/src/api.ts +++ b/packages/web/src/api.ts @@ -64,6 +64,7 @@ export interface TraceSummary { durationMs: number; notation: string; stepCount: number; + usage?: { inputTokens: number; outputTokens: number }; } export interface QueryTrace extends TraceSummary { diff --git a/packages/web/src/components/GraphView.tsx b/packages/web/src/components/GraphView.tsx index 41cfec2..62b39e1 100644 --- a/packages/web/src/components/GraphView.tsx +++ b/packages/web/src/components/GraphView.tsx @@ -36,6 +36,16 @@ const KIND_COLOR: Record = { const radius = (n: SimNode) => 5 + Math.sqrt(n.links) * 3.5; +/** 3241 → "3.2k" — compact token counts for the traces panel. */ +function fmtTokens(n: number): string { + return n >= 1000 ? `${(n / 1000).toFixed(1)}k` : String(n); +} + +function usageLabel(usage?: { inputTokens: number; outputTokens: number }): string | null { + if (!usage) return null; + return `${fmtTokens(usage.inputTokens)}→${fmtTokens(usage.outputTokens)} tok`; +} + /** The traversal chain: concept visits in step order (reads + writes), deduped consecutively. */ function traceVisits(trace: QueryTrace): { path: string; seq: number; write: boolean }[] { const visits: { path: string; seq: number; write: boolean }[] = []; @@ -319,7 +329,12 @@ export function GraphView({ /> {t.input} -
{t.notation}
+
+ {t.notation} + {usageLabel(t.usage) && ( + {usageLabel(t.usage)} + )} +
))} @@ -337,6 +352,11 @@ export function GraphView({ {activeTrace.kind} {activeTrace.input} + {usageLabel(activeTrace.usage) && ( + + {usageLabel(activeTrace.usage)} + + )}