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

ksadk/a2a/task_event_outbox.py

@@ -0,0 +1,326 @@
+"""Durable local outbox for AgentEngine A2A task event batches."""
+
+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 Any
+
+ENV_A2A_EVENT_OUTBOX_PATH = "KSADK_A2A_EVENT_OUTBOX_PATH"
+DEFAULT_A2A_EVENT_OUTBOX_PATH = ".agentengine/a2a_event_outbox.sqlite3"
+
+
+@dataclass(frozen=True)
+class A2ATaskEventBatch:
+    sequence: int
+    batch_id: str
+    platform_task_id: str
+    events: list[dict[str, Any]]
+    attempt_count: int = 0
+
+
+class A2ATaskEventOutbox(ABC):
+    """Persists task-sink batches until AgentEngine acknowledges them."""
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def ensure_writable(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def enqueue(
+        self,
+        *,
+        platform_task_id: str,
+        events: list[dict[str, Any]],
+    ) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def pending(self, *, limit: int = 100) -> list[A2ATaskEventBatch]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def acknowledge(self, batch_id: str) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def record_failure(self, batch_id: str, error: str) -> None:
+        raise NotImplementedError
+
+
+def _serialize_events(events: list[dict[str, Any]]) -> str:
+    return json.dumps(events, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+
+
+def _batch_id(platform_task_id: str, events_json: str) -> str:
+    digest = hashlib.sha256(f"{platform_task_id}:{events_json}".encode("utf-8")).hexdigest()
+    return f"a2a-event-batch-{digest}"
+
+
+class InMemoryA2ATaskEventOutbox(A2ATaskEventOutbox):
+    """Test/local adapter; production Runtime deployments should use SQLite."""
+
+    def __init__(self) -> None:
+        self._batches: dict[str, A2ATaskEventBatch] = {}
+        self._lock = asyncio.Lock()
+        self._next_sequence = 0
+
+    async def initialize(self) -> None:
+        return None
+
+    async def ensure_writable(self) -> None:
+        return None
+
+    async def enqueue(
+        self,
+        *,
+        platform_task_id: str,
+        events: list[dict[str, Any]],
+    ) -> str:
+        events_json = _serialize_events(events)
+        batch_id = _batch_id(platform_task_id, events_json)
+        async with self._lock:
+            if batch_id not in self._batches:
+                self._next_sequence += 1
+                self._batches[batch_id] = A2ATaskEventBatch(
+                    sequence=self._next_sequence,
+                    batch_id=batch_id,
+                    platform_task_id=platform_task_id,
+                    events=json.loads(events_json),
+                )
+        return batch_id
+
+    async def pending(self, *, limit: int = 100) -> list[A2ATaskEventBatch]:
+        async with self._lock:
+            return sorted(self._batches.values(), key=lambda batch: batch.sequence)[:limit]
+
+    async def acknowledge(self, batch_id: str) -> None:
+        async with self._lock:
+            self._batches.pop(batch_id, None)
+
+    async def record_failure(self, batch_id: str, error: str) -> None:
+        async with self._lock:
+            batch = self._batches.get(batch_id)
+            if batch is not None:
+                self._batches[batch_id] = A2ATaskEventBatch(
+                    sequence=batch.sequence,
+                    batch_id=batch.batch_id,
+                    platform_task_id=batch.platform_task_id,
+                    events=batch.events,
+                    attempt_count=batch.attempt_count + 1,
+                )
+
+
+class SQLiteA2ATaskEventOutbox(A2ATaskEventOutbox):
+    """SQLite adapter with stable batch IDs and crash-safe acknowledgement."""
+
+    def __init__(self, path: str | os.PathLike[str] | None = None) -> None:
+        configured = path or os.getenv(ENV_A2A_EVENT_OUTBOX_PATH) or DEFAULT_A2A_EVENT_OUTBOX_PATH
+        self._path = Path(configured).expanduser()
+        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
+
+    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 a2a_task_event_outbox (
+                    sequence INTEGER NOT NULL,
+                    batch_id TEXT PRIMARY KEY,
+                    platform_task_id TEXT NOT NULL,
+                    events_json TEXT NOT NULL,
+                    created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                    attempt_count INTEGER NOT NULL DEFAULT 0,
+                    last_error TEXT
+                )
+                """
+            )
+            columns = {
+                str(row[1])
+                for row in connection.execute("PRAGMA table_info(a2a_task_event_outbox)").fetchall()
+            }
+            if "sequence" not in columns:
+                connection.execute("ALTER TABLE a2a_task_event_outbox ADD COLUMN sequence INTEGER")
+                connection.execute(
+                    "UPDATE a2a_task_event_outbox SET sequence = rowid WHERE sequence IS NULL"
+                )
+            connection.execute(
+                """
+                CREATE TABLE IF NOT EXISTS a2a_task_event_outbox_sequence (
+                    sequence INTEGER PRIMARY KEY AUTOINCREMENT
+                )
+                """
+            )
+            connection.execute(
+                """
+                INSERT OR IGNORE INTO a2a_task_event_outbox_sequence(sequence)
+                SELECT sequence FROM a2a_task_event_outbox
+                """
+            )
+            connection.execute(
+                """
+                CREATE INDEX IF NOT EXISTS ix_a2a_task_event_outbox_created
+                ON a2a_task_event_outbox(sequence)
+                """
+            )
+            connection.commit()
+        finally:
+            connection.close()
+        os.chmod(self._path, 0o600)
+
+    async def enqueue(
+        self,
+        *,
+        platform_task_id: str,
+        events: list[dict[str, Any]],
+    ) -> str:
+        await self.initialize()
+        events_json = _serialize_events(events)
+        batch_id = _batch_id(platform_task_id, events_json)
+        await asyncio.to_thread(
+            self._enqueue_sync,
+            batch_id,
+            platform_task_id,
+            events_json,
+        )
+        return batch_id
+
+    async def ensure_writable(self) -> None:
+        await self.initialize()
+        await asyncio.to_thread(self._ensure_writable_sync)
+
+    async def pending(self, *, limit: int = 100) -> list[A2ATaskEventBatch]:
+        await self.initialize()
+        rows = await asyncio.to_thread(self._pending_sync, limit)
+        return [
+            A2ATaskEventBatch(
+                sequence=int(row[0]),
+                batch_id=str(row[1]),
+                platform_task_id=str(row[2]),
+                events=list(json.loads(str(row[3]))),
+                attempt_count=int(row[4]),
+            )
+            for row in rows
+        ]
+
+    async def acknowledge(self, batch_id: str) -> None:
+        await self.initialize()
+        await asyncio.to_thread(
+            self._execute,
+            "DELETE FROM a2a_task_event_outbox WHERE batch_id = ?",
+            (batch_id,),
+        )
+
+    async def record_failure(self, batch_id: str, error: str) -> None:
+        await self.initialize()
+        await asyncio.to_thread(
+            self._execute,
+            """
+            UPDATE a2a_task_event_outbox
+            SET attempt_count = attempt_count + 1, last_error = ?
+            WHERE batch_id = ?
+            """,
+            (error[:512], batch_id),
+        )
+
+    def _execute(self, statement: str, params: tuple[Any, ...]) -> None:
+        connection = sqlite3.connect(self._path, timeout=5)
+        try:
+            connection.execute("PRAGMA synchronous=FULL")
+            connection.execute(statement, params)
+            connection.commit()
+        finally:
+            connection.close()
+
+    def _enqueue_sync(self, batch_id: str, platform_task_id: str, events_json: str) -> None:
+        connection = sqlite3.connect(self._path, timeout=5, isolation_level=None)
+        try:
+            connection.execute("PRAGMA synchronous=FULL")
+            connection.execute("BEGIN IMMEDIATE")
+            existing = connection.execute(
+                "SELECT 1 FROM a2a_task_event_outbox WHERE batch_id = ?",
+                (batch_id,),
+            ).fetchone()
+            if existing is None:
+                cursor = connection.execute(
+                    "INSERT INTO a2a_task_event_outbox_sequence DEFAULT VALUES"
+                )
+                if cursor.lastrowid is None:
+                    raise RuntimeError("SQLite did not return an outbox sequence")
+                sequence = int(cursor.lastrowid)
+                connection.execute(
+                    """
+                    INSERT INTO a2a_task_event_outbox(
+                        sequence, batch_id, platform_task_id, events_json
+                    ) VALUES (?, ?, ?, ?)
+                    """,
+                    (sequence, batch_id, platform_task_id, events_json),
+                )
+            connection.execute("COMMIT")
+        except BaseException:
+            connection.rollback()
+            raise
+        finally:
+            connection.close()
+
+    def _ensure_writable_sync(self) -> None:
+        connection = sqlite3.connect(self._path, timeout=5, isolation_level=None)
+        try:
+            connection.execute("BEGIN IMMEDIATE")
+            connection.execute("ROLLBACK")
+        finally:
+            connection.close()
+
+    def _pending_sync(self, limit: int) -> list[tuple[Any, ...]]:
+        connection = sqlite3.connect(self._path, timeout=5)
+        try:
+            return list(
+                connection.execute(
+                    """
+                    SELECT sequence, batch_id, platform_task_id, events_json, attempt_count
+                    FROM a2a_task_event_outbox
+                    ORDER BY sequence
+                    LIMIT ?
+                    """,
+                    (limit,),
+                ).fetchall()
+            )
+        finally:
+            connection.close()
+
+
+__all__ = [
+    "A2ATaskEventBatch",
+    "A2ATaskEventOutbox",
+    "DEFAULT_A2A_EVENT_OUTBOX_PATH",
+    "ENV_A2A_EVENT_OUTBOX_PATH",
+    "InMemoryA2ATaskEventOutbox",
+    "SQLiteA2ATaskEventOutbox",
+]

ksadk/a2a/task_store.py

@@ -0,0 +1,206 @@
+"""A2A durable TaskStore (goal-05)。
+
+契约(§4.6):托管被调 task 使用 SDK ``DatabaseTaskStore``,表名隔离为
+``ksadk_a2a_tasks``,owner resolver 必须包含 account/runtime identity。
+**不用 InMemoryTaskStore**——7/31 重启恢复发布门禁依赖 durable store。
+
+注意:TaskStore 只持久化**协议 Task**,不等于 runner checkpoint;两者必须共同成功
+才能声明 durable async(§7.2)。SQLite 只用于快速单测;进程崩溃恢复门禁必须使用
+真实 PostgreSQL 和独立 OS 进程。
+"""
+
+from __future__ import annotations
+
+import warnings
+from threading import Lock
+from typing import Any, Callable, Mapping, Optional, Sequence
+
+from a2a.server import models as _sdk_models
+from a2a.server.context import ServerCallContext
+from a2a.server.routes.common import DefaultServerCallContextBuilder
+from a2a.server.tasks import DatabaseTaskStore
+from sqlalchemy.exc import SAWarning
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from starlette.requests import Request
+
+from ksadk.a2a.identity import A2AIngressIdentity, A2ATrustedIdentityResolver
+
+#: 契约 §4.6 规定的隔离表名。
+A2A_TASK_TABLE = "ksadk_a2a_tasks"
+
+DEFAULT_ACCOUNT_HEADERS = ("x-ksc-account-id", "x-account-id")
+DEFAULT_RUNTIME_HEADERS = ("x-auth-agent-id", "x-ksc-agent-id", "x-runtime-id")
+
+
+# a2a-sdk 1.1.0 dynamically declares a new ORM model for every custom-table
+# store. Repeating the same table in one process raises an SQLAlchemy duplicate
+# table error. Cache only KsADK's model instead of monkeypatching the SDK module.
+_task_model_cache: dict[str, type[Any]] = {}
+_task_model_lock = Lock()
+
+
+def _task_model(table_name: str) -> type[Any]:
+    with _task_model_lock:
+        model = _task_model_cache.get(table_name)
+        if model is None:
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    message=(
+                        "This declarative base already contains a class with the same class name"
+                    ),
+                    category=SAWarning,
+                )
+                model = _sdk_models.create_task_model(table_name)
+            _task_model_cache[table_name] = model
+        return model
+
+
+class _KsADKDatabaseTaskStore(DatabaseTaskStore):
+    """DatabaseTaskStore with a process-local, non-global custom-table model."""
+
+    def __init__(
+        self,
+        *,
+        engine: AsyncEngine,
+        create_table: bool,
+        table_name: str,
+        owner_resolver: Callable[[ServerCallContext], str],
+    ) -> None:
+        super().__init__(
+            engine=engine,
+            create_table=create_table,
+            owner_resolver=owner_resolver,
+        )
+        self.task_model = _task_model(table_name)
+
+
+def _first(state: dict[str, Any], *keys: str) -> Optional[str]:
+    for key in keys:
+        value = state.get(key)
+        if value:
+            return str(value)
+    return None
+
+
+def default_owner_resolver(context: ServerCallContext) -> str:
+    """默认 owner resolver:组合 account + runtime identity(契约 §4.6 要求两者)。
+
+    owner 形如 ``"{account}/{runtime}"``:account 取 ``account_id``/``tenant_id``,
+    runtime 取 ``runtime_id``/``agent_id``。**同一 account 下不同 runtime 形成独立任务
+    权限域**(不会共用);只凑齐一侧时用一侧;两侧都缺时退回 ``user``/``anonymous``。
+    生产环境应由鉴权/STS 在 context.state 注入这些键;本地无鉴权时退回匿名。
+    """
+    state = getattr(context, "state", None) or {}
+    verified_identity = state.get(A2ATrustedIdentityResolver.state_key)
+    if isinstance(verified_identity, A2AIngressIdentity):
+        return verified_identity.owner_key()
+    account = _first(state, "account_id", "tenant_id")
+    runtime = _first(state, "runtime_id", "agent_id")
+    parts = [part for part in (account, runtime) if part]
+    if parts:
+        return "/".join(parts)
+    return _first(state, "user") or "anonymous"
+
+
+class A2AOwnerContextBuilder(DefaultServerCallContextBuilder):
+    """Build A2A call contexts from platform-authenticated HTTP identity.
+
+    The SDK default stores headers only under ``state['headers']`` while the
+    task store intentionally reads normalized top-level identity fields.  This
+    builder is shared by JSON-RPC and REST so both transports enforce the same
+    account/runtime owner boundary.  Deployments must strip/replace these
+    platform headers at the public edge; message metadata is never trusted.
+    """
+
+    def __init__(
+        self,
+        *,
+        account_headers: Sequence[str] = DEFAULT_ACCOUNT_HEADERS,
+        runtime_headers: Sequence[str] = DEFAULT_RUNTIME_HEADERS,
+        identity_resolver: A2ATrustedIdentityResolver | None = None,
+        allow_unverified_identity: bool = True,
+    ) -> None:
+        self._account_headers = tuple(header.lower() for header in account_headers)
+        self._runtime_headers = tuple(header.lower() for header in runtime_headers)
+        self._identity_resolver = identity_resolver
+        self._allow_unverified_identity = allow_unverified_identity
+
+    @staticmethod
+    def _header(headers: Mapping[str, Any], names: Sequence[str]) -> str | None:
+        for name in names:
+            value = str(headers.get(name) or "").strip()
+            if value:
+                return value
+        return None
+
+    def build(self, request: Request) -> ServerCallContext:
+        context = super().build(request)
+        state = dict(context.state or {})
+        headers = {str(key).lower(): value for key, value in request.headers.items()}
+        scope_state = request.scope.get("state")
+        trusted_state = dict(scope_state) if isinstance(scope_state, Mapping) else {}
+        if self._identity_resolver is not None:
+            identity = self._identity_resolver.resolve(request)
+            state[A2ATrustedIdentityResolver.state_key] = identity
+            state["account_id"] = identity.account_id
+            state["tenant_id"] = identity.tenant_id
+            state["caller_principal_type"] = identity.caller_principal_type
+            state["caller_principal_id"] = identity.caller_principal_id
+        elif not self._allow_unverified_identity:
+            raise PermissionError("verified Gateway identity is required for inbound A2A")
+        else:
+            account = _first(trusted_state, "account_id", "tenant_id") or self._header(
+                headers, self._account_headers
+            )
+            runtime = _first(trusted_state, "runtime_id", "agent_id") or self._header(
+                headers, self._runtime_headers
+            )
+            if account:
+                state["account_id"] = account
+            if runtime:
+                state["runtime_id"] = runtime
+        context.state = state
+        context.tenant = default_owner_resolver(context)
+        return context
+
+
+def build_a2a_task_store(
+    dsn: str | None = None,
+    *,
+    engine: Optional[AsyncEngine] = None,
+    table_name: str = A2A_TASK_TABLE,
+    owner_resolver: Optional[Callable[[ServerCallContext], str]] = None,
+    create_table: bool = True,
+) -> DatabaseTaskStore:
+    """构建 durable ``DatabaseTaskStore``。
+
+    参数:
+        dsn: 数据库 DSN(如 ``sqlite+aiosqlite:///path.db`` 或
+            ``postgresql+asyncpg://...``)。与 ``engine`` 二选一。
+        engine: 已存在的 SQLAlchemy AsyncEngine(优先于 dsn)。
+        table_name: 隔离表名,默认 ``ksadk_a2a_tasks``。
+        owner_resolver: owner 解析器,默认 :func:`default_owner_resolver`。
+        create_table: 是否在 initialize 时建表(测试/本地默认 True;生产可 False
+            由迁移管理)。
+    """
+    if engine is None:
+        if not dsn:
+            raise ValueError("build_a2a_task_store 需要 dsn 或 engine 之一")
+        engine = create_async_engine(dsn)
+    return _KsADKDatabaseTaskStore(
+        engine=engine,
+        create_table=create_table,
+        table_name=table_name,
+        owner_resolver=owner_resolver or default_owner_resolver,
+    )
+
+
+__all__ = [
+    "A2A_TASK_TABLE",
+    "A2AOwnerContextBuilder",
+    "DEFAULT_ACCOUNT_HEADERS",
+    "DEFAULT_RUNTIME_HEADERS",
+    "build_a2a_task_store",
+    "default_owner_resolver",
+]

ksadk/a2ui/__init__.py

@@ -0,0 +1,27 @@
+"""A2UI(Agent 驱动 UI)— canonical:docs/A2UI-agent驱动UI技术方案.md。"""
+
+from ksadk.a2ui.core import A2UICore
+from ksadk.a2ui.models import (
+    BASIC_CATALOG,
+    BASIC_CATALOG_ID,
+    A2UIValidationError,
+    ActionReceipt,
+    Component,
+    PendingInteraction,
+    Surface,
+)
+from ksadk.a2ui.renderer import A2UIRenderer, RenderedNode, submit_a2ui_action
+
+__all__ = [
+    "A2UICore",
+    "A2UIRenderer",
+    "RenderedNode",
+    "submit_a2ui_action",
+    "A2UIValidationError",
+    "ActionReceipt",
+    "BASIC_CATALOG",
+    "BASIC_CATALOG_ID",
+    "Component",
+    "PendingInteraction",
+    "Surface",
+]