fix(acp): send provider-resolved context limit in usage updates - #11014
fix(acp): send provider-resolved context limit in usage updates#11014matt2e wants to merge 9 commits into
Conversation
Berd (a goose serve client) shows a 128K token window for ACP-backed agents (codex, claude) and we need to trace where that value comes from at runtime. Add greppable one-line tracing::info! logs, no behavior change: - model_config: log provider, model, resolved context limit, and its source (explicit_override / config_override / model_mapping / default_fallback) whenever a model config is materialized - acp/provider: log the model info (if any) the ACP agent reported at session setup, log when the agent reports a context window size via usage_update, and log which source get_context_limit resolved from - acp/server: log session id, provider, model, used tokens, and context limit for every usage update sent to the client Verified with cargo check -p goose and cargo clippy -p goose --lib. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
Berd showed a 128K token window for ACP-backed agents (codex, claude) even after the agent reported its real window. A confirmed repro showed AcpProvider capturing the agent's usage_update size (e.g. 258,400 tokens) while build_usage_updates — the single funnel for every usage notification sent to clients — read session.model_config.context_limit() directly, so the 128K default fallback went out on the wire. Resolve the limit through the session provider's get_context_limit() (which prefers the ACP agent's reported window) and pass it into build_usage_updates, keeping model_config.context_limit() as the fallback when no provider is available or resolution fails: - add resolve_provider_context_limit(), used by the prompt-completion path (agent in hand) and notify_session_setup (peeks the active session map without triggering activation) - thread the resolved limit through send_session_setup_notifications - report the actual resolution source (provider / session_model_config_fallback / default_fallback) in the existing "usage update sent to client" diagnostic log - add a unit test that a provider-resolved limit overrides the model config default Verified with cargo check -p goose, cargo clippy -p goose --lib, and cargo test -p goose --lib acp:: (221 passed). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
This reverts commit 3112758. The tracing::info! diagnostics served their purpose: they traced the 128K token window shown in Berd to build_usage_updates reading the model-config default, which 6fdf728 fixed by threading the provider-resolved context limit through. Remove the logs now that the root cause is fixed. The revert conflicted in build_usage_updates because the fix landed on top of the diagnostics; resolved by keeping the provider_context_limit resolution (falling back to the session model config) and dropping only the log statement and its source labels. Verified with cargo check -p goose, cargo clippy -p goose --lib -D warnings, and cargo test -p goose --lib acp:: (221 passed). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
6fdf728 made notify_session_setup await resolve_provider_context_limit before session/new, session/load, and session/fork return. That call is not guaranteed cheap: the OpenAI-compatible provider probes /v1/models (5s timeout) on a cold cache and ollama-cloud fetches /api/show, so a slow endpoint added seconds to session setup. For the ACP case that motivated the fix the await was pure cost anyway — AcpProvider only learns its window when the child agent sends a UsageUpdate, which a fresh session hasn't received yet at setup time, so the fallback went out regardless. Send the model-config fallback immediately and push a correction out-of-band; usage updates are already a push channel and Berd applies contextLimit from any usage_update, in or out of a turn: - extract send_usage_update_notifications from send_session_setup_notifications, and use it for the duplicated block on the prompt-completion path - notify_session_setup sends setup notifications with no provider limit and spawns spawn_context_limit_refresh, modelled on spawn_session_name_update_notifier - the background task resolves through the provider, sends nothing when the limit is unresolved or matches the fallback, and otherwise re-fetches session and usage totals (Berd overwrites accumulatedTotal from usage.used, so a stale snapshot could regress the displayed count) before sending the corrected update - the end-of-prompt resolution stays inline; it covers the ACP agent reporting its window mid-first-prompt, and the setup task has warmed the per-model cache for probing providers Verified with cargo fmt, cargo check -p goose --lib, cargo clippy -p goose --lib --all-targets -D warnings, and cargo test -p goose --lib acp:: (224 passed, including the new skip-decision cases). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
![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: f88fb1a425
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| &cx, | ||
| &session, | ||
| &totals, | ||
| Some(context_limit), |
There was a problem hiding this comment.
If the setup-time context-limit probe is slow and the client changes the model or provider before it finishes, this task reloads the current session but still sends context_limit resolved from the captured model_config. The model-switch path only sends a config update, so this late usage update can overwrite the client's contextLimit with the previous model's window until another prompt completes; re-resolve against the reloaded session or skip when the session's current model/provider no longer matches the captured one.
Useful? React with 👍 / 👎.
The context-limit work on this branch (6fdf728, f88fb1a) was tested only at the unit level. Usage updates were invisible to the integration harness — `Notification` had no `UsageUpdate` variant and `to_notifications` dropped them — so no integration test could observe a context limit at all, and the original 128K bug would have sailed through the suite. `spawn_context_limit_refresh` had no coverage beyond its embedded pure function. Surface usage updates in the harness and assert the out-of-band correction end to end: - add `Notification::UsageUpdate { used, context_limit }` and map `SessionUpdate::UsageUpdate` in `to_notifications`; every prompt expectation now pins the token count and TEST_MODEL's window, so a regression to the provider default breaks 14 assertions across both the server and provider connections - add `OpenAiFixture::with_n_ctx`, serving a `/v1/models` payload that advertises `meta.n_ctx` for a model the way llama.cpp and Ollama do, behind a 200ms delay standing in for a slow endpoint - add `run_context_limit_correction_notification`, registered for the server connection alongside the session-name notification test: a session on a model absent from the canonical registry falls back to DEFAULT_CONTEXT_LIMIT (128K, the reported symptom) while the probe reports 258,400, and the corrected update must arrive after `new_session` returns Because the fixture drops every notification sent before `new_session` returns, the test also pins resolution to the background path. Verified both regressions fail it: stubbing out `spawn_context_limit_refresh`, and re-inlining resolution into `notify_session_setup`, each produce "expected a corrected usage update, got []". Verified with cargo fmt, cargo clippy -p goose --all-targets -D warnings, and the ACP integration suites (acp_server_test 49 passed, acp_provider_test 12 passed, plus custom_requests, custom_provider_methods, fork_session, secret_cache_invalidation and transport_auth). Ran acp_server_test three times and the new test five times for flakiness. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
The context-limit work on this branch annotated test constants and helpers with multi-sentence doc comments, out of step with the sparse single-line notes used in the ACP test files, and several drifted into reviewer-talk (the 128K symptom history, what a regression would break) rather than constraints the next reader needs. Cut each comment down to the part the code cannot say, matching the style around it: - acp_common_tests: drop the docs on TEST_MODEL_CONTEXT_LIMIT and usage_update (both self-documenting), reduce UNKNOWN_MODEL to the one-line registry-fallback fact, and keep only the load-bearing test-mechanics constraint on run_context_limit_correction_notification (pre-return notifications are dropped, so only the out-of-band correction can pass) - acp_fixtures: shrink N_CTX_PROBE_DELAY and with_n_ctx docs to match the two-line style of OpenAiFixture::new - server.rs: tighten the spawn_context_limit_refresh doc to three lines while keeping the network-probe rationale The backwards-compatibility comment in response_builder.rs was moved from server.rs, not added here, so it stays. Comment-only change, verified with cargo fmt (no further reflow). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 46321d6aa7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| SessionUpdate::UsageUpdate(updates.standard), | ||
| ))?; | ||
| } | ||
| let provider_context_limit = resolve_provider_context_limit(&agent, &session).await; |
There was a problem hiding this comment.
When the current provider is OpenAI-compatible with an uncatalogued model, resolve_provider_context_limit calls Provider::get_context_limit, and OpenAiProvider::get_context_limit performs a /v1/models probe with a 5s timeout on a cold cache (crates/goose-providers/src/openai.rs lines 711-716). Because this await runs after the agent stream or cancellation has finished but before returning the PromptResponse, a slow or hung local/Ollama/llama.cpp endpoint keeps the prompt request open for seconds even though all output has already been sent; prefer emitting the fallback here and relying on the background correction/cache, or skip the probe for canceled/cold-cache turns.
Useful? React with 👍 / 👎.
spawn_context_limit_refresh captures the session's model config, resolves
a window from it over the network, then sends the correction — but never
rechecks that the session is still on that model. The switch path
(dispatch.rs "model"/"provider") persists a new config and sends no usage
update of its own, so a switch landing inside the probe (up to the
OpenAI-compatible provider's 5s /v1/models timeout) leaves the late
correction free to clobber the client's window. Worse than merely stale:
resolve_model_context_limit calls agent.provider() at resolve time, so
after a provider swap it asks the new provider about the old model config
and the answer matches neither configuration.
Guard the send instead of re-resolving, following the precedent at
dispatch.rs:206 ("provider changed before inventory refresh completed") —
re-resolving restarts the same race, and the correction baseline is the
fallback setup already sent for the captured model, which is wrong for
the new one:
- add context_limit_refresh_applies, comparing model_name plus the
computed context_limit() rather than the whole struct: ModelConfig has
no PartialEq, and fields irrelevant to the window (request_params after
a thinking-effort switch) shouldn't suppress a valid correction
- check it in the spawned task after the session reload; a switch landing
between the reload and the send still escapes, but the window shrinks
from seconds to microseconds and self-corrects at the next prompt
Provider switches are caught indirectly, since update_provider replaces
the session's model config with the new provider's model.
The switch path still never updates the client's contextLimit at all —
after switching, the client shows the old window until the next prompt
completes. That is pre-existing and left alone; sending an update on the
switch path is new client-visible behavior and belongs in its own change.
Tests: four #[test_case]s over the decision (unchanged / different model /
changed override / no config), and an integration test asserting the
probe's 258,400 never reaches a client that switched models right after
session setup. The switch is awaited rather than raced: on_set_model
persists before build_config_update, the slower step of the same request.
Confirmed the integration test fails without the guard ("expected no
usage update for the pre-switch model, got [ConfigOption, UsageUpdate {
used: 0, context_limit: 258400 }]") and passes five runs in a row with it.
Verified with cargo fmt --check, cargo clippy -p goose --all-targets
-D warnings, cargo test -p goose --lib acp:: (228 passed), and the ACP
integration suites (acp_server_test 50, acp_provider_test 12, plus
custom_requests, custom_provider_methods, fork_session,
secret_cache_invalidation and transport_auth).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Matt Toohey <contact@matttoohey.com>
The guard added in 9e6b9ab carried a three-line doc comment whose first sentence and a half restated what the function name and call site already say — that a correction resolved for one model config isn't sent once the session has moved on. Every sibling free function in that run of server.rs (to_nonnegative_u64, prompt_stop_reason, context_limit_correction, build_usage_updates, validate_absolute_cwd) carries no comment at all. Keep only the part the code cannot say: - context_limit_refresh_applies: two lines on why the comparison is name plus computed window rather than whole-struct equality — no PartialEq on ModelConfig, and an unrelated field change shouldn't suppress a still-valid correction - run_context_limit_refresh_skipped_after_model_switch: tightened to two lines, matching the sibling correction test's note, and kept because the persist-before-inventory ordering is load-bearing test mechanics — it is why awaiting the switch doesn't defeat the race the test sets up Comment-only change. Verified with cargo fmt --check, cargo clippy -p goose --all-targets -D warnings, cargo test -p goose --lib acp:: (228 passed), and the context-limit integration tests (2 passed). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9e6b9abfe7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| current.model_name == captured.model_name | ||
| && current.context_limit() == captured.context_limit() |
There was a problem hiding this comment.
Fresh evidence: this new guard treats the same model name/window as enough to apply the correction, but the corrected value was resolved by the captured provider. If the setup probe is still running and the session switches to another provider with the same uncatalogued model/default context, this check still passes and sends the old provider’s probed context limit to the new provider’s session; capture and compare the provider name before sending the correction.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f8982325c4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| SessionUpdate::UsageUpdate(updates.standard), | ||
| ))?; | ||
| } | ||
| let provider_context_limit = resolve_provider_context_limit(&agent, &session).await; |
There was a problem hiding this comment.
When this runs against a provider whose context limit lookup performs a cold network probe, the session has already been unlocked and activeRunId has been sent as None above, so start_active_run can accept another prompt while this await is still pending. Because session and totals were captured before the await, the first request can later send an older UsageUpdate and overwrite the client's token count/context limit from the newer turn; resolve before releasing the active run, or reload the session and totals after the await before sending.
Useful? React with 👍 / 👎.
The guard added in 9e6b9ab compares model_name plus the computed context_limit(), and neither changes when a session switches to a different provider serving the same uncatalogued model. ModelConfig carries no provider identity at all, so the guard had nothing provider-shaped to compare: 9e6b9ab's claim that provider switches are caught indirectly holds only when the replacement changes the model name or the window, and several of the switch path's model-resolution fallbacks preserve the name. The leak is exactly the population this feature targets. A session on an OpenAI-compatible endpoint A with a model absent from the canonical registry sends the 128K fallback and starts A's /v1/models probe; a switch to endpoint B serving the same name lands mid-probe, still uncatalogued so still 128K; A's probed window is then pushed to a session on B, whose real n_ctx may differ. Two endpoints serving the same local model with different context configurations is mundane in the llama.cpp/Ollama world. Exposure is the full probe duration (up to the provider's 5s timeout), not the microseconds-wide reload-to-send residue. Session already persists provider_name and the switch path updates it in the same atomic session update as the model config, so the fix follows the guard's existing shape: - capture session.provider_name alongside the model config in spawn_context_limit_refresh - compare it against the reloaded session's in context_limit_refresh_applies, as plain Option<&str> equality — None == None passes, since a switch would have populated the field - extend the #[test_case] table with provider-changed and provider-unrecorded rows Residual imprecision, in the fail-safe direction: a switch landing between the spawn-time capture and the probe's agent.provider() call resolves through the new provider, and the guard now skips a correction that was in fact valid; the client keeps the honest fallback and the next prompt completion resolves it inline. Closing that too would mean recording the provider the probe actually used, restructuring the task to fetch the provider before probing rather than reusing resolve_model_context_limit — not worth it for a microsecond window with a safe failure mode. No integration test: the fixture's provider factory ignores the provider name, and reaching the same-name case would need a second provider whose registry default_model is the uncatalogued model, so a fixture switch changes the model name too and is already caught by the existing check. The guard is a pure function and the model-switch integration test already pins the send-path plumbing. Verified with cargo fmt, cargo clippy -p goose --all-targets -D warnings, cargo test -p goose --lib acp:: (230 passed), and the ACP integration suites (acp_server_test 50, acp_provider_test 12), including the context-limit correction test which still receives its correction — so the captured and reloaded provider names agree on the normal path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Matt Toohey <contact@matttoohey.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a4f1d0d920
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if let Err(error) = send_usage_update_notifications( | ||
| &cx, | ||
| &session, | ||
| &totals, | ||
| Some(context_limit), |
There was a problem hiding this comment.
When a client closes a session while this spawned context-limit probe is still pending, on_close_session only marks the id closed/removes it from self.sessions and leaves the stored session row, so the reload above still succeeds and this send emits a UsageUpdate after close_session has completed. ACP clients that drop closed sessions can receive an unexpected update for a resource that subsequent requests treat as not found; check closed_session_ids or that the session is still loaded before sending the correction.
Useful? React with 👍 / 👎.
|
Thanks @matt2e for the detailed investigation and for building integration coverage around ACP usage updates. This work uncovered useful constraints, especially that provider context-limit discovery may involve network I/O and that delayed results must account for model, provider, usage, and session lifecycle changes. We’re going to close this implementation in favor of #10966. That issue is intended to solve this class of bug by making This PR instead reconciles those two sources inside the ACP server by sending a model-config fallback and later pushing a provider-resolved correction. The resulting background task, timing guards, and lifecycle races are the kind of path-specific behavior #10966 is meant to avoid rather than extend. Please add any findings that are not already captured—particularly the network-probe latency and invalidation cases—to #10966. The ACP usage-update integration coverage here will also be useful when an implementation is agreed. Once #10966 reaches Ready, implementation can proceed against that design. Thank you again for tracing the regression and documenting it so thoroughly. |
2 participants
ACP clients showed a 128K token window for ACP-backed agents (codex, claude) even after the child agent reported its real window.
AcpProvidercaptured the agent'susage_updatesize (e.g. 258,400 tokens), butbuild_usage_updates— the single funnel for every usage notification sent to clients — readsession.model_config.context_limit()directly, so the default fallback went out on the wire.This is a regression fix: #9455 originally wired the agent-reported window through to clients (merged May 29), and the ACP session-setup refactor in #9488 (merged June 4) dropped that wiring. The refactor consolidated usage notifications into
build_usage_updates(session: &Session), whose signature takes only the session, and deleted theprovider.get_model_config().context_limit()hop at the old call site — so from then on the persisted model config's 128K default was sent instead of the provider's value. Nothing caught it because only the capture side was tested: the #9455 test that the provider stores and surfaces the reported size kept passing while the value went unconsumed. This PR adds the missing consumption-side test (provider-resolved limit overrides the model-config default).Changes
build_usage_updatestakes an optional provider-resolved context limit and prefers it overmodel_config.context_limit(), which stays the fallback.resolve_provider_context_limitresolves the window through the session provider'sget_context_limit(), which prefers the ACP agent's reported size.send_usage_update_notificationsinstead of duplicating the custom/standard notification block./v1/models(5s timeout) on a cold cache, ollama-cloud fetches/api/show— so a slow endpoint would add seconds tosession/new,session/load, andsession/fork. Instead setup sends the model-config fallback immediately and spawns a background refresh (modelled onspawn_session_name_update_notifier) that pushes a corrected usage update only when the resolved limit differs from the fallback. Usage updates are already a push channel, and clients applycontextLimitfrom any usage update, in or out of a turn.usage.usedand a stale snapshot could regress the displayed count.The branch also reverts the temporary info-level diagnostics that were used to trace the root cause.
Verification
cargo fmt,cargo check -p goose --lib,cargo clippy -p goose --lib --all-targets -- -D warnings, andcargo test -p goose --lib acp::(224 passed, including new cases covering provider-limit preference and the correction skip decision).🤖 Generated with Claude Code