fix(#4251): guard provider ownership under the session lock

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

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"