"""HTTP tests for the AI Chat scratchpad (power users only). Covers the access gate, the streamed reply, and the demo-key model pinning. This pinning must not let a public visitor reach paid models through this page. python -m pytest tests/test_chat.py -v """ import os import tempfile _tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False) _tmp.close() os.environ["AIDND_DB_PATH"] = _tmp.name os.environ.pop("AIDND_DATABASE_URL", None) os.environ.pop("DATABASE_URL", None) import pytest from fastapi import Depends from fastapi.testclient import TestClient from app import auth, limits, models from app.database import Base, SessionLocal, engine, get_db from app.main import app from app.routers import chat class FakeProvider: """Records what it was constructed with, then streams a fixed reply. Stands in for the real egress point, so asserting on last_key/last_model is asserting on exactly what would have gone over the wire.""" last_usage = None last_model = None last_key = None last_endpoint = None last_messages = None def __init__(self, endpoint_url, api_key, model, api_mode="chat", reasoning_max_tokens=0): FakeProvider.last_model = model FakeProvider.last_key = api_key FakeProvider.last_endpoint = endpoint_url async def chat(self, messages, *, temperature, max_tokens): FakeProvider.last_messages = messages yield ("reasoning", "hmm") yield ("text", "Hello back.") @pytest.fixture() def client(monkeypatch): Base.metadata.create_all(bind=engine) setup = SessionLocal() user = models.User(is_guest=False, email="power@example.com") setup.add(user) setup.flush() setup.add(models.Settings(user_id=user.id, api_key="enc:dummy", model="test-model")) setup.commit() user_id = user.id setup.close() monkeypatch.setattr(chat, "OpenAICompatibleProvider", FakeProvider) monkeypatch.setattr(limits, "rate_limit", lambda *a, **k: None) # Multi-user mode is what makes the power-user gate meaningful, because # local mode trusts everyone. The allowlist is set per test. monkeypatch.setattr(auth, "MULTI_USER", True) monkeypatch.setattr(auth, "POWER_USERS", {"power@example.com"}) # These tests deliberately do not stub resolve_provider_config. The # point is to exercise the real BYOK-vs-demo decision, since that # decision is what keeps the shared key off paid models. Each test # picks a mode with _byok/_demo below. def _current_user(db=Depends(get_db)): return db.get(models.User, user_id) app.dependency_overrides[auth.get_current_user] = _current_user try: yield TestClient(app) finally: app.dependency_overrides.clear() Base.metadata.drop_all(bind=engine) def _send(client, **extra): return client.post("/api/chat/stream", json={"messages": [{"role": "user", "content": "hi"}], **extra}) def _byok(monkeypatch): """The user brought their own key: no demo key in play, any model allowed.""" monkeypatch.setattr(auth, "demo_enabled", lambda: False) db = SessionLocal() try: settings = db.query(models.Settings).first() settings.api_key = "sk-my-own-key" # legacy-plaintext path: used as-is db.commit() finally: db.close() def _demo(monkeypatch, whitelist=("free/allowed",)): """The user has no key, so turns run on the server-funded demo key.""" monkeypatch.setattr(auth, "demo_enabled", lambda: True) monkeypatch.setattr(auth, "DEMO_API_KEY", "demo-key") monkeypatch.setattr(auth, "DEMO_ENDPOINT_URL", "http://demo") monkeypatch.setattr(auth, "DEMO_MODELS", list(whitelist)) def test_non_power_user_gets_404(client, monkeypatch): monkeypatch.setattr(auth, "POWER_USERS", set()) assert _send(client).status_code == 404 assert client.get("/api/chat/config").status_code == 404 def test_power_user_streams_a_reply(client, monkeypatch): _byok(monkeypatch) resp = _send(client) assert resp.status_code == 200, resp.text assert '"type": "reasoning"' in resp.text assert "Hello back." in resp.text assert '"type": "done"' in resp.text assert FakeProvider.last_messages == [{"role": "user", "content": "hi"}] def test_system_prompt_and_model_override_are_honoured(client, monkeypatch): _byok(monkeypatch) resp = client.post("/api/chat/stream", json={ "messages": [ {"role": "system", "content": "Be terse."}, {"role": "user", "content": "hi"}, ], "model": "some/other-model", }) assert resp.status_code == 200, resp.text # BYOK: any model the user names is passed straight through, on their key. assert FakeProvider.last_model == "some/other-model" assert FakeProvider.last_key == "sk-my-own-key" assert FakeProvider.last_messages[0] == {"role": "system", "content": "Be terse."} def test_demo_key_pins_model_to_whitelist(client, monkeypatch): _demo(monkeypatch) resp = _send(client, model="expensive/paid-model") assert resp.status_code == 200, resp.text # Refused visibly: the whitelisted model runs instead, with a note. The # paid slug must never reach the wire alongside the server-funded key. assert FakeProvider.last_model == "free/allowed" assert FakeProvider.last_key == "demo-key" assert '"type": "note"' in resp.text # A whitelisted model is still selectable on the demo key. _demo(monkeypatch, ["free/allowed", "free/second"]) resp = _send(client, model="free/second") assert resp.status_code == 200, resp.text assert FakeProvider.last_model == "free/second" def test_demo_key_ignores_an_off_whitelist_settings_model(client, monkeypatch): """The override is not the only untrusted input. `Settings.model` is also user-set, and it must be pinned the same way when there is no BYOK key.""" _demo(monkeypatch) db = SessionLocal() try: db.query(models.Settings).first().model = "expensive/paid-model" db.commit() finally: db.close() resp = _send(client) assert resp.status_code == 200, resp.text assert FakeProvider.last_model == "free/allowed" def test_demo_key_endpoint_cannot_be_redirected(client, monkeypatch): """A user-controlled `endpoint_url` would leak the key itself, which is worse than spending it. The demo branch pins the URL too.""" _demo(monkeypatch) db = SessionLocal() try: db.query(models.Settings).first().endpoint_url = "http://attacker.example/v1" db.commit() finally: db.close() assert _send(client).status_code == 200 assert FakeProvider.last_endpoint == "http://demo" assert FakeProvider.last_key == "demo-key" def test_provider_config_refuses_server_funded_paid_model(monkeypatch): """The structural backstop: a hand-built config (a future code path that forgets to go through resolve_provider_config) cannot run a server-funded turn on an off-whitelist model.""" monkeypatch.setattr(auth, "DEMO_API_KEY", "demo-key") monkeypatch.setattr(auth, "DEMO_MODELS", ["free/allowed"]) with pytest.raises(ValueError): auth.ProviderConfig("http://demo", "demo-key", "expensive/paid-model", True) auth.ProviderConfig("http://demo", "demo-key", "free/allowed", True) # whitelisted: fine # The user's own key with any model stays fine. auth.ProviderConfig("http://any", "sk-mine", "expensive/paid-model", False) def test_byok_user_may_reuse_the_demo_keys_value(client, monkeypatch): """Regression: the demo key is just an OpenRouter key, so a user can paste that same value into their own Settings. That is still BYOK, because the user is paying, and it must not trip the guard. It used to raise on every resolution, which returned a 500 from `GET /auth/me` and broke the entire SPA (no nav, no chat).""" monkeypatch.setattr(auth, "demo_enabled", lambda: True) monkeypatch.setattr(auth, "DEMO_API_KEY", "shared-key") monkeypatch.setattr(auth, "DEMO_ENDPOINT_URL", "http://demo") monkeypatch.setattr(auth, "DEMO_MODELS", ["free/allowed"]) db = SessionLocal() try: settings = db.query(models.Settings).first() settings.api_key = "shared-key" # same value, but supplied by the user settings.model = "expensive/paid-model" # their spend, their choice db.commit() finally: db.close() assert client.get("/api/auth/me").status_code == 200 assert client.get("/api/chat/config").status_code == 200 resp = _send(client) assert resp.status_code == 200, resp.text assert FakeProvider.last_model == "expensive/paid-model" assert FakeProvider.last_key == "shared-key" def test_resolve_provider_config_is_the_single_choke_point(monkeypatch): """Turns, AI Chat, and the connection test all resolve through this one function, so pinning it here pins every caller. No DB or HTTP needed.""" monkeypatch.setattr(auth, "demo_enabled", lambda: True) monkeypatch.setattr(auth, "DEMO_API_KEY", "demo-key") monkeypatch.setattr(auth, "DEMO_ENDPOINT_URL", "http://demo") monkeypatch.setattr(auth, "DEMO_MODELS", ["free/allowed"]) # No key of their own: both endpoint and model are pinned, regardless # of what they set. no_key = models.Settings(endpoint_url="http://mine/v1", api_key="", model="expensive/paid") assert auth.resolve_provider_config(no_key) == auth.ProviderConfig( "http://demo", "demo-key", "free/allowed", True) assert auth.resolve_provider_config( no_key, model_override="expensive/paid").model == "free/allowed" assert auth.resolve_provider_config( no_key, model_override="free/allowed").model == "free/allowed" # Their own key: their endpoint, their key, their choice of model. byok = models.Settings(endpoint_url="http://mine/v1", api_key="sk-mine", model="expensive/paid") assert auth.resolve_provider_config(byok) == auth.ProviderConfig( "http://mine/v1", "sk-mine", "expensive/paid", False) def test_oversized_conversation_is_refused(client, monkeypatch): _byok(monkeypatch) huge = "x" * 90_000 resp = client.post("/api/chat/stream", json={ "messages": [{"role": "user", "content": huge} for _ in range(5)], }) assert resp.status_code == 413, resp.text