fix(agent): roll back switch_model state when client rebuild fails - #33199
Closed
briandevans wants to merge 1 commit into
Closed
fix(agent): roll back switch_model state when client rebuild fails#33199briandevans wants to merge 1 commit into
briandevans wants to merge 1 commit into
Conversation
`switch_model()` swapped `agent.model`/`agent.provider`/`agent.base_url`/ `agent.api_mode` several lines before the new client was built. Any exception during the rebuild (bad API key, network error, MiniMax OAuth failure, or import failure) left the agent in a torn state: the new model/provider names against the *old* client. The next turn then hit a model/provider mismatch (e.g. HTTP 400 "claude-sonnet-4-6 is not supported on openai-codex") and the user only recovered via fallback activation. Snapshot every field this function mutates before the swap, wrap the swap + client rebuild + prompt-caching evaluation in `try`/`finally`, and restore the snapshot on any uncommitted failure. The exception continues to propagate so callers (`model_switch.switch_model`, the TUI/CLI `/model` handlers) keep seeing failure and can surface it to the user — the only change is that the agent is now internally consistent when they do. Tests cover both failure paths (chat_completions and anthropic_messages) and confirm the happy path still commits the new state.
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a rollback/transaction-style guard to switch_model so that if client rebuild fails mid-switch, the agent’s runtime fields are restored to the pre-switch state (regression for #33175).
Changes:
- Wrap
switch_modelcore field updates + client rebuild in atry/finallywith a pre-swap snapshot and rollback on exceptions. - Add regression tests asserting rollback on
chat_completionsandanthropic_messagesrebuild failures. - Add a “happy path” test ensuring successful switches still persist new state.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| tests/run_agent/test_switch_model_rollback.py | New regression tests for rollback-on-failure and commit-on-success behavior. |
| agent/agent_runtime_helpers.py | Introduces rollback snapshot/restore in switch_model to prevent torn agent state on rebuild exceptions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1374
to
+1389
| _rollback = { | ||
| "model": agent.model, | ||
| "provider": agent.provider, | ||
| "base_url": agent.base_url, | ||
| "api_mode": agent.api_mode, | ||
| "api_key": agent.api_key, | ||
| "_config_context_length": getattr(agent, "_config_context_length", None), | ||
| "client": getattr(agent, "client", None), | ||
| "_client_kwargs": dict(getattr(agent, "_client_kwargs", {}) or {}), | ||
| "_anthropic_api_key": getattr(agent, "_anthropic_api_key", None), | ||
| "_anthropic_base_url": getattr(agent, "_anthropic_base_url", None), | ||
| "_anthropic_client": getattr(agent, "_anthropic_client", None), | ||
| "_is_anthropic_oauth": getattr(agent, "_is_anthropic_oauth", False), | ||
| "_use_prompt_caching": getattr(agent, "_use_prompt_caching", False), | ||
| "_use_native_cache_layout": getattr(agent, "_use_native_cache_layout", False), | ||
| } |
Comment on lines
+1409
to
+1410
| # Invalidate transport cache — new api_mode may need a different transport | ||
| if hasattr(agent, "_transport_cache"): |
Comment on lines
+1480
to
+1482
| if not _committed: | ||
| for _attr, _val in _rollback.items(): | ||
| setattr(agent, _attr, _val) |
Comment on lines
+1470
to
1485
| agent._use_prompt_caching, agent._use_native_cache_layout = ( | ||
| agent._anthropic_prompt_cache_policy( | ||
| provider=new_provider, | ||
| base_url=agent.base_url, | ||
| api_mode=api_mode, | ||
| model=new_model, | ||
| ) | ||
| ) | ||
| ) | ||
| _committed = True | ||
| finally: | ||
| if not _committed: | ||
| for _attr, _val in _rollback.items(): | ||
| setattr(agent, _attr, _val) | ||
|
|
||
| # ── LM Studio: preload before probing context length ── | ||
| agent._ensure_lmstudio_runtime_loaded() |
| f"failure; diverged fields: {[k for k in before if before[k] != after[k]]}" | ||
| ) | ||
|
|
||
|
|
Collaborator
|
Superseded by #33228 (teknium salvage onto current main with additional test coverage). Same fix for switch_model() rollback on client rebuild failure. |
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
agent_runtime_helpers.switch_model()swappedagent.model,agent.provider,agent.base_url, andagent.api_modeseveral linesbefore the new client was built. Any exception during the rebuild —
bad API key, network error, MiniMax OAuth failure, or import failure —
left the agent in a torn state: the new model/provider against the
old client. The next turn then hit a model/provider mismatch (e.g.
HTTP 400
"claude-sonnet-4-6 is not supported on openai-codex").The fix snapshots every field this function mutates before the swap
(
model,provider,base_url,api_mode,api_key,client,_client_kwargs,_config_context_length, the four_anthropic_*fields, and the two prompt-caching flags), wraps the swap + client
rebuild + prompt-cache policy evaluation in
try/finally, andrestores the snapshot on any uncommitted failure. The exception
continues to propagate so callers (
model_switch.switch_model, theTUI/CLI
/modelhandlers) keep seeing the failure — the only change isthat the agent is now internally consistent when they do.
Related Issue
Fixes #33175
Type of Change
Changes Made
agent/agent_runtime_helpers.py— snapshot pre-swap state, wrap themutating section in
try/finallywith a_committedflag,restore the snapshot when the body raises before completing. No
behaviour change on the happy path.
tests/run_agent/test_switch_model_rollback.py— three new tests:chat_completions rebuild failure restores all state, anthropic
rebuild failure restores all state, happy-path commit still applies
the new state.
How to Test
uv run --with pytest --with pytest-xdist --with pytest-asyncio python3 -m pytest tests/run_agent/test_switch_model_rollback.py tests/run_agent/test_switch_model_context.py tests/run_agent/test_switch_model_fallback_prune.py tests/hermes_cli/test_model_switch_variant_tags.py tests/hermes_cli/test_model_switch_custom_providers.py tests/hermes_cli/test_model_switch_opencode_anthropic.py tests/gateway/test_model_command_flat_string_config.py -v— 49 passed locally.tests/run_agent/test_switch_model_rollback.py; the two failure-pathtests fail with exactly the torn-state symptom described in switch_model() leaves agent in torn state if client rebuild fails #33175
(
model,provider,api_mode,base_url,_anthropic_*allstuck on the new values).
Checklist
Code
fix(scope):,feat(scope):, etc.)Documentation & Housekeeping
docs/, docstrings) — N/A, behaviour is unchanged on the happy pathcli-config.yaml.exampleif I added/changed config keys — N/ACONTRIBUTING.mdorAGENTS.mdif I changed architecture or workflows — N/AContract Protected
Invariant: after
switch_model()returns, either (a) the new model,provider, base_url, api_mode, api_key, and matching client are all
committed together, or (b) all of them are still on the old values.
There is no observable intermediate state where the model name and the
client object disagree.
Known-bad inputs covered by tests:
_create_openai_clientraises (badcredentials, network error, transport build failure).
build_anthropic_clientraises(
ConnectionError,AuthenticationError, import failure).Negative case: the happy-path test confirms the new state is
committed when no exception is raised — the rollback wrapper does not
silently swallow successful swaps.
Sibling code paths:
_try_activate_fallback()inrun_agent.pymirrors the same client-swap logic for the fallback case (see the
"mirrors
_try_activate_fallback" note in this function's docstring).That path is turn-scoped and resets each turn, so torn state is less
visible there, but it would benefit from the same rollback discipline
in a follow-up — intentionally left out of this PR's scope to keep the
diff small and focused on the reported symptom. Happy to widen if
preferred.