qzl
504 lines
16 KiB
Python
504 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
from types import SimpleNamespace
|
|
from uuid import uuid4
|
|
|
|
from ag_ui.core import RunAgentInput
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app import app
|
|
from core.auth.models import CurrentUser
|
|
from v1.agent import router as agent_router
|
|
from v1.agent.dependencies import get_agent_service
|
|
from v1.users.dependencies import get_current_user
|
|
|
|
|
|
class _FakeAgentService:
|
|
def __init__(self) -> None:
|
|
self._stream_called = False
|
|
|
|
async def enqueue_run(self, *, run_input: RunAgentInput, current_user: CurrentUser):
|
|
del current_user
|
|
return SimpleNamespace(
|
|
task_id="task-run-1",
|
|
thread_id=run_input.thread_id,
|
|
run_id=run_input.run_id,
|
|
created=False,
|
|
)
|
|
|
|
async def enqueue_resume(
|
|
self,
|
|
*,
|
|
thread_id: str,
|
|
run_input: RunAgentInput,
|
|
current_user: CurrentUser,
|
|
):
|
|
del thread_id, current_user
|
|
return SimpleNamespace(
|
|
task_id="task-resume-1",
|
|
thread_id=run_input.thread_id,
|
|
run_id=run_input.run_id,
|
|
created=False,
|
|
)
|
|
|
|
async def stream_events(
|
|
self,
|
|
*,
|
|
thread_id: str,
|
|
last_event_id: str | None,
|
|
current_user: CurrentUser,
|
|
) -> list[dict[str, object]]:
|
|
del thread_id, current_user
|
|
if self._stream_called:
|
|
return []
|
|
self._stream_called = True
|
|
return [
|
|
{
|
|
"id": "2-0",
|
|
"event": {
|
|
"type": "RUN_STARTED",
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-1",
|
|
},
|
|
"cursor": last_event_id,
|
|
}
|
|
]
|
|
|
|
async def get_history_snapshot(
|
|
self,
|
|
*,
|
|
thread_id: str,
|
|
before: str | None,
|
|
current_user: CurrentUser,
|
|
) -> dict[str, object]:
|
|
del current_user
|
|
return {
|
|
"type": "STATE_SNAPSHOT",
|
|
"threadId": thread_id,
|
|
"snapshot": {
|
|
"scope": "history_day",
|
|
"day": before or "2026-03-07",
|
|
"hasMore": False,
|
|
"messages": [
|
|
{
|
|
"id": "msg-h1",
|
|
"role": "assistant",
|
|
"content": "history-message",
|
|
}
|
|
],
|
|
},
|
|
}
|
|
|
|
async def get_user_history_snapshot(
|
|
self,
|
|
*,
|
|
current_user: CurrentUser,
|
|
thread_id: str | None,
|
|
before: str | None,
|
|
) -> dict[str, object]:
|
|
del current_user, before
|
|
return {
|
|
"type": "STATE_SNAPSHOT",
|
|
"threadId": thread_id or "00000000-0000-0000-0000-000000000001",
|
|
"snapshot": {
|
|
"scope": "history_day",
|
|
"day": "2026-03-07",
|
|
"hasMore": False,
|
|
"messages": [],
|
|
},
|
|
}
|
|
|
|
|
|
def test_run_requires_auth_and_returns_202_task_id() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
client = TestClient(app)
|
|
original_allow_run = agent_router._allow_run_request
|
|
|
|
async def _allow_run(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
agent_router._allow_run_request = _allow_run # type: ignore[assignment]
|
|
|
|
try:
|
|
unauthorized = client.post(
|
|
"/api/v1/agent/runs",
|
|
json={
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-1",
|
|
"state": {},
|
|
"messages": [{"id": "u1", "role": "user", "content": "hello"}],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
},
|
|
)
|
|
assert unauthorized.status_code == 401
|
|
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
authorized = client.post(
|
|
"/api/v1/agent/runs",
|
|
json={
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-1",
|
|
"state": {},
|
|
"messages": [{"id": "u1", "role": "user", "content": "hello"}],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
},
|
|
)
|
|
assert authorized.status_code == 202
|
|
assert authorized.json()["taskId"] == "task-run-1"
|
|
assert authorized.json()["threadId"] == "00000000-0000-0000-0000-000000000001"
|
|
assert authorized.json()["runId"] == "run-1"
|
|
assert authorized.json()["created"] is False
|
|
finally:
|
|
agent_router._allow_run_request = original_allow_run # type: ignore[assignment]
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_stream_reads_from_last_event_id() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
client = TestClient(app)
|
|
original_acquire = agent_router._acquire_sse_slot
|
|
original_release = agent_router._release_sse_slot
|
|
|
|
async def _allow_slot(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
async def _noop_release(*, user_id: str) -> None:
|
|
del user_id
|
|
return None
|
|
|
|
agent_router._acquire_sse_slot = _allow_slot # type: ignore[assignment]
|
|
agent_router._release_sse_slot = _noop_release # type: ignore[assignment]
|
|
|
|
try:
|
|
response = client.get(
|
|
"/api/v1/agent/runs/00000000-0000-0000-0000-000000000001/events?idle_limit=1",
|
|
headers={"Last-Event-ID": "1-0"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.headers["content-type"].startswith("text/event-stream")
|
|
assert "id: 2-0" in response.text
|
|
assert "event: RUN_STARTED" in response.text
|
|
assert '"threadId":"00000000-0000-0000-0000-000000000001"' in response.text
|
|
finally:
|
|
agent_router._acquire_sse_slot = original_acquire # type: ignore[assignment]
|
|
agent_router._release_sse_slot = original_release # type: ignore[assignment]
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_stream_rejects_invalid_last_event_id() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
client = TestClient(app)
|
|
|
|
try:
|
|
response = client.get(
|
|
"/api/v1/agent/runs/00000000-0000-0000-0000-000000000001/events",
|
|
headers={"Last-Event-ID": "bad-id"},
|
|
)
|
|
assert response.status_code == 422
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_history_returns_state_snapshot() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
client = TestClient(app)
|
|
|
|
try:
|
|
unauthorized = client.get(
|
|
"/api/v1/agent/runs/00000000-0000-0000-0000-000000000001/history"
|
|
)
|
|
assert unauthorized.status_code == 401
|
|
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
authorized = client.get(
|
|
"/api/v1/agent/runs/00000000-0000-0000-0000-000000000001/history",
|
|
params={"before": "2026-03-07"},
|
|
)
|
|
assert authorized.status_code == 200
|
|
payload = authorized.json()
|
|
assert payload["type"] == "STATE_SNAPSHOT"
|
|
assert payload["threadId"] == "00000000-0000-0000-0000-000000000001"
|
|
assert payload["snapshot"]["scope"] == "history_day"
|
|
assert payload["snapshot"]["day"] == "2026-03-07"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_user_history_returns_latest_snapshot() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
client = TestClient(app)
|
|
try:
|
|
response = client.get("/api/v1/agent/history")
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["type"] == "STATE_SNAPSHOT"
|
|
assert body["threadId"] == "00000000-0000-0000-0000-000000000001"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_run_rejects_oversized_user_text_payload() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
client = TestClient(app)
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/runs",
|
|
json={
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-oversize",
|
|
"state": {},
|
|
"messages": [
|
|
{
|
|
"id": "u1",
|
|
"role": "user",
|
|
"content": "x" * 11000,
|
|
}
|
|
],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
},
|
|
)
|
|
assert response.status_code == 422
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_run_rejects_client_supplied_history_messages() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
client = TestClient(app)
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/runs",
|
|
json={
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-history",
|
|
"state": {},
|
|
"messages": [
|
|
{"id": "a1", "role": "assistant", "content": "old"},
|
|
{"id": "u1", "role": "user", "content": "new"},
|
|
],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
},
|
|
)
|
|
assert response.status_code == 422
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_resume_accepts_tool_message_without_user_message() -> None:
|
|
app.dependency_overrides[get_agent_service] = lambda: _FakeAgentService()
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
client = TestClient(app)
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/runs/00000000-0000-0000-0000-000000000001/resume",
|
|
json={
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-resume-1",
|
|
"state": {},
|
|
"messages": [
|
|
{
|
|
"id": "tool-1",
|
|
"role": "tool",
|
|
"toolCallId": "call-1",
|
|
"content": '{"toolName":"navigate_to_route","toolArgs":{"target":"/calendar/dayweek"},"nonce":"n1","result":{"ok":true}}',
|
|
}
|
|
],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
},
|
|
)
|
|
assert response.status_code == 202
|
|
assert response.json()["taskId"] == "task-resume-1"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_asr_transcribe_returns_sync_transcript(monkeypatch) -> None:
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
|
|
async def _allow_transcribe(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_transcribe_request", _allow_transcribe)
|
|
|
|
async def mock_transcribe_file(file_path: str, filename: str) -> str:
|
|
assert file_path.endswith(".wav")
|
|
assert filename == "test.wav"
|
|
return "这是测试转写结果"
|
|
|
|
monkeypatch.setattr(
|
|
"v1.agent.service.asr_service.transcribe_file",
|
|
mock_transcribe_file,
|
|
)
|
|
|
|
client = TestClient(app)
|
|
|
|
wav_content = b"RIFF\x24\x80\x00\x00WAVEfmt "
|
|
wav_file = BytesIO(wav_content)
|
|
wav_file.name = "test.wav"
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/transcribe",
|
|
files={"audio": ("test.wav", wav_file, "audio/wav")},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "transcript" in data
|
|
assert data["transcript"] == "这是测试转写结果"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_asr_transcribe_rejects_oversized_audio(monkeypatch) -> None:
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
|
|
monkeypatch.setattr(agent_router, "_MAX_TRANSCRIBE_AUDIO_BYTES", 4)
|
|
|
|
async def _allow_transcribe(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_transcribe_request", _allow_transcribe)
|
|
|
|
client = TestClient(app)
|
|
oversized = BytesIO(b"12345")
|
|
oversized.name = "test.wav"
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/transcribe",
|
|
files={"audio": ("test.wav", oversized, "audio/wav")},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert response.json()["detail"] == "Audio file too large"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_asr_transcribe_rejects_non_wav_audio(monkeypatch) -> None:
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
|
|
async def _allow_transcribe(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_transcribe_request", _allow_transcribe)
|
|
|
|
client = TestClient(app)
|
|
fake_mp3 = BytesIO(b"fake-mp3")
|
|
fake_mp3.name = "test.mp3"
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/transcribe",
|
|
files={"audio": ("test.mp3", fake_mp3, "audio/mpeg")},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert response.json()["detail"] == "Unsupported audio format"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_asr_transcribe_rejects_invalid_wav_payload(monkeypatch) -> None:
|
|
app.dependency_overrides[get_current_user] = lambda: CurrentUser(
|
|
id=uuid4(), email="user@example.com"
|
|
)
|
|
|
|
async def _allow_transcribe(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_transcribe_request", _allow_transcribe)
|
|
|
|
client = TestClient(app)
|
|
fake_payload = BytesIO(b"not-a-wav")
|
|
fake_payload.name = "test.wav"
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/transcribe",
|
|
files={"audio": ("test.wav", fake_payload, "audio/wav")},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert response.json()["detail"] == "Unsupported audio format"
|
|
finally:
|
|
app.dependency_overrides = {}
|
|
|
|
|
|
def test_asr_transcribe_rejects_when_rate_limited_for_current_user(monkeypatch) -> None:
|
|
known_user = CurrentUser(id=uuid4(), email="user@example.com")
|
|
app.dependency_overrides[get_current_user] = lambda: known_user
|
|
|
|
captured_user_ids: list[str] = []
|
|
|
|
async def _deny_transcribe(*, user_id: str) -> bool:
|
|
captured_user_ids.append(user_id)
|
|
return False
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_transcribe_request", _deny_transcribe)
|
|
|
|
client = TestClient(app)
|
|
wav_content = b"RIFF\x24\x80\x00\x00WAVEfmt "
|
|
wav_file = BytesIO(wav_content)
|
|
wav_file.name = "test.wav"
|
|
|
|
try:
|
|
response = client.post(
|
|
"/api/v1/agent/transcribe",
|
|
files={"audio": ("test.wav", wav_file, "audio/wav")},
|
|
)
|
|
|
|
assert response.status_code == 429
|
|
assert response.json()["detail"] == "Too many transcribe requests"
|
|
assert captured_user_ids == [str(known_user.id)]
|
|
finally:
|
|
app.dependency_overrides = {}
|