fix(coding-agent): expose rpc context usage closes #2550

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

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.

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(),
 		};
 	}
 

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();
+		}
+	});
+});