diff --git a/routes/chat_routes.py b/routes/chat_routes.py index 1fb2bd6ad..403ec9539 100644 --- a/routes/chat_routes.py +++ b/routes/chat_routes.py @@ -527,7 +527,7 @@ def _first_image_attachment(chat_handler, att_ids: List[str], owner: str | None return None -def _recover_empty_session_model(sess, session_id: str, owner: str | None = None) -> bool: +def _recover_empty_session_model(sess, session_id: str, owner: str | None = None, *, allow_live_probes: bool = True) -> bool: """Re-populate sess.model from the matching endpoint's cached models. Covers the window between endpoint setup and the first chat send: the @@ -535,6 +535,11 @@ def _recover_empty_session_model(sess, session_id: str, owner: str | None = None written (Issue #587 — UI uses the cached endpoint list, not s.model). For ChatGPT Subscription, also repairs stale OpenAI API model names such as ``gpt-5`` that are not accepted by the Codex-backed ChatGPT account route. + + Bearer chat callers set ``allow_live_probes`` to false. They may use the + already-persisted visible cache for this request, but recovery must not + resolve provider credentials, refresh the catalog, or persist a model/cache + change as a side effect. """ current_model = (getattr(sess, "model", "") or "").strip() endpoint_url = (getattr(sess, "endpoint_url", "") or "").strip() @@ -585,7 +590,7 @@ def _recover_empty_session_model(sess, session_id: str, owner: str | None = None visible = cached if current_model and current_model in {str(item).strip() for item in visible}: return False - if is_chatgpt_subscription: + if is_chatgpt_subscription and allow_live_probes: live_models = [] if getattr(ep, "provider_auth_id", None): try: @@ -617,6 +622,16 @@ def _recover_empty_session_model(sess, session_id: str, owner: str | None = None if not isinstance(model, str) or not model.strip(): return False model = model.strip() + if not allow_live_probes: + # Keep this request usable without turning cache-based recovery + # into a durable session mutation. The normal chat save path will + # persist user/assistant messages, not this transient selection. + sess.model = model + logger.info( + "Recovered session model for %s from cached endpoint model %r (no persistence)", + session_id, model, + ) + return True # Persist so the next request, websocket reconnect, or page reload # picks up the same model (we'd otherwise re-pick on every send # and silently switch on the user if the cached order shifts). @@ -779,13 +794,19 @@ def setup_chat_routes( except KeyError: raise HTTPException(404, f"Session '{session}' not found") owner = effective_user(request) + request_capability = build_request_capability(request) if _clear_orphaned_session_endpoint(sess, owner=owner): raise HTTPException(400, "Selected model endpoint was removed. Pick another model in Settings.") # Empty model + live endpoint = setup race (Issue #587). Repair from # the endpoint's cached model list before privilege checks, which # otherwise see "" and behave inconsistently with the allowlist. - _recover_empty_session_model(sess, session, owner=owner) + _recover_empty_session_model( + sess, + session, + owner=owner, + allow_live_probes=request_capability.allow_live_probes, + ) if not getattr(sess, "model", "").strip(): raise HTTPException( 400, @@ -798,7 +819,6 @@ def setup_chat_routes( # non-streaming path can't be used to bypass). _enforce_chat_privileges(request, sess) - request_capability = build_request_capability(request) api_token_request = request_capability.is_bearer tool_policy = build_effective_tool_policy(last_user_message=message) allow_tool_preprocessing = ( @@ -1311,7 +1331,12 @@ def setup_chat_routes( # the first cached model off the matching endpoint so the # upstream isn't called with model="" (which surfaces as a # generic 401/503). - _recover_empty_session_model(sess, session, owner=owner) + _recover_empty_session_model( + sess, + session, + owner=owner, + allow_live_probes=request_capability.allow_live_probes, + ) if not getattr(sess, "model", "").strip(): raise HTTPException( 400, diff --git a/tests/test_api_token_forward_fixes.py b/tests/test_api_token_forward_fixes.py index b02f5e506..5239cfccd 100644 --- a/tests/test_api_token_forward_fixes.py +++ b/tests/test_api_token_forward_fixes.py @@ -395,7 +395,7 @@ def test_bearer_context_preprocessing_does_not_fetch_embedded_urls(monkeypatch): async def test_sync_bearer_chat_cannot_use_research_memory_or_background_extraction(monkeypatch): from routes import chat_routes - calls = {"memory": 0, "research": 0, "post": []} + calls = {"memory": 0, "research": 0, "post": [], "recovery": []} class _ChatHandler: async def handle_memory_command(self, _session, _message): @@ -435,7 +435,11 @@ async def test_sync_bearer_chat_cannot_use_research_memory_or_background_extract monkeypatch.setattr(chat_routes, "_verify_session_owner", lambda *args, **kwargs: None) monkeypatch.setattr(chat_routes, "_clear_orphaned_session_endpoint", lambda *args, **kwargs: False) - monkeypatch.setattr(chat_routes, "_recover_empty_session_model", lambda *args, **kwargs: False) + def recover(*args, **kwargs): + calls["recovery"].append(kwargs) + return False + + monkeypatch.setattr(chat_routes, "_recover_empty_session_model", recover) monkeypatch.setattr(chat_routes, "_enforce_chat_privileges", lambda *args, **kwargs: None) monkeypatch.setattr(chat_routes, "resolve_session_auth", lambda *args, **kwargs: None) monkeypatch.setattr(chat_routes, "build_chat_context", build_context) @@ -483,6 +487,7 @@ async def test_sync_bearer_chat_cannot_use_research_memory_or_background_extract assert calls["memory"] == 0 assert calls["research"] == 0 assert calls["post"] and calls["post"][0]["allow_background_extraction"] is False + assert calls["recovery"] == [{"owner": "alice", "allow_live_probes": False}] @pytest.mark.asyncio @@ -497,6 +502,13 @@ async def test_stream_bearer_chat_disables_deferred_memory_extraction(monkeypatc captured, capture_completion=True, ) + recovery_calls = [] + + def recover(*args, **kwargs): + recovery_calls.append(kwargs) + return False + + monkeypatch.setattr(chat_routes, "_recover_empty_session_model", recover) request = SimpleNamespace( headers={}, app=SimpleNamespace(state=SimpleNamespace(auth_manager=None)), @@ -520,6 +532,184 @@ async def test_stream_bearer_chat_disables_deferred_memory_extraction(monkeypatc assert captured["post_processed"] assert captured["post_processed"][0][1]["allow_background_extraction"] is False + assert recovery_calls == [{"owner": "alice", "allow_live_probes": False}] + + +class _RecoveryPredicate: + def __or__(self, _other): + return self + + +class _RecoveryColumn: + def __eq__(self, _value): + return _RecoveryPredicate() + + +class _RecoveryEndpointModel: + is_enabled = _RecoveryColumn() + owner = _RecoveryColumn() + + +class _RecoverySessionModel: + id = _RecoveryColumn() + owner = _RecoveryColumn() + + +class _RecoveryQuery: + def __init__(self, db, model): + self.db = db + self.model = model + + def filter(self, *args, **kwargs): + return self + + def all(self): + if self.model is _RecoveryEndpointModel: + return [self.db.endpoint] + return [] + + def first(self): + if self.model is _RecoverySessionModel: + return self.db.session_row + return None + + +class _RecoveryDb: + def __init__(self, endpoint, session_row): + self.endpoint = endpoint + self.session_row = session_row + self.commits = 0 + self.rollbacks = 0 + + def query(self, model): + return _RecoveryQuery(self, model) + + def commit(self): + self.commits += 1 + + def rollback(self): + self.rollbacks += 1 + + def close(self): + return None + + +def _recovery_harness(monkeypatch, cached_models): + from routes import chat_routes + from src import chatgpt_subscription + + endpoint = SimpleNamespace( + id="endpoint-1", + base_url="https://chatgpt.com", + cached_models=json.dumps(cached_models), + hidden_models=None, + provider_auth_id="provider-auth-1", + ) + session_row = SimpleNamespace( + id="session-1", + owner="alice", + model="", + updated_at=None, + ) + db = _RecoveryDb(endpoint, session_row) + sess = SimpleNamespace( + id="session-1", + endpoint_url="https://chatgpt.com/backend-api/codex", + model="", + headers={}, + ) + + monkeypatch.setattr(chat_routes, "SessionLocal", lambda: db) + monkeypatch.setattr(chat_routes, "ModelEndpoint", _RecoveryEndpointModel) + monkeypatch.setattr(chat_routes, "DBSession", _RecoverySessionModel) + monkeypatch.setattr(chat_routes, "_session_url_matches_endpoint", lambda *args: True) + monkeypatch.setattr( + chatgpt_subscription, + "is_chatgpt_subscription_base", + lambda _url: True, + ) + return chat_routes, db, endpoint, session_row, sess + + +def test_bearer_empty_model_recovery_fails_without_cache_or_live_probe(monkeypatch): + chat_routes, db, endpoint, session_row, sess = _recovery_harness(monkeypatch, []) + from src import chatgpt_subscription, endpoint_resolver + + def forbidden(*args, **kwargs): + raise AssertionError("bearer recovery must not resolve credentials or fetch models") + + monkeypatch.setattr(chatgpt_subscription, "fetch_available_models", forbidden) + monkeypatch.setattr(endpoint_resolver, "resolve_endpoint_runtime", forbidden) + + assert chat_routes._recover_empty_session_model( + sess, + "session-1", + owner="alice", + allow_live_probes=False, + ) is False + assert sess.model == "" + assert session_row.model == "" + assert endpoint.cached_models == "[]" + assert db.commits == 0 + assert db.rollbacks == 0 + + +def test_bearer_model_recovery_uses_cache_without_endpoint_or_session_writes(monkeypatch): + chat_routes, db, endpoint, session_row, sess = _recovery_harness( + monkeypatch, + ["cached-model"], + ) + from src import chatgpt_subscription, endpoint_resolver + + def forbidden(*args, **kwargs): + raise AssertionError("bearer recovery must not resolve credentials or fetch models") + + monkeypatch.setattr(chatgpt_subscription, "fetch_available_models", forbidden) + monkeypatch.setattr(endpoint_resolver, "resolve_endpoint_runtime", forbidden) + + assert chat_routes._recover_empty_session_model( + sess, + "session-1", + owner="alice", + allow_live_probes=False, + ) is True + assert sess.model == "cached-model" + assert session_row.model == "" + assert endpoint.cached_models == '["cached-model"]' + assert db.commits == 0 + assert db.rollbacks == 0 + + +def test_interactive_model_recovery_retains_live_catalog_and_persistence(monkeypatch): + chat_routes, db, endpoint, session_row, sess = _recovery_harness(monkeypatch, []) + from src import chatgpt_subscription, endpoint_resolver + + seen = {} + + def resolve(ep, owner=None): + seen["resolve"] = (ep, owner) + return ep.base_url, "owner-secret" + + def fetch(api_key): + seen["fetch"] = api_key + return ["gpt-live"] + + monkeypatch.setattr(endpoint_resolver, "resolve_endpoint_runtime", resolve) + monkeypatch.setattr(chatgpt_subscription, "fetch_available_models", fetch) + + # Interactive callers retain the helper's live-probe default. + assert chat_routes._recover_empty_session_model( + sess, + "session-1", + owner="alice", + ) is True + assert seen["resolve"] == (endpoint, "alice") + assert seen["fetch"] == "owner-secret" + assert sess.model == "gpt-live" + assert session_row.model == "gpt-live" + assert json.loads(endpoint.cached_models) == ["gpt-live"] + assert db.commits == 2 + assert db.rollbacks == 0 @pytest.mark.asyncio