agent-runtime-sdk 0.1.0__py3-none-any.whl

agent_runtime/runtime_engine/context.py

@@ -0,0 +1,469 @@
+"""运行时任务上下文与等待状态模型。
+
+这个文件只放 runtime 共享的状态定义，不放具体业务流程：
+
+- `current_task_id` / `current_task_pool`: 让运行中的 agent 能拿到当前任务上下文
+- `TaskContext`: 一次任务在执行期需要保存的状态
+- `WaitState`: runtime 暂停等待用户输入或授权时的统一模型
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import threading
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Any, Callable, Protocol
+
+
+current_task_id: contextvars.ContextVar[str] = contextvars.ContextVar("current_task_id")
+current_task_pool: contextvars.ContextVar["TaskPool"] = contextvars.ContextVar(
+    "current_task_pool"
+)
+
+WAIT_TYPE_INPUT_REQUIRED = "input_required"
+WAIT_TYPE_AUTH_REQUIRED = "auth_required"
+
+
+class TaskPhase(StrEnum):
+    """任务内部生命周期阶段。
+
+    这里的 phase 是 runtime 内部状态，不直接等价于 A2A 的 TaskState。
+    它的目标是让 task_flow / tool_flow 用一套明确状态，而不是依赖多个布尔位组合。
+    """
+
+    CREATED = "created"
+    SUBMITTED = "submitted"
+    RUNNING = "running"
+    WAITING_INPUT = "waiting_input"
+    WAITING_AUTH = "waiting_auth"
+    CANCELLING = "cancelling"
+    CANCELLED = "cancelled"
+    TIMED_OUT = "timed_out"
+    FAILED = "failed"
+    COMPLETED = "completed"
+    FINALIZED = "finalized"
+
+
+class TaskUpdaterProtocol(Protocol):
+    """task updater 的最小行为约束。"""
+
+    event_queue: Any
+
+    def new_agent_message(self, *, parts: list[Any], metadata: dict[str, Any]) -> Any:
+        ...
+
+    async def update_status(
+        self, status: Any, message: Any = None, metadata: dict[str, Any] | None = None
+    ) -> None:
+        ...
+
+    async def requires_input(self, message: Any = None) -> None:
+        ...
+
+    async def complete(self, message: Any = None) -> None:
+        ...
+
+    async def failed(self, message: Any = None) -> None:
+        ...
+
+    async def add_artifact(
+        self, *, parts: list[Any], name: str, metadata: dict[str, Any]
+    ) -> None:
+        ...
+
+
+@dataclass(frozen=True)
+class InputRequiredWaitState:
+    """任务因缺少用户输入而暂停时的等待态。"""
+
+    prompt: str
+
+    @property
+    def type(self) -> str:
+        return WAIT_TYPE_INPUT_REQUIRED
+
+    def to_payload(self) -> dict[str, Any]:
+        return {"type": self.type, "data": {"prompt": self.prompt}}
+
+
+@dataclass(frozen=True)
+class AuthRequiredWaitState:
+    """任务因需要用户授权确认而暂停时的等待态。"""
+
+    tool_name: str | None = None
+    args: dict[str, Any] | None = None
+
+    @property
+    def type(self) -> str:
+        return WAIT_TYPE_AUTH_REQUIRED
+
+    def to_payload(self) -> dict[str, Any]:
+        data: dict[str, Any] = {}
+        if self.tool_name is not None:
+            data["tool_name"] = self.tool_name
+        if self.args is not None:
+            data["args"] = self.args
+        return {"type": self.type, "data": data}
+
+
+WaitState = InputRequiredWaitState | AuthRequiredWaitState
+
+
+def wait_state_type(wait_state: WaitState | None) -> str | None:
+    if wait_state is None:
+        return None
+    return wait_state.type
+
+
+def wait_state_payload(wait_state: WaitState) -> dict[str, Any]:
+    return wait_state.to_payload()
+
+
+@dataclass
+class TaskRuntimeRefs:
+    """任务运行期持有的资源引用。"""
+
+    agent: object | None = None
+    agent_task: Future[Any] | None = None
+    task_executor: ThreadPoolExecutor | None = None
+    event: asyncio.Event | None = None
+    loop: asyncio.AbstractEventLoop | None = None
+    updater: TaskUpdaterProtocol | Any | None = None
+
+
+@dataclass
+class TaskWaitContext:
+    """等待用户恢复时的状态。"""
+
+    item: WaitState | None = None
+    emitted: bool = False
+    resume_payload: str | dict[str, Any] | None = None
+    timeout_handle: asyncio.Handle | None = None
+
+
+@dataclass
+class TaskOverrideContext:
+    """等待授权恢复时的附加状态。"""
+
+    auth_denied: bool = False
+    tool_args: Any = None
+    tool_name: str | None = None
+
+
+@dataclass
+class TaskControlContext:
+    """任务的生命周期与停止控制状态。"""
+
+    phase: TaskPhase = TaskPhase.CREATED
+    stop_requested: bool = False
+    finalized: bool = False
+    timed_out: bool = False
+    stop_reason: str | None = None
+    stop_error: Exception | None = None
+    task_timeout_handle: asyncio.Handle | None = None
+
+
+@dataclass
+class ActiveToolCall:
+    """当前正在执行的工具调用。"""
+
+    tool_name: str
+    mcp_name: str | None = None
+    cancel_hook: Callable[[], None] | None = None
+
+    def cancel(self) -> None:
+        if self.cancel_hook is not None:
+            self.cancel_hook()
+
+
+@dataclass(init=False)
+class TaskContext:
+    """一次任务在运行期的完整上下文。
+
+    当前实现把 task 拆成几块：
+
+    - `runtime`: 线程池、future、loop、updater 这些运行时引用
+    - `wait`: 等待用户输入/授权时的状态
+    - `overrides`: 授权确认恢复时的参数覆盖状态
+    - `control`: phase、取消、超时、finalize 这些生命周期状态
+    - `active_tool`: 当前执行中的工具，供超时/取消主动打断
+
+    为了降低迁移成本，这里保留了旧字段的 property 兼容访问。
+    """
+
+    runtime: TaskRuntimeRefs
+    wait: TaskWaitContext
+    overrides: TaskOverrideContext
+    control: TaskControlContext
+    active_tool: ActiveToolCall | None
+    auth_required: Any
+
+    def __init__(
+        self,
+        *,
+        agent: object | None = None,
+        agent_task: Future[Any] | None = None,
+        task_executor: ThreadPoolExecutor | None = None,
+        event: asyncio.Event | None = None,
+        wait_item: WaitState | None = None,
+        wait_item_emitted: bool = False,
+        user_input: str | dict[str, Any] | None = None,
+        auth_required: Any = None,
+        auth_denied: bool = False,
+        terminate: bool = False,
+        loop: asyncio.AbstractEventLoop | None = None,
+        updater: TaskUpdaterProtocol | Any | None = None,
+        finalized: bool = False,
+        tool_args_override: Any = None,
+        tool_name_override: str | None = None,
+        timed_out: bool = False,
+        timeout_reason: str | None = None,
+        timeout_handle: asyncio.Handle | None = None,
+        phase: TaskPhase = TaskPhase.CREATED,
+        runtime: TaskRuntimeRefs | None = None,
+        wait: TaskWaitContext | None = None,
+        overrides: TaskOverrideContext | None = None,
+        control: TaskControlContext | None = None,
+        active_tool: ActiveToolCall | None = None,
+    ) -> None:
+        self.runtime = runtime or TaskRuntimeRefs(
+            agent=agent,
+            agent_task=agent_task,
+            task_executor=task_executor,
+            event=event,
+            loop=loop,
+            updater=updater,
+        )
+        self.wait = wait or TaskWaitContext(
+            item=wait_item,
+            emitted=wait_item_emitted,
+            resume_payload=user_input,
+            timeout_handle=timeout_handle,
+        )
+        self.overrides = overrides or TaskOverrideContext(
+            auth_denied=auth_denied,
+            tool_args=tool_args_override,
+            tool_name=tool_name_override,
+        )
+        self.control = control or TaskControlContext(
+            phase=phase,
+            stop_requested=terminate,
+            finalized=finalized,
+            timed_out=timed_out,
+            stop_reason=timeout_reason,
+        )
+        self.active_tool = active_tool
+        self.auth_required = auth_required
+
+    def set_phase(self, phase: TaskPhase) -> None:
+        self.control.phase = phase
+
+    def request_stop(self, *, error: Exception, reason: str, timed_out: bool) -> None:
+        self.control.stop_requested = True
+        self.control.timed_out = timed_out
+        self.control.stop_reason = reason
+        self.control.stop_error = error
+
+    def clear_wait(self) -> None:
+        self.wait.item = None
+        self.wait.emitted = False
+        self.wait.resume_payload = None
+
+    def clear_wait_item(self) -> None:
+        self.wait.item = None
+        self.wait.emitted = False
+
+    def set_wait(self, wait_state: WaitState, phase: TaskPhase) -> None:
+        self.wait.item = wait_state
+        self.wait.emitted = False
+        self.wait.resume_payload = None
+        self.control.timed_out = False
+        self.control.stop_reason = None
+        self.set_phase(phase)
+
+    def set_active_tool(self, active_tool: ActiveToolCall | None) -> None:
+        self.active_tool = active_tool
+
+    def clear_active_tool(self) -> None:
+        self.active_tool = None
+
+    @property
+    def is_waiting(self) -> bool:
+        return self.wait.item is not None
+
+    @property
+    def is_terminal(self) -> bool:
+        return self.control.phase in {
+            TaskPhase.CANCELLED,
+            TaskPhase.COMPLETED,
+            TaskPhase.FAILED,
+            TaskPhase.TIMED_OUT,
+            TaskPhase.FINALIZED,
+        }
+
+    @property
+    def agent(self):
+        return self.runtime.agent
+
+    @agent.setter
+    def agent(self, value: object | None) -> None:
+        self.runtime.agent = value
+
+    @property
+    def agent_task(self):
+        return self.runtime.agent_task
+
+    @agent_task.setter
+    def agent_task(self, value: Future[Any] | None) -> None:
+        self.runtime.agent_task = value
+
+    @property
+    def task_executor(self):
+        return self.runtime.task_executor
+
+    @task_executor.setter
+    def task_executor(self, value: ThreadPoolExecutor | None) -> None:
+        self.runtime.task_executor = value
+
+    @property
+    def event(self):
+        return self.runtime.event
+
+    @event.setter
+    def event(self, value: asyncio.Event | None) -> None:
+        self.runtime.event = value
+
+    @property
+    def wait_item(self):
+        return self.wait.item
+
+    @wait_item.setter
+    def wait_item(self, value: WaitState | None) -> None:
+        self.wait.item = value
+
+    @property
+    def wait_item_emitted(self) -> bool:
+        return self.wait.emitted
+
+    @wait_item_emitted.setter
+    def wait_item_emitted(self, value: bool) -> None:
+        self.wait.emitted = value
+
+    @property
+    def user_input(self):
+        return self.wait.resume_payload
+
+    @user_input.setter
+    def user_input(self, value: str | dict[str, Any] | None) -> None:
+        self.wait.resume_payload = value
+
+    @property
+    def auth_denied(self) -> bool:
+        return self.overrides.auth_denied
+
+    @auth_denied.setter
+    def auth_denied(self, value: bool) -> None:
+        self.overrides.auth_denied = value
+
+    @property
+    def terminate(self) -> bool:
+        return self.control.stop_requested
+
+    @terminate.setter
+    def terminate(self, value: bool) -> None:
+        self.control.stop_requested = value
+
+    @property
+    def loop(self):
+        return self.runtime.loop
+
+    @loop.setter
+    def loop(self, value: asyncio.AbstractEventLoop | None) -> None:
+        self.runtime.loop = value
+
+    @property
+    def updater(self):
+        return self.runtime.updater
+
+    @updater.setter
+    def updater(self, value: TaskUpdaterProtocol | Any | None) -> None:
+        self.runtime.updater = value
+
+    @property
+    def finalized(self) -> bool:
+        return self.control.finalized
+
+    @finalized.setter
+    def finalized(self, value: bool) -> None:
+        self.control.finalized = value
+
+    @property
+    def tool_args_override(self):
+        return self.overrides.tool_args
+
+    @tool_args_override.setter
+    def tool_args_override(self, value: Any) -> None:
+        self.overrides.tool_args = value
+
+    @property
+    def tool_name_override(self) -> str | None:
+        return self.overrides.tool_name
+
+    @tool_name_override.setter
+    def tool_name_override(self, value: str | None) -> None:
+        self.overrides.tool_name = value
+
+    @property
+    def timed_out(self) -> bool:
+        return self.control.timed_out
+
+    @timed_out.setter
+    def timed_out(self, value: bool) -> None:
+        self.control.timed_out = value
+
+    @property
+    def timeout_reason(self) -> str | None:
+        return self.control.stop_reason
+
+    @timeout_reason.setter
+    def timeout_reason(self, value: str | None) -> None:
+        self.control.stop_reason = value
+
+    @property
+    def timeout_handle(self):
+        return self.wait.timeout_handle
+
+    @timeout_handle.setter
+    def timeout_handle(self, value: asyncio.Handle | None) -> None:
+        self.wait.timeout_handle = value
+
+
+class TaskPool:
+    """线程安全的任务上下文池。"""
+
+    def __init__(self):
+        self._tasks: dict[str, TaskContext] = {}
+        self._lock = threading.RLock()
+
+    def __contains__(self, task_id: str) -> bool:
+        with self._lock:
+            return task_id in self._tasks
+
+    def __getitem__(self, task_id: str) -> TaskContext:
+        with self._lock:
+            return self._tasks[task_id]
+
+    def __setitem__(self, task_id: str, ctx: TaskContext) -> None:
+        with self._lock:
+            self._tasks[task_id] = ctx
+
+    def get(self, task_id: str, default: Any = None) -> Any:
+        with self._lock:
+            return self._tasks.get(task_id, default)
+
+    def pop(self, task_id: str, default: Any = None) -> Any:
+        with self._lock:
+            return self._tasks.pop(task_id, default)

agent_runtime/runtime_engine/loading.py

@@ -0,0 +1,170 @@
+"""runtime reload 阶段的内部引擎。
+
+职责很聚焦：
+
+- 加载或重载 plugin
+- 发现所有可用 MCP 工具
+- 建立 `tool_name -> source_mcp` 索引
+
+这一步只准备运行时元数据，不负责构造单次请求 agent。
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from ..config.definitions import (
+    DEFAULT_PLUGIN_CLASS_NAME,
+    DEFAULT_PLUGIN_MODULE_PATH,
+    PluginSettings,
+)
+from ..mcp.metadata import DiscoveredTool, normalize_discovered_tools
+from ..mcp.utils import index_tool_sources, merge_mcp_headers
+from ..plugins import BaseAgentPlugin, load_plugin
+
+if TYPE_CHECKING:
+    from ..runtime import ManagedAgentRuntime
+
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeLoading:
+    """负责 runtime reload 阶段的 plugin 与工具元数据准备。"""
+
+    def reload(
+        self, discover: bool = True, skip_plugin_load: bool = False
+    ) -> "ManagedAgentRuntime":
+        """重新加载插件和工具描述。"""
+        self._reset_reload_state()
+        logger.info(
+            "Reloading runtime agent_id=%s discover=%s skip_plugin_load=%s",
+            self._agent.agent_id,
+            discover,
+            skip_plugin_load,
+        )
+        try:
+            plugin = self._resolve_plugin_for_reload(skip_plugin_load=skip_plugin_load)
+            discovered_tools = self._discover_tools_metadata() if discover else []
+            tool_sources = self._index_tool_sources(discovered_tools)
+            self._apply_reload_result(
+                plugin=plugin,
+                discovered_tools=discovered_tools,
+                tool_sources=tool_sources,
+            )
+            logger.info(
+                "Runtime reload finished agent_id=%s discovered_tools=%s",
+                self._agent.agent_id,
+                len(self.discovered_tools),
+            )
+        except Exception as exc:
+            self._apply_reload_failure(exc)
+        return self
+
+    def _reset_reload_state(self) -> None:
+        self._asgi_app = None
+        self.load_error = None
+        self.load_exception = None
+        self.discovered_tools = []
+        self._tool_source_by_name = {}
+
+    def _resolve_plugin_for_reload(
+        self, *, skip_plugin_load: bool
+    ) -> BaseAgentPlugin | None:
+        if skip_plugin_load:
+            return self.plugin
+        return self._load_plugin_instance()
+
+    def _load_plugin_instance(self) -> BaseAgentPlugin:
+        plugin_settings = self._effective_plugin_settings()
+        if plugin_settings is not None:
+            plugin = load_plugin(
+                self._resolved_plugin_module_path(plugin_settings),
+                plugin_settings.class_name,
+                plugin_settings.config,
+            )
+        else:
+            from ..plugins.default import DefaultPlugin
+
+            plugin = DefaultPlugin()
+        logger.debug(
+            "Loaded plugin for agent_id=%s plugin=%s",
+            self._agent.agent_id,
+            type(plugin).__name__,
+        )
+        return plugin
+
+    def _discover_tools_metadata(self) -> list[DiscoveredTool]:
+        """发现阶段只收集工具摘要，不保留可执行 tool。"""
+
+        discovered_tools: list[DiscoveredTool] = []
+        for mcp in self._mcps:
+            if not mcp.enabled:
+                logger.debug(
+                    "Skipping disabled MCP during reload agent_id=%s mcp=%s",
+                    self._agent.agent_id,
+                    mcp.name,
+                )
+                continue
+            discovered_tools.extend(self._discover_mcp_tools_metadata(mcp))
+        return discovered_tools
+
+    def _discover_mcp_tools_metadata(self, mcp) -> list[DiscoveredTool]:
+        logger.debug(
+            "Discovering MCP tools agent_id=%s mcp=%s url=%s",
+            self._agent.agent_id,
+            mcp.name,
+            mcp.url,
+        )
+        return normalize_discovered_tools(
+            self.discoverer(mcp, headers=merge_mcp_headers(mcp)), mcp.name
+        )
+
+    def _index_tool_sources(
+        self, discovered_tools: list[DiscoveredTool]
+    ) -> dict[str, str]:
+        return index_tool_sources(discovered_tools)
+
+    def _apply_reload_result(
+        self,
+        *,
+        plugin: BaseAgentPlugin | None,
+        discovered_tools: list[DiscoveredTool],
+        tool_sources: dict[str, str],
+    ) -> None:
+        if plugin is not None:
+            self.plugin = plugin
+        self.discovered_tools = discovered_tools
+        self._tool_source_by_name = tool_sources
+
+    def _apply_reload_failure(self, exc: Exception) -> None:
+        self.discovered_tools = []
+        self.load_exception = exc
+        self.load_error = f"{type(exc).__name__}: {exc}"
+        logger.exception("Runtime reload failed agent_id=%s", self._agent.agent_id)
+
+    def _effective_plugin_settings(self) -> PluginSettings | None:
+        if self._runtime_config.plugin is not None:
+            return self._runtime_config.plugin
+
+        default_plugin_path = (
+            Path(self._runtime_config.base_path).resolve() / DEFAULT_PLUGIN_MODULE_PATH
+        )
+        if default_plugin_path.exists():
+            return PluginSettings(
+                module_path=DEFAULT_PLUGIN_MODULE_PATH,
+                class_name=DEFAULT_PLUGIN_CLASS_NAME,
+                config={},
+            )
+        return None
+
+    def _resolved_plugin_module_path(self, plugin_settings: PluginSettings) -> str:
+        module_path = plugin_settings.module_path
+        if module_path.endswith(".py") or "/" in module_path:
+            path = Path(module_path)
+            if not path.is_absolute():
+                path = Path(self._runtime_config.base_path).resolve() / path
+            return path.as_posix()
+        return module_path