Repository navigation
📐 feat: AgentContext.projectContextUsage: pre-send context snapshot - #244
Conversation
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 585888511e
ℹ️ 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".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| effectiveInstructionTokens, | ||
| calibrationRatio, | ||
| } = prune({ | ||
| messages, |
There was a problem hiding this comment.
When projectContextUsage is called on a branch under high context pressure, passing the caller's messages array directly into createPruneMessages lets the pruner perform its normal in-place mutations (for example maskConsumedToolResults and pre-flight truncation replace entries in messages[i]). In the intended page-load/pre-send projection path, this can permanently truncate or mask the session history before any real model call or summarization sees it, even though the method is advertised as a safe projection; pass cloned messages (or add a read-only pruner mode) before invoking it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 8a3651d. projectContextUsage now passes a shallow copy ([...messages]) to the pruner. Confirmed both maskConsumedToolResults and preFlightTruncateToolResults mutate via slot replacement (messages[i] = new ToolMessage(...)), not in-place content edits, so a shallow array copy fully protects the caller's history. Added a test asserting the caller's tool-message slot + content are unchanged after a projection under context pressure.
| tokenCounter: this.tokenCounter, | ||
| maxTokens: this.maxContextTokens, | ||
| thinkingEnabled: isThinkingEnabled(this.provider, this.clientOptions), | ||
| indexTokenCountMap: this.indexTokenCountMap, |
There was a problem hiding this comment.
Because projectContextUsage accepts an arbitrary messages array but seeds the pruner with this.indexTokenCountMap, any call on an existing context whose token map was built for a different branch/model state will reuse counts by numeric index without checking that they belong to these message objects. In that scenario the projection can prune or report usage based on stale entries (the pruner skips recounting defined indices), so page-load/window-switch estimates can be wrong; build a fresh per-call token map for the supplied messages or require a matching map from the caller.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 8a3651d. The method now recounts a fresh indexTokenCountMap for the supplied messages via this.tokenCounter (the context's own map is keyed to the live run's branch). A caller that already has matching counts (e.g. stored per-message tokenCounts) may pass them via the new opts.indexTokenCountMap. Tests: a stale/empty context map still prunes correctly (recounted), and a caller-supplied map is honored.
| usageMetadata: opts?.usageMetadata ?? this.currentUsage, | ||
| lastCallUsage: this.lastCallUsage, | ||
| totalTokensFresh: this.totalTokensFresh, |
There was a problem hiding this comment.
When this context has just completed a live call, currentUsage/lastCallUsage are still fresh, but projectContextUsage creates a brand-new pruner whose lastTurnStartIndex starts at 0. Passing that previous call's usage into the fresh pruner makes calibration compare the prior provider input against the entire projected branch (often including the previous AI output and later messages), which can drive the projected calibration/instruction overhead and remaining budget far away from what the live stateful pruner would report; default projections should skip live usage or mark it stale unless the caller supplies matching pruner state.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 8a3651d. The projection no longer pulls this.currentUsage/this.lastCallUsage/totalTokensFresh — it passes lastCallUsage: undefined, totalTokensFresh: false, and only an explicitly-supplied opts.usageMetadata. So a default projection is uncalibrated (the honest pre-send view) and a context with fresh live usage projects identically to a clean one. Calibration is opt-in. Added a test asserting equal remainingContextTokens/calibrationRatio with vs without live usage set.
|
@codex review |
danny-avila
changed the title
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8a3651db60
ℹ️ 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".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| effectiveInstructionTokens, | ||
| calibrationRatio, | ||
| } = prune({ | ||
| messages: [...messages], |
There was a problem hiding this comment.
When thinking is enabled and pruning needs to reattach a reasoning block, createPruneMessages can reach getMessagesWithinTokenLimit → addThinkingBlock, which unshifts into the existing AIMessage.content array. Since this projection only shallow-copies the outer array here, the caller's original message objects can still be mutated despite the method being documented as safe/non-mutating; a page-load or window-switch projection could therefore corrupt the branch it is only supposed to inspect.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in f8fd752. Confirmed addThinkingBlock does content.unshift(...) on the original AIMessage's content array (in place, before returning a new message), so the shallow array copy wasn't enough. The projection now clones array-content messages via the existing cloneMessage (now exported) before pruning, giving the pruner its own content arrays. Added a test (thinking enabled) asserting the caller's AI content array reference + length are unchanged after a projection.
| messages: [...messages], | ||
| usageMetadata: opts?.usageMetadata, | ||
| lastCallUsage: undefined, | ||
| totalTokensFresh: false, |
There was a problem hiding this comment.
opts.usageMetadata is exposed for callers that have a provider snapshot, but this call always passes totalTokensFresh: false; createPruneMessages skips its calibration block whenever that flag is false, so complete metadata cannot adjust the ratio used for remainingContextTokens. In reload or model-switch flows that pass persisted usage to avoid raw-token estimates, the projection remains uncalibrated and can disagree with the live ON_CONTEXT_USAGE snapshot for providers whose accounting differs from the local tokenizer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — totalTokensFresh: false made the calibration block (if (currentUsage && totalTokensFresh !== false)) unreachable, so usageMetadata was dead. Rather than re-enable re-derivation (the recalibration hazard from the previous round — a fresh pruner at startIndex 0 comparing prior input against the whole branch), I replaced the opt with a calibrationRatio seed. The seed ratio is applied throughout the pruner's budget math regardless of totalTokensFresh, so callers pass the persisted snapshot's calibrationRatio (which the live path already computed) to get a calibrated projection — no re-derivation. Test: an explicit ratio is reflected in the result and changes remainingContextTokens.
| summarizationEnabled: this.summarizationEnabled, | ||
| reserveRatio: this.summarizationConfig?.reserveRatio, | ||
| calibrationRatio: this.calibrationRatio, | ||
| getInstructionTokens: () => this.instructionTokens, |
There was a problem hiding this comment.
When the context has been marked systemRunnableStale (for example after setHandoffContext, setInitialSummary, or clearing/changing summary state), the live graph refreshes agentContext.systemRunnable before pruning, which recounts systemMessageTokens; this projection path never does that, so getInstructionTokens can use stale overhead and under- or over-estimate pruning for transferred agents or updated system prompts. Callers using the projection for those branches can get a snapshot that no longer matches the prompt a real call would send.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in f8fd752. projectContextUsage now calls this.initializeSystemRunnable() first (idempotent — recomputes only when stale), mirroring the live graph's refresh before pruning, so handoff/summary changes recount systemMessageTokens instead of projecting with stale instruction overhead. Test: setHandoffContext (marks stale, defers the recount) followed by a projection recomputes systemMessageTokens higher.
|
@codex review |
|
Codex Review: Didn't find any major issues. Hooray! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Add a pure, off-the-hot-path projection of the context-usage snapshot for an arbitrary message set, without invoking the model. It runs the same pruner + budget math the live path uses (createPruneMessages → getTokenBudgetBreakdown → syncBudgetDerivedFields), so a "what would I send" projection matches a real call's ON_CONTEXT_USAGE snapshot. Returns null when no tokenizer/window is configured; uses a local pruner and never mutates the context. This backs a host-side (e.g. LibreChat) context gauge in the states the live snapshot can't cover — page load of a snapshot-less branch, and model/window switches where a persisted snapshot is stale — so the gauge always reflects the SDK's authoritative view instead of a client-side estimate. syncBudgetDerivedFields moves from Graph.ts to a shared src/messages/budget.ts so the live and projected paths share one derivation (no drift). Behavior of the live snapshot is unchanged.
…ection Address three findings on the pre-send projection's reuse of the stateful pruner: - P1: pass a shallow copy of `messages` to the pruner. Its mask/pre-flight truncation replaces tool-result slots in place, which would otherwise permanently truncate the caller's session history during a projection. - Recount a fresh per-message token map for the supplied messages (the context's own map is keyed to the live run's branch and would missum an arbitrary branch). Callers may pass a map they guarantee matches. - Do not pull this context's live `currentUsage`/`lastCallUsage`: a fresh pruner (startIndex 0) would recalibrate the prior call's provider input against the whole projected branch, skewing the estimate. Calibration is opt-in via an explicit `usageMetadata`.
…refresh - Clone array-content messages before projecting. The pruner unshifts reasoning blocks into AI content arrays in place (addThinkingBlock), so a shallow array copy alone left the caller's history mutable; reuse the existing `cloneMessage` (export it) to give the pruner its own arrays. - Replace the (dead) `usageMetadata` opt with a `calibrationRatio` seed. The pruner skips calibration when `totalTokensFresh` is false, so passing usage never adjusted the ratio. Re-deriving from usage is the recalibration hazard flagged earlier anyway — instead apply the persisted ratio as a static seed (the same value the live snapshot carries), no re-derivation. - Refresh a stale system runnable (`initializeSystemRunnable`) before projecting, as the live graph does, so handoff/summary changes recount system overhead instead of projecting with stale instruction tokens.
Expose a thin async entry that builds a throwaway AgentContext from the same AgentInputs a run uses, awaits its instruction/tool token accounting, then runs the shared projection. Hosts (e.g. LibreChat) don't hold an AgentContext — this lets them project a branch's pre-send context snapshot from the config they already assemble for createRun, without exposing the AgentContext class.
0914956 to
1e1c691
Compare
1 participant
Summary
Adds
AgentContext.projectContextUsage(messages, opts?)— a pure, off-the-hot-path projection of the context-usage snapshot for an arbitrary message set, without invoking the model.It runs the same pruner + budget math the live path uses (
createPruneMessages→getTokenBudgetBreakdown→syncBudgetDerivedFields), so a "what would I send" projection produces numbers identical to a real call'sON_CONTEXT_USAGEsnapshot. It returnsnullwhen the context has no tokenizer/window, uses a local pruner, and never mutates the context — safe to call repeatedly off the run loop.Why
A host-side context gauge (e.g. LibreChat) is accurate while streaming and on reload of a turn that emitted a live snapshot. But two states have no SDK data and currently fall back to a client-side estimate that sums the whole branch with no window/prune awareness:
projectContextUsagelets the host ask the SDK "what would the context be for this branch under this config" so the gauge always reflects the SDK's authoritative pruning, not a re-derivation.Changes
AgentContext.projectContextUsage(...)(new). Pre-send counterpart to the live snapshot; omits the live post-format reconciliation (provider-specific, invoke-time) — a small, acceptable delta for a pre-send estimate.syncBudgetDerivedFieldsmoves fromGraph.tsto a sharedsrc/messages/budget.tsso the live and projected paths share one derivation (no drift). Live snapshot behavior is unchanged — the function body and its two call sites are identical; only the location moved.Tests
projectContextUsagecases inAgentContext.test.ts: null without tokenizer/window; keeps the whole branch + reports headroom when it fits; prunes older messages when the branch exceeds the window; does not mutate the context.AgentContextsuite (89) andmessagessuite (349) green; the livecontext-usage-eventemission spec (2) still passes, confirming the relocation didn't change the hot path. Typecheck + lint clean.(The two
src/llm/{anthropic,google}/llm.spec.tssuites fail locally with "API key not found" — pre-existing, env-only, unrelated.)Follow-up (separate, LibreChat)
A projection endpoint that builds the agent context for a branch + resolved config and calls this method, plus client wiring so the gauge prefers live snapshot → fresh persisted snapshot → projection → estimate.