zl-q
7b8865e256
- 新增 stepStarted/stepFinished 事件类型支持 - 前端实现图片附件上传和预览功能 - 后端增强工具结果存储和事件处理 - 完善相关单元测试和集成测试
269 lines
8.3 KiB
Python
269 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
from uuid import uuid4
|
|
|
|
from ag_ui.core import RunAgentInput
|
|
from fastapi import HTTPException
|
|
import pytest
|
|
|
|
from core.auth.models import CurrentUser
|
|
from v1.agent import router as agent_router
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_allow_run_request_fails_closed_when_redis_unavailable(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def _raise_redis_error():
|
|
raise RuntimeError("redis unavailable")
|
|
|
|
monkeypatch.setattr(agent_router, "get_or_init_redis_client", _raise_redis_error)
|
|
|
|
allowed = await agent_router._allow_run_request(user_id="user-1")
|
|
|
|
assert allowed is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_sse_slot_fails_closed_when_redis_unavailable(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def _raise_redis_error():
|
|
raise RuntimeError("redis unavailable")
|
|
|
|
monkeypatch.setattr(agent_router, "get_or_init_redis_client", _raise_redis_error)
|
|
|
|
allowed = await agent_router._acquire_sse_slot(user_id="user-1")
|
|
|
|
assert allowed is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_allow_transcribe_request_fails_closed_when_redis_unavailable(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def _raise_redis_error():
|
|
raise RuntimeError("redis unavailable")
|
|
|
|
monkeypatch.setattr(agent_router, "get_or_init_redis_client", _raise_redis_error)
|
|
|
|
allowed = await agent_router._allow_transcribe_request(user_id="user-1")
|
|
|
|
assert allowed is False
|
|
|
|
|
|
def _resume_input_with_tool_message() -> RunAgentInput:
|
|
return RunAgentInput.model_validate(
|
|
{
|
|
"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","result":{"ok":true}}',
|
|
}
|
|
],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enqueue_resume_rejects_without_tool_contract() -> None:
|
|
request = RunAgentInput.model_validate(
|
|
{
|
|
"threadId": "00000000-0000-0000-0000-000000000001",
|
|
"runId": "run-resume-invalid",
|
|
"state": {},
|
|
"messages": [
|
|
{
|
|
"id": "u1",
|
|
"role": "user",
|
|
"content": "continue",
|
|
}
|
|
],
|
|
"tools": [],
|
|
"context": [],
|
|
"forwardedProps": {},
|
|
}
|
|
)
|
|
|
|
class _Service:
|
|
async def enqueue_resume(self, **kwargs): # noqa: ANN003
|
|
del kwargs
|
|
raise AssertionError("enqueue_resume should not be called")
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await agent_router.enqueue_resume(
|
|
thread_id="00000000-0000-0000-0000-000000000001",
|
|
request=request,
|
|
service=cast(Any, _Service()),
|
|
current_user=CurrentUser(id=uuid4(), email="user@example.com"),
|
|
)
|
|
|
|
assert exc_info.value.status_code == 422
|
|
assert (
|
|
exc_info.value.detail
|
|
== "RunAgentInput.messages requires a tool message with toolCallId for resume"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enqueue_resume_rejects_when_rate_limited(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
request = _resume_input_with_tool_message()
|
|
|
|
async def _deny_run(*, user_id: str) -> bool:
|
|
del user_id
|
|
return False
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_run_request", _deny_run)
|
|
|
|
class _Service:
|
|
async def enqueue_resume(self, **kwargs): # noqa: ANN003
|
|
del kwargs
|
|
raise AssertionError("enqueue_resume should not be called")
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await agent_router.enqueue_resume(
|
|
thread_id="00000000-0000-0000-0000-000000000001",
|
|
request=request,
|
|
service=cast(Any, _Service()),
|
|
current_user=CurrentUser(id=uuid4(), email="user@example.com"),
|
|
)
|
|
|
|
assert exc_info.value.status_code == 429
|
|
assert exc_info.value.detail == "Too many run requests"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enqueue_resume_accepts_valid_tool_contract(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
request = _resume_input_with_tool_message()
|
|
|
|
async def _allow_run(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
monkeypatch.setattr(agent_router, "_allow_run_request", _allow_run)
|
|
|
|
class _Service:
|
|
async def enqueue_resume(self, **kwargs): # noqa: ANN003
|
|
return SimpleNamespace(
|
|
task_id="task-resume-1",
|
|
thread_id=kwargs["thread_id"],
|
|
run_id=kwargs["run_input"].run_id,
|
|
created=False,
|
|
)
|
|
|
|
result = await agent_router.enqueue_resume(
|
|
thread_id="00000000-0000-0000-0000-000000000001",
|
|
request=request,
|
|
service=cast(Any, _Service()),
|
|
current_user=CurrentUser(id=uuid4(), email="user@example.com"),
|
|
)
|
|
|
|
assert result.task_id == "task-resume-1"
|
|
assert result.thread_id == "00000000-0000-0000-0000-000000000001"
|
|
assert result.run_id == "run-resume-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_events_retries_on_redis_timeout(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def _acquire(*, user_id: str) -> bool:
|
|
del user_id
|
|
return True
|
|
|
|
async def _release(*, user_id: str) -> None:
|
|
del user_id
|
|
|
|
monkeypatch.setattr(agent_router, "_acquire_sse_slot", _acquire)
|
|
monkeypatch.setattr(agent_router, "_release_sse_slot", _release)
|
|
|
|
class _Request:
|
|
async def is_disconnected(self) -> bool:
|
|
return False
|
|
|
|
class _Service:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
async def stream_events(self, **kwargs): # noqa: ANN003
|
|
del kwargs
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
raise RuntimeError("Timeout reading from localhost:6379")
|
|
if self.calls == 2:
|
|
return [{"id": "1-0", "event": {"type": "RUN_FINISHED"}}]
|
|
return []
|
|
|
|
response = await agent_router.stream_events(
|
|
request=cast(Any, _Request()),
|
|
thread_id="00000000-0000-0000-0000-000000000001",
|
|
service=cast(Any, _Service()),
|
|
current_user=CurrentUser(id=uuid4(), email="user@example.com"),
|
|
last_event_id=None,
|
|
idle_limit=2,
|
|
)
|
|
|
|
chunks: list[str] = []
|
|
async for chunk in response.body_iterator:
|
|
chunks.append(str(chunk))
|
|
if any("RUN_FINISHED" in item for item in chunks):
|
|
break
|
|
|
|
merged = "".join(chunks)
|
|
assert "event: RUN_FINISHED" in merged
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_attachment_preview_rejects_negative_index() -> None:
|
|
class _Service:
|
|
async def get_attachment_preview(self, **kwargs): # noqa: ANN003
|
|
del kwargs
|
|
raise AssertionError("get_attachment_preview should not be called")
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await agent_router.get_attachment_preview(
|
|
thread_id="00000000-0000-0000-0000-000000000001",
|
|
message_id="00000000-0000-0000-0000-000000000010",
|
|
attachment_index=-1,
|
|
service=cast(Any, _Service()),
|
|
current_user=CurrentUser(id=uuid4(), email="user@example.com"),
|
|
)
|
|
|
|
assert exc_info.value.status_code == 422
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_attachment_preview_returns_streaming_response() -> None:
|
|
class _Service:
|
|
async def get_attachment_preview(self, **kwargs): # noqa: ANN003
|
|
del kwargs
|
|
return b"png-bytes", "image/png"
|
|
|
|
response = await agent_router.get_attachment_preview(
|
|
thread_id="00000000-0000-0000-0000-000000000001",
|
|
message_id="00000000-0000-0000-0000-000000000010",
|
|
attachment_index=0,
|
|
service=cast(Any, _Service()),
|
|
current_user=CurrentUser(id=uuid4(), email="user@example.com"),
|
|
)
|
|
chunks: list[bytes] = []
|
|
async for chunk in response.body_iterator:
|
|
chunks.append(cast(bytes, chunk))
|
|
|
|
assert response.media_type == "image/png"
|
|
assert b"".join(chunks) == b"png-bytes"
|