🤖 refactor: move token budget orchestration behind context management

[ThomasK33]

Summary

Move Token Budget orchestration (admission, warnings, rollover, legacy flush recovery) out of AgentSession into a controller-owned TokenBudgetStrategy, then remove the forwarding shim so AgentSession holds no strategy-owned state. Behavior is unchanged. This finishes the extraction that #4266 started for Summarize and Continuous.

Layer 3 of 3. Most of the diff is relocation: strategies/tokenBudget.ts is +1,081 lines and agentSession.ts is +139/−1,037.

Background

AgentSession owned admission, the stream pipeline, and history, and it also decided which context-management strategy ran and kept each strategy's state. #4266 moved Summarize and Continuous behind a per-session SessionContextController (src/node/services/contextManagement/). Token Budget stayed in AgentSession behind a transitional forwarding shim.

Implementation

• contextManagement/strategies/tokenBudget.ts: the Token Budget strategy, moved from AgentSession with its admission, advisory, rollover, and recovery logic. It keeps the synchronous generation/admission receipts: validatePreparation(receipt) still runs immediately before dispatch with no intervening await.
• SessionContextController owns the strategy per session. AgentSession calls the controller's token-budget query at the existing evaluation points. The shim and the strategy-owned fields (isTokenBudgetActive, onContextBudgetStepSettled, budget state) are gone.
• SessionContextHost gains only the members the strategy needs; the session still owns publication, queue mutation, and the stream lifecycle.
• Automatic compaction request construction stays on the host because it reads live workspace settings and the shared provider/auth view.
• di/appRuntime.ts documents lifetimes: ContextManagement is app-lifetime, controllers are session-owned, and in-flight work holds leases through enterExecution.
• Characterization tests move to controller-private strategy seams with unchanged behavioral assertions, plus coverage for stale preparations and synchronous continuation ownership. The recovery-fence witness now queues its marker as a promise reaction, so a suite that stubs queueMicrotask cannot fake a failure.

Validation

• Full make test-unit on the validated change set (600fd38b3: these three layers on 266625dd, before the stack was rebuilt onto current main): 20,013 pass, 9 skip, 0 fail across 967 files; Storybook/DOM 42/42. BUN_JSC_useWasmIPInt=0 applied to that invocation only (Bun 1.3.5 Wasm crash, Bun #17841).
• make static-check-full passed on the same change set. make static-check, the context-management and agentSession.* suites, and the compaction IPC mock suites pass on this head.
• Dogfood (D4). Real Anthropic adapter against an owned loopback SSE stub, in a clean make dev-server-sandbox, with screenshots and video. Every row below passed; the stub recorded each request.
  • Summarize: automatic boundary plus follow-up, then manual /compact.
  • Continuous: eager summary while the stream is live, then a consumed mid-stream prefix swap. The swapped request dropped the old seed and carried the summary, and the journal was written. A durable fold follows at settle.
  • Continuous lifecycle: setting the threshold to 100% discards staged work, yet a consumed swap still folds exactly once. Stop clears the observation latch.
  • Two workspaces on one backend: A was held under Continuous while the global default switched to Summarize, and B compacted through Summarize. Zero cross-lane markers across more than 300 requests.
  • Token Budget, agent-led per ADR 0006: warning, handoff advisory, then a new_context or forced rollover with Continue. Manual /compact still summarizes. With PTC+RLM on, settings show Token Budget as inactive.
  • Mobile at 375 px, and a graceful backend restart with viewer reconnect.

Risks

Medium surface, low expected behavior change. The strategy governs when a context window rolls over. The main regression risk is an await inserted before a synchronous admission point or reset. Deferred-barrier tests guard that, along with the unchanged characterization suites.

Follow-ups, pre-existing and not changed here:

1. If a provider error races the Continuous fallback, applyDurably can assert "partial to be committed first" (continuousCompactor.ts:750, since 🤖 feat: add continuous compaction experiment #4086). The assertion is caught and recovery state is preserved.
2. applyJsoncEdits re-parses the document once per edit.

📋 Implementation Plan (excerpt: summary and phases 4–5; the full plan exceeds GitHub's body limit)

0. Summary

Extract the three context-management strategies out of AgentSession into one app-provided Effect service, ContextManagement, that opens a session-owned controller per AgentSession. Strategies become internal adapters behind the controller. AgentSession keeps admission, the stream pipeline, follow-up dispatch, completion publication and history authority; it stops knowing which automatic strategy is active or how its orchestration works.

Six increments: Phase 0 tests · 1 selection · 2 seam/DI/shared collaborators · 3 automatic arbitration (Summarize + Continuous, which are interleaved today) · 4 Token Budget · 5 cleanup. Each leaves a usable system with current behavior preserved; unmigrated code stays exactly where it is until its owning phase. Semantics changes (per-turn pinning, idle-timer relocation) are out of scope.

Baseline fbbea2b16403ceb6f52f71af33a13a06c9c466fa, effect@4.0.0-rc.112. Research artifact (outside repo): /home/coder/.cache/compaction-fp4j-research/compaction-effect-research.md.

Net product LoC (estimate, all phases): ≈ 0 to +200 overall; agentSession.ts −900 to −1000. Per phase below. Estimates, not gates.

Phase 4 — Token Budget behind the seam

Work: TokenBudgetStrategy owns fields :1038–1049, prepareContextBudgetSend, batch prefix/rollover (:4521–4584) → proceed{prefixRows, assemblySnapshot, receipt}; checkContextBudgetBeforeSend → reject; stepSettlement(stream) reproduces :8253 predicate (budgetActive(options) || stream.contextBudgetFlushTurn); onContextBudgetStepSettled enqueues flush/continue via host.enqueueContinuation (same dedupe keys, thinking override, 1-step bound); onSendFailure for context_budget_exceeded (:8278–8285), budget retry :8700; degradeFlushEntryToContinuation; budget clear folded into the budget-clearing reasons of §4.1 and to the front of onUserInterrupt; transitionalCompactionHandler deleted. Session calls validatePreparation(receipt) synchronously where :4572–4585/:6128–6136 check today.
Transaction boundary: The strategy prepares rollover rows, snapshots and continuation policy. The session retains full request assembly, candidate disposal/transfer, accepted snapshot-ID remapping and file-baseline association, context-reset side effects, durable publication/rollback, and retry dispatch. Capture receipts at the original generation checkpoints, not during earlier prefix preparation. The final generation/admission check and invocation of its protected operation stay in one synchronous segment; the entire recovery transaction is not atomic. Persisted flush restoration yields only in branches that already awaited admission. Prepared rollover requests always install settlement; ordinary stream starts retain the active-mode-or-admitted-flush predicate, with live stream capture when the callback runs.

Acceptance: agentSession.tokenBudget.test.ts, streamManager.contextBudget.test.ts, agentSession.memoryContext.test.ts unchanged (agentSession.budgetGate.test.ts covers goal monetary budgets; runs as a regression gate only); flush turn hidden/1-step/restricted tools; new_context rollover; unknown limit skips tokenization; Deferred-barrier test proving no await between validatePreparation and dispatch.
Gate: suites above + all previous; make static-check; dogfood D3. LoC: ≈ +430 moved+glue / −420 (net +10 to +60).

Phase 5 — Cleanup

Work: remove strategy-specific fields/methods left in AgentSession (isTokenBudgetActive, getContinuousCompactionContext, onContextBudgetStepSettled, continuousCompactor, budget fields); move buildAutoCompactionRequest into the controller if all its inputs are strategy-local (otherwise it stays a host method — decided by inspection, not assumed); doc note in di/appRuntime.ts (ContextManagement app-lifetime; controllers session-owned; leases via enterExecution). No Effect-native conversions in this phase (§2, §3.5).
Acceptance: AgentSession declares none of the listed strategy-owned fields or methods; calls to controller queries may retain their names. No transitional* member remains; all gates green. (The ≈ −900 to −1000 line reduction is an expectation, not a criterion.)
Gate: make test-unit; make static-check-full; dogfood D4. LoC: ≈ −150 to −250.

---

Generated with xum • Model: anthropic:claude-opus-5-5 • Thinking: xhigh • Cost: $3484.72

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-23T06:39:25.896353Z	6c64594	New commits
🔒 Security Review	✅ Completed 2026-09-23T06:44:20.736823Z	6c64594	New commits

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.