NotesAgentic/backend/app/agent/runtime.py

"""Agent 运行时：负责模型轮次、工具调用、权限确认与事件发布。"""

from __future__ import annotations

import asyncio
import json
from app.agent.async_trace import AsyncTraceWriter
from app.operation_logs import log_event, agent_run_id
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from datetime import datetime, timezone
from time import perf_counter
from typing import TYPE_CHECKING
from uuid import uuid4

from app.agent.permissions import PermissionManager, PermissionMode
from app.agent.tools import ToolExecutionContext, ToolNotFoundError, ToolRegistry
from app.agent.trace_repository import AgentTraceRepository, sanitize_trace_value
from app.contracts import (
    AgentEvent,
    AgentEventType,
    AgentRun,
    AgentRunCreateRequest,
    AgentRunStatus,
    AgentTraceResponse,
    Citation,
    Message,
    MessageRole,
    ModelRequest,
    ToolCall,
    ToolResult,
)
from app.providers.registry import ProviderRegistry
from app.providers.base import ProviderError

if TYPE_CHECKING:
    from app.extensions import AgentConfiguration, SkillRuntime


class AgentRunNotFoundError(LookupError):
    pass


class AgentCapacityError(RuntimeError):
    pass


TERMINAL_STATUSES = {
    AgentRunStatus.completed,
    AgentRunStatus.failed,
    AgentRunStatus.cancelled,
}
MAX_RUN_RECORDS = 200
MAX_EVENTS_PER_RUN = 2_000
MAX_TOOL_CALLS_PER_TURN = 50


@dataclass(slots=True)
class RunRecord:
    """单次运行的可变上下文，仅由 AgentRuntime 持有。"""

    run: AgentRun
    request: AgentRunCreateRequest
    skill_config: AgentConfiguration | None = None
    allowed_tools: list[str] = field(default_factory=list)
    events: list[AgentEvent] = field(default_factory=list)
    subscribers: set[asyncio.Queue[AgentEvent]] = field(default_factory=set)
    task: asyncio.Task[None] | None = None
    next_sequence: int = 0
    publish_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    cancel_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    persisted_run: AgentRun | None = None


class AgentRuntime:
    """进程内 Agent 编排器；对外返回深拷贝，避免调用方修改运行状态。"""

    def __init__(
        self,
        providers: ProviderRegistry,
        tools: ToolRegistry,
        permissions: PermissionManager,
        skills: SkillRuntime | None = None,
        trace_repository: AgentTraceRepository | None = None,
    ) -> None:
        self.providers = providers
        self.tools = tools
        self.permissions = permissions
        self.skills = skills
        self.trace_repository = trace_repository or AgentTraceRepository()
        self._records: dict[str, RunRecord] = {}
        self._writer = AsyncTraceWriter(self.trace_repository)

    async def create_run(self, request: AgentRunCreateRequest) -> AgentRun:
        self._prune_records()
        provider = self.providers.get(request.provider_id)
        skill_config = None
        if request.skill_id:
            if request.skill_id.startswith("user_skill_"):
                from app.services.user_skills import build_agent_configuration

                skill_config = await asyncio.to_thread(
                    build_agent_configuration,
                    request.skill_id,
                    provider.config.capabilities,
                    self.tools,
                )
            else:
                if self.skills is None:
                    raise RuntimeError("Skill Runtime is not configured.")
                skill_config = self.skills.build_agent_configuration(
                    request.skill_id, provider.config.capabilities
                )
        now = datetime.now(timezone.utc)
        run = AgentRun(
            run_id=f"run_{uuid4().hex}",
            status=AgentRunStatus.queued,
            input=request.input,
            provider_id=request.provider_id,
            model=request.model,
            skill_id=request.skill_id,
            max_steps=request.max_steps,
            token_budget=request.token_budget,
            created_at=now,
            updated_at=now,
        )
        allowed_tools = list(request.allowed_tools)
        if skill_config is not None:
            # 同时指定 Skill 与工具白名单时取交集，避免 Skill 扩大调用权限。
            allowed_tools = (
                [name for name in skill_config.allowed_tools if name in allowed_tools]
                if allowed_tools
                else list(skill_config.allowed_tools)
            )
        record = RunRecord(
            run=run,
            request=request,
            skill_config=skill_config,
            allowed_tools=allowed_tools,
        )
        # 在让渡给并发创建者之前保留容量。
        self._records[run.run_id] = record
        try:
            cancelled = await self._writer.submit('create', run.model_copy(deep=True), request.model_copy(deep=True), self._config_snapshot(record))
        except BaseException:
            self._records.pop(run.run_id, None)
            raise
        record.persisted_run = run.model_copy(deep=True)
        log_event('agent', 'run.created', run_id=run.run_id, provider_id=run.provider_id, model=run.model)
        if cancelled:
            await self._finish_cancelled(record)
            raise asyncio.CancelledError
        record.task = asyncio.create_task(self._execute(record), name=run.run_id)
        return run.model_copy(deep=True)

    def get_run(self, run_id: str) -> AgentRun:
        record = self._records.get(run_id)
        if record is not None:
            return (record.persisted_run or record.run).model_copy(deep=True)
        run = self.trace_repository.recover_interrupted(run_id)
        if run is None:
            raise AgentRunNotFoundError(run_id)
        return run.model_copy(deep=True)

    def list_runs(self, limit: int, offset: int) -> tuple[list[AgentRun], int]:
        items, total = self.trace_repository.list_runs(limit=limit, offset=offset)
        recovered = [
            self.trace_repository.recover_interrupted(item.run_id) or item
            if item.run_id not in self._records
            else (self._records[item.run_id].persisted_run or self._records[item.run_id].run).model_copy(deep=True)
            for item in items
        ]
        return recovered, total

    async def cancel(self, run_id: str) -> AgentRun:
        record = self._records.get(run_id)
        if record is None:
            return self.get_run(run_id)
        async with record.cancel_lock:
            if record.task and not record.task.done():
                if record.run.status not in TERMINAL_STATUSES:
                    record.task.cancel()
                    self.permissions.cancel_run(run_id)
                await asyncio.gather(record.task, return_exceptions=True)
            if record.run.status not in TERMINAL_STATUSES:
                await self._finish_cancelled(record)
            return (record.persisted_run or record.run).model_copy(deep=True)

    async def resolve_permission(self, run_id: str, request_id: str, decision: str) -> bool:
        record = self._records.get(run_id)
        if record is None:
            return False
        ticket = self.permissions.get_ticket(run_id, request_id)
        resolved = self.permissions.resolve(run_id, request_id, decision)
        if resolved:
            await self._publish(
                record,
                AgentEventType.permission_resolved,
                {
                    "request_id": request_id,
                    "permission": ticket.permission if ticket else None,
                    "decision": decision,
                },
            )
        return resolved

    async def events(
        self, run_id: str, *, after_sequence: int = -1
    ) -> AsyncIterator[AgentEvent]:
        record = self._records.get(run_id)
        run = self.get_run(run_id)
        if record is None:
            for event in await asyncio.to_thread(self.trace_repository.list_events,
                run_id, after_sequence=after_sequence
            ):
                yield event
            return

        # 先注册再异步读取历史；历史与实时队列的交界用 sequence 去重。
        queue: asyncio.Queue[AgentEvent] = asyncio.Queue()
        record.subscribers.add(queue)
        last_sequence = after_sequence
        try:
            history = await asyncio.to_thread(self.trace_repository.list_events,
                run_id, after_sequence=after_sequence)
            for event in history:
                last_sequence = event.sequence
                yield event
                if event.event in {AgentEventType.run_completed, AgentEventType.run_failed, AgentEventType.run_cancelled}:
                    return
            if run.status in TERMINAL_STATUSES:
                return
            while True:
                event = await queue.get()
                if event.sequence <= last_sequence:
                    continue
                last_sequence = event.sequence
                yield event.model_copy(deep=True)
                if event.event in {
                    AgentEventType.run_completed,
                    AgentEventType.run_failed,
                    AgentEventType.run_cancelled,
                }:
                    return
        finally:
            record.subscribers.discard(queue)

    async def wait(self, run_id: str) -> AgentRun:
        record = self._records.get(run_id)
        if record is None:
            return self.get_run(run_id)
        if record.task:
            try:
                await asyncio.shield(record.task)
            except asyncio.CancelledError:
                pass
        return (record.persisted_run or record.run).model_copy(deep=True)

    def get_trace(
        self, run_id: str, *, after_sequence: int, limit: int
    ) -> AgentTraceResponse:
        self.get_run(run_id)
        trace = self.trace_repository.get_trace(
            run_id, after_sequence=after_sequence, limit=limit
        )
        if trace is None:
            raise AgentRunNotFoundError(run_id)
        return trace

    async def _execute(self, record: RunRecord) -> None:
        token = agent_run_id.set(record.run.run_id)
        try:
            async with asyncio.timeout(record.request.run_timeout_seconds):
                await self._run_loop(record)
        except asyncio.CancelledError:
            if record.run.status != AgentRunStatus.cancelled:
                await self._finish_cancelled(record)
        except TimeoutError:
            await self._fail(record, "AGENT_TIMEOUT", "Agent run exceeded its timeout.")
        except ProviderError as exc:
            await self._fail(record, exc.code, exc.message)
        except Exception as exc:
            log_event('agent', 'execution.failed', level='ERROR', error=exc, run_id=record.run.run_id)
            await self._fail(record, "AGENT_FAILED", str(exc))
        finally:
            self.permissions.cancel_run(record.run.run_id)
            agent_run_id.reset(token)

    async def shutdown(self) -> None:
        results = await asyncio.gather(*(self.cancel(run_id) for run_id in list(self._records)), return_exceptions=True)
        for result in results:
            if isinstance(result, BaseException):
                log_event('agent', 'shutdown.failed', level='ERROR', error=result)
        await self._writer.queue.join()

    async def _run_loop(self, record: RunRecord) -> None:
        record.run.status = AgentRunStatus.running
        record.run.updated_at = datetime.now(timezone.utc)
        await self._publish(
            record,
            AgentEventType.run_started,
            {"provider_id": record.request.provider_id, "model": record.request.model},
        )

        messages = [Message(role=MessageRole.user, content=record.request.input)]
        allowed_tools = self.tools.definitions(record.allowed_tools)
        provider = self.providers.get(record.request.provider_id).adapter

        for step in range(1, record.request.max_steps + 1):
            record.run.current_step = step
            record.run.updated_at = datetime.now(timezone.utc)
            model_call_id = f"model_call_{uuid4().hex}"
            started_at = perf_counter()
            await self._publish(
                record,
                AgentEventType.model_call_started,
                {
                    "model_call_id": model_call_id,
                    "step": step,
                    "provider_id": record.request.provider_id,
                    "model": record.request.model,
                },
            )
            try:
                turn = await provider.complete(
                    ModelRequest(
                        provider_id=record.request.provider_id,
                        model=record.request.model,
                        system=(record.skill_config.system_prompt if record.skill_config else None),
                        messages=messages,
                        tools=allowed_tools,
                        metadata=self._request_metadata(record),
                    )
                )
            except Exception as exc:
                await self._publish(
                    record,
                    AgentEventType.model_call_failed,
                    {
                        "model_call_id": model_call_id,
                        "duration_ms": int((perf_counter() - started_at) * 1000),
                        "error_code": getattr(exc, "code", type(exc).__name__),
                    },
                )
                raise
            await self._publish(
                record,
                AgentEventType.model_call_completed,
                {
                    "model_call_id": model_call_id,
                    "duration_ms": int((perf_counter() - started_at) * 1000),
                    "finish_reason": "tool_calls" if turn.tool_calls else "stop",
                    "input_tokens": turn.input_tokens,
                    "output_tokens": turn.output_tokens,
                    "tool_call_count": len(turn.tool_calls),
                },
            )
            record.run.token_usage += turn.input_tokens + turn.output_tokens
            await self._publish(
                record,
                AgentEventType.usage,
                {"token_usage": record.run.token_usage},
            )
            if (
                record.request.token_budget is not None
                and record.run.token_usage > record.request.token_budget
            ):
                await self._fail(record, "TOKEN_BUDGET_EXCEEDED", "Agent token budget exceeded.")
                return

            if turn.tool_calls:
                if len(turn.tool_calls) > MAX_TOOL_CALLS_PER_TURN:
                    await self._fail(
                        record,
                        "TOO_MANY_TOOL_CALLS",
                        f"Provider requested more than {MAX_TOOL_CALLS_PER_TURN} tools in one turn.",
                    )
                    return
                calls = [
                    ToolCall(
                        tool_call_id=item.tool_call_id,
                        name=item.name,
                        arguments=item.arguments,
                    )
                    for item in turn.tool_calls
                ]
                messages.append(
                    Message(role=MessageRole.assistant, content=turn.text or "", reasoning_content=turn.reasoning_content, tool_calls=calls)
                )
                # 工具可以并发执行，但结果按模型原始调用顺序写回上下文，保证轮次可复现。
                semaphore = asyncio.Semaphore(record.request.max_concurrent_tools)

                async def execute(call: ToolCall) -> ToolResult:
                    async with semaphore:
                        return await self._execute_tool(record, call, model_call_id)

                executions = [asyncio.create_task(execute(call)) for call in calls]
                try:
                    results = await asyncio.gather(*executions)
                finally:
                    for execution in executions:
                        if not execution.done():
                            execution.cancel()
                    await asyncio.gather(*executions, return_exceptions=True)
                for call, result in zip(calls, results):
                    record.run.tool_results.append(result)
                    await self._collect_citations(record, result)
                    messages.append(
                        Message(
                            role=MessageRole.tool,
                            name=call.name,
                            tool_call_id=call.tool_call_id,
                            content=json.dumps(result.model_dump(mode="json"), ensure_ascii=False),
                        )
                    )
                continue

            if turn.text is not None:
                record.run.output = turn.text
                await self._publish(record, AgentEventType.text_delta, {"text": turn.text})
                record.run.status = AgentRunStatus.completed
                record.run.updated_at = datetime.now(timezone.utc)
                await self._publish(
                    record,
                    AgentEventType.run_completed,
                    {"output": turn.text, "token_usage": record.run.token_usage},
                )
                return

            await self._fail(record, "EMPTY_MODEL_RESPONSE", "Provider returned no text or tool call.")
            return

        await self._fail(record, "MAX_STEPS_EXCEEDED", "Agent reached its maximum step count.")

    async def _execute_tool(
        self, record: RunRecord, call: ToolCall, parent_model_call_id: str
    ) -> ToolResult:
        started_at = perf_counter()
        call_data = call.model_dump(mode="json")
        call_data["parent_model_call_id"] = parent_model_call_id
        await self._publish(record, AgentEventType.tool_call, call_data)
        try:
            registered = self.tools.get(call.name)
        except ToolNotFoundError:
            registered = None

        if registered is not None and call.name not in record.allowed_tools:
            result = ToolResult(
                tool_call_id=call.tool_call_id,
                name=call.name,
                success=False,
                error_code="TOOL_NOT_ALLOWED",
                error_message="Tool is not included in allowed_tools.",
            )
            await self._publish_tool_result(
                record, result, parent_model_call_id, started_at
            )
            return result

        permission = registered.definition.permission if registered else None
        if permission == "network.request" and not record.request.allow_network:
            result = ToolResult(
                tool_call_id=call.tool_call_id,
                name=call.name,
                success=False,
                error_code="NETWORK_NOT_ALLOWED",
                error_message="Agent run does not allow network tools.",
            )
            await self._publish_tool_result(
                record, result, parent_model_call_id, started_at
            )
            return result
        mode = self.permissions.mode_for(permission)
        if mode == PermissionMode.deny:
            result = self._permission_denied(call)
        elif mode == PermissionMode.confirm and permission:
            # 运行状态必须在等待期间可见，前端才能展示并处理权限确认卡片。
            ticket = self.permissions.create_ticket(record.run.run_id, permission)
            record.run.status = AgentRunStatus.waiting_permission
            await self._publish(
                record,
                AgentEventType.permission_required,
                {
                    "request_id": ticket.request_id,
                    "permission": permission,
                    "tool_call": call.model_dump(mode="json"),
                },
            )
            try:
                decision = await self.permissions.wait(
                    ticket, timeout=record.request.tool_timeout_seconds
                )
            except TimeoutError:
                record.run.status = AgentRunStatus.running
                result = ToolResult(
                    tool_call_id=call.tool_call_id,
                    name=call.name,
                    success=False,
                    error_code="PERMISSION_TIMEOUT",
                    error_message="Tool permission confirmation timed out.",
                )
                await self._publish_tool_result(
                    record, result, parent_model_call_id, started_at
                )
                return result
            record.run.status = AgentRunStatus.running
            record.run.updated_at = datetime.now(timezone.utc)
            async with record.publish_lock:
                snapshot = record.run.model_copy(deep=True)
                cancelled = await self._writer.submit('save', snapshot)
                record.persisted_run = snapshot
                if cancelled:
                    raise asyncio.CancelledError
            result = (
                await self._invoke_tool(record, call)
                if decision in {"allow_once", "allow_session"}
                else self._permission_denied(call)
            )
        else:
            result = await self._invoke_tool(record, call)

        await self._publish_tool_result(record, result, parent_model_call_id, started_at)
        return result

    async def _publish_tool_result(
        self,
        record: RunRecord,
        result: ToolResult,
        parent_model_call_id: str,
        started_at: float,
    ) -> None:
        data = result.model_dump(mode="json")
        data["parent_model_call_id"] = parent_model_call_id
        data["duration_ms"] = int((perf_counter() - started_at) * 1000)
        await self._publish(record, AgentEventType.tool_result, data)

    async def _invoke_tool(self, record: RunRecord, call: ToolCall) -> ToolResult:
        try:
            return await asyncio.wait_for(
                self.tools.execute(
                    call,
                    ToolExecutionContext(
                        run_id=record.run.run_id,
                        tool_call_id=call.tool_call_id,
                    ),
                ),
                timeout=record.request.tool_timeout_seconds,
            )
        except TimeoutError:
            return ToolResult(
                tool_call_id=call.tool_call_id,
                name=call.name,
                success=False,
                error_code="TOOL_TIMEOUT",
                error_message="Tool execution timed out.",
            )

    @staticmethod
    def _permission_denied(call: ToolCall) -> ToolResult:
        return ToolResult(
            tool_call_id=call.tool_call_id,
            name=call.name,
            success=False,
            error_code="PERMISSION_DENIED",
            error_message="Tool permission was denied.",
        )

    async def _finish_cancelled(self, record: RunRecord) -> None:
        record.run.cancelled = True
        record.run.status = AgentRunStatus.cancelled
        record.run.updated_at = datetime.now(timezone.utc)
        await self._publish(record, AgentEventType.run_cancelled, {})

    async def _fail(self, record: RunRecord, code: str, message: str) -> None:
        log_event('agent', 'run.error', level='ERROR', run_id=record.run.run_id, error_code=code)
        if record.persisted_run and record.persisted_run.status in TERMINAL_STATUSES:
            return
        record.run.status = AgentRunStatus.failed
        record.run.error_code = code
        record.run.error_message = message
        record.run.updated_at = datetime.now(timezone.utc)
        await self._publish(
            record,
            AgentEventType.run_failed,
            {"code": code, "message": message},
        )

    async def _publish(
        self, record: RunRecord, event_type: AgentEventType, data: dict[str, object]
    ) -> None:
        async with record.publish_lock:
            sanitized = sanitize_trace_value(data)
            assert isinstance(sanitized, dict)
            event = AgentEvent(
                event=event_type,
                run_id=record.run.run_id,
                sequence=record.next_sequence,
                data=sanitized,
                timestamp=datetime.now(timezone.utc),
            )
            snapshot = record.run.model_copy(deep=True)
            try:
                cancelled = await self._writer.submit('event', snapshot, event)
            except Exception as exc:
                log_event('agent', 'trace.write_failed', level='ERROR', error=exc, run_id=record.run.run_id)
                raise
            record.next_sequence += 1
            record.persisted_run = snapshot
            record.events.append(event)
            log_event('agent', event_type.value,
                      level='ERROR' if event_type.value.endswith('Failed') or data.get('success') is False else 'INFO',
                      run_id=record.run.run_id, provider_id=record.run.provider_id, model=record.run.model,
                      sequence=event.sequence, step=record.run.current_step, status=snapshot.status.value,
                      tool=data.get('name'), error_code=data.get('code') or data.get('error_code'))
            # 内存只保留实时订阅窗口；完整审计轨迹由 SQLite 保存。
            if len(record.events) > MAX_EVENTS_PER_RUN:
                del record.events[: len(record.events) - MAX_EVENTS_PER_RUN]
            for queue in record.subscribers:
                queue.put_nowait(event)
            if cancelled:
                raise asyncio.CancelledError

    @staticmethod
    def _request_metadata(record: RunRecord) -> dict[str, object]:
        metadata = dict(record.request.metadata)
        metadata["run_id"] = record.run.run_id
        if record.skill_config is not None:
            metadata["skill_id"] = record.skill_config.skill_id
            metadata["retrieval"] = record.skill_config.retrieval.model_dump(mode="json")
        return metadata

    def _config_snapshot(self, record: RunRecord) -> dict[str, object]:
        provider = self.providers.get(record.request.provider_id).config
        return {
            "provider_id": record.request.provider_id,
            "provider_type": provider.provider_type.value,
            "model": record.request.model,
            "capabilities": [item.value for item in provider.capabilities],
            "skill_id": record.request.skill_id,
            "allowed_tools": list(record.allowed_tools),
            "max_steps": record.request.max_steps,
            "token_budget": record.request.token_budget,
            "allow_network": record.request.allow_network,
            "metadata": record.request.metadata,
        }

    async def _collect_citations(self, record: RunRecord, result: ToolResult) -> None:
        if not result.success or not isinstance(result.output, dict):
            return
        items = result.output.get("items")
        if not isinstance(items, list):
            return
        known = {citation.citation_id for citation in record.run.citations}
        for item in items:
            if not isinstance(item, dict) or not isinstance(item.get("citation"), dict):
                continue
            try:
                citation = Citation.model_validate(item["citation"])
            except ValueError:
                continue
            if citation.citation_id in known:
                continue
            known.add(citation.citation_id)
            record.run.citations.append(citation)
            await self._publish(record, AgentEventType.citation, citation.model_dump(mode="json"))

    def _get_record(self, run_id: str) -> RunRecord:
        try:
            return self._records[run_id]
        except KeyError as exc:
            raise AgentRunNotFoundError(run_id) from exc

    def _prune_records(self) -> None:
        # 只清理终态记录，绝不为了容量取消仍在执行或等待授权的任务。
        overflow = len(self._records) - MAX_RUN_RECORDS + 1
        if overflow <= 0:
            return
        terminal = sorted(
            (
                record
                for record in self._records.values()
                if record.run.status in TERMINAL_STATUSES and (record.task is None or record.task.done())
            ),
            key=lambda record: record.run.updated_at,
        )
        for record in terminal[:overflow]:
            self._records.pop(record.run.run_id, None)
        if len(self._records) >= MAX_RUN_RECORDS:
            raise AgentCapacityError("Too many active Agent runs.")