Parth
e7d75c3b05
* Rewrite comments in Google developer documentation style Rewrite the comments and docstrings across the backend core modules so they read plainly. The previous prose was accurate but dense and figurative, which made it slow to skim. Applies the Google developer documentation style guide: short sentences, active voice, present tense, American spelling, and no metaphors, idioms, or rhetorical asides. Replaces em-dash chains with separate sentences.
258 lines
10 KiB
Python
258 lines
10 KiB
Python
"""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
|