#14913: fix: update context pruning to notify session metadata after pruning & fix Usage type #14833
agents
stale
size: S
Cluster:
Session Pruning Enhancements
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
#14879: fix: persist session metadata to sessions.json after context pruning
by skylarkoo7 · 2026-02-12
89.6%
#12581: feat(hooks): emit session prune lifecycle event
by vincentkoc · 2026-02-09
80.1%
#19878: fix: Handle compaction when fallback model has smaller context window
by gaurav10gg · 2026-02-18
80.0%
#5360: fix(compaction): add emergency pruning for context overflow
by sgwannabe · 2026-01-31
79.3%
#10567: feat(agents): add configurable startup session pruning
by abutlabs · 2026-02-06
79.1%
#10465: Context pruning: strip image blocks instead of skipping
by quentintou · 2026-02-06
79.1%
#11999: fix: add session-growth guard to prevent unbounded session store gr...
by reverendrewind · 2026-02-08
78.7%
#5057: fix: add per-tool-result hard cap to prevent context overflow
by hanxiao · 2026-01-31
76.9%
#10997: fix: enable cache-ttl pruning on first load after restart
by anotb · 2026-02-07
76.8%
#15726: fix(sessions): use model contextWindow instead of agent contextToke...
by lailoo · 2026-02-13
76.7%