jiang
6a4478bfe4
Remove retired query paths and repository compatibility branches, and share task test fixtures. Cancel pending questions and resolve their notifications when recovering expired runs.
247 lines
12 KiB
Python
247 lines
12 KiB
Python
from contextlib import asynccontextmanager
|
|
from datetime import datetime, timedelta, timezone
|
|
from unittest.mock import AsyncMock
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from app.auth.project_dependencies import ProjectContext
|
|
from app.domain.domain_errors import DomainError
|
|
from app.domain.schemas.agent_messages import AgentMessageCreate, AgentMessageDelivery, AgentMessageReply
|
|
from app.domain.schemas.operational_task import AgentRunCommand
|
|
from app.services.agent_messages import AgentMessageService
|
|
from app.services.operational_tasks import OperationalTaskService
|
|
|
|
|
|
@pytest.fixture
|
|
def anyio_backend():
|
|
return "asyncio"
|
|
|
|
|
|
NOW = datetime(2026, 10, 13, 0, tzinfo=timezone.utc)
|
|
|
|
|
|
class Repository:
|
|
def __init__(self):
|
|
self.project_id, self.task_id, self.run_id = uuid4(), uuid4(), uuid4()
|
|
self.run = {"run_id": str(self.run_id), "task_id": str(self.task_id),
|
|
"status": "running", "agent_session_id": "task-session",
|
|
"lease_expires_at": NOW + timedelta(minutes=1), "revision": 2,
|
|
"generation": 1, "execution_request_id": str(uuid4()),
|
|
"next_dispatch_at": NOW, "created_at": NOW, "updated_at": NOW}
|
|
self.messages = {}
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self):
|
|
yield self
|
|
|
|
async def lock(self, key):
|
|
pass
|
|
|
|
async def get_run(self, run_id):
|
|
return self.run if run_id == str(self.run_id) else None
|
|
|
|
async def get_task(self, task_id):
|
|
return {"task_type": "event_investigation"} if task_id == str(self.task_id) else None
|
|
|
|
async def list_questions(self, task_id):
|
|
return []
|
|
|
|
async def pending_calculation_jobs(self, task_id, run_id, execution_request_id):
|
|
return []
|
|
|
|
async def get_message(self, message_id):
|
|
return self.messages.get(message_id)
|
|
|
|
async def get_message_request(self, owner_id, session_id, request_id):
|
|
return next((row for row in self.messages.values()
|
|
if row["source_owner_id"] == owner_id and row["source_session_id"] == session_id
|
|
and row["request_id"] == request_id), None)
|
|
|
|
async def insert_message(self, message):
|
|
self.messages[message["message_id"]] = message
|
|
|
|
async def update_message_delivery(self, message_id, update):
|
|
self.messages[message_id].update(update)
|
|
|
|
|
|
def setup():
|
|
repo = Repository()
|
|
actor = ProjectContext(repo.project_id, "test", uuid4(), "member")
|
|
worker = ProjectContext(repo.project_id, "test", uuid4(), "member", system_role="service")
|
|
service = AgentMessageService(repo, actor, clock=lambda: NOW)
|
|
trusted = AgentMessageService(repo, worker, clock=lambda: NOW)
|
|
command = AgentMessageCreate(request_id=uuid4(), task_id=repo.task_id, run_id=repo.run_id,
|
|
source_session_id="main-session", text="核对当前压力")
|
|
return repo, actor, service, trusted, command
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_messages_are_bound_to_exact_run_and_idempotent_even_after_finish():
|
|
repo, actor, service, trusted, command = setup()
|
|
sent = await service.create(command)
|
|
assert sent.source_owner_id == actor.user_id == sent.conversation_owner_id
|
|
assert sent.native_message_id == f"msg_{sent.message_id}"
|
|
repo.run["status"] = "succeeded"
|
|
assert (await service.create(command)).message_id == sent.message_id
|
|
with pytest.raises(DomainError, match="request_id") as error:
|
|
await service.create(command.model_copy(update={"text": "different"}))
|
|
assert error.value.code == "IDEMPOTENCY_CONFLICT"
|
|
with pytest.raises(DomainError) as error:
|
|
await service.create(command.model_copy(update={"request_id": uuid4()}))
|
|
assert error.value.code == "RUN_TERMINAL"
|
|
assert len(repo.messages) == 1
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_wrong_task_or_expired_lease_cannot_accept_message():
|
|
repo, actor, service, trusted, command = setup()
|
|
with pytest.raises(DomainError) as error:
|
|
await service.create(command.model_copy(update={"task_id": uuid4()}))
|
|
assert error.value.code == "RUN_NOT_FOUND"
|
|
repo.run["lease_expires_at"] = NOW
|
|
with pytest.raises(DomainError) as error:
|
|
await service.create(command)
|
|
assert error.value.code == "LEASE_EXPIRED"
|
|
repo.run.update(status="queued", agent_session_id=None)
|
|
assert (await service.create(command)).target_session_id is None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_reply_has_original_conversation_owner_and_immutable_channel():
|
|
repo, actor, service, trusted, command = setup()
|
|
incoming = await service.create(command)
|
|
reply_command = AgentMessageReply(request_id=uuid4(), source_session_id="task-session", text="压力正常")
|
|
with pytest.raises(DomainError) as error:
|
|
await service.reply(incoming.message_id, reply_command)
|
|
assert error.value.code == "SERVICE_REQUIRED"
|
|
with pytest.raises(DomainError) as error:
|
|
await trusted.reply(incoming.message_id, reply_command.model_copy(update={"source_session_id": "other"}))
|
|
assert error.value.code == "SESSION_MISMATCH"
|
|
reply = await trusted.reply(incoming.message_id, reply_command)
|
|
assert reply.source_owner_id == trusted.context.user_id
|
|
assert reply.conversation_owner_id == actor.user_id
|
|
assert reply.target_session_id == command.source_session_id
|
|
followup = await service.reply(reply.message_id, AgentMessageReply(
|
|
request_id=uuid4(), source_session_id="main-session", text="请继续核对"))
|
|
assert followup.direction == "to_task" and followup.in_reply_to == reply.message_id
|
|
assert followup.target_session_id == "task-session"
|
|
with pytest.raises(DomainError) as error:
|
|
await service.reply(reply.message_id, AgentMessageReply(
|
|
request_id=uuid4(), source_session_id="wrong-main", text="伪造来源"))
|
|
assert error.value.code == "SESSION_MISMATCH"
|
|
repo.run["status"] = "succeeded"
|
|
assert (await trusted.reply(incoming.message_id, reply_command)).message_id == reply.message_id
|
|
with pytest.raises(DomainError) as error:
|
|
await service.delivery(reply.message_id, AgentMessageDelivery(status="delivered", target_session_id="other"))
|
|
assert error.value.code == "SESSION_MISMATCH"
|
|
received = await service.delivery(reply.message_id, AgentMessageDelivery(status="enqueued"))
|
|
assert received.status == "enqueued" and received.delivered_at is None
|
|
assert (await service.delivery(reply.message_id, AgentMessageDelivery(status="delivered"))).delivered_at == NOW
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_inbound_delivery_requires_service_run_session_and_fixed_native_identity():
|
|
repo, actor, service, trusted, command = setup()
|
|
incoming = await service.create(command)
|
|
with pytest.raises(DomainError) as error:
|
|
await service.delivery(incoming.message_id, AgentMessageDelivery(status="delivered"))
|
|
assert error.value.code == "SERVICE_REQUIRED"
|
|
with pytest.raises(DomainError) as error:
|
|
await trusted.delivery(incoming.message_id, AgentMessageDelivery(status="enqueued", target_session_id="wrong"))
|
|
assert error.value.code == "SESSION_MISMATCH"
|
|
with pytest.raises(DomainError) as error:
|
|
await trusted.delivery(incoming.message_id, AgentMessageDelivery(status="delivered", target_session_id="task-session", native_message_id="changed"))
|
|
assert error.value.code == "MESSAGE_ID_MISMATCH"
|
|
await trusted.delivery(incoming.message_id, AgentMessageDelivery(status="enqueued", target_session_id="task-session"))
|
|
assert (await service.get(incoming.message_id)).status == "enqueued"
|
|
await trusted.delivery(incoming.message_id, AgentMessageDelivery(status="delivered", target_session_id="task-session"))
|
|
with pytest.raises(DomainError) as error:
|
|
await trusted.delivery(incoming.message_id, AgentMessageDelivery(status="failed"))
|
|
assert error.value.code == "MESSAGE_TERMINAL"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_other_project_member_cannot_read_private_message_or_pending_dispatch():
|
|
repo, actor, service, trusted, command = setup()
|
|
incoming = await service.create(command)
|
|
outsider = AgentMessageService(repo, ProjectContext(repo.project_id, "test", uuid4(), "member"))
|
|
with pytest.raises(DomainError) as error:
|
|
await outsider.get(incoming.message_id)
|
|
assert error.value.code == "MESSAGE_NOT_FOUND"
|
|
with pytest.raises(DomainError) as error:
|
|
await outsider.pending(20)
|
|
assert error.value.code == "SERVICE_REQUIRED"
|
|
viewer_service = AgentMessageService(repo, ProjectContext(
|
|
repo.project_id, "test", uuid4(), "viewer", system_role="service"))
|
|
with pytest.raises(DomainError) as error:
|
|
await viewer_service.get(incoming.message_id)
|
|
assert error.value.code == "SERVICE_REQUIRED"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_own_session_listing_uses_owner_and_keyset_cursor():
|
|
repo, actor, service, trusted, command = setup()
|
|
sent = await service.create(command)
|
|
other = await service.create(command.model_copy(update={"request_id": uuid4(), "text": "后续问题"}))
|
|
repo.list_session_messages = AsyncMock(return_value=[sent.model_dump(), other.model_dump()])
|
|
page = await service.list_session("main-session", 1, None, pending_only=True)
|
|
assert len(page.items) == 1 and page.next_cursor
|
|
repo.list_session_messages.assert_awaited_once_with(
|
|
str(actor.user_id), "main-session", 2,
|
|
cursor_at=None, cursor_id=None, pending_only=True)
|
|
repo.list_session_messages.reset_mock()
|
|
repo.list_session_messages.return_value = []
|
|
await service.list_session("main-session", 1, page.next_cursor)
|
|
assert repo.list_session_messages.await_args.kwargs["cursor_id"] == sent.message_id
|
|
with pytest.raises(DomainError) as error:
|
|
await service.list_session("main-session", 20, "invalid")
|
|
assert error.value.code == "INVALID_CURSOR"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_pending_dispatch_page_returns_forward_keyset_cursor():
|
|
repo, actor, service, trusted, command = setup()
|
|
first = await service.create(command)
|
|
second = await service.create(command.model_copy(update={"request_id": uuid4()}))
|
|
repo.list_pending_messages = AsyncMock(return_value=[first.model_dump(), second.model_dump()])
|
|
page = await trusted.pending(1)
|
|
assert page.items[0].message_id == first.message_id and page.next_cursor
|
|
await trusted.pending(1, page.next_cursor)
|
|
assert repo.list_pending_messages.await_args.kwargs == {
|
|
"cursor_at": first.created_at, "cursor_id": first.message_id}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_requeue_locks_and_rechecks_run_before_cancelling_it():
|
|
repo, actor, service, trusted, command = setup()
|
|
run = repo.run.copy()
|
|
replacement = {**run, "run_id": str(uuid4())}
|
|
repo.get_active_run = AsyncMock(side_effect=[run, replacement, replacement])
|
|
repo.lock = AsyncMock()
|
|
repo.update_run, repo.insert_run, repo.fail_run_messages = AsyncMock(), AsyncMock(), AsyncMock()
|
|
task_service = OperationalTaskService(repo, actor, [], clock=lambda: NOW)
|
|
task_service._run_event = AsyncMock()
|
|
task = {"task_id": str(repo.task_id), "execution_generation": 1, "revision": 2}
|
|
await task_service._requeue_task(task, NOW)
|
|
assert [call.args[0] for call in repo.lock.await_args_list] == [
|
|
f"run:{run['run_id']}", f"run:{replacement['run_id']}"]
|
|
assert repo.update_run.await_args.args[0]["run_id"] == replacement["run_id"]
|
|
repo.fail_run_messages.assert_awaited_once_with(replacement["run_id"], "Run superseded by requeue", NOW)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_finish_refuses_messages_enqueued_but_not_entered_in_model_context():
|
|
repo, actor, service, trusted, command = setup()
|
|
repo.get_request = AsyncMock(return_value=None)
|
|
repo.pending_run_messages = AsyncMock(return_value=["pending-message"])
|
|
repo.get_latest_report = AsyncMock()
|
|
task_service = OperationalTaskService(repo, actor, [], clock=lambda: NOW)
|
|
task_service._task = AsyncMock(return_value={"task_type": "event_investigation"})
|
|
with pytest.raises(DomainError) as error:
|
|
await task_service.complete_run(repo.task_id, repo.run_id, AgentRunCommand(
|
|
request_id=uuid4(), expected_revision=2, source_session_id="task-session"))
|
|
assert error.value.code == "PENDING_AGENT_MESSAGES" and error.value.status == 409
|
|
assert error.value.refs == ["pending-message"]
|
|
repo.get_latest_report.assert_not_awaited()
|