From c60f6a8ab3b36155d572fe9c10df61f7613281ee Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Thu, 18 Jun 2026 23:01:10 +0200 Subject: [PATCH] feat(coding-agent): expose post-compaction token estimates closes #5877 --- packages/coding-agent/CHANGELOG.md | 4 ++++ packages/coding-agent/docs/rpc.md | 4 ++++ packages/coding-agent/src/core/agent-session.ts | 15 ++++++++++++++- .../src/core/compaction/compaction.ts | 1 + .../test/suite/agent-session-compaction.test.ts | 5 +++++ 5 files changed, 28 insertions(+), 1 deletion(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 0fcc303a..e58cd877 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added estimated post-compaction token counts to compact results and compaction events ([#5877](https://github.com/earendil-works/pi/issues/5877)). + ### Fixed - Fixed compaction to refuse sessions with no eligible messages instead of producing empty summaries ([#4811](https://github.com/earendil-works/pi/issues/4811)). diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 9aa16ffc..a9942409 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -374,11 +374,14 @@ Response: "summary": "Summary of conversation...", "firstKeptEntryId": "abc123", "tokensBefore": 150000, + "estimatedTokensAfter": 32000, "details": {} } } ``` +`estimatedTokensAfter` is a heuristic estimate over the rebuilt message context immediately after compaction, not a provider-exact token count. + #### set_auto_compaction Enable or disable automatic compaction when context is nearly full. @@ -924,6 +927,7 @@ The `reason` field is `"manual"`, `"threshold"`, or `"overflow"`. "summary": "Summary of conversation...", "firstKeptEntryId": "abc123", "tokensBefore": 150000, + "estimatedTokensAfter": 32000, "details": {} }, "aborted": false, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index cae330f3..f8af4320 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -45,6 +45,7 @@ import { collectEntriesForBranchSummary, compact, estimateContextTokens, + estimateTokens, generateBranchSummary, prepareCompaction, shouldCompact, @@ -242,6 +243,14 @@ interface ToolDefinitionEntry { sourceInfo: SourceInfo; } +function estimateMessagesTokens(messages: AgentMessage[]): number { + let tokens = 0; + for (const message of messages) { + tokens += estimateTokens(message); + } + return tokens; +} + // ============================================================================ // Constants // ============================================================================ @@ -1726,6 +1735,7 @@ export class AgentSession { const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; + const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); // Get the saved compaction entry for the extension event const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as @@ -1740,10 +1750,11 @@ export class AgentSession { }); } - const compactionResult = { + const compactionResult: CompactionResult = { summary, firstKeptEntryId, tokensBefore, + estimatedTokensAfter, details, }; this._emit({ @@ -1999,6 +2010,7 @@ export class AgentSession { const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; + const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); // Get the saved compaction entry for the extension event const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as @@ -2017,6 +2029,7 @@ export class AgentSession { summary, firstKeptEntryId, tokensBefore, + estimatedTokensAfter, details, }; this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index f3f3d59f..83315a84 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -104,6 +104,7 @@ export interface CompactionResult { summary: string; firstKeptEntryId: string; tokensBefore: number; + estimatedTokensAfter?: number; /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; } diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index a3302adf..ba5ae5c7 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -5,6 +5,7 @@ import { type Model, } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; +import { estimateTokens } from "../../src/core/compaction/index.ts"; import { createHarness, type Harness } from "./harness.ts"; type SessionWithCompactionInternals = { @@ -118,8 +119,10 @@ describe("AgentSession compaction characterization", () => { const result = await harness.session.compact(); const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + const estimatedTokensAfter = harness.session.messages.reduce((sum, message) => sum + estimateTokens(message), 0); expect(result.summary).toBe("summary from extension"); + expect(result.estimatedTokensAfter).toBe(estimatedTokensAfter); expect(compactionEntries).toHaveLength(1); expect(harness.session.messages[0]?.role).toBe("compactionSummary"); }); @@ -161,7 +164,9 @@ describe("AgentSession compaction characterization", () => { await sessionInternals._runAutoCompaction("threshold", false); const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + const compactionEnd = harness.eventsOfType("compaction_end").at(-1); expect(compactionEntries).toHaveLength(1); + expect(compactionEnd?.result?.estimatedTokensAfter).toBeGreaterThan(0); expect(getStreamCallCount()).toBe(1); });