Agent/cyber_agent/core/task_protocol_service.py

"""TaskProtocol 命令的单进程原子执行边界。"""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from copy import deepcopy
from collections.abc import AsyncIterator, Callable
from typing import Any
from weakref import WeakValueDictionary

from cyber_agent.core.agent_mode import require_mutable_trace_policy
from cyber_agent.core.context_policy import get_authorized_context_snapshot
from cyber_agent.core.task_protocol import (
    TaskProgress,
    TaskProgressRevision,
    append_task_progress_revision,
    ensure_task_protocol,
    task_contract_ref,
)
from cyber_agent.trace.models import Trace
from cyber_agent.trace.protocols import TraceStore


_TRACE_LOCKS: WeakValueDictionary[
    tuple[asyncio.AbstractEventLoop, int, str], asyncio.Lock
] = WeakValueDictionary()


class TaskProtocolStateError(RuntimeError):
    """TaskProtocol 命令不能在当前持久化状态执行。"""


class TaskProtocolService:
    """将协议领域纯函数与 TraceStore 持久化组合起来。"""

    def __init__(self, trace_store: TraceStore, event_service=None) -> None:
        self.trace_store = trace_store
        self.event_service = event_service

    def _lock(self, trace_id: str) -> asyncio.Lock:
        key = (asyncio.get_running_loop(), id(self.trace_store), trace_id)
        lock = _TRACE_LOCKS.get(key)
        if lock is None:
            lock = asyncio.Lock()
            _TRACE_LOCKS[key] = lock
        return lock

    @asynccontextmanager
    async def transaction(self, trace_id: str) -> AsyncIterator[None]:
        """Serialize one complete TaskProtocol read-modify-write command."""
        async with self._lock(trace_id):
            yield

    @asynccontextmanager
    async def transaction_many(
        self,
        *trace_ids: str,
    ) -> AsyncIterator[None]:
        """Lock multiple protocol records in stable order to avoid deadlocks."""
        locks = [self._lock(item) for item in sorted(set(trace_ids))]
        for lock in locks:
            await lock.acquire()
        try:
            yield
        finally:
            for lock in reversed(locks):
                lock.release()

    @asynccontextmanager
    async def locked_trace(self, trace_id: str) -> AsyncIterator[Trace]:
        """Load one fresh Trace while holding its protocol command lock."""
        async with self.transaction(trace_id):
            trace = await self.trace_store.get_trace(trace_id)
            if trace is None:
                raise TaskProtocolStateError(f"Trace not found: {trace_id}")
            yield trace

    async def mutate_state(
        self,
        trace_id: str,
        mutation: Callable[[Trace, dict[str, Any]], Any],
    ) -> Any:
        """Apply one synchronous protocol mutation to a fresh persisted Trace."""
        async with self.locked_trace(trace_id) as trace:
            state = ensure_task_protocol(trace.context)
            result = mutation(trace, state)
            await self.trace_store.update_trace(trace_id, context=trace.context)
            return result

    async def update_progress(
        self,
        trace_id: str,
        *,
        expected_revision: int,
        progress: TaskProgress | dict,
        effective_at_sequence: int,
    ) -> TaskProgressRevision:
        async with self.transaction(trace_id):
            trace = await self.trace_store.get_trace(trace_id)
            if trace is None:
                raise TaskProtocolStateError(f"Trace not found: {trace_id}")
            policy = require_mutable_trace_policy(trace.context)
            if not policy.requires_task_progress:
                raise TaskProtocolStateError(
                    "update_task_progress requires Recursive revision 3"
                )
            if (
                trace.agent_type == "validator"
                or trace.context.get("created_by_tool") == "validator"
            ):
                raise TaskProtocolStateError(
                    "Validator traces cannot update TaskProgress"
                )

            parsed = TaskProgress.model_validate(progress)
            root_trace_id = trace.context.get("root_trace_id") or trace.trace_id
            for item in (
                *parsed.questions,
                *parsed.blockers,
                *parsed.findings,
                *parsed.hypotheses,
                *parsed.work_items,
            ):
                for ref in item.context_refs:
                    get_authorized_context_snapshot(
                        trace.context,
                        ref_id=ref.ref_id,
                        version=ref.version,
                        root_trace_id=root_trace_id,
                        uid=trace.uid,
                    )

            next_context = deepcopy(trace.context)
            state = ensure_task_protocol(next_context)
            if state["pending_reviews"]:
                raise TaskProtocolStateError(
                    "Review all child TaskReports before updating TaskProgress"
                )
            if state["next_actions"]:
                raise TaskProtocolStateError(
                    "Execute every approved next action before updating TaskProgress"
                )
            if state["task_report"] is not None:
                raise TaskProtocolStateError(
                    "TaskProgress cannot change after TaskReport submission"
                )
            revision = append_task_progress_revision(
                state,
                parsed,
                expected_revision=expected_revision,
                effective_at_sequence=effective_at_sequence,
                contract_ref=task_contract_ref(
                    state,
                    root_task_anchor_hash=next_context.get(
                        "root_task_anchor_hash"
                    ),
                ),
            )
            await self.trace_store.update_trace(trace_id, context=next_context)
            if self.event_service is not None:
                await self.event_service.emit_after_commit(
                    source_trace_id=trace_id,
                    event_type="task.progress.revised",
                    event_key=(
                        f"task.progress.revised:{trace_id}:{revision.revision}"
                    ),
                    effective_at_sequence=revision.effective_at_sequence,
                    payload={"revision": revision.model_dump(mode="json")},
                )
            return revision