agentengine-sdk-python 0.8.5__py3-none-any.whl → 0.8.6__py3-none-any.whl

ksadk/a2a/task_event_outbox_owned.py

@@ -0,0 +1,130 @@
+"""Runtime-owned PostgreSQL event delivery, isolated from legacy consumers."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+
+from sqlalchemy import text
+
+from ksadk.a2a.task_event_outbox import (
+    A2ATaskEventBatch,
+    A2ATaskEventOutbox,
+    _batch_id,
+    _serialize_events,
+)
+
+
+class OwnedPostgresA2ATaskEventOutbox(A2ATaskEventOutbox):
+    """Replicas share their runtime's queue, never another workload's queue.
+
+    Separate tables are intentional: an older unscoped dispatcher must not
+    consume new owned batches during a rolling upgrade. Legacy rows are kept
+    untouched; their owner cannot safely be inferred from their event payload.
+    """
+
+    def __init__(self, engine, *, account_id: str, tenant_id: str, runtime_id: str):
+        identity = (account_id, tenant_id, runtime_id)
+        if any(not isinstance(value, str) or not value.strip() for value in identity):
+            raise ValueError("A2A_OUTBOX_OWNER_REQUIRED")
+        self._engine = engine
+        self._owner = hashlib.sha256(
+            json.dumps(identity, separators=(",", ":")).encode()
+        ).hexdigest()
+
+    async def initialize(self):
+        async with self._engine.begin() as connection:
+            await connection.execute(text("""
+                CREATE TABLE IF NOT EXISTS ksadk_a2a_owned_task_event_outbox (
+                    sequence BIGSERIAL PRIMARY KEY,
+                    owner_key TEXT NOT NULL,
+                    batch_id TEXT NOT NULL UNIQUE,
+                    platform_task_id TEXT NOT NULL,
+                    events_json TEXT NOT NULL,
+                    created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                    attempt_count INTEGER NOT NULL DEFAULT 0,
+                    last_error TEXT
+                )
+            """))
+            await connection.execute(text("""
+                CREATE TABLE IF NOT EXISTS ksadk_a2a_owned_task_event_dead_letters (
+                    sequence BIGINT NOT NULL,
+                    owner_key TEXT NOT NULL,
+                    batch_id TEXT PRIMARY KEY,
+                    platform_task_id TEXT NOT NULL,
+                    events_json TEXT NOT NULL,
+                    created_at TIMESTAMPTZ NOT NULL,
+                    attempt_count INTEGER NOT NULL,
+                    last_error TEXT NOT NULL,
+                    dead_lettered_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+                )
+            """))
+            await connection.execute(text("""
+                CREATE INDEX IF NOT EXISTS ix_ksadk_a2a_owned_outbox_owner_sequence
+                ON ksadk_a2a_owned_task_event_outbox(owner_key, sequence)
+            """))
+
+    async def ensure_writable(self):
+        await self.initialize()
+        async with self._engine.begin() as connection:
+            await connection.execute(text(
+                "LOCK TABLE ksadk_a2a_owned_task_event_outbox IN ROW EXCLUSIVE MODE"
+            ))
+
+    async def enqueue(self, *, platform_task_id, events):
+        payload = _serialize_events(events)
+        batch_id = _batch_id(self._owner + ":" + platform_task_id, payload)
+        async with self._engine.begin() as connection:
+            await connection.execute(text("""
+                INSERT INTO ksadk_a2a_owned_task_event_outbox(
+                    owner_key, batch_id, platform_task_id, events_json
+                ) VALUES (:owner, :batch, :task, :events)
+                ON CONFLICT(batch_id) DO NOTHING
+            """), {"owner": self._owner, "batch": batch_id,
+                   "task": platform_task_id, "events": payload})
+        return batch_id
+
+    async def pending(self, *, limit=100):
+        async with self._engine.connect() as connection:
+            result = await connection.execute(text("""
+                SELECT sequence, batch_id, platform_task_id, events_json, attempt_count
+                FROM ksadk_a2a_owned_task_event_outbox
+                WHERE owner_key = :owner ORDER BY sequence LIMIT :limit
+            """), {"owner": self._owner, "limit": limit})
+            return [A2ATaskEventBatch(
+                sequence=int(row[0]), batch_id=str(row[1]), platform_task_id=str(row[2]),
+                events=list(json.loads(row[3])), attempt_count=int(row[4]),
+            ) for row in result]
+
+    async def acknowledge(self, batch_id):
+        async with self._engine.begin() as connection:
+            await connection.execute(text("""
+                DELETE FROM ksadk_a2a_owned_task_event_outbox
+                WHERE owner_key = :owner AND batch_id = :batch
+            """), {"owner": self._owner, "batch": batch_id})
+
+    async def record_failure(self, batch_id, error):
+        async with self._engine.begin() as connection:
+            await connection.execute(text("""
+                UPDATE ksadk_a2a_owned_task_event_outbox
+                SET attempt_count = attempt_count + 1, last_error = :error
+                WHERE owner_key = :owner AND batch_id = :batch
+            """), {"owner": self._owner, "batch": batch_id, "error": error[:512]})
+
+    async def dead_letter(self, batch_id, error):
+        async with self._engine.begin() as connection:
+            params = {"owner": self._owner, "batch": batch_id, "error": error[:512]}
+            await connection.execute(text("""
+                INSERT INTO ksadk_a2a_owned_task_event_dead_letters(
+                    sequence, owner_key, batch_id, platform_task_id, events_json,
+                    created_at, attempt_count, last_error
+                ) SELECT sequence, owner_key, batch_id, platform_task_id, events_json,
+                         created_at, attempt_count, :error
+                  FROM ksadk_a2a_owned_task_event_outbox
+                  WHERE owner_key = :owner AND batch_id = :batch
+                ON CONFLICT(batch_id) DO NOTHING
+            """), params)
+            await connection.execute(text("""
+                DELETE FROM ksadk_a2a_owned_task_event_outbox
+                WHERE owner_key = :owner AND batch_id = :batch
+            """), {"owner": self._owner, "batch": batch_id})

ksadk/a2a/task_store.py

@@ -19,6 +19,8 @@ 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 a2a.types import Task, TaskState, TaskStatus
+from sqlalchemy import update
 from sqlalchemy.exc import SAWarning
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from starlette.requests import Request
@@ -75,6 +77,41 @@ class _KsADKDatabaseTaskStore(DatabaseTaskStore):
         self.task_model = _task_model(table_name)
 
 
+    async def save_confirmed_cancellation(
+        self, task: Task, status: TaskStatus, context: ServerCallContext,
+    ) -> Task | None:
+        """Repair a lost cancel event without replacing a concurrently committed terminal."""
+        if status.state != TaskState.TASK_STATE_CANCELED:
+            raise ValueError("Only confirmed cancellation can use this repair")
+        return await self.save_confirmed_terminal(task, status, context)
+
+    async def save_confirmed_terminal(
+        self, task: Task, status: TaskStatus, context: ServerCallContext,
+    ) -> Task | None:
+        """Apply an authoritative terminal status without replacing a newer terminal."""
+        if status.state not in {TaskState.TASK_STATE_FAILED, TaskState.TASK_STATE_CANCELED,
+                                TaskState.TASK_STATE_COMPLETED}:
+            raise ValueError("Only confirmed terminal states can use this repair")
+        await self._ensure_initialized()
+        owner = self.owner_resolver(context)
+        canceled = Task()
+        canceled.CopyFrom(task)
+        canceled.status.CopyFrom(status)
+        record = self._to_orm(canceled, owner)
+        # Guard inside UPDATE, rather than a separate read/check/write. Update only
+        # status fields so concurrent artifacts/history cannot be replaced by a cache.
+        active = ["TASK_STATE_SUBMITTED", "TASK_STATE_WORKING",
+                  "TASK_STATE_INPUT_REQUIRED", "TASK_STATE_AUTH_REQUIRED"]
+        async with self.async_session_maker.begin() as session:
+            await session.execute(update(self.task_model).where(
+                self.task_model.id == task.id,
+                self.task_model.owner == owner,
+                self.task_model.protocol_version == "1.0",
+                self.task_model.status["state"].as_string().in_(active),
+            ).values(status=record.status, last_updated=record.last_updated))
+        return await self.get(task.id, context)
+
+
 def _first(state: dict[str, Any], *keys: str) -> Optional[str]:
     for key in keys:
         value = state.get(key)

ksadk/a2a/tool_activity.py

@@ -0,0 +1,176 @@
+"""Explicit public tool disclosure at the callee wire boundary."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+from a2a.types import Part
+from google.protobuf.json_format import ParseDict
+from google.protobuf.struct_pb2 import Value
+from pydantic import BaseModel, ConfigDict, Field
+
+from ksadk.a2a.presentation import (
+    A2A_TOOL_ACTIVITY_EXTENSION_URI,
+    A2A_TOOL_ACTIVITY_MAX_EVENT_BYTES,
+    A2A_TOOL_ACTIVITY_MAX_INVOCATION_BYTES,
+    build_public_tool_activity,
+)
+from ksadk.events.canonical import ItemCompleted, ItemSnapshotReplaced, ItemStarted, ItemUpdated
+from ksadk.events.content import ToolCallContent, ToolResultContent
+from ksadk.events.identity import stable_item_id
+
+
+@dataclass(frozen=True)
+class PublicToolRule:
+    """Only named top-level fields cross the boundary; everything else is private."""
+
+    name: str
+    arguments: frozenset[str] = frozenset()
+    results: frozenset[str] = frozenset()
+
+
+@dataclass(frozen=True)
+class PublicToolActivityPolicy:
+    tools: Mapping[str, PublicToolRule] = field(default_factory=dict)
+
+    @classmethod
+    def from_json(cls, payload: str) -> PublicToolActivityPolicy:
+        """Load a platform-owned declaration, never a request/model override."""
+        declaration = _PublicToolDeclaration.model_validate_json(payload)
+        return cls(
+            {
+                name: PublicToolRule(rule.name, frozenset(rule.arguments), frozenset(rule.results))
+                for name, rule in declaration.tools.items()
+            }
+        )
+
+
+class _PublicToolFields(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+    name: str = Field(pattern=r"^[A-Za-z_][A-Za-z0-9_.:-]{0,127}$")
+    arguments: list[str] = Field(default_factory=list, max_length=100)
+    results: list[str] = Field(default_factory=list, max_length=100)
+
+
+class _PublicToolDeclaration(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+    tools: dict[str, _PublicToolFields] = Field(max_length=100)
+
+
+class PublicToolActivityProducer:
+    def __init__(self, policy: PublicToolActivityPolicy) -> None:
+        self.policy = policy
+        self._calls: dict[tuple[str, str], PublicToolRule] = {}
+        self._bytes = 0
+        self._seen: set[str] = set()
+        self._budget_exhausted = False
+
+    async def _publish(
+        self,
+        *,
+        activity_payload: dict[str, Any],
+        public_item: str,
+        public_event: str,
+        updater: Any,
+    ) -> None:
+        if public_event in self._seen or self._budget_exhausted:
+            return
+        activity = build_public_tool_activity(activity_payload)
+        size = len(activity.model_dump_json().encode())
+        # Reserve room for an explicit, sanitized truncation observation.
+        # Never consume the entire invocation budget and silently drop output.
+        exhausted = self._bytes + size > (
+            A2A_TOOL_ACTIVITY_MAX_INVOCATION_BYTES - A2A_TOOL_ACTIVITY_MAX_EVENT_BYTES
+        )
+        if exhausted:
+            activity = build_public_tool_activity(
+                {**activity_payload, "arguments": None, "result": None, "truncated": True}
+            )
+            size = len(activity.model_dump_json().encode())
+        terminal = activity.operation in {"completed", "failed"}
+        await updater.add_artifact(
+            artifact_id=public_item,
+            name="public-tool-activity",
+            parts=[
+                Part(
+                    data=ParseDict(activity.model_dump(mode="json"), Value()),
+                    metadata={"extension_uri": A2A_TOOL_ACTIVITY_EXTENSION_URI},
+                )
+            ],
+            metadata={
+                "event_id": public_event,
+                "extension_uri": A2A_TOOL_ACTIVITY_EXTENSION_URI,
+            },
+            append=False,
+            last_chunk=terminal,
+        )
+        # A failed store/write is retryable with the same public event identity.
+        self._bytes += size
+        self._seen.add(public_event)
+        self._budget_exhausted = exhausted
+
+    async def emit(self, event: Any, updater: Any) -> None:
+        if not isinstance(event, (ItemStarted, ItemUpdated, ItemSnapshotReplaced, ItemCompleted)):
+            return
+        if event.item_kind not in {"tool_call", "tool_result"}:
+            return
+        snapshot = getattr(event, "snapshot", None) or getattr(event, "initial", None)
+        parts = snapshot.parts if snapshot is not None else (getattr(event, "update", None),)
+        for part in parts:
+            if isinstance(part, ToolCallContent):
+                rule = self.policy.tools.get(part.name)
+                if rule is None:
+                    continue
+                self._calls[(event.scope_id, part.call_id)] = rule
+            elif isinstance(part, ToolResultContent):
+                rule = self._calls.get((event.scope_id, part.call_id))
+                if rule is None:
+                    continue
+            else:
+                continue
+            public_item = stable_item_id("a2a-tool", event.scope_id, event.item_id)
+            public_call = stable_item_id("a2a-tool", event.scope_id, part.call_id)
+            is_call = isinstance(part, ToolCallContent)
+            is_error = isinstance(part, ToolResultContent) and part.is_error
+            complete = isinstance(event, ItemCompleted)
+            operation = "failed" if is_error else "completed" if complete else "updated"
+            common = {
+                "item_id": public_item,
+                "call_id": public_call,
+                "item_kind": "tool_call" if is_call else "tool_result",
+                "execution_status": (
+                    "failed" if is_error else "completed" if complete and not is_call else "running"
+                ),
+                "name": rule.name if is_call else None,
+                "arguments": self._select(part.arguments, rule.arguments) if is_call else None,
+                "result": None if is_call else self._select(part.result, rule.results),
+                "is_error": is_error,
+            }
+            if is_call and complete:
+                started_event = stable_item_id("a2a-tool", event.event_id, part.part_id, "started")
+                await self._publish(
+                    activity_payload={
+                        **common,
+                        "event_id": started_event,
+                        "operation": "started",
+                        "execution_status": "running",
+                    },
+                    public_item=public_item,
+                    public_event=started_event,
+                    updater=updater,
+                )
+            public_event = stable_item_id("a2a-tool", event.event_id, part.part_id, operation)
+            await self._publish(
+                activity_payload={**common, "event_id": public_event, "operation": operation},
+                public_item=public_item,
+                public_event=public_event,
+                updater=updater,
+            )
+
+    @staticmethod
+    def _select(value: Any, keys: frozenset[str]) -> dict[str, Any]:
+        if not isinstance(value, Mapping):
+            return {}
+        return {key: value[key] for key in keys if key in value}

ksadk/a2a/workload_tokens.py

@@ -0,0 +1,134 @@
+"""Audience-specific projected workload-token loading and validation."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import errno
+import json
+import os
+import stat
+import time
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+ENV_A2A_TOKEN_DIR = "KSADK_A2A_TOKEN_DIR"
+
+DEFAULT_A2A_TOKEN_DIR = "/var/run/secrets/agentengine/a2a"
+
+MAX_WORKLOAD_TOKEN_BYTES = 16 * 1024
+
+
+class WorkloadTokenProvider(ABC):
+    """Returns a short-lived workload token for one exact audience."""
+
+    @abstractmethod
+    def get_token(self, audience: str) -> str:
+        raise NotImplementedError
+
+
+class FileWorkloadTokenProvider(WorkloadTokenProvider):
+    """Reads audience-specific projected JWT files before every request."""
+
+    def __init__(self, token_dir: str | os.PathLike[str] | None = None) -> None:
+        configured = token_dir or os.getenv(ENV_A2A_TOKEN_DIR) or DEFAULT_A2A_TOKEN_DIR
+        self._token_dir = Path(configured)
+
+    def get_token(self, audience: str) -> str:
+        if not audience or "/" in audience or audience in {".", ".."}:
+            raise ValueError(f"invalid workload token audience: {audience!r}")
+        path = self._token_dir / f"{audience}.jwt"
+        if path.is_symlink():
+            # AtomicWriter rotates ..data atomically. Only accept its exact key
+            # layout, and open the resolved generation without following links.
+            if os.readlink(path) != f"..data/{path.name}":
+                raise RuntimeError(f"invalid workload token symlink for {audience!r}")
+            try:
+                volume = self._token_dir.resolve(strict=True)
+                relative = path.resolve(strict=True).relative_to(volume)
+                flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
+                directory = os.open(volume, flags | os.O_DIRECTORY)
+                try:
+                    for component in relative.parts[:-1]:
+                        child = os.open(component, flags | os.O_DIRECTORY, dir_fd=directory)
+                        os.close(directory)
+                        directory = child
+                    descriptor = os.open(relative.name, flags | os.O_NONBLOCK, dir_fd=directory)
+                finally:
+                    os.close(directory)
+            except (OSError, ValueError, RuntimeError) as exc:
+                raise RuntimeError(f"unsafe workload token projection for {audience!r}") from exc
+        else:
+            descriptor = self._open_regular_token(path, audience)
+        return self._read_token(descriptor, path, audience)
+
+    @staticmethod
+    def _open_regular_token(path: Path, audience: str) -> int:
+        flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
+        try:
+            return os.open(path, flags | os.O_NONBLOCK)
+        except OSError as exc:
+            if exc.errno == errno.ELOOP:
+                raise RuntimeError(
+                    f"workload token for audience {audience!r} must not be a symlink: {path}"
+                ) from exc
+            raise RuntimeError(f"missing workload token for audience {audience!r}: {path}") from exc
+
+    def _read_token(self, descriptor: int, path: Path, audience: str) -> str:
+        try:
+            token_stat = os.fstat(descriptor)
+            if not stat.S_ISREG(token_stat.st_mode):
+                raise RuntimeError(
+                    f"workload token for audience {audience!r} is not a regular file: {path}"
+                )
+            if token_stat.st_size > MAX_WORKLOAD_TOKEN_BYTES:
+                raise RuntimeError(
+                    f"workload token for audience {audience!r} exceeds 16 KiB: {path}"
+                )
+            # Kubernetes adds group-read for a configured pod fsGroup. Keep
+            # credentials read-only and never permit access by other users.
+            if stat.S_IMODE(token_stat.st_mode) not in {0o400, 0o440}:
+                raise RuntimeError(
+                    f"workload token for audience {audience!r} must have mode 0400 or 0440: {path}"
+                )
+            with os.fdopen(descriptor, "rb", closefd=False) as token_file:
+                raw_token = token_file.read(MAX_WORKLOAD_TOKEN_BYTES + 1)
+            if len(raw_token) > MAX_WORKLOAD_TOKEN_BYTES:
+                raise RuntimeError(
+                    f"workload token for audience {audience!r} exceeds 16 KiB: {path}"
+                )
+            token = raw_token.decode("ascii").strip()
+        except UnicodeError as exc:
+            raise RuntimeError(
+                f"workload token for audience {audience!r} must be ASCII: {path}"
+            ) from exc
+        finally:
+            os.close(descriptor)
+        if not token:
+            raise RuntimeError(f"empty workload token for audience {audience!r}: {path}")
+        self._validate_expiry(token, audience=audience, path=path)
+        return token
+
+    @staticmethod
+    def _validate_expiry(token: str, *, audience: str, path: Path) -> None:
+        parts = token.split(".")
+        if len(parts) != 3:
+            raise RuntimeError(f"workload token for audience {audience!r} is not a JWT: {path}")
+        try:
+            payload_segment = parts[1] + "=" * (-len(parts[1]) % 4)
+            payload = json.loads(base64.urlsafe_b64decode(payload_segment).decode("utf-8"))
+            expires_at = int(payload["exp"])
+        except (binascii.Error, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
+            raise RuntimeError(
+                f"workload token for audience {audience!r} has no valid exp claim: {path}"
+            ) from exc
+        token_audience = payload.get("aud")
+        accepted_audiences = (
+            {str(item) for item in token_audience}
+            if isinstance(token_audience, list)
+            else {str(token_audience)}
+        )
+        if audience not in accepted_audiences:
+            raise RuntimeError(f"workload token audience mismatch for {audience!r}: {path}")
+        if expires_at <= int(time.time()):
+            raise RuntimeError(f"workload token for audience {audience!r} is expired: {path}")