CyberGril/tests/test_api.py

"""R1 invariants through HTTP and real loopback streaming, with isolated databases."""

import asyncio
import concurrent.futures
import json
import socket
import threading
import time
from contextlib import contextmanager
from uuid import uuid4

import httpx
import pytest
import uvicorn
from conftest import engineering_budget
from fastapi.testclient import TestClient

from cybergril.api.app import create_app
from cybergril.contracts import HealthStatus, LLMDelta
from cybergril.providers.llm.echo import EchoLLM
from cybergril.think.router import build_router
from cybergril.think.scenes import Tier

CREATE = {"name": "星河", "birth": "2000-10-18", "hour": "亥时", "relation": "朋友"}


class ControlledLLM(EchoLLM):
    def __init__(self):
        super().__init__()
        self.calls = []
        self.closed = 0
        self.mode = "normal"
        self.health_calls = 0

    async def health(self):
        self.health_calls += 1
        if self.mode == "health_timeout":
            await asyncio.sleep(30)
        return HealthStatus(ok=self.mode != "unreachable")

    async def stream(self, messages, **kwargs):
        self.calls.append(messages)
        try:
            if self.mode == "error":
                raise OSError("secret upstream configuration must never be exposed")
            yield LLMDelta(text="答")
            if self.mode == "slow":
                await asyncio.sleep(30)
            if self.mode == "truncated":
                return
            yield LLMDelta(text="复")
            yield LLMDelta(done=True)
        finally:
            self.closed += 1


def make_app(path, provider=None, **kwargs):
    provider = provider or ControlledLLM()
    router = build_router(
        providers={Tier.HIGH: provider, Tier.STANDARD: provider, Tier.LOW: provider},
        budget=engineering_budget(),
    )
    return create_app(db_path=path, router=router, **kwargs), provider


def turn(client, text="你好", turn_id=None, epoch=None):
    return client.post(
        "/api/chat",
        json={
            "text": text,
            "turn_id": turn_id or str(uuid4()),
            "epoch": epoch or client.get("/api/health").json()["epoch"],
        },
    )


def events(response):
    return [
        json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: ")
    ]


def test_preview_validation_singleton_and_readonly(tmp_path):
    app, _ = make_app(tmp_path / "test.db")
    with TestClient(app) as client:
        assert client.get("/api/companion").status_code == 404
        preview = client.post(
            "/api/companion/preview", json={k: CREATE[k] for k in ("name", "birth", "hour")}
        )
        assert preview.status_code == 200
        assert len(preview.json()["pillars"]["day"]) == 2
        assert client.get("/api/companion").status_code == 404
        assert app.state.store.query_one("SELECT count(*) n FROM bazi_core")["n"] == 0
        for bad in [
            dict(CREATE, name="<x>"),
            dict(CREATE, birth="2000-02-30"),
            dict(CREATE, birth="2020-01-01"),
            dict(CREATE, birth="1899-01-01"),
            dict(CREATE, hour="未知"),
            dict(CREATE, relation=""),
            dict(CREATE, nickname="x\x01"),
            {k: v for k, v in CREATE.items() if k != "relation"},
            dict(CREATE, derived={}),
        ]:
            response = client.post("/api/companion", json=bad)
            assert response.status_code == 422, response.text
            assert response.json()["fields"]
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as pool:
            codes = list(
                pool.map(lambda _: client.post("/api/companion", json=CREATE).status_code, range(5))
            )
        assert sorted(codes) == [201, 409, 409, 409, 409]
        identity = client.get("/api/companion").json()
        assert identity["color"] == preview.json()
        assert "derived" not in identity and "strength" not in identity
        assert (
            client.patch(
                "/api/relation", json={"relation": "恋人", "birth": "1990-01-01"}
            ).status_code
            == 422
        )
        assert app.state.store.query_one("PRAGMA journal_mode")[0] == "wal"


def test_history_context_relation_retry_and_restart(tmp_path):
    path = tmp_path / "test.db"
    app, provider = make_app(path)
    with TestClient(app) as client:
        client.post("/api/companion", json=CREATE)
        epoch = client.get("/api/health").json()["epoch"]
        tid = str(uuid4())
        result = events(turn(client, "我的暗号是蓝星", tid))
        assert result[-1]["t"] == "chat.final"
        assert all(e["turn_id"] == tid and e["epoch"] == epoch for e in result)
        assert result[0]["t"] == "chat.delta"
        turn(client, "记得暗号吗")
        assert any(m.content == "我的暗号是蓝星" for m in provider.calls[-1])
        assert (
            "星河" in provider.calls[-1][0].content
            and "尚未接通" not in provider.calls[-1][-1].content
        )
        assert len(client.get("/api/messages").json()["messages"]) == 4
        assert events(turn(client, "我的暗号是蓝星", tid))[-1]["t"] == "chat.final"
        assert len(provider.calls) == 2  # A lost final can be replayed without regenerating.
        assert turn(client, "tampered", tid).status_code == 409
        first = client.get("/api/messages?limit=2").json()
        second = client.get(f"/api/messages?limit=2&before={first['next_before']}").json()
        assert second["messages"][-1]["id"] < first["messages"][0]["id"]
        assert client.get("/api/messages?limit=201").status_code == 422
        core = dict(app.state.store.query_one("SELECT * FROM bazi_core"))
        updated = client.patch(
            "/api/relation", json={"relation": "恋人", "nickname": "小星", "boundary": "慢慢来"}
        ).json()
        assert updated["relation"] == "恋人"
        assert dict(app.state.store.query_one("SELECT * FROM bazi_core")) == core
        assert len(client.get("/api/messages").json()["messages"]) == 4
    app2, _ = make_app(path)
    with TestClient(app2) as client:
        assert client.get("/api/health").json()["epoch"] == epoch + 1
        assert client.get("/api/companion").json() == updated
        assert len(client.get("/api/messages").json()["messages"]) == 4
        assert turn(client, epoch=epoch).json()["epoch"] == epoch + 1
        assert events(turn(client))[-1]["t"] == "chat.final"


def test_errors_retry_context_excludes_failed_and_health_cache(tmp_path):
    app, provider = make_app(tmp_path / "test.db", stream_timeout=0.03)
    with TestClient(app) as client:
        client.post("/api/companion", json=CREATE)
        provider.mode = "unreachable"
        health = client.get("/api/health").json()
        assert not health["llm"]["ok"]
        assert not health["capabilities"]["life"]
        client.get("/api/health")
        assert provider.health_calls == 1
        for mode in ("error", "slow", "truncated"):
            provider.mode = mode
            tid = str(uuid4())
            response = turn(client, text="失败的消息" + mode, turn_id=tid)
            assert events(response)[-1]["t"] == "chat.error"
            assert "secret" not in response.text
            assert client.get("/api/messages").json()["messages"][-1]["status"] == "failed"
        provider.mode = "normal"
        response = turn(client, text="重试成功")
        assert events(response)[-1]["t"] == "chat.final"
        assert len(provider.calls[-1]) == 2  # system + current; no failed turns
        provider.mode = "error"
        tid = str(uuid4())
        turn(client, text="同一轮", turn_id=tid)
        count = len(client.get("/api/messages").json()["messages"])
        provider.mode = "normal"
        turn(client, text="同一轮", turn_id=tid)
        assert len(client.get("/api/messages").json()["messages"]) == count + 1
        assert provider.closed == len(provider.calls)


@contextmanager
def live_server(app, port=0):
    sock = socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("127.0.0.1", port))
    sock.listen()
    server = uvicorn.Server(uvicorn.Config(app, log_level="error", lifespan="on"))
    thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]})
    thread.start()
    try:
        deadline = time.monotonic() + 10
        while not server.started:
            if not thread.is_alive() or time.monotonic() > deadline:
                raise RuntimeError("local test server failed to start")
            time.sleep(0.01)
        yield f"http://127.0.0.1:{sock.getsockname()[1]}"
    finally:
        server.should_exit = True
        thread.join(timeout=10)
        sock.close()
        assert not thread.is_alive()


@pytest.mark.asyncio
async def test_real_disconnect_and_immediate_new_turn(tmp_path):
    app, provider = make_app(tmp_path / "test.db")
    with live_server(app) as url:
        async with httpx.AsyncClient(base_url=url) as client:
            await client.post("/api/companion", json=CREATE)
            epoch = (await client.get("/api/health")).json()["epoch"]
            provider.mode = "slow"
            tid = str(uuid4())
            async with client.stream(
                "POST", "/api/chat", json={"text": "取消我", "turn_id": tid, "epoch": epoch}
            ) as response:
                async for line in response.aiter_lines():
                    if line.startswith("data:"):
                        break
            # Closing the browser response must cancel upstream while its next token is blocked.
            deadline = time.monotonic() + 1
            while provider.closed < 1 and time.monotonic() < deadline:
                await asyncio.sleep(0.01)
            assert provider.closed == 1
            provider.mode = "normal"
            new_tid = str(uuid4())
            response = await client.post(
                "/api/chat", json={"text": "新问题", "turn_id": new_tid, "epoch": epoch}
            )
            assert events(response)[-1]["t"] == "chat.final"
            assert all(e["turn_id"] == new_tid for e in events(response))
            messages = (await client.get("/api/messages")).json()["messages"]
            assert messages[0]["status"] == "cancelled"
            assert len(messages) == 3  # discarded partial reply is never a complete answer
            assert len(provider.calls[-1]) == 2  # cancelled turn excluded from context


def test_concurrent_identical_retry_generates_once(tmp_path):
    app, provider = make_app(tmp_path / "test.db")
    with TestClient(app) as client:
        client.post("/api/companion", json=CREATE)
        epoch = client.get("/api/health").json()["epoch"]
        tid = str(uuid4())
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
            responses = list(pool.map(lambda _: turn(client, "同一轮", tid, epoch), range(4)))
        assert all(events(response)[-1]["t"] == "chat.final" for response in responses)
        assert len(provider.calls) == 1
        assert len(client.get("/api/messages").json()["messages"]) == 2


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "status,data,expected",
    [
        (401, {}, False),
        (404, {}, False),
        (503, {}, False),
        (200, {"data": [{"id": "other"}]}, False),
        (200, {"data": [{"id": "test-model"}]}, True),
    ],
)
async def test_model_health_requires_success_and_requested_model(status, data, expected):
    from cybergril.providers.llm.openai_compat import OpenAICompatLLM

    provider = OpenAICompatLLM(base_url="http://isolated/v1", model="test-model")
    await provider._client.aclose()
    provider._client = httpx.AsyncClient(
        base_url="http://isolated/v1",
        transport=httpx.MockTransport(lambda request: httpx.Response(status, json=data)),
    )
    try:
        assert (await provider.health()).ok is expected
    finally:
        await provider.aclose()


def test_health_timeout_is_bounded_and_cached(tmp_path):
    app, provider = make_app(tmp_path / "test.db")
    provider.mode = "health_timeout"
    with TestClient(app) as client:
        start = time.monotonic()
        assert client.get("/api/health").json()["llm"]["ok"] is False
        assert time.monotonic() - start < 6
        assert client.get("/api/health").json()["llm"]["ok"] is False
        assert provider.health_calls == 1


@pytest.mark.asyncio
async def test_disconnect_closes_real_openai_http_stream(tmp_path):
    from fastapi import FastAPI
    from fastapi.responses import StreamingResponse

    from cybergril.providers.llm.openai_compat import OpenAICompatLLM

    closed = threading.Event()
    gateway = FastAPI()

    @gateway.post("/v1/chat/completions")
    async def completion():
        async def chunks():
            try:
                # ends with a sentence terminator so the HEL-908 reasoning
                # separator flushes it immediately (this test is about upstream
                # closure on disconnect, not about hold semantics)
                yield 'data: {"choices":[{"delta":{"content":"部分。"}}]}\n\n'
                await asyncio.sleep(30)
            finally:
                closed.set()

        return StreamingResponse(chunks(), media_type="text/event-stream")

    with live_server(gateway) as gateway_url:
        provider = OpenAICompatLLM(base_url=gateway_url + "/v1", model="controlled")
        app, _ = make_app(tmp_path / "test.db", provider)
        with live_server(app) as url:
            async with httpx.AsyncClient(base_url=url) as client:
                await client.post("/api/companion", json=CREATE)
                async with client.stream(
                    "POST",
                    "/api/chat",
                    json={"text": "断开上游", "turn_id": str(uuid4()), "epoch": 1},
                ) as response:
                    async for line in response.aiter_lines():
                        if line.startswith("data:"):
                            break
                deadline = time.monotonic() + 1
                while not closed.is_set() and time.monotonic() < deadline:
                    await asyncio.sleep(0.01)
                assert closed.is_set()  # Real OpenAICompatLLM's HTTP response was actually closed.


def test_retry_old_turn_preserves_user_assistant_context_pairs(tmp_path):
    app, provider = make_app(tmp_path / "test.db")
    with TestClient(app) as client:
        client.post("/api/companion", json=CREATE)
        tid = str(uuid4())
        provider.mode = "error"
        turn(client, "旧轮失败", tid)
        provider.mode = "normal"
        turn(client, "中间的新轮")
        turn(client, "旧轮失败", tid)
        turn(client, "最后一轮")
        context = provider.calls[-1]
        assert [m.role for m in context] == [
            "system",
            "user",
            "assistant",
            "user",
            "assistant",
            "user",
        ]
        assert [m.content for m in context[1::2]] == ["中间的新轮", "旧轮失败", "最后一轮"]


def test_upstream_http_503_maps_to_chinese_sse_error(tmp_path):
    from cybergril.providers.llm.openai_compat import OpenAICompatLLM

    provider = OpenAICompatLLM(base_url="http://isolated/v1", model="controlled")
    provider._client = httpx.AsyncClient(
        base_url="http://isolated/v1",
        transport=httpx.MockTransport(
            lambda request: httpx.Response(503, json={"error": "secret gateway path"})
        ),
    )
    app, _ = make_app(tmp_path / "test.db", provider)
    with TestClient(app) as client:
        client.post("/api/companion", json=CREATE)
        response = turn(client, epoch=1)
        assert events(response)[-1]["t"] == "chat.error"
        assert "稍后再试" in response.text and "secret" not in response.text
        assert client.get("/api/messages").json()["messages"][0]["status"] == "failed"