feat(server, web): add context usage tracking and display in chat component; enhance event handling and tests

Author: TKassis

server/src/agent/events.ts

@@ -3,14 +3,25 @@
  * frontend consumes. We deliberately flatten the streaming deltas and drop
  * Pi-internal lifecycle noise so the client contract stays small.
  */
-import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent";
+import type { AgentSessionEvent, ContextUsage } from "@earendil-works/pi-coding-agent";
 import { skillLabelForRead } from "./skill-label.ts";
 
 export interface ClientFrame {
   type: string;
   [k: string]: unknown;
 }
 
+/** Keep Pi's context-utilization shape explicit on the wire. */
+export function contextUsageFrame(usage: ContextUsage | undefined): ClientFrame | null {
+  if (!usage) return null;
+  return {
+    type: "context_usage",
+    tokens: usage.tokens,
+    contextWindow: usage.contextWindow,
+    percent: usage.percent,
+  };
+}
+
 /** Frontmatter skill name when a `read` call is a skill activation. */
 export function skillFieldFor(
   toolName: string,

server/src/api/sessions.ts

@@ -10,7 +10,7 @@ import type { FastifyInstance } from "fastify";
 import { activePaths, getProject, touchProject } from "../projects.ts";
 import { corsResponseHeaders } from "../cors.ts";
 import { currentProjectId } from "../scope.ts";
-import { toClientFrame, type ClientFrame } from "../agent/events.ts";
+import { contextUsageFrame, toClientFrame, type ClientFrame } from "../agent/events.ts";
 import { setFusionConfig } from "../agent/fusion-bridge.ts";
 import {
   pendingInterviewFor,
@@ -113,7 +113,11 @@ export async function registerSessionRoutes(app: FastifyInstance): Promise<void>
         reply.code(404);
         return { detail: "No such session" };
       }
-      return { messages: toHistory(file, paths.sandbox) };
+      const session = await getSession(currentProjectId(), paths, req.params.id);
+      return {
+        messages: toHistory(file, paths.sandbox),
+        contextUsage: session?.getContextUsage() ?? null,
+      };
     } catch (err) {
       reply.code(400);
       return { detail: (err as Error).message };
@@ -484,10 +488,15 @@ export async function registerSessionRoutes(app: FastifyInstance): Promise<void>
         const write = (frame: ClientFrame) => {
           if (!raw.writableEnded) raw.write(`data: ${JSON.stringify(frame)}\n\n`);
         };
+        const writeContextUsage = () => {
+          const frame = contextUsageFrame(session.getContextUsage());
+          if (frame) write(frame);
+        };
         // Synthetic route-level frame (Pi events carry no run id): lets the
         // client stamp provisional notebook entries with this run before the
         // authoritative refetch.
         write({ type: "run_start", runId });
+        writeContextUsage();
 
         // Hard budget cap: refuse to run if the project has reached its limit.
         const budget = isBudgetExceeded(projectId);
@@ -518,6 +527,7 @@ export async function registerSessionRoutes(app: FastifyInstance): Promise<void>
           }
           const frame = toClientFrame(ev, sandboxRoot);
           if (frame) write(frame);
+          if (ev.type === "turn_end") writeContextUsage();
         });
 
         req.raw.on("close", () => {
@@ -559,6 +569,7 @@ export async function registerSessionRoutes(app: FastifyInstance): Promise<void>
             // restart/compaction-proof); `tokens` is Pi's in-context cumulative;
             // `runCost`/`runTokens` are the delta for THIS turn, so the UI can
             // attribute a price to the message that just completed.
+            writeContextUsage();
             write({
               type: "cost",
               cost: sessionCostSummary(req.params.id, projectId).totalUsd,

server/test/backend.test.ts

@@ -31,7 +31,11 @@ import {
 } from "../src/agent/web-access-bridge.ts";
 import { guessMime, isUserVisible } from "../src/sandbox-fs.ts";
 import { listProjectSkills, seedProjectSkills } from "../src/agent/skills.ts";
-import { toClientFrame, relativizeSandboxPaths } from "../src/agent/events.ts";
+import {
+  contextUsageFrame,
+  toClientFrame,
+  relativizeSandboxPaths,
+} from "../src/agent/events.ts";
 import { helperPython, HELPERS_DIR } from "../src/helpers-env.ts";
 import { sciHelperFor } from "../src/api/sci-helpers.ts";
 
@@ -255,6 +259,22 @@ describe("skills", () => {
 });
 
 describe("events → client frames", () => {
+  it("maps Pi context utilization onto the client wire shape", () => {
+    expect(
+      contextUsageFrame({
+        tokens: 42_000,
+        contextWindow: 200_000,
+        percent: 21,
+      }),
+    ).toEqual({
+      type: "context_usage",
+      tokens: 42_000,
+      contextWindow: 200_000,
+      percent: 21,
+    });
+    expect(contextUsageFrame(undefined)).toBeNull();
+  });
+
   it("maps text/thinking deltas and tool/lifecycle events", () => {
     expect(toClientFrame({ type: "agent_start" } as never)).toEqual({ type: "agent_start" });
     expect(

web/src/components/chat-tab.tsx

@@ -41,6 +41,7 @@ import { onChatPrefill } from "@/lib/chat-prefill";
 import { buildSkillsContext, type Skill } from "@/components/skills-selector";
 import { AddContextMenu } from "@/components/add-context-menu";
 import { ContextChipsBar } from "@/components/context-chips";
+import { ContextUsageIndicator } from "@/components/context-usage-indicator";
 import { CitationBadge } from "@/components/citation-badge";
 import { NotebookEntryChip, ReasoningBlock, ToolActivityList } from "@/components/tool-activity";
 import { InterviewCard } from "@/components/interview-form";
@@ -54,7 +55,12 @@ import {
   type PromptImage,
 } from "@/lib/image-attachments";
 import { suggestSkillsForFiles } from "@/lib/skill-suggestions";
-import { useAgent, type ActivityItem, type ChatMessage } from "@/lib/use-agent";
+import {
+  useAgent,
+  type ActivityItem,
+  type ChatMessage,
+  type ContextUsage,
+} from "@/lib/use-agent";
 import type { NotebookEntry } from "@/lib/notebook";
 import { routeSubmit, steerNotStreamingFallback, type SendIntent } from "@/lib/chat-routing";
 import { SpeechInput } from "@/components/ai-elements/speech-input";
@@ -462,6 +468,7 @@ function ChatInput({
   onDbsChange,
   selectedModel,
   onModelChange,
+  contextUsage,
   selectedComputeTarget,
   onComputeTargetChange,
   thinkingLevel,
@@ -495,6 +502,7 @@ function ChatInput({
   onDbsChange: (dbs: Database[]) => void;
   selectedModel: Model;
   onModelChange: (model: Model) => void;
+  contextUsage: ContextUsage | null;
   selectedComputeTarget: ModalInstance | null;
   onComputeTargetChange: (instance: ModalInstance | null) => void;
   thinkingLevel: ThinkingLevel;
@@ -823,6 +831,7 @@ function ChatInput({
                 onChange={onComputeTargetChange}
                 modalConfigured={modalConfigured}
               />
+              <ContextUsageIndicator usage={contextUsage} />
             </div>
             <div className="flex items-center gap-1.5 shrink-0">
               <InfoTooltip
@@ -1057,7 +1066,19 @@ export const ChatTab = forwardRef<ChatTabHandle, ChatTabProps>(function ChatTab(
   },
   ref,
 ) {
-  const { messages, status, send, stop, steer, pendingSteers, getSessionId, loadSession, notebookEntries, subagentCompletions } = useAgent();
+  const {
+    messages,
+    contextUsage,
+    status,
+    send,
+    stop,
+    steer,
+    pendingSteers,
+    getSessionId,
+    loadSession,
+    notebookEntries,
+    subagentCompletions,
+  } = useAgent();
   const isStreaming = status === "streaming" || status === "submitted";
   // Scopes the deep-link querySelector to THIS tab's transcript.
   const rootRef = useRef<HTMLDivElement>(null);
@@ -1468,6 +1489,7 @@ export const ChatTab = forwardRef<ChatTabHandle, ChatTabProps>(function ChatTab(
             onDbsChange={setSelectedDbs}
             selectedModel={selectedModel}
             onModelChange={setSelectedModel}
+            contextUsage={contextUsage}
             selectedComputeTarget={selectedComputeTarget}
             onComputeTargetChange={setSelectedComputeTarget}
             thinkingLevel={thinkingLevel}

web/src/components/context-usage-indicator.test.tsx

@@ -0,0 +1,40 @@
+import { render, screen } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { ContextUsageIndicator } from "./context-usage-indicator";
+
+describe("ContextUsageIndicator", () => {
+  it("shows Pi context utilization accessibly", () => {
+    render(
+      <TooltipProvider>
+        <ContextUsageIndicator
+          usage={{ tokens: 42_000, contextWindow: 200_000, percent: 21 }}
+        />
+      </TooltipProvider>,
+    );
+
+    expect(screen.getByRole("status")).toHaveTextContent("21%");
+    expect(screen.getByRole("status")).toHaveAccessibleName(
+      "Model context 21.0 percent, 42,000 of 200,000 tokens",
+    );
+  });
+
+  it("shows the post-compaction unknown state and hides without usage", () => {
+    const { rerender } = render(
+      <TooltipProvider>
+        <ContextUsageIndicator
+          usage={{ tokens: null, contextWindow: 200_000, percent: null }}
+        />
+      </TooltipProvider>,
+    );
+    expect(screen.getByRole("status")).toHaveTextContent("?%");
+
+    rerender(
+      <TooltipProvider>
+        <ContextUsageIndicator usage={null} />
+      </TooltipProvider>,
+    );
+    expect(screen.queryByRole("status")).not.toBeInTheDocument();
+  });
+});

web/src/components/context-usage-indicator.tsx

@@ -0,0 +1,61 @@
+"use client";
+
+import { GaugeIcon } from "lucide-react";
+
+import { InfoTooltip } from "@/components/ui/info-tooltip";
+import { cn, formatCompactTokens } from "@/lib/utils";
+import type { ContextUsage } from "@/lib/use-agent";
+
+export function ContextUsageIndicator({ usage }: { usage: ContextUsage | null }) {
+  if (!usage) return null;
+
+  const known = usage.tokens !== null && usage.percent !== null;
+  const tokens = usage.tokens ?? 0;
+  const percent = usage.percent ?? 0;
+  const fill = Math.min(100, Math.max(0, percent));
+  const critical = known && percent > 90;
+  const warning = known && percent > 70;
+  const value = known ? `${Math.round(percent)}%` : "?%";
+  const ariaLabel = known
+    ? `Model context ${percent.toFixed(1)} percent, ${tokens.toLocaleString()} of ${usage.contextWindow.toLocaleString()} tokens`
+    : `Model context usage recalculating, ${usage.contextWindow.toLocaleString()} token window`;
+
+  return (
+    <InfoTooltip
+      content={
+        <>
+          <b>Model context</b>
+          <br />
+          {known
+            ? `${formatCompactTokens(tokens)} of ${formatCompactTokens(usage.contextWindow)} tokens (${percent.toFixed(1)}%)`
+            : `Recalculating after compaction · ${formatCompactTokens(usage.contextWindow)} token window`}
+          <br />
+          Pi estimates how much of the selected model&apos;s context window this
+          conversation currently uses.
+        </>
+      }
+    >
+      <span
+        role="status"
+        aria-label={ariaLabel}
+        className={cn(
+          "inline-flex h-7 cursor-help items-center gap-1.5 rounded-md px-2 font-mono text-[11px] tabular-nums text-muted-foreground",
+          critical && "bg-destructive/10 text-destructive",
+          warning && !critical && "bg-amber-500/10 text-amber-700 dark:text-amber-400",
+        )}
+      >
+        <GaugeIcon className="size-3.5 shrink-0" aria-hidden />
+        <span>{value}</span>
+        <span className="h-1 w-6 overflow-hidden rounded-full bg-muted" aria-hidden>
+          <span
+            className={cn(
+              "block h-full rounded-full transition-[width]",
+              critical ? "bg-destructive" : warning ? "bg-amber-500" : "bg-primary",
+            )}
+            style={{ width: `${fill}%` }}
+          />
+        </span>
+      </span>
+    </InfoTooltip>
+  );
+}

web/src/lib/use-agent.test.ts

@@ -1,7 +1,7 @@
 import { act, renderHook } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import * as projects from "@/lib/projects";
-import { buildRunBody, useAgent } from "@/lib/use-agent";
+import { buildRunBody, parseContextUsage, useAgent } from "@/lib/use-agent";
 
 /** Build an SSE response body streaming one `data: <json>\n\n` frame per entry. */
 function sseStream(frames: unknown[]): ReadableStream<Uint8Array> {
@@ -56,6 +56,22 @@ describe("buildRunBody", () => {
   });
 });
 
+describe("parseContextUsage", () => {
+  it("accepts known and post-compaction Pi usage", () => {
+    expect(
+      parseContextUsage({ tokens: 42_000, contextWindow: 200_000, percent: 21 }),
+    ).toEqual({ tokens: 42_000, contextWindow: 200_000, percent: 21 });
+    expect(
+      parseContextUsage({ tokens: null, contextWindow: 200_000, percent: null }),
+    ).toEqual({ tokens: null, contextWindow: 200_000, percent: null });
+  });
+
+  it("rejects malformed usage", () => {
+    expect(parseContextUsage({ tokens: -1, contextWindow: 200_000, percent: -1 })).toBeNull();
+    expect(parseContextUsage({ tokens: 1, contextWindow: 0, percent: 1 })).toBeNull();
+  });
+});
+
 describe("useAgent notebook accumulation", () => {
   afterEach(() => vi.restoreAllMocks());
 
@@ -87,4 +103,37 @@ describe("useAgent notebook accumulation", () => {
 
     expect(result.current.notebookEntries.map((e) => e.id)).toEqual(["tc_1"]);
   });
+
+  it("tracks Pi context utilization from the SSE stream", async () => {
+    vi.spyOn(projects, "apiFetch").mockImplementation(async (path: string) => {
+      if (path === "/sessions") {
+        return new Response(JSON.stringify({ id: "s1" }), { status: 200 });
+      }
+      if (path === "/sessions/s1/run") {
+        return new Response(
+          sseStream([
+            {
+              type: "context_usage",
+              tokens: 42_000,
+              contextWindow: 200_000,
+              percent: 21,
+            },
+          ]),
+          { status: 200 },
+        );
+      }
+      throw new Error(`unexpected apiFetch path: ${path}`);
+    });
+
+    const { result } = renderHook(() => useAgent());
+    await act(async () => {
+      await result.current.send("hi");
+    });
+
+    expect(result.current.contextUsage).toEqual({
+      tokens: 42_000,
+      contextWindow: 200_000,
+      percent: 21,
+    });
+  });
 });