agentengine-sdk-python 0.6.9__py3-none-any.whl → 0.8.0__py3-none-any.whl

ksadk/a2a/resume_store.py

@@ -0,0 +1,201 @@
+"""Runtime-local storage for opaque A2A input-required recovery state."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import os
+import sqlite3
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, TypeAlias
+
+from ksadk.runtime.adapter import ResumePayload, ResumeTarget, RunHandle
+
+A2AResumePayloadKind: TypeAlias = Literal[
+    "tool_result",
+    "approval_decision",
+    "hitl_answer",
+    "free_text",
+]
+
+
+@dataclass(frozen=True)
+class A2AResumeState:
+    """Private Runtime state needed to resume one protocol Task."""
+
+    handle: RunHandle
+    target: ResumeTarget
+    payload_kind: A2AResumePayloadKind
+    call_id: str | None = None
+
+
+class A2AResumeStateStore(ABC):
+    """Persists recovery state outside public A2A Task metadata."""
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def put(self, key: str, state: A2AResumeState) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get(self, key: str) -> A2AResumeState | None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete(self, key: str) -> None:
+        raise NotImplementedError
+
+
+class InMemoryA2AResumeStateStore(A2AResumeStateStore):
+    """Local-development store; it intentionally does not survive a restart."""
+
+    def __init__(self) -> None:
+        self._states: dict[str, A2AResumeState] = {}
+
+    async def initialize(self) -> None:
+        return None
+
+    async def put(self, key: str, state: A2AResumeState) -> None:
+        self._states[key] = state
+
+    async def get(self, key: str) -> A2AResumeState | None:
+        return self._states.get(key)
+
+    async def delete(self, key: str) -> None:
+        self._states.pop(key, None)
+
+
+class SQLiteA2AResumeStateStore(A2AResumeStateStore):
+    """Crash-safe Runtime-local store for single-replica/persistent-volume deployments."""
+
+    def __init__(self, path: str | os.PathLike[str]) -> None:
+        self._path = Path(path)
+        self._initialized = False
+        self._init_lock = asyncio.Lock()
+
+    @property
+    def path(self) -> Path:
+        return self._path
+
+    async def initialize(self) -> None:
+        if self._initialized:
+            return
+        async with self._init_lock:
+            if self._initialized:
+                return
+            await asyncio.to_thread(self._initialize_sync)
+            self._initialized = True
+
+    async def put(self, key: str, state: A2AResumeState) -> None:
+        await self.initialize()
+        payload = json.dumps(
+            {
+                "handle": state.handle.model_dump(mode="json"),
+                "target": state.target.model_dump(mode="json"),
+                "payload_kind": state.payload_kind,
+                "call_id": state.call_id,
+            },
+            separators=(",", ":"),
+            sort_keys=True,
+        )
+        await asyncio.to_thread(self._put_sync, key, payload)
+
+    async def get(self, key: str) -> A2AResumeState | None:
+        await self.initialize()
+        payload = await asyncio.to_thread(self._get_sync, key)
+        if payload is None:
+            return None
+        try:
+            raw = json.loads(payload)
+            return A2AResumeState(
+                handle=RunHandle.model_validate(raw["handle"]),
+                target=ResumeTarget.model_validate(raw["target"]),
+                payload_kind=ResumePayload.model_validate(
+                    {"kind": raw["payload_kind"]}
+                ).kind,
+                call_id=str(raw["call_id"]) if raw.get("call_id") is not None else None,
+            )
+        except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
+            raise RuntimeError("A2A resume state is invalid") from exc
+
+    async def delete(self, key: str) -> None:
+        await self.initialize()
+        await asyncio.to_thread(self._delete_sync, key)
+
+    def _initialize_sync(self) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        connection = sqlite3.connect(self._path, timeout=5)
+        try:
+            connection.execute("PRAGMA journal_mode=WAL")
+            connection.execute("PRAGMA synchronous=FULL")
+            connection.execute(
+                """
+                CREATE TABLE IF NOT EXISTS ksadk_a2a_resume_states (
+                    resume_key_sha256 TEXT PRIMARY KEY,
+                    state_json TEXT NOT NULL,
+                    updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
+                )
+                """
+            )
+            connection.commit()
+        finally:
+            connection.close()
+        os.chmod(self._path, 0o600)
+
+    @staticmethod
+    def _hash(key: str) -> str:
+        return hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+    def _put_sync(self, key: str, payload: str) -> None:
+        connection = sqlite3.connect(self._path, timeout=5)
+        try:
+            connection.execute(
+                """
+                INSERT INTO ksadk_a2a_resume_states(resume_key_sha256, state_json)
+                VALUES (?, ?)
+                ON CONFLICT(resume_key_sha256) DO UPDATE SET
+                    state_json = excluded.state_json,
+                    updated_at = CURRENT_TIMESTAMP
+                """,
+                (self._hash(key), payload),
+            )
+            connection.commit()
+        finally:
+            connection.close()
+
+    def _get_sync(self, key: str) -> str | None:
+        connection = sqlite3.connect(self._path, timeout=5)
+        try:
+            row = connection.execute(
+                "SELECT state_json FROM ksadk_a2a_resume_states WHERE resume_key_sha256 = ?",
+                (self._hash(key),),
+            ).fetchone()
+            return str(row[0]) if row is not None else None
+        finally:
+            connection.close()
+
+    def _delete_sync(self, key: str) -> None:
+        connection = sqlite3.connect(self._path, timeout=5)
+        try:
+            connection.execute(
+                "DELETE FROM ksadk_a2a_resume_states WHERE resume_key_sha256 = ?",
+                (self._hash(key),),
+            )
+            connection.commit()
+        finally:
+            connection.close()
+
+
+__all__ = [
+    "A2AResumeState",
+    "A2AResumePayloadKind",
+    "A2AResumeStateStore",
+    "InMemoryA2AResumeStateStore",
+    "SQLiteA2AResumeStateStore",
+]

ksadk/a2a/routes.py

@@ -0,0 +1,105 @@
+"""A2A 协议路由装配 — create_runtime_app 的 A2A seam (goal-05 §8)。
+
+契约 §8:
+
+```python
+def create_runtime_app(config):
+    app = create_base_data_plane_app(config)
+    if config.a2a.enabled:
+        add_a2a_protocol_routes(app, config.a2a)
+    return app
+```
+
+A2A route、TaskStore、task adapter 和 card 构造**只有一份实现**,普通 runtime app
+与 HarnessApp 共用;A2A 是数据面 route group,不承载注册副作用。``A2AConfig``
+保留给本地开发和协议一致性测试；AgentEngine 产品 Runtime 必须通过
+``AgentEngineA2ABootstrap`` 装配 durable storage、trusted ingress 和 egress guard。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Optional, Sequence
+
+from a2a.server.routes import add_a2a_routes_to_fastapi
+from a2a.server.tasks import TaskStore
+from a2a.types import AgentSkill
+from fastapi import FastAPI
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from ksadk.a2a.card import build_agent_card
+from ksadk.a2a.server import A2AProtocolServer
+from ksadk.a2a.task_adapter import A2ARuntimeTaskAdapter
+from ksadk.a2a.task_store import A2A_TASK_TABLE, A2AOwnerContextBuilder, build_a2a_task_store
+
+
+@dataclass
+class A2AConfig:
+    """本地/一致性测试使用的 A2A 协议装配配置(数据面 route group)。"""
+
+    enabled: bool = False
+    base_url: str = "http://127.0.0.1:8000"
+    agent_name: str = "agent"
+    description: str = ""
+    version: str = "1.0.0"
+    skills: Sequence[str | AgentSkill] = field(default_factory=tuple)
+    streaming: bool = True
+    prefer_stream: bool = True
+    # 本地调试可显式开放 reasoning artifact；托管 Runtime 默认关闭，避免把
+    # 模型内部推理写入公开 A2A Task。
+    include_reasoning: bool = False
+    # durable task store:dsn(如 sqlite+aiosqlite:///.agentengine/a2a_tasks.db 或
+    # postgresql+asyncpg://...)或外部传入 engine/task_store。
+    task_store_dsn: Optional[str] = None
+    task_table: str = A2A_TASK_TABLE
+    create_table: bool = True
+
+
+def add_a2a_protocol_routes(
+    app: FastAPI,
+    runner: Any,
+    config: A2AConfig,
+    *,
+    task_adapter: A2ARuntimeTaskAdapter,
+    task_store: Optional[TaskStore] = None,
+    engine: Optional[AsyncEngine] = None,
+    context_builder: A2AOwnerContextBuilder | None = None,
+) -> A2AProtocolServer:
+    """把 A2A 协议路由(JSONRPC + HTTP+JSON + AgentCard)装配到 app。
+
+    返回构建好的 :class:`A2AProtocolServer`(便于调用方取 request_handler /
+    task_store 做后续 subscribe/cancel 或测试)。
+    """
+    card = build_agent_card(
+        name=config.agent_name,
+        base_url=config.base_url,
+        description=config.description,
+        version=config.version,
+        skills=config.skills,
+        streaming=config.streaming,
+    )
+    store = task_store or build_a2a_task_store(
+        dsn=config.task_store_dsn,
+        engine=engine,
+        table_name=config.task_table,
+        create_table=config.create_table,
+    )
+    server = A2AProtocolServer(
+        runner,
+        agent_card=card,
+        task_store=store,
+        task_adapter=task_adapter,
+        context_builder=context_builder,
+        prefer_stream=config.prefer_stream,
+        include_reasoning=config.include_reasoning,
+    )
+    add_a2a_routes_to_fastapi(
+        app,
+        agent_card_routes=server.agent_card_routes(),
+        jsonrpc_routes=server.jsonrpc_routes(),
+        rest_routes=server.rest_routes(),
+    )
+    return server
+
+
+__all__ = ["A2AConfig", "add_a2a_protocol_routes"]

ksadk/a2a/server.py

@@ -1,79 +1,126 @@
-"""High-level server helpers for exposing a KsADK runner over A2A."""
+"""A2AProtocolServer — 托管 Agent 的 A2A 协议数据面 (goal-05)。
+
+契约 §3.2:``A2AProtocolServer`` = latest SDK route factories + AgentCard +
+request handler + durable task store。装配进 create_runtime_app(§8),
+**不再是独立 Starlette app**(旧 demo 的 ``A2AStarletteApplication`` 已废弃)。
+
+路由由 ``ksadk.a2a.routes.add_a2a_protocol_routes`` 统一装配,一份实现,
+普通 runtime app 与 HarnessApp 共用。
+"""
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import Any
+from typing import Any, cast
 
-from a2a.server.apps import A2AStarletteApplication
+from a2a.server.agent_execution import RequestContext, SimpleRequestContextBuilder
+from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers import DefaultRequestHandler
-from a2a.server.tasks import InMemoryTaskStore
-from starlette.applications import Starlette
+from a2a.server.routes import (
+    create_agent_card_routes,
+    create_jsonrpc_routes,
+    create_rest_routes,
+)
+from a2a.server.tasks import TaskStore
+from a2a.types import AgentCard, SendMessageRequest, Task
+
+from ksadk.a2a.card import JSONRPC_PATH, REST_PATH_PREFIX
+from ksadk.a2a.executor import A2ARuntimeExecutor
+from ksadk.a2a.task_adapter import A2ARuntimeTaskAdapter
+from ksadk.a2a.task_store import A2AOwnerContextBuilder
+
 
-from ksadk.a2a.card_builder import AgentCardBuilder
-from ksadk.a2a.executor import KsAgentExecutor
+class _DurableRequestContextBuilder(SimpleRequestContextBuilder):
+    """Recover the persisted context when a follow-up only carries task_id."""
+
+    def __init__(self, task_store: TaskStore) -> None:
+        super().__init__(should_populate_referred_tasks=False, task_store=task_store)
+        self._durable_task_store = task_store
+
+    async def build(
+        self,
+        context: ServerCallContext,
+        params: SendMessageRequest | None = None,
+        task_id: str | None = None,
+        context_id: str | None = None,
+        task: Task | None = None,
+    ) -> RequestContext:
+        if task_id and (task is None or not context_id):
+            persisted = await self._durable_task_store.get(task_id, context)
+            if persisted is not None:
+                task = persisted
+                context_id = context_id or persisted.context_id
+        return await super().build(
+            context=context,
+            params=params,
+            task_id=task_id,
+            context_id=context_id,
+            task=task,
+        )
 
 
-class KsA2AServer:
-    """Expose a KsADK runner as an A2A-compatible ASGI app."""
+class A2AProtocolServer:
+    """把一个 ksadk runner 暴露为 A2A 协议数据面。
+
+    参数:
+        runner: 当前 runtime 的 runner。
+        agent_card: 符合 wire 1.0 的 AgentCard(``ksadk.a2a.card.build_agent_card``)。
+        task_store: durable ``DatabaseTaskStore``(``ksadk.a2a.task_store``)。
+        task_adapter: 可选 ``A2ARuntimeTaskAdapter``(提供则 cancel 走 RuntimeAdapter.cancel)。
+        prefer_stream: 是否优先用 runner.stream(默认 True)。
+        include_reasoning: 是否把 reasoning 输出为 ``adk_thought`` artifact。
+    """
 
     def __init__(
         self,
         runner: Any,
-        app_name: str,
-        url: str,
-        description: str = "",
-        skills: Sequence[str] | None = None,
-        version: str = "1.0.0",
+        *,
+        agent_card: AgentCard,
+        task_store: TaskStore,
+        task_adapter: A2ARuntimeTaskAdapter,
+        context_builder: A2AOwnerContextBuilder | None = None,
+        prefer_stream: bool = True,
+        include_reasoning: bool = False,
     ) -> None:
         self.runner = runner
-        self.app_name = app_name
-        self.url = url
-        self.description = description
-        self.skills = list(skills or [])
-        self.version = version
-
-        self.agent_card = AgentCardBuilder(
-            name=app_name,
-            url=url,
-            description=description,
-            skills=self.skills,
-            version=version,
-        ).build()
-        self.executor = KsAgentExecutor(runner=runner)
-        self.task_store = InMemoryTaskStore()
+        self.agent_card = agent_card
+        self.task_store = task_store
+        self.task_adapter = task_adapter
+        self.context_builder = context_builder or A2AOwnerContextBuilder()
+        self.executor = A2ARuntimeExecutor(
+            runner,
+            task_adapter=task_adapter,
+            prefer_stream=prefer_stream,
+            include_reasoning=include_reasoning,
+        )
         self.request_handler = DefaultRequestHandler(
             agent_executor=self.executor,
-            task_store=self.task_store,
+            task_store=task_store,
+            agent_card=agent_card,
+            request_context_builder=_DurableRequestContextBuilder(task_store),
         )
-        self.application = A2AStarletteApplication(
-            agent_card=self.agent_card,
-            http_handler=self.request_handler,
+
+    def agent_card_routes(self) -> list[Any]:
+        return cast(list[Any], create_agent_card_routes(self.agent_card))
+
+    def jsonrpc_routes(self) -> list[Any]:
+        return cast(
+            list[Any],
+            create_jsonrpc_routes(
+                self.request_handler,
+                rpc_url=JSONRPC_PATH,
+                context_builder=self.context_builder,
+            ),
         )
 
-    def build(self, **kwargs: Any) -> Starlette:
-        """Build the Starlette application with the standard A2A routes."""
-        return self.application.build(**kwargs)
-
-    def add_routes_to_app(self, app: Starlette) -> None:
-        """Attach the A2A routes to an existing Starlette application."""
-        self.application.add_routes_to_app(app)
-
-
-def to_a2a(
-    runner: Any,
-    app_name: str,
-    url: str,
-    description: str = "",
-    skills: Sequence[str] | None = None,
-    version: str = "1.0.0",
-) -> KsA2AServer:
-    """Convenience helper for building a KsA2AServer."""
-    return KsA2AServer(
-        runner=runner,
-        app_name=app_name,
-        url=url,
-        description=description,
-        skills=skills,
-        version=version,
-    )
+    def rest_routes(self) -> list[Any]:
+        return cast(
+            list[Any],
+            create_rest_routes(
+                self.request_handler,
+                path_prefix=REST_PATH_PREFIX,
+                context_builder=self.context_builder,
+            ),
+        )
+
+
+__all__ = ["A2AProtocolServer"]