382 lines
15 KiB
Python
382 lines
15 KiB
Python
"""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 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}
|
|
)
|
|
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:
|
|
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"
|