feat: harden workflow gates and expose token usage

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