Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
benvinegar merged 1 commit into from
Feb 23, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CONFIGURATION.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
Expand Down
100 changes: 85 additions & 15 deletions pi/extensions/context.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot Feb 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inconsistent async handling - tool_result uses void (fire-and-forget) while turn_end and session_start use await

Suggested change
void persistSnapshotFromContext(ctx);
await persistSnapshotFromContext(ctx);

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Prompt To Fix With AI 
This is a comment left during a code review.
Path: pi/extensions/context.ts
Line: 523

Comment:
Inconsistent async handling - `tool_result` uses `void` (fire-and-forget) while `turn_end` and `session_start` use `await`

```suggestion
		await persistSnapshotFromContext(ctx);
```

<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>

How can I resolve this? If you propose a fix, please make it concise.


// Only count successful reads.
if ((event as any).toolName !== "read") return;
if ((event as any).isError) return;
Expand All @@ -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) => {
Expand Down Expand Up @@ -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;
Expand All @@ -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)"}`);
Expand All @@ -559,7 +629,7 @@ export default function contextExtension(pi: ExtensionAPI) {
systemPromptTokens,
agentTokens,
toolsTokens,
activeTools: activeToolNames.length,
activeTools,
}
: null,
agentFiles: agentFilePaths,
Expand Down
25 changes: 25 additions & 0 deletions slack-bridge/broker-bridge.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [];
Expand Down Expand Up @@ -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,
Expand All @@ -246,6 +270,7 @@ function buildPullMeta(maxMessages, waitSeconds) {
poll_count: brokerPollCount + 1,
max_messages: maxMessages,
wait_seconds: waitSeconds,
...(contextUsage ? contextUsage : {}),
};
}

Expand Down
25 changes: 24 additions & 1 deletion test/broker-bridge.integration.test.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 = "";
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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,
Expand Down