"""Messages never change task generation or silently target a replacement run.""" from datetime import datetime, timezone from uuid import UUID, uuid4 from app.auth.permissions import OPERATIONAL_TASK_EXECUTE, permissions_for_context from app.domain.domain_errors import DomainError, require from app.domain.schemas.agent_messages import ( AgentMessage, AgentMessageCreate, AgentMessageDelivery, AgentMessagePage, AgentMessageReply, ) def _message_cursor(cursor: str | None): if not cursor: return None, None try: stamp, identifier = cursor.rsplit("|", 1) cursor_at, cursor_id = datetime.fromisoformat(stamp), UUID(identifier) require(cursor_at.tzinfo is not None, "INVALID_CURSOR", "Cursor needs a timezone", status=400) return cursor_at, cursor_id except (ValueError, TypeError) as exc: raise DomainError("INVALID_CURSOR", "Invalid message cursor", status=400) from exc class AgentMessageService: def __init__(self, repository, context, *, clock=None): self.repo = repository self.context = context self.clock = clock or (lambda: datetime.now(timezone.utc)) def _require_service(self) -> None: require(self.context.system_role == "service" and OPERATIONAL_TASK_EXECUTE in permissions_for_context(self.context), "SERVICE_REQUIRED", "Message dispatch requires a project service account", status=403) async def get(self, message_id: UUID) -> AgentMessage: if self.context.system_role == "service": self._require_service() row = await self.repo.get_message(str(message_id)) require(row is not None, "MESSAGE_NOT_FOUND", "Message not found", status=404) require(self.context.system_role == "service" or str(row["conversation_owner_id"]) == str(self.context.user_id), "MESSAGE_NOT_FOUND", "Message not found", status=404) return AgentMessage.model_validate(row) async def _active_run(self, task_id: UUID, run_id: UUID) -> dict: run = await self.repo.get_run(str(run_id)) require(run is not None and str(run["task_id"]) == str(task_id), "RUN_NOT_FOUND", "Run does not belong to the task", status=404) task = await self.repo.get_task(str(task_id)) require(task is not None and task["task_type"] == "event_investigation", "INVALID_TASK_TYPE", "Messages require an investigation task", status=409) require(run["status"] in {"queued", "claimed", "running"}, "RUN_TERMINAL", "Run no longer accepts messages", status=409) if run["status"] != "queued": expires = run.get("lease_expires_at") if isinstance(expires, str): expires = datetime.fromisoformat(expires.replace("Z", "+00:00")) require(expires is not None and expires > self.clock(), "LEASE_EXPIRED", "Run lease expired", status=409) return run async def _insert(self, *, command, owner_id, conversation_owner_id, direction, task_id, run_id, target_session_id, in_reply_to=None) -> AgentMessage: prior = await self.repo.get_message_request( str(owner_id), command.source_session_id, str(command.request_id), ) expected = { "task_id": str(task_id), "run_id": str(run_id), "direction": direction, "in_reply_to": str(in_reply_to) if in_reply_to else None, "text": command.text, } if prior is not None: require(all((str(prior[key]) if prior[key] is not None else None) == value for key, value in expected.items()), "IDEMPOTENCY_CONFLICT", "request_id was used with different content", status=409) return AgentMessage.model_validate(prior) # The run lock serializes message creation with completion/cancellation. await self._active_run(task_id, run_id) message_id = uuid4() now = self.clock() message = AgentMessage( project_id=self.context.project_id, message_id=message_id, request_id=command.request_id, task_id=task_id, run_id=run_id, source_session_id=command.source_session_id, source_owner_id=owner_id, conversation_owner_id=conversation_owner_id, target_session_id=target_session_id, direction=direction, in_reply_to=in_reply_to, text=command.text, status="pending", native_message_id=f"msg_{message_id}", created_at=now, updated_at=now, ) await self.repo.insert_message(message.model_dump(mode="json")) return message async def create(self, command: AgentMessageCreate) -> AgentMessage: async with self.repo.transaction(): await self.repo.lock(f"run:{command.run_id}") run = await self.repo.get_run(str(command.run_id)) return await self._insert( command=command, owner_id=self.context.user_id, conversation_owner_id=self.context.user_id, direction="to_task", task_id=command.task_id, run_id=command.run_id, target_session_id=run.get("agent_session_id") if run else None, ) async def reply(self, message_id: UUID, command: AgentMessageReply) -> AgentMessage: async with self.repo.transaction(): original = await self.get(message_id) await self.repo.lock(f"run:{original.run_id}") run = await self.repo.get_run(str(original.run_id)) if original.direction == "to_task": self._require_service() require(run is not None and run.get("agent_session_id") == command.source_session_id, "SESSION_MISMATCH", "Session does not own the message run", status=409) direction, target = "to_chat", original.source_session_id else: require(original.conversation_owner_id == self.context.user_id, "MESSAGE_NOT_FOUND", "Message not found", status=404) require(command.source_session_id == original.target_session_id, "SESSION_MISMATCH", "Only the original chat can reply", status=409) direction = "to_task" target = run.get("agent_session_id") if run else None return await self._insert( command=command, owner_id=self.context.user_id, conversation_owner_id=original.conversation_owner_id, direction=direction, task_id=original.task_id, run_id=original.run_id, target_session_id=target, in_reply_to=message_id, ) async def list_session( self, session_id: str, limit: int, cursor: str | None, pending_only: bool = False, ) -> AgentMessagePage: cursor_at, cursor_id = _message_cursor(cursor) kwargs = {"cursor_at": cursor_at, "cursor_id": cursor_id, "pending_only": pending_only} if self.context.system_role == "service": self._require_service() rows = await self.repo.list_run_session_messages(session_id, limit + 1, **kwargs) if not rows: rows = await self.repo.list_session_messages(str(self.context.user_id), session_id, limit + 1, **kwargs) else: rows = await self.repo.list_session_messages(str(self.context.user_id), session_id, limit + 1, **kwargs) items = [AgentMessage.model_validate(row) for row in rows[:limit]] next_cursor = f"{items[-1].created_at.isoformat()}|{items[-1].message_id}" if len(rows) > limit else None return AgentMessagePage(items=items, next_cursor=next_cursor) async def pending(self, limit: int, cursor: str | None = None) -> AgentMessagePage: self._require_service() cursor_at, cursor_id = _message_cursor(cursor) rows = await self.repo.list_pending_messages(limit + 1, cursor_at=cursor_at, cursor_id=cursor_id) items = [AgentMessage.model_validate(row) for row in rows[:limit]] next_cursor = f"{items[-1].created_at.isoformat()}|{items[-1].message_id}" if len(rows) > limit else None return AgentMessagePage(items=items, next_cursor=next_cursor) async def delivery(self, message_id: UUID, command: AgentMessageDelivery) -> AgentMessage: async with self.repo.transaction(): message = await self.get(message_id) if self.context.system_role == "service": self._require_service() if self.context.system_role != "service": require(message.direction == "to_chat", "SERVICE_REQUIRED", "Only the receiver can acknowledge a chat reply", status=403) await self.repo.lock(f"run:{message.run_id}") await self.repo.lock(f"message:{message_id}") message = await self.get(message_id) if message.status in {"delivered", "failed"}: require(message.status == command.status, "MESSAGE_TERMINAL", "Message already has a terminal delivery result", status=409) return message if message.direction == "to_task" and command.status != "failed": run = await self._active_run(message.task_id, message.run_id) target = run.get("agent_session_id") require(target is not None and command.target_session_id == target, "SESSION_MISMATCH", "Delivery must target the owning run session", status=409) else: target = message.target_session_id require(command.target_session_id in {None, target}, "SESSION_MISMATCH", "Delivery cannot change the reply channel", status=409) require(command.native_message_id in {None, message.native_message_id}, "MESSAGE_ID_MISMATCH", "Native message identity is immutable", status=409) now = self.clock() update = { "status": command.status, "target_session_id": target, "failure_detail": command.failure_detail, "updated_at": now, "delivered_at": now if command.status == "delivered" else None, } await self.repo.update_message_delivery(str(message_id), update) return message.model_copy(update=update)