← Back to PRs

#14913: fix: update context pruning to notify session metadata after pruning & fix Usage type #14833

by ScreenTechnicals open 2026-02-12 19:22 View on GitHub →
agents stale size: S
Updates the context pruning extension to report estimated token counts after pruning, allowing session metadata (e.g. sessions.json) to reflect accurate contextTokens immediately. ### Changes - context-pruning/runtime.ts — Added optional onPruned - callback to ContextPruningRuntimeValue, invoked with the estimated token count after successful pruning. - context-pruning/pruner.ts — Exported estimateContextChars and CHARS_PER_TOKEN_ESTIMATE so the extension can compute post-prune token estimates. - context-pruning/extension.ts — Call runtime.onPruned(estimatedTokens) after pruning succeeds, before resetting the TTL window. - pi-embedded-runner/extensions.ts — Wire up onPruned to update session context token metadata. - pi-embedded-runner/run/attempt.ts — Pass sessionKey through to the extension setup. context-pruning.test.ts - Fixed Usage type (total → totalTokens + cost object) to match the updated @mariozechner/pi-ai interface. Added tests for onPruned - callback behavior (called on prune, not called when no pruning occurs). <!-- greptile_comment --> <h2>Greptile Overview</h2> <h3>Greptile Summary</h3> This PR extends the context-pruning runtime with an `onPruned` callback, exports character/token estimation helpers from the pruner, and wires the embedded runner to persist updated token metadata to the session store immediately after pruning. It also updates tests to the new `Usage` shape (`totalTokens` + `cost`) and adds coverage ensuring `onPruned` fires only when pruning actually changes the message list. <h3>Confidence Score: 3/5</h3> - This PR is close to mergeable but has a high-likelihood correctness issue in how session token metadata is persisted after pruning. - Core pruning/onPruned wiring and tests look coherent, but the embedded runner writes the post-prune *context size estimate* into `totalTokens`/`inputTokens`/`outputTokens`, which are typically used for usage/billing totals; that will produce incorrect session metadata and can break consumers of usage fields. Fixing the field mapping should make the change low-risk. - src/agents/pi-embedded-runner/extensions.ts <!-- greptile_other_comments_section --> <sub>(2/5) Greptile learns from your feedback when you react with thumbs up/down!</sub> <!-- /greptile_comment -->

Most Similar PRs