feat(coding-agent): expose post-compaction token estimates

closes #5877
This commit is contained in:
Armin Ronacher
2026-06-18 23:01:10 +02:00
parent b09fbde00e
commit c60f6a8ab3
5 changed files with 28 additions and 1 deletions

View File

@@ -2,6 +2,10 @@
## [Unreleased] ## [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
- Fixed compaction to refuse sessions with no eligible messages instead of producing empty summaries ([#4811](https://github.com/earendil-works/pi/issues/4811)). - Fixed compaction to refuse sessions with no eligible messages instead of producing empty summaries ([#4811](https://github.com/earendil-works/pi/issues/4811)).

View File

@@ -374,11 +374,14 @@ Response:
"summary": "Summary of conversation...", "summary": "Summary of conversation...",
"firstKeptEntryId": "abc123", "firstKeptEntryId": "abc123",
"tokensBefore": 150000, "tokensBefore": 150000,
"estimatedTokensAfter": 32000,
"details": {} "details": {}
} }
} }
``` ```
`estimatedTokensAfter` is a heuristic estimate over the rebuilt message context immediately after compaction, not a provider-exact token count.
#### set_auto_compaction #### set_auto_compaction
Enable or disable automatic compaction when context is nearly full. 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...", "summary": "Summary of conversation...",
"firstKeptEntryId": "abc123", "firstKeptEntryId": "abc123",
"tokensBefore": 150000, "tokensBefore": 150000,
"estimatedTokensAfter": 32000,
"details": {} "details": {}
}, },
"aborted": false, "aborted": false,

View File

@@ -45,6 +45,7 @@ import {
collectEntriesForBranchSummary, collectEntriesForBranchSummary,
compact, compact,
estimateContextTokens, estimateContextTokens,
estimateTokens,
generateBranchSummary, generateBranchSummary,
prepareCompaction, prepareCompaction,
shouldCompact, shouldCompact,
@@ -242,6 +243,14 @@ interface ToolDefinitionEntry {
sourceInfo: SourceInfo; sourceInfo: SourceInfo;
} }
function estimateMessagesTokens(messages: AgentMessage[]): number {
let tokens = 0;
for (const message of messages) {
tokens += estimateTokens(message);
}
return tokens;
}
// ============================================================================ // ============================================================================
// Constants // Constants
// ============================================================================ // ============================================================================
@@ -1726,6 +1735,7 @@ export class AgentSession {
const newEntries = this.sessionManager.getEntries(); const newEntries = this.sessionManager.getEntries();
const sessionContext = this.sessionManager.buildSessionContext(); const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages; this.agent.state.messages = sessionContext.messages;
const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages);
// Get the saved compaction entry for the extension event // Get the saved compaction entry for the extension event
const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as 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, summary,
firstKeptEntryId, firstKeptEntryId,
tokensBefore, tokensBefore,
estimatedTokensAfter,
details, details,
}; };
this._emit({ this._emit({
@@ -1999,6 +2010,7 @@ export class AgentSession {
const newEntries = this.sessionManager.getEntries(); const newEntries = this.sessionManager.getEntries();
const sessionContext = this.sessionManager.buildSessionContext(); const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages; this.agent.state.messages = sessionContext.messages;
const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages);
// Get the saved compaction entry for the extension event // Get the saved compaction entry for the extension event
const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
@@ -2017,6 +2029,7 @@ export class AgentSession {
summary, summary,
firstKeptEntryId, firstKeptEntryId,
tokensBefore, tokensBefore,
estimatedTokensAfter,
details, details,
}; };
this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });

View File

@@ -104,6 +104,7 @@ export interface CompactionResult<T = unknown> {
summary: string; summary: string;
firstKeptEntryId: string; firstKeptEntryId: string;
tokensBefore: number; tokensBefore: number;
estimatedTokensAfter?: number;
/** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */
details?: T; details?: T;
} }

View File

@@ -5,6 +5,7 @@ import {
type Model, type Model,
} from "@earendil-works/pi-ai"; } from "@earendil-works/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest"; import { afterEach, describe, expect, it, vi } from "vitest";
import { estimateTokens } from "../../src/core/compaction/index.ts";
import { createHarness, type Harness } from "./harness.ts"; import { createHarness, type Harness } from "./harness.ts";
type SessionWithCompactionInternals = { type SessionWithCompactionInternals = {
@@ -118,8 +119,10 @@ describe("AgentSession compaction characterization", () => {
const result = await harness.session.compact(); const result = await harness.session.compact();
const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); 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.summary).toBe("summary from extension");
expect(result.estimatedTokensAfter).toBe(estimatedTokensAfter);
expect(compactionEntries).toHaveLength(1); expect(compactionEntries).toHaveLength(1);
expect(harness.session.messages[0]?.role).toBe("compactionSummary"); expect(harness.session.messages[0]?.role).toBe("compactionSummary");
}); });
@@ -161,7 +164,9 @@ describe("AgentSession compaction characterization", () => {
await sessionInternals._runAutoCompaction("threshold", false); await sessionInternals._runAutoCompaction("threshold", false);
const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction");
const compactionEnd = harness.eventsOfType("compaction_end").at(-1);
expect(compactionEntries).toHaveLength(1); expect(compactionEntries).toHaveLength(1);
expect(compactionEnd?.result?.estimatedTokensAfter).toBeGreaterThan(0);
expect(getStreamCallCount()).toBe(1); expect(getStreamCallCount()).toBe(1);
}); });