From 835296b1d7d999100ee1c465b4958031d0bb5adc Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 24 Mar 2026 03:08:51 +0100 Subject: [PATCH] fix(coding-agent): expose rpc context usage closes #2550 --- packages/coding-agent/CHANGELOG.md | 4 + packages/coding-agent/docs/rpc.md | 13 +- .../coding-agent/src/core/agent-session.ts | 2 + .../test/agent-session-stats.test.ts | 143 ++++++++++++++++++ 4 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 packages/coding-agent/test/agent-session-stats.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 612414a9..3b0d3d96 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed RPC `get_session_stats` to expose `contextUsage`, so headless clients can read actual current context-window usage instead of deriving it from token totals ([#2550](https://github.com/badlogic/pi-mono/issues/2550)) + ## [0.62.0] - 2026-03-23 ### New Features diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 53350cdd..162ffd94 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -494,7 +494,7 @@ Response: #### get_session_stats -Get token usage and cost statistics. +Get token usage, cost statistics, and current context window usage. ```json {"type": "get_session_stats"} @@ -521,11 +521,20 @@ Response: "cacheWrite": 5000, "total": 105000 }, - "cost": 0.45 + "cost": 0.45, + "contextUsage": { + "tokens": 60000, + "contextWindow": 200000, + "percent": 30 + } } } ``` +`tokens` contains assistant usage totals for the current session state. `contextUsage` contains the actual current context-window estimate used for compaction and footer display. + +`contextUsage` is omitted when no model or context window is available. `contextUsage.tokens` and `contextUsage.percent` are `null` immediately after compaction until a fresh post-compaction assistant response provides valid usage data. + #### export_html Export session to an HTML file. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 31534b25..09271db4 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -200,6 +200,7 @@ export interface SessionStats { total: number; }; cost: number; + contextUsage?: ContextUsage; } interface ToolDefinitionEntry { @@ -3029,6 +3030,7 @@ export class AgentSession { total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, }, cost: totalCost, + contextUsage: this.getContextUsage(), }; } diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts new file mode 100644 index 00000000..d9cd192b --- /dev/null +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -0,0 +1,143 @@ +import { Agent } from "@mariozechner/pi-agent-core"; +import { type AssistantMessage, getModel, type Usage } from "@mariozechner/pi-ai"; +import { describe, expect, it } from "vitest"; +import { AgentSession } from "../src/core/agent-session.js"; +import { AuthStorage } from "../src/core/auth-storage.js"; +import { ModelRegistry } from "../src/core/model-registry.js"; +import { SessionManager } from "../src/core/session-manager.js"; +import { SettingsManager } from "../src/core/settings-manager.js"; +import { createTestResourceLoader } from "./utilities.js"; + +const model = getModel("anthropic", "claude-sonnet-4-5")!; + +function createUsage(totalTokens: number): Usage { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }; +} + +function createAssistantMessage(text: string, totalTokens: number, timestamp: number): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(totalTokens), + stopReason: "stop", + timestamp, + }; +} + +function createUserMessage(text: string, timestamp: number) { + return { + role: "user" as const, + content: text, + timestamp, + }; +} + +function createSession() { + const settingsManager = SettingsManager.inMemory(); + const sessionManager = SessionManager.inMemory(); + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + const session = new AgentSession({ + agent: new Agent({ + getApiKey: () => "test-key", + initialState: { + model, + systemPrompt: "You are a helpful assistant.", + tools: [], + thinkingLevel: "high", + }, + }), + sessionManager, + settingsManager, + cwd: process.cwd(), + modelRegistry: new ModelRegistry(authStorage, undefined), + resourceLoader: createTestResourceLoader(), + }); + + return { session, sessionManager }; +} + +function syncAgentMessages(session: AgentSession, sessionManager: SessionManager): void { + session.agent.replaceMessages(sessionManager.buildSessionContext().messages); +} + +describe("AgentSession.getSessionStats", () => { + it("exposes the current context usage alongside token totals", () => { + const { session, sessionManager } = createSession(); + + try { + sessionManager.appendMessage(createUserMessage("hello", 1)); + sessionManager.appendMessage(createAssistantMessage("hi", 200, 2)); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.contextUsage).toEqual(session.getContextUsage()); + expect(stats.contextUsage?.tokens).toBe(200); + expect(stats.contextUsage?.contextWindow).toBe(model.contextWindow); + expect(stats.contextUsage?.percent).toBe((200 / model.contextWindow) * 100); + } finally { + session.dispose(); + } + }); + + it("reports unknown current context usage immediately after compaction", () => { + const { session, sessionManager } = createSession(); + + try { + sessionManager.appendMessage(createUserMessage("first", 1)); + sessionManager.appendMessage(createAssistantMessage("response1", 180_000, 2)); + const keptUserId = sessionManager.appendMessage(createUserMessage("second", 3)); + sessionManager.appendMessage(createAssistantMessage("response2", 195_000, 4)); + sessionManager.appendCompaction("summary", keptUserId, 195_000); + sessionManager.appendMessage(createUserMessage("third", 5)); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.tokens.input).toBe(195_000); + expect(stats.contextUsage).toBeDefined(); + expect(stats.contextUsage?.tokens).toBeNull(); + expect(stats.contextUsage?.percent).toBeNull(); + } finally { + session.dispose(); + } + }); + + it("uses post-compaction usage for current context instead of stale kept usage", () => { + const { session, sessionManager } = createSession(); + + try { + sessionManager.appendMessage(createUserMessage("first", 1)); + sessionManager.appendMessage(createAssistantMessage("response1", 180_000, 2)); + const keptUserId = sessionManager.appendMessage(createUserMessage("second", 3)); + sessionManager.appendMessage(createAssistantMessage("response2", 195_000, 4)); + sessionManager.appendCompaction("summary", keptUserId, 195_000); + sessionManager.appendMessage(createUserMessage("third", 5)); + sessionManager.appendMessage(createAssistantMessage("response3", 25_000, 6)); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.tokens.input).toBe(220_000); + expect(stats.contextUsage).toBeDefined(); + expect(stats.contextUsage?.tokens).toBe(25_000); + expect(stats.contextUsage?.percent).toBe((25_000 / model.contextWindow) * 100); + } finally { + session.dispose(); + } + }); +});