interactive-story/backend/tests/test_chat.py

"""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