feat: 支持 agent 运行取消功能
This commit is contained in:
@@ -100,6 +100,7 @@ class _FakeRepository:
|
||||
class _FakeQueue:
|
||||
def __init__(self) -> None:
|
||||
self.commands: list[dict[str, object]] = []
|
||||
self.cancel_requests: list[dict[str, str]] = []
|
||||
|
||||
async def enqueue(
|
||||
self, *, command: dict[str, object], dedup_key: str | None
|
||||
@@ -108,6 +109,21 @@ class _FakeQueue:
|
||||
self.commands.append(command)
|
||||
return "task-1"
|
||||
|
||||
async def request_cancel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
requested_by: str,
|
||||
) -> None:
|
||||
self.cancel_requests.append(
|
||||
{
|
||||
"thread_id": thread_id,
|
||||
"run_id": run_id,
|
||||
"requested_by": requested_by,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _FakeStream:
|
||||
async def read(
|
||||
@@ -469,3 +485,56 @@ async def test_get_history_snapshot_filters_out_tool_messages() -> None:
|
||||
)
|
||||
|
||||
assert [message.role for message in snapshot.messages] == ["user", "assistant"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_run_requests_queue_cancel_for_owner() -> None:
|
||||
queue = _FakeQueue()
|
||||
service = AgentService(
|
||||
repository=_FakeRepository(),
|
||||
queue=queue,
|
||||
stream=_FakeStream(),
|
||||
attachment_storage=_FakeAttachmentStorage(),
|
||||
)
|
||||
|
||||
result = await service.cancel_run(
|
||||
thread_id="00000000-0000-0000-0000-000000000001",
|
||||
run_id="run-cancel-1",
|
||||
current_user=_user(),
|
||||
)
|
||||
|
||||
assert result.accepted is True
|
||||
assert result.thread_id == "00000000-0000-0000-0000-000000000001"
|
||||
assert result.run_id == "run-cancel-1"
|
||||
assert queue.cancel_requests == [
|
||||
{
|
||||
"thread_id": "00000000-0000-0000-0000-000000000001",
|
||||
"run_id": "run-cancel-1",
|
||||
"requested_by": "00000000-0000-0000-0000-000000000001",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_run_rejects_non_owner() -> None:
|
||||
queue = _FakeQueue()
|
||||
service = AgentService(
|
||||
repository=_FakeRepository(),
|
||||
queue=queue,
|
||||
stream=_FakeStream(),
|
||||
attachment_storage=_FakeAttachmentStorage(),
|
||||
)
|
||||
other_user = CurrentUser(
|
||||
id=UUID("00000000-0000-0000-0000-000000000099"),
|
||||
phone="+8613812340000",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await service.cancel_run(
|
||||
thread_id="00000000-0000-0000-0000-000000000001",
|
||||
run_id="run-cancel-2",
|
||||
current_user=other_user,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert queue.cancel_requests == []
|
||||
|
||||
Reference in New Issue
Block a user