Agent/cyber_agent/core/task_protocol_service.py

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

from __future__ import annotations

import asyncio
from copy import deepcopy
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.protocols import TraceStore


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


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


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

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

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

    async def update_progress(
        self,
        trace_id: str,
        *,
        expected_revision: int,
        progress: TaskProgress | dict,
        effective_at_sequence: int,
    ) -> TaskProgressRevision:
        async with self._lock(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)
            return revision