feat: per-query context usage in traces and UI

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));
         }
       },
     });

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,
     };
   }
 }

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();
+  });
+});

packages/server/src/api/browse.ts

@@ -53,14 +53,15 @@ 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,
         startedAt,
         durationMs,
         notation,
         stepCount: steps.length,
+        usage,
       }))
     );
   });

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 {

packages/web/src/components/GraphView.tsx

@@ -36,6 +36,16 @@ const KIND_COLOR: Record<TraceSummary["kind"], string> = {
 
 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({
                   />
                   <span className="truncate text-zinc-200">{t.input}</span>
                 </div>
-                <div className="mt-0.5 truncate font-mono text-[10px] text-zinc-500">{t.notation}</div>
+                <div className="mt-0.5 flex items-baseline gap-2">
+                  <span className="truncate font-mono text-[10px] text-zinc-500">{t.notation}</span>
+                  {usageLabel(t.usage) && (
+                    <span className="ml-auto shrink-0 font-mono text-[10px] text-zinc-600">{usageLabel(t.usage)}</span>
+                  )}
+                </div>
               </button>
             ))}
           </div>
@@ -337,6 +352,11 @@ export function GraphView({
               {activeTrace.kind}
             </span>
             <span className="truncate text-zinc-200">{activeTrace.input}</span>
+            {usageLabel(activeTrace.usage) && (
+              <span className="shrink-0 rounded bg-zinc-800 px-1.5 py-0.5 font-mono text-[10px] text-zinc-400">
+                {usageLabel(activeTrace.usage)}
+              </span>
+            )}
             <button
               onClick={closeTrace}
               className="ml-auto shrink-0 rounded px-1.5 text-zinc-400 hover:bg-zinc-800 hover:text-zinc-100"