diff --git a/api/streaming.py b/api/streaming.py index 8a3d756a..161fe5da 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -6694,23 +6694,35 @@ def _run_agent_streaming( s = get_session(session_id) update_active_run(stream_id, phase="running", session_id=session_id) s.workspace = str(Path(workspace).expanduser().resolve()) - _last_persisted_model = getattr(s, "model", None) + _last_persisted_model = None + _last_persisted_provider = None _turn_owns_persisted_model = False provider_context = ( str(model_provider).strip().lower() if model_provider is not None else getattr(s, "model_provider", None) ) + provider_context = str(provider_context).strip().lower() if provider_context else None + _agent_lock = _get_session_agent_lock(session_id) # #4251: the route layer already persisted this turn's model under the # session lock before dispatch, so a mismatch here means a newer picker # write won the race and must not be clobbered by the worker thread. - if _last_persisted_model in (None, "", model): - s.model = model - s.model_provider = provider_context or None - _last_persisted_model = model - _turn_owns_persisted_model = True + with _agent_lock: + _last_persisted_model = getattr(s, "model", None) + _last_persisted_provider = getattr(s, "model_provider", None) + if _last_persisted_provider is not None: + _last_persisted_provider = str(_last_persisted_provider).strip().lower() or None + _persisted_model_is_empty = _last_persisted_model in (None, "") + _provider_matches = _last_persisted_provider in (None, provider_context) + if _persisted_model_is_empty or ( + _last_persisted_model == model and _provider_matches + ): + s.model = model + s.model_provider = provider_context + _last_persisted_model = model + _last_persisted_provider = provider_context + _turn_owns_persisted_model = True - _agent_lock = _get_session_agent_lock(session_id) # TD1: set thread-local env context so concurrent sessions don't clobber globals # Check for pre-flight cancel (user cancelled before agent even started) if cancel_event.is_set(): @@ -6751,12 +6763,22 @@ def _run_agent_streaming( has_profile=bool(getattr(s, "profile", None)), ) # #4251: only apply the profile-repair persistence if this turn still - # owns the session model value it last wrote. - if _turn_owns_persisted_model and getattr(s, "model", None) == _last_persisted_model: - s.model_provider = provider_context - if _repaired and model != (s.model or ""): - s.model = model - _last_persisted_model = model + # owns the session model/provider pair it last wrote. + provider_context = str(provider_context).strip().lower() if provider_context else None + with _agent_lock: + _current_provider = getattr(s, "model_provider", None) + if _current_provider is not None: + _current_provider = str(_current_provider).strip().lower() or None + if ( + _turn_owns_persisted_model + and getattr(s, "model", None) == _last_persisted_model + and _current_provider == _last_persisted_provider + ): + s.model_provider = provider_context + _last_persisted_provider = provider_context + if _repaired and model != (s.model or ""): + s.model = model + _last_persisted_model = model # Capture the resolved profile name now, while profile context is # reliable. Used in the compression migration block to stamp s.profile diff --git a/tests/test_issue4251_wakeup_model_picker_race.py b/tests/test_issue4251_wakeup_model_picker_race.py index dfb1385f..f5481de8 100644 --- a/tests/test_issue4251_wakeup_model_picker_race.py +++ b/tests/test_issue4251_wakeup_model_picker_race.py @@ -173,6 +173,19 @@ def test_dispatch_stamp_does_not_clobber_newer_picker_model(monkeypatch): assert fake_session.model_provider == "openrouter" +def test_dispatch_stamp_does_not_clobber_newer_picker_provider_only_choice(monkeypatch): + fake_session = FakeSession(model="haiku-4-5", model_provider="openrouter") + + _run_streaming_turn( + monkeypatch, + fake_session, + stream_id="stream-4251-provider-race", + ) + + assert fake_session.model == "haiku-4-5" + assert fake_session.model_provider == "openrouter" + + def test_dispatch_stamp_persists_resolved_model_when_no_race(monkeypatch): fake_session = FakeSession(model="haiku-4-5", model_provider=None) @@ -216,3 +229,31 @@ def test_profile_repair_skips_persistence_when_newer_picker_choice_already_won(m assert fake_session.model == "opus-4-8" assert fake_session.model_provider == "openrouter" + + +def test_dispatch_stamp_snapshots_provider_under_agent_lock(monkeypatch): + fake_session = FakeSession(model="haiku-4-5", model_provider="anthropic") + + class PickerUpdateLock: + def __init__(self): + self._entered = 0 + + def __enter__(self): + self._entered += 1 + if self._entered == 1: + fake_session.model_provider = "openrouter" + return self + + def __exit__(self, exc_type, exc, tb): + return False + + monkeypatch.setattr(streaming, "_get_session_agent_lock", lambda _session_id: PickerUpdateLock()) + + _run_streaming_turn( + monkeypatch, + fake_session, + stream_id="stream-4251-lock-race", + ) + + assert fake_session.model == "haiku-4-5" + assert fake_session.model_provider == "openrouter"