feat: harden workflow gates and expose token usage

backend/src/bridge/session-bridge.ts

@@ -7,11 +7,20 @@ export type ToolActivity = {
   status: "running" | "completed" | "failed";
 };
 
+export type TokenUsage = {
+  input: number;
+  cacheRead: number;
+  output: number;
+  totalTokens: number;
+  contextWindow: number;
+};
+
 export type ClientEvent =
   | { type: "ui_request"; ui_request: any }
   | { type: "text_delta"; text: string }
   | { type: "activity_delta"; text: string }
   | { type: "activity_event"; activity: ToolActivity }
+  | { type: "usage"; usage: TokenUsage }
   | { type: "info"; [k: string]: any }
   | { type: "system_event"; event: string };
 
@@ -25,6 +34,7 @@ export class SessionBridge {
   // del descriptor (che viaggia opaco nel `title`). Va memorizzato e rimandato indietro,
   // altrimenti Pi scarta la risposta e ctx.ui.input non si risolve mai (stuck senza output).
   private pendingPiId: string | null = null;
+  private contextWindow = 0;
   private cbs = new Set<(e: ClientEvent) => void>();
 
   constructor(private rpc: RpcClient) {
@@ -36,17 +46,16 @@ export class SessionBridge {
         this.pendingPiId = m.id as string | null;
         this.state = "waiting";
         this.fan({ type: "ui_request", ui_request: descriptor });
-      } else if (
-        m.type === "message_end" &&
-        m.message?.role === "assistant" &&
-        m.message?.stopReason === "error"
-      ) {
-        this.markFailed();
-        this.fan({
-          type: "info",
-          level: "error",
-          text: "Model request failed. Check provider connectivity, then Resume the session.",
-        });
+      } else if (m.type === "message_end" && m.message?.role === "assistant") {
+        this.emitUsage(m.message.usage);
+        if (m.message.stopReason === "error") {
+          this.markFailed();
+          this.fan({
+            type: "info",
+            level: "error",
+            text: "Model request failed. Check provider connectivity, then Resume the session.",
+          });
+        }
       } else if (m.type === "extension_ui_request" && m.method === "notify") {
         this.fan({ type: "info", level: m.notifyType ?? "info", text: m.message ?? "" });
       } else if (m.type === "message_update" && m.assistantMessageEvent?.type === "text_delta") {
@@ -93,12 +102,38 @@ export class SessionBridge {
     });
   }
 
+  private emitUsage(usage: any): void {
+    if (!usage || this.contextWindow <= 0) return;
+    const number = (value: unknown): number =>
+      typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : 0;
+    const input = number(usage.input);
+    const cacheRead = number(usage.cacheRead);
+    const output = number(usage.output);
+    const reportedTotal = number(usage.totalTokens);
+    this.fan({
+      type: "usage",
+      usage: {
+        input,
+        cacheRead,
+        output,
+        totalTokens: reportedTotal || input + cacheRead + output,
+        contextWindow: this.contextWindow,
+      },
+    });
+  }
+
   private fan(e: ClientEvent) { for (const cb of this.cbs) cb(e); }
 
   turnState(): TurnState { return this.state; }
 
   beginTurn(): void { this.state = "running"; }
 
+  setContextWindow(value: unknown): void {
+    if (typeof value === "number" && Number.isFinite(value) && value > 0) {
+      this.contextWindow = value;
+    }
+  }
+
   /** Record a backend-detected failure; callers own any sanitized client message. */
   markFailed(): void { this.state = "failed"; }
 

backend/src/pi/pi-process-manager.ts

@@ -166,7 +166,10 @@ export class PiProcessManager {
     const thinking = o.thinking ?? this.cfg.defaults.thinking;
 
     if (provider && model) {
-      await rt.rpc.request({ type: "set_model", provider, modelId: model } as object & { type: string });
+      const response = await rt.rpc.request(
+        { type: "set_model", provider, modelId: model } as object & { type: string },
+      );
+      rt.bridge.setContextWindow(response?.data?.contextWindow);
     }
     if (thinking) {
       await rt.rpc.request({ type: "set_thinking_level", level: thinking } as object & { type: string });

backend/test/pi-process-manager.test.ts

@@ -207,6 +207,43 @@ test("createFor does not prompt until start is called", async () => {
   mgr.teardown("sid-deferred");
 });
 
+test("configure gives the bridge the context window returned by set_model", async () => {
+  const child = recordingChild();
+  child.stdin.write = (data: unknown) => {
+    const request = JSON.parse(String(data));
+    child._writes.push(String(data));
+    if (request.type === "set_model") {
+      queueMicrotask(() => child.stdout.emit("data", `${JSON.stringify({
+        type: "response",
+        id: request.id,
+        success: true,
+        data: { contextWindow: 200_000 },
+      })}\n`));
+    }
+    return true;
+  };
+  const mgr = new PiProcessManager(loadConfig({}), { spawnFn: () => child as any });
+  const rt = mgr.createFor("context-window", {});
+  const seen: any[] = [];
+  rt.bridge.onClientEvent((event) => seen.push(event));
+
+  await mgr.configure(rt, { provider: "zai", model: "glm-5.2" });
+  child.stdout.emit("data", `${JSON.stringify({
+    type: "message_end",
+    message: {
+      role: "assistant",
+      stopReason: "stop",
+      usage: { input: 10, cacheRead: 20, output: 5, totalTokens: 35 },
+    },
+  })}\n`);
+
+  expect(seen).toContainEqual({
+    type: "usage",
+    usage: { input: 10, cacheRead: 20, output: 5, totalTokens: 35, contextWindow: 200_000 },
+  });
+  mgr.teardown("context-window");
+});
+
 test("a created runtime is active during configure/bootstrap", () => {
   const child = recordingChild();
   const mgr = new PiProcessManager(loadConfig({}), { spawnFn: () => child as any });

backend/test/session-bridge.test.ts

@@ -42,6 +42,43 @@ test("real Pi thinking_delta becomes a dedicated activity_delta to the FE", () =
   expect(seen).toEqual([{ type: "activity_delta", text: "Valuto le ambiguità" }]);
 });
 
+test("assistant message_end exposes sanitized token usage with the configured context window", () => {
+  const { rpc, fire } = fakeRpc();
+  const bridge = new SessionBridge(rpc);
+  const seen: any[] = [];
+  bridge.onClientEvent((event) => seen.push(event));
+  bridge.setContextWindow(200_000);
+
+  fire({
+    type: "message_end",
+    message: {
+      role: "assistant",
+      stopReason: "stop",
+      usage: {
+        input: 110_348,
+        cacheRead: 633_344,
+        cacheWrite: 12,
+        output: 33_652,
+        totalTokens: 777_356,
+        cost: { total: 99 },
+      },
+      content: "DO_NOT_FORWARD",
+    },
+  });
+
+  expect(seen).toEqual([{
+    type: "usage",
+    usage: {
+      input: 110_348,
+      cacheRead: 633_344,
+      output: 33_652,
+      totalTokens: 777_356,
+      contextWindow: 200_000,
+    },
+  }]);
+  expect(JSON.stringify(seen)).not.toMatch(/DO_NOT_FORWARD|cost|cacheWrite/);
+});
+
 test("assistant provider errors are sanitized and leave the turn failed", () => {
   const { rpc, fire } = fakeRpc();
   const bridge = new SessionBridge(rpc);