zero-agent-core 0.2.0__py3-none-any.whl

agent_core/runtime/service.py

@@ -0,0 +1,557 @@
+"""通用 Agent Runtime 生命周期服务。
+
+Runtime 负责运行实例和生命周期，业务编排仍由调用方注入到工具或上层服务中。
+这样 API 层不需要重复实现 start、resume、cancel 和事件发布逻辑。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Sequence
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from agent_core.access import AccessContext, enforce_access
+from agent_core.callbacks import CallbackHandler, CallbackManager
+from agent_core.ports import EventSink, RunStore
+from agent_core.protocol.runtime import RunContext, RunEvent
+from agent_core.runtime.react import ReActAgent
+from agent_core.runtime.store import MemoryRunStore
+from agent_core.storage.lease import RunLeaseStore
+from agent_core.storage.runtime import RuntimeConcurrencyError
+
+logger = logging.getLogger(__name__)
+
+
+class RunExecutor(Protocol):
+    """自定义 Run 执行器：在 Run 生命周期内替代默认的 ``agent.run`` 调用。
+
+    Runtime 会在挂接执行器前完成状态置位、``run_started`` 事件、checkpoint
+    回调绑定和持久化；执行器只需驱动具体的 Agent 实现并返回最终回答。
+    应用可用它接入按模型快照选择 Agent、业务编排分支等自定义执行策略。
+    """
+
+    async def __call__(
+        self,
+        context: RunContext,
+        user_input: str,
+        *,
+        memory_context: str | None = None,
+        checkpoint_state: dict[str, Any] | None = None,
+    ) -> str: ...
+
+
+class MemoryEventSink(EventSink):
+    """测试和开发环境使用的事件接收器。"""
+
+    def __init__(self) -> None:
+        self.events: list[RunEvent] = []
+
+    async def publish(self, event: RunEvent) -> None:
+        self.events.append(event)
+
+
+@dataclass
+class RuntimeRun:
+    """运行句柄，供需要后台任务的 API 层使用。"""
+
+    context: RunContext
+    task: asyncio.Task[str] | None = None
+
+
+class AgentRuntime:
+    """围绕 ``ReActAgent`` 的通用生命周期门面。"""
+
+    def __init__(
+        self,
+        agent: ReActAgent,
+        *,
+        run_store: RunStore | None = None,
+        event_sink: EventSink | None = None,
+        callbacks: Sequence[CallbackHandler] | None = None,
+        callback_manager: CallbackManager | None = None,
+        lease_store: RunLeaseStore | None = None,
+        worker_id: str | None = None,
+        lease_seconds: float = 30.0,
+        heartbeat_seconds: float = 10.0,
+        require_access: bool = False,
+        run_executor: RunExecutor | None = None,
+    ) -> None:
+        if lease_seconds <= 0:
+            raise ValueError("lease_seconds 必须大于 0")
+        if heartbeat_seconds <= 0 or heartbeat_seconds >= lease_seconds:
+            raise ValueError("heartbeat_seconds 必须大于 0 且小于 lease_seconds")
+        if callbacks is not None and callback_manager is not None:
+            raise ValueError("callbacks 和 callback_manager 不能同时传入")
+        self._agent = agent
+        self._run_store = run_store or MemoryRunStore()
+        self._event_sink = event_sink
+        self._callback_manager = callback_manager or CallbackManager(callbacks)
+        self._lease_store = lease_store
+        self._worker_id = worker_id or f"worker-{uuid.uuid4()}"
+        self._lease_seconds = lease_seconds
+        self._heartbeat_seconds = heartbeat_seconds
+        self._require_access = require_access
+        self._run_executor = run_executor
+        self._tasks: dict[str, asyncio.Task[str]] = {}
+        self._event_cursors: dict[str, int] = {}
+        self._lock = asyncio.Lock()
+
+    async def start(
+        self,
+        thread_id: str,
+        user_input: str,
+        *,
+        user_id: str | None = None,
+        tenant_id: str | None = None,
+        parent_run_id: str | None = None,
+        tags: Sequence[str] | None = None,
+        idempotency_key: str | None = None,
+        memory_context: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        checkpoint_state: dict[str, Any] | None = None,
+    ) -> RuntimeRun:
+        """创建后台运行实例并立即返回句柄。"""
+        self._require_owner(user_id, tenant_id)
+        async with self._lock:
+            if idempotency_key:
+                existing = await self._run_store.find_run_by_idempotency(tenant_id, user_id, idempotency_key)
+                if existing is not None:
+                    task = self._tasks.get(existing.run_id)
+                    return RuntimeRun(context=existing, task=task)
+
+            context = RunContext(
+                thread_id=thread_id,
+                run_id=str(uuid.uuid4()),
+                user_id=user_id,
+                tenant_id=tenant_id,
+                parent_run_id=parent_run_id,
+                tags=list(tags or ()),
+                idempotency_key=idempotency_key,
+                metadata={
+                    **(metadata or {}),
+                    "_input": user_input,
+                },
+                checkpoint=checkpoint_state or {},
+            )
+            await self._run_store.save_run(context)
+            await self._claim(context)
+            task = asyncio.create_task(
+                self._run_with_lease(
+                    context,
+                    self._execute(
+                        context,
+                        user_input,
+                        memory_context=memory_context,
+                        checkpoint_state=checkpoint_state,
+                    ),
+                )
+            )
+            self._tasks[context.run_id] = task
+            task.add_done_callback(lambda _: self._cleanup_run(context.run_id))
+            return RuntimeRun(context=context, task=task)
+
+    async def run(
+        self,
+        thread_id: str,
+        user_input: str,
+        **kwargs: Any,
+    ) -> RunContext:
+        """同步等待一次完整运行并返回最终上下文。"""
+        handle = await self.start(thread_id, user_input, **kwargs)
+        access = self._context_access(handle.context)
+        await self.wait(
+            handle.context.thread_id,
+            handle.context.run_id,
+            access=access,
+        )
+        return await self.get(
+            handle.context.thread_id,
+            handle.context.run_id,
+            access=access,
+        )
+
+    async def wait(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        access: AccessContext | None = None,
+    ) -> RunContext:
+        """等待进程内任务；进程重启后可读取持久化状态。"""
+        context = await self._run_store.load_run(thread_id, run_id)
+        if context is None:
+            raise KeyError(f"run 不存在：{run_id}")
+        self._check_access(context, access)
+        task = self._tasks.get(run_id)
+        if task is not None:
+            await asyncio.shield(task)
+        context = await self._run_store.load_run(thread_id, run_id)
+        if context is None:
+            raise KeyError(f"run 不存在：{run_id}")
+        self._check_access(context, access)
+        return context
+
+    async def get(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        access: AccessContext | None = None,
+    ) -> RunContext:
+        """读取运行上下文。"""
+        context = await self._run_store.load_run(thread_id, run_id)
+        if context is None:
+            raise KeyError(f"run 不存在：{run_id}")
+        self._check_access(context, access)
+        return context
+
+    async def resume(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        access: AccessContext | None = None,
+    ) -> RuntimeRun:
+        """从已有 checkpoint 继续运行同一个 run。"""
+        async with self._lock:
+            if run_id in self._tasks:
+                return RuntimeRun(
+                    context=await self.get(thread_id, run_id, access=access),
+                    task=self._tasks[run_id],
+                )
+            context = await self.get(thread_id, run_id, access=access)
+            if not context.checkpoint:
+                raise ValueError("该 run 没有可恢复的 checkpoint")
+            if context.status in {"completed", "cancelled"}:
+                raise ValueError(f"状态为 {context.status} 的 run 不能恢复")
+            self._event_cursors.setdefault(context.run_id, len(context.events))
+            await self._claim(context)
+            try:
+                context.status = "queued"
+                context.emit("run_resumed")
+                await self._persist(context)
+            except Exception:
+                await self._release(context)
+                raise
+            task = asyncio.create_task(
+                self._run_with_lease(
+                    context,
+                    self._execute(
+                        context,
+                        "",
+                        checkpoint_state=context.checkpoint,
+                    ),
+                )
+            )
+            self._tasks[run_id] = task
+            task.add_done_callback(lambda _: self._cleanup_run(run_id))
+            return RuntimeRun(context=context, task=task)
+
+    async def cancel(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        access: AccessContext | None = None,
+    ) -> bool:
+        """取消进程内运行；已完成或只存在于数据库中的 run 不会被伪造修改。"""
+        context = await self.get(thread_id, run_id, access=access)
+        task = self._tasks.get(run_id)
+        if task is None or task.done():
+            return False
+        task.cancel()
+        gather_result = await asyncio.gather(task, return_exceptions=True)
+        for outcome in gather_result:
+            # 取消方不再吞掉业务异常；这里记录日志，避免取消失败与取消成功不可区分。
+            if isinstance(outcome, BaseException) and not isinstance(outcome, asyncio.CancelledError):
+                logger.warning("被取消的 run %s 在结束前抛出异常", run_id, exc_info=outcome)
+        context = await self.get(thread_id, run_id, access=access)
+        return context.status == "cancelled"
+
+    async def stream(
+        self,
+        thread_id: str,
+        user_input: str,
+        **kwargs: Any,
+    ) -> AsyncIterator[str]:
+        """以流式方式执行一次运行，并复用同一生命周期持久化逻辑。"""
+        context = RunContext(
+            thread_id=thread_id,
+            run_id=str(uuid.uuid4()),
+            user_id=kwargs.pop("user_id", None),
+            tenant_id=kwargs.pop("tenant_id", None),
+            parent_run_id=kwargs.pop("parent_run_id", None),
+            tags=list(kwargs.pop("tags", ())),
+            idempotency_key=kwargs.pop("idempotency_key", None),
+            metadata={**kwargs.pop("metadata", {}), "_input": user_input},
+        )
+        self._require_owner(context.user_id, context.tenant_id)
+        await self._run_store.save_run(context)
+        await self._claim(context)
+        context.status = "running"
+        context.emit("run_started", input=user_input)
+        context.checkpoint_callback = self._checkpoint_callback(context)
+        await self._persist(context)
+        lease_errors: list[BaseException] = []
+        heartbeat = self._start_stream_heartbeat(context, lease_errors)
+        try:
+            async for chunk in self._agent.stream(
+                user_input,
+                memory_context=kwargs.pop("memory_context", None),
+                thread_id=thread_id,
+                run_context=context,
+                checkpoint_state=kwargs.pop("checkpoint_state", None),
+            ):
+                context.emit("model_stream_chunk", component="model", delta=chunk)
+                await self._publish_pending_events(context)
+                yield chunk
+        except asyncio.CancelledError:
+            if not lease_errors:
+                raise
+            context.status = "interrupted"
+            context.emit(
+                "run_interrupted",
+                reason="worker_lease_lost",
+                error=str(lease_errors[0]),
+            )
+            raise RuntimeConcurrencyError(f"run {context.run_id} 的 Worker 租约已丢失") from lease_errors[0]
+        finally:
+            if heartbeat is not None:
+                heartbeat.cancel()
+                await asyncio.gather(heartbeat, return_exceptions=True)
+            await self._release(context)
+            await self._persist(context)
+            self._event_cursors.pop(context.run_id, None)
+
+    async def recover_expired_runs(self, *, auto_resume: bool = False) -> list[RunContext]:
+        """标记租约过期的 Run，并可从 checkpoint 自动恢复。"""
+        if self._lease_store is None:
+            return []
+        recovered: list[RunContext] = []
+        for lease in await self._lease_store.list_expired():
+            context = await self._run_store.load_run(lease.thread_id, lease.run_id)
+            await self._lease_store.release(lease.thread_id, lease.run_id, lease.worker_id)
+            if context is None or context.status not in {"queued", "running"}:
+                continue
+            self._event_cursors.setdefault(context.run_id, len(context.events))
+            context.status = "interrupted"
+            context.emit("run_interrupted", reason="worker_lease_expired")
+            await self._persist(context)
+            recovered.append(context)
+            if auto_resume and context.checkpoint:
+                await self.resume(
+                    context.thread_id,
+                    context.run_id,
+                    access=self._context_access(context),
+                )
+        return recovered
+
+    async def _claim(self, context: RunContext) -> None:
+        if self._lease_store is None:
+            return
+        claimed = await self._lease_store.claim(
+            context.thread_id,
+            context.run_id,
+            self._worker_id,
+            ttl_seconds=self._lease_seconds,
+        )
+        if not claimed:
+            raise RuntimeConcurrencyError(f"run {context.run_id} 已被其他 Worker 认领")
+
+    async def _release(self, context: RunContext) -> None:
+        if self._lease_store is not None:
+            await self._lease_store.release(
+                context.thread_id,
+                context.run_id,
+                self._worker_id,
+            )
+
+    async def _heartbeat(self, context: RunContext) -> None:
+        assert self._lease_store is not None
+        while True:
+            await asyncio.sleep(self._heartbeat_seconds)
+            renewed = await self._lease_store.renew(
+                context.thread_id,
+                context.run_id,
+                self._worker_id,
+                ttl_seconds=self._lease_seconds,
+            )
+            if not renewed:
+                raise RuntimeConcurrencyError(f"run {context.run_id} 的 Worker 租约已丢失")
+
+    async def _run_with_lease(
+        self,
+        context: RunContext,
+        execution: Coroutine[Any, Any, str],
+    ) -> str:
+        if self._lease_store is None:
+            return await execution
+        execution_task = asyncio.create_task(execution)
+        heartbeat_task = asyncio.create_task(self._heartbeat(context))
+        try:
+            done, _ = await asyncio.wait(
+                {execution_task, heartbeat_task},
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            # 两个任务可能在同一轮事件循环同时完成；执行结果优先，避免丢失成功终态。
+            if execution_task in done:
+                return await execution_task
+            if heartbeat_task in done:
+                error = heartbeat_task.exception()
+                execution_task.cancel()
+                await asyncio.gather(execution_task, return_exceptions=True)
+                context.status = "interrupted"
+                context.emit(
+                    "run_interrupted",
+                    reason="worker_lease_lost",
+                    error=str(error) if error is not None else None,
+                )
+                await self._persist(context)
+                if error is not None:
+                    raise error
+                raise RuntimeConcurrencyError(f"run {context.run_id} 的 Worker 心跳意外结束")
+            return await execution_task
+        finally:
+            if not execution_task.done():
+                execution_task.cancel()
+                await asyncio.gather(execution_task, return_exceptions=True)
+            heartbeat_task.cancel()
+            await asyncio.gather(heartbeat_task, return_exceptions=True)
+            await self._release(context)
+
+    def _start_stream_heartbeat(
+        self,
+        context: RunContext,
+        errors: list[BaseException],
+    ) -> asyncio.Task[None] | None:
+        if self._lease_store is None:
+            return None
+        owner = asyncio.current_task()
+
+        async def monitor() -> None:
+            try:
+                await self._heartbeat(context)
+            except Exception as exc:
+                errors.append(exc)
+                context.emit("run_lease_lost", worker_id=self._worker_id)
+                if owner is not None:
+                    owner.cancel()
+                raise
+
+        return asyncio.create_task(monitor())
+
+    def _check_access(
+        self,
+        context: RunContext,
+        access: AccessContext | None,
+    ) -> None:
+        enforce_access(
+            strict=self._require_access,
+            accessor="AgentRuntime",
+            subject="run",
+            owner_user_id=context.user_id,
+            owner_tenant_id=context.tenant_id,
+            access=access,
+        )
+
+    def _require_owner(
+        self,
+        user_id: str | None,
+        tenant_id: str | None,
+    ) -> None:
+        if self._require_access and (user_id is None or tenant_id is None):
+            raise PermissionError("严格访问模式要求同时提供 user_id 和 tenant_id")
+
+    @staticmethod
+    def _context_access(context: RunContext) -> AccessContext | None:
+        if context.user_id is None or context.tenant_id is None:
+            return None
+        return AccessContext(user_id=context.user_id, tenant_id=context.tenant_id)
+
+    async def _execute(
+        self,
+        context: RunContext,
+        user_input: str,
+        *,
+        memory_context: str | None = None,
+        checkpoint_state: dict[str, Any] | None = None,
+    ) -> str:
+        context.status = "running"
+        context.emit("run_started", input=user_input)
+        context.checkpoint_callback = self._checkpoint_callback(context)
+        await self._persist(context)
+        try:
+            if self._run_executor is not None:
+                # 自定义执行器：生命周期（状态/事件/持久化/租约）仍由 Runtime 收口。
+                answer = await self._run_executor(
+                    context,
+                    user_input,
+                    memory_context=memory_context,
+                    checkpoint_state=checkpoint_state,
+                )
+            else:
+                answer = await self._agent.run(
+                    user_input,
+                    memory_context=memory_context,
+                    thread_id=context.thread_id,
+                    run_context=context,
+                    checkpoint_state=checkpoint_state,
+                )
+        except asyncio.CancelledError:
+            if context.status == "running":
+                await self._finish_terminal(context, "cancelled", "run_cancelled")
+            raise
+        except Exception as exc:
+            if context.status == "running":
+                await self._finish_terminal(context, "failed", "run_failed", error=str(exc))
+            raise
+        # 内置路径由 Agent Loop 自己收口；执行器路径在此统一落到 completed。
+        if context.status == "running":
+            await self._finish_terminal(context, "completed", "run_completed", answer=answer)
+        return answer
+
+    async def _finish_terminal(
+        self,
+        context: RunContext,
+        status: str,
+        event_type: str,
+        **payload: Any,
+    ) -> None:
+        context.status = status  # type: ignore[assignment]
+        context.emit(event_type, **payload)
+        await self._persist(context)
+
+    def _checkpoint_callback(
+        self,
+        context: RunContext,
+    ) -> Callable[[dict[str, Any]], Awaitable[None]]:
+        async def persist(_state: dict[str, Any]) -> None:
+            await self._persist(context)
+
+        return persist
+
+    async def _persist(self, context: RunContext) -> None:
+        await self._run_store.save_run(context)
+        await self._publish_pending_events(context)
+
+    async def _publish_pending_events(self, context: RunContext) -> None:
+        """把尚未发布的事件依次交给进程内 Callback 和进程外 EventSink。"""
+
+        cursor = self._event_cursors.get(context.run_id, 0)
+        for event in context.events[cursor:]:
+            await self._callback_manager.dispatch(event)
+            if self._event_sink is not None:
+                try:
+                    await self._event_sink.publish(event)
+                except Exception:
+                    # 事件订阅者故障不能覆盖 Agent 的最终结果。
+                    logger.exception("发布 Agent 运行事件失败：%s", event.event_type)
+        self._event_cursors[context.run_id] = len(context.events)
+
+    def _cleanup_run(self, run_id: str) -> None:
+        """后台任务完全结束后释放进程内句柄和事件游标。"""
+        self._tasks.pop(run_id, None)
+        self._event_cursors.pop(run_id, None)

agent_core/runtime/store.py

@@ -0,0 +1,55 @@
+"""Runtime 的默认内存实现。
+
+它只适合测试和单进程开发环境。生产应用应通过 ``RunStore`` 注入 SQLite、
+PostgreSQL 或其他持久化实现。
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+from agent_core.ports import RunStore
+from agent_core.protocol.runtime import RunContext
+
+
+class MemoryRunStore(RunStore):
+    """带基本版本控制和幂等查询的内存运行存储。"""
+
+    def __init__(self) -> None:
+        self._runs: dict[tuple[str, str], dict[str, Any]] = {}
+
+    async def save_run(self, context: RunContext) -> None:
+        key = (context.thread_id, context.run_id)
+        previous = self._runs.get(key)
+        previous_version = int(previous.get("version", 0)) if previous else 0
+        if previous is not None and context.version != previous_version:
+            raise RuntimeError(f"run {context.run_id} 版本冲突")
+        context.version = previous_version + 1
+        self._runs[key] = copy.deepcopy(context.to_dict())
+
+    async def load_run(self, thread_id: str, run_id: str) -> RunContext | None:
+        value = self._runs.get((thread_id, run_id))
+        if value is None:
+            return None
+        return _context_from_dict(value)
+
+    async def find_run_by_idempotency(
+        self,
+        tenant_id: str | None,
+        user_id: str | None,
+        idempotency_key: str,
+    ) -> RunContext | None:
+        for value in self._runs.values():
+            if (
+                value.get("tenant_id") == tenant_id
+                and value.get("user_id") == user_id
+                and value.get("idempotency_key") == idempotency_key
+            ):
+                return _context_from_dict(value)
+        return None
+
+
+def _context_from_dict(value: dict[str, Any]) -> RunContext:
+    """从持久化字典恢复 Core 上下文，忽略不可序列化的回调。"""
+    return RunContext.from_dict(copy.deepcopy(value))

agent_core/serialization/__init__.py

@@ -0,0 +1,21 @@
+"""安全、版本化的 Core 对象序列化。"""
+
+from agent_core.serialization.registry import (
+    SerializedEnvelope,
+    SerializerRegistry,
+    default_serializer_registry,
+    dumpd,
+    dumps,
+    load,
+    loads,
+)
+
+__all__ = [
+    "SerializedEnvelope",
+    "SerializerRegistry",
+    "default_serializer_registry",
+    "dumpd",
+    "dumps",
+    "load",
+    "loads",
+]