bridge: emit context-window stats in inbox pull meta (#150)

benvinegar · web-flow · commit 94790e3df233 · 2026-02-23T16:47:22.000-05:00
diff --git a/CONFIGURATION.md b/CONFIGURATION.md
@@ -113,6 +113,8 @@ Set by `sudo baudbot broker register` when using brokered Slack OAuth flow.
 | `SLACK_BROKER_DEDUPE_TTL_MS` | Dedupe cache TTL in milliseconds (default: `1200000`) |
 | `BAUDBOT_AGENT_VERSION` | Optional override for broker observability `meta.agent_version` (otherwise read from `~/.pi/agent/baudbot-version.json` when available) |
 
+Broker mode also emits best-effort context usage telemetry in inbox pull `meta` by reading `~/.pi/agent/context-usage.json` (written by the `context` extension on session start/turn end/tool results).
+
 ### Kernel (Cloud Browsers)
 
 | Variable | Description | How to get it |
diff --git a/pi/extensions/context.ts b/pi/extensions/context.ts
@@ -169,7 +169,7 @@ function extractCostTotal(usage: any): number {
 	return 0;
 }
 
-function sumSessionUsage(ctx: ExtensionCommandContext): {
+function sumSessionUsage(ctx: ExtensionContext | ExtensionCommandContext): {
 	input: number;
 	output: number;
 	cacheRead: number;
@@ -206,6 +206,71 @@ function sumSessionUsage(ctx: ExtensionCommandContext): {
 	};
 }
 
+function estimateToolDefinitionTokens(pi: ExtensionAPI): { toolsTokens: number; activeTools: number } {
+	const TOOL_FUDGE = 1.5;
+	const activeToolNames = pi.getActiveTools();
+	const toolInfoByName = new Map(pi.getAllTools().map((t) => [t.name, t] as const));
+	let toolsTokens = 0;
+	for (const name of activeToolNames) {
+		const info = toolInfoByName.get(name);
+		const blob = `${name}\n${info?.description ?? ""}`;
+		toolsTokens += estimateTokens(blob);
+	}
+	toolsTokens = Math.round(toolsTokens * TOOL_FUDGE);
+	return { toolsTokens, activeTools: activeToolNames.length };
+}
+
+type ContextUsageSnapshot = {
+	generated_at: string;
+	session_id: string;
+	context_window_used_tokens?: number;
+	context_window_limit_tokens?: number;
+	context_window_used_pct?: number;
+	session_total_tokens: number;
+	session_total_cost_usd: number;
+	message_tokens?: number;
+	tools_tokens?: number;
+};
+
+function buildContextUsageSnapshot(pi: ExtensionAPI, ctx: ExtensionContext): ContextUsageSnapshot {
+	const usage = ctx.getContextUsage();
+	const { toolsTokens } = estimateToolDefinitionTokens(pi);
+	const messageTokens = usage?.tokens ?? 0;
+	const contextWindow = usage?.contextWindow ?? 0;
+	const effectiveTokens = Math.max(0, messageTokens + toolsTokens);
+	const percent = contextWindow > 0 ? (effectiveTokens / contextWindow) * 100 : 0;
+	const sessionUsage = sumSessionUsage(ctx);
+
+	const snapshot: ContextUsageSnapshot = {
+		generated_at: new Date().toISOString(),
+		session_id: ctx.sessionManager.getSessionId(),
+		session_total_tokens: Math.max(0, sessionUsage.totalTokens),
+		session_total_cost_usd: Math.max(0, sessionUsage.totalCost),
+	};
+
+	if (contextWindow > 0) {
+		snapshot.context_window_used_tokens = effectiveTokens;
+		snapshot.context_window_limit_tokens = contextWindow;
+		snapshot.context_window_used_pct = percent;
+		snapshot.message_tokens = messageTokens;
+		snapshot.tools_tokens = toolsTokens;
+	}
+
+	return snapshot;
+}
+
+async function persistContextUsageSnapshot(snapshot: ContextUsageSnapshot): Promise<void> {
+	const snapshotPath = path.join(os.homedir(), ".pi", "agent", "context-usage.json");
+	const tmpPath = `${snapshotPath}.tmp`;
+	try {
+		await fs.mkdir(path.dirname(snapshotPath), { recursive: true });
+		await fs.writeFile(tmpPath, `${JSON.stringify(snapshot, null, 2)}\n`, { mode: 0o600 });
+		await fs.rename(tmpPath, snapshotPath);
+	} catch {
+		// Best-effort only. Observability should not disrupt agent runtime.
+	}
+}
+
 function shortenPath(p: string, cwd: string): string {
 	const rp = path.resolve(p);
 	const rc = path.resolve(cwd);
@@ -449,7 +514,14 @@ export default function contextExtension(pi: ExtensionAPI) {
 		return best?.name ?? null;
 	};
 
+	const persistSnapshotFromContext = async (ctx: ExtensionContext): Promise<void> => {
+		const snapshot = buildContextUsageSnapshot(pi, ctx);
+		await persistContextUsageSnapshot(snapshot);
+	};
+
 	pi.on("tool_result", (event: ToolResultEvent, ctx: ExtensionContext) => {
+		void persistSnapshotFromContext(ctx);
+
 		// Only count successful reads.
 		if ((event as any).toolName !== "read") return;
 		if ((event as any).isError) return;
@@ -469,6 +541,14 @@ export default function contextExtension(pi: ExtensionAPI) {
 		}
 	});
 
+	pi.on("turn_end", async (_event, ctx: ExtensionContext) => {
+		await persistSnapshotFromContext(ctx);
+	});
+
+	pi.on("session_start", async (_event, ctx: ExtensionContext) => {
+		await persistSnapshotFromContext(ctx);
+	});
+
 	pi.registerCommand("context", {
 		description: "Show loaded context overview",
 		handler: async (_args, ctx: ExtensionCommandContext) => {
@@ -503,18 +583,8 @@ export default function contextExtension(pi: ExtensionAPI) {
 			const ctxWindow = usage?.contextWindow ?? 0;
 
 			// Tool definitions are not part of ctx.getContextUsage() (it estimates message tokens).
-			// We approximate their token impact from tool name + description, and apply a fudge
-			// factor to account for parameters/schema/formatting.
-			const TOOL_FUDGE = 1.5;
-			const activeToolNames = pi.getActiveTools();
-			const toolInfoByName = new Map(pi.getAllTools().map((t) => [t.name, t] as const));
-			let toolsTokens = 0;
-			for (const name of activeToolNames) {
-				const info = toolInfoByName.get(name);
-				const blob = `${name}\n${info?.description ?? ""}`;
-				toolsTokens += estimateTokens(blob);
-			}
-			toolsTokens = Math.round(toolsTokens * TOOL_FUDGE);
+			// We approximate their token impact from tool name + description.
+			const { toolsTokens, activeTools } = estimateToolDefinitionTokens(pi);
 
 			const effectiveTokens = messageTokens + toolsTokens;
 			const percent = ctxWindow > 0 ? (effectiveTokens / ctxWindow) * 100 : 0;
@@ -533,7 +603,7 @@ export default function contextExtension(pi: ExtensionAPI) {
 					lines.push("Window: (unknown)");
 				}
 				lines.push(`System: ~${systemPromptTokens.toLocaleString()} tok (AGENTS ~${agentTokens.toLocaleString()})`);
-				lines.push(`Tools: ~${toolsTokens.toLocaleString()} tok (${activeToolNames.length} active)`);
+				lines.push(`Tools: ~${toolsTokens.toLocaleString()} tok (${activeTools} active)`);
 				lines.push(`AGENTS: ${agentFilePaths.length ? joinComma(agentFilePaths) : "(none)"}`);
 				lines.push(`Extensions (${extensionFiles.length}): ${extensionFiles.length ? joinComma(extensionFiles) : "(none)"}`);
 				lines.push(`Skills (${skills.length}): ${skills.length ? joinComma(skills) : "(none)"}`);
@@ -559,7 +629,7 @@ export default function contextExtension(pi: ExtensionAPI) {
 						systemPromptTokens,
 						agentTokens,
 						toolsTokens,
-						activeTools: activeToolNames.length,
+						activeTools,
 					}
 					: null,
 				agentFiles: agentFilePaths,
diff --git a/slack-bridge/broker-bridge.mjs b/slack-bridge/broker-bridge.mjs
@@ -53,6 +53,7 @@ const MAX_BACKOFF_MS = 30_000;
 const INBOX_PROTOCOL_VERSION = "2026-02-1";
 const BROKER_HEALTH_PATH = path.join(homedir(), ".pi", "agent", "broker-health.json");
 const BAUDBOT_VERSION_PATH = path.join(homedir(), ".pi", "agent", "baudbot-version.json");
+const CONTEXT_USAGE_PATH = path.join(homedir(), ".pi", "agent", "context-usage.json");
 const LOG_BUFFER_MAX_LINES = 1000;
 
 const logLineBuffer = [];
@@ -228,10 +229,33 @@ function countActivePiSessions() {
   }
 }
 
+function readContextUsageSnapshot() {
+  try {
+    const raw = fs.readFileSync(CONTEXT_USAGE_PATH, "utf8");
+    const parsed = JSON.parse(raw);
+    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return null;
+
+    const readFiniteNumber = (value) => (typeof value === "number" && Number.isFinite(value) ? value : null);
+    const snapshot = {
+      context_window_used_tokens: readFiniteNumber(parsed.context_window_used_tokens),
+      context_window_limit_tokens: readFiniteNumber(parsed.context_window_limit_tokens),
+      context_window_used_pct: readFiniteNumber(parsed.context_window_used_pct),
+      session_total_tokens: readFiniteNumber(parsed.session_total_tokens),
+      session_total_cost_usd: readFiniteNumber(parsed.session_total_cost_usd),
+    };
+
+    const hasAny = Object.values(snapshot).some((value) => value !== null);
+    return hasAny ? snapshot : null;
+  } catch {
+    return null;
+  }
+}
+
 function buildPullMeta(maxMessages, waitSeconds) {
   const { activeSessions, activeDevAgents } = countActivePiSessions();
   const bridgeUptimeHours = Math.max(0, (Date.now() - bridgeStartedAtMs) / (1000 * 60 * 60));
   const systemUptimeHours = Math.max(0, getSystemUptimeSeconds() / (60 * 60));
+  const contextUsage = readContextUsageSnapshot();
 
   return {
     agent_version: agentVersion,
@@ -246,6 +270,7 @@ function buildPullMeta(maxMessages, waitSeconds) {
     poll_count: brokerPollCount + 1,
     max_messages: maxMessages,
     wait_seconds: waitSeconds,
+    ...(contextUsage ? contextUsage : {}),
   };
 }
 
diff --git a/test/broker-bridge.integration.test.mjs b/test/broker-bridge.integration.test.mjs
@@ -4,7 +4,7 @@ import { spawn } from "node:child_process";
 import net from "node:net";
 import path from "node:path";
 import { fileURLToPath } from "node:url";
-import { mkdtempSync, mkdirSync, rmSync } from "node:fs";
+import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import sodium from "libsodium-wrappers-sumo";
 import {
@@ -502,6 +502,23 @@ describe("broker pull bridge semi-integration", () => {
     const signKeypair = sodium.crypto_sign_seed_keypair(new Uint8Array(signingSeed));
     let pullPayload = null;
 
+    const tempHome = mkdtempSync(path.join(tmpdir(), "baudbot-broker-test-"));
+    tempDirs.push(tempHome);
+    const contextUsageDir = path.join(tempHome, ".pi", "agent");
+    mkdirSync(contextUsageDir, { recursive: true });
+    writeFileSync(
+      path.join(contextUsageDir, "context-usage.json"),
+      `${JSON.stringify({
+        generated_at: "2026-02-23T00:00:00.000Z",
+        session_id: "session-test",
+        context_window_used_tokens: 12345,
+        context_window_limit_tokens: 200000,
+        context_window_used_pct: 6.1725,
+        session_total_tokens: 54321,
+        session_total_cost_usd: 1.25,
+      }, null, 2)}\n`,
+    );
+
     const broker = createServer(async (req, res) => {
       if (req.method === "POST" && req.url === "/api/inbox/pull") {
         let raw = "";
@@ -547,6 +564,7 @@ describe("broker pull bridge semi-integration", () => {
       cwd: bridgeCwd,
       env: {
         ...cleanEnv(),
+        HOME: tempHome,
         SLACK_BROKER_URL: brokerUrl,
         SLACK_BROKER_WORKSPACE_ID: workspaceId,
         SLACK_BROKER_SERVER_PRIVATE_KEY: b64(32, 11),
@@ -581,6 +599,11 @@ describe("broker pull bridge semi-integration", () => {
     expect(typeof pullPayload.meta.agent_version).toBe("string");
     expect(pullPayload.meta.heartbeat_runs).toBeGreaterThanOrEqual(0);
     expect(pullPayload.meta.heartbeat_consecutive_errors).toBeGreaterThanOrEqual(0);
+    expect(pullPayload.meta.context_window_used_tokens).toBe(12345);
+    expect(pullPayload.meta.context_window_limit_tokens).toBe(200000);
+    expect(pullPayload.meta.context_window_used_pct).toBe(6.1725);
+    expect(pullPayload.meta.session_total_tokens).toBe(54321);
+    expect(pullPayload.meta.session_total_cost_usd).toBe(1.25);
 
     const canonical = canonicalizeProtocolRequest(workspaceId, "2026-02-1", "inbox.pull", pullPayload.timestamp, {
       max_messages: 10,
