diff --git a/CONFIGURATION.md b/CONFIGURATION.md index df95dab..27c190a 100644 --- 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 index b29c035..ef890f5 100644 --- 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 { + 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 => { + 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 index 55268d8..61d8c89 100755 --- 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 index 798a508..549d2b5 100644 --- 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,