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

ksadk/a2a/invocation_checkpoint.py

@@ -0,0 +1,274 @@
+"""Private caller checkpoints, separate from the public canonical event log."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import time
+from contextlib import asynccontextmanager
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator
+from sqlalchemy import text
+
+from ksadk.a2a.invocation_types import A2AInvocationHandle, A2AInvocationRequest
+
+
+class A2AInvocationCheckpoint(BaseModel):
+    """Private recovery input and protocol state; never a public Task payload.
+
+    ``state`` is a versioned adapter/service checkpoint, not DI objects, permits
+    or transport credentials. Canonical events remain in their existing store.
+    """
+
+    model_config = ConfigDict(extra="forbid", frozen=True, hide_input_in_errors=True)
+    version: Literal[1] = 1
+    request: A2AInvocationRequest
+    handle: A2AInvocationHandle
+    send_entered: bool = False
+    state: dict[str, JsonValue] = Field(default_factory=dict, repr=False)
+
+    @model_validator(mode="after")
+    def identity_matches(self):
+        excluded = {"platform_task_id", "remote_task_id", "remote_context_id"}
+        if self.request.handle.model_dump(exclude=excluded) != self.handle.model_dump(
+            exclude=excluded
+        ):
+            raise ValueError("A2A_CHECKPOINT_IDENTITY_CONFLICT")
+        return self
+
+
+class InvocationCheckpointBusy(RuntimeError):
+    pass
+
+
+class _CheckpointLease:
+    def __init__(self, connection, key, scope_id):
+        self._connection, self._key, self._scope_id = connection, key, scope_id
+        self._active = True
+
+    def _require_active(self):
+        if not self._active or self._connection.closed or self._connection.invalidated:
+            raise RuntimeError("A2A_CHECKPOINT_LEASE_CLOSED")
+
+    async def load(self) -> A2AInvocationCheckpoint | None:
+        self._require_active()
+        result = await self._connection.execute(
+            text(
+                "SELECT checkpoint_json FROM ksadk_a2a_invocation_checkpoints "
+                "WHERE owner_scope_key=:key"
+            ),
+            {"key": self._key},
+        )
+        payload = result.scalar_one_or_none()
+        await self._connection.commit()
+        if payload is None:
+            return None
+        checkpoint = A2AInvocationCheckpoint.model_validate_json(payload)
+        if checkpoint.handle.scope_id != self._scope_id:
+            raise ValueError("A2A_CHECKPOINT_IDENTITY_CONFLICT")
+        return checkpoint
+
+    async def cancel_requested(self) -> bool:
+        self._require_active()
+        result = await self._connection.scalar(
+            text(
+                "SELECT cancel_requested FROM ksadk_a2a_invocation_checkpoints "
+                "WHERE owner_scope_key=:key"
+            ),
+            {"key": self._key},
+        )
+        await self._connection.commit()
+        return bool(result)
+
+    async def clear_cancel_request(self):
+        self._require_active()
+        await self._connection.execute(
+            text(
+                "UPDATE ksadk_a2a_invocation_checkpoints SET cancel_requested=FALSE "
+                "WHERE owner_scope_key=:key"
+            ),
+            {"key": self._key},
+        )
+        await self._connection.commit()
+
+    async def save(self, checkpoint: A2AInvocationCheckpoint) -> None:
+        self._require_active()
+        # Revalidate even a model_copy from a caller. JSON encoding disallows
+        # nonfinite numbers and arbitrary Python object serialization.
+        checkpoint = A2AInvocationCheckpoint.model_validate(checkpoint.model_dump())
+        if checkpoint.handle.scope_id != self._scope_id:
+            raise ValueError("A2A_CHECKPOINT_IDENTITY_CONFLICT")
+        previous = await self.load()
+        if previous is not None:
+            if previous.request != checkpoint.request:
+                raise ValueError("A2A_CHECKPOINT_REQUEST_CONFLICT")
+            if previous.send_entered and not checkpoint.send_entered:
+                raise ValueError("A2A_CHECKPOINT_SEND_REGRESSION")
+            for locator in ("platform_task_id", "remote_task_id", "remote_context_id"):
+                before = getattr(previous.handle, locator)
+                if before is not None and before != getattr(checkpoint.handle, locator):
+                    raise ValueError("A2A_CHECKPOINT_BINDING_CONFLICT")
+        payload = json.dumps(
+            checkpoint.model_dump(mode="json"), allow_nan=False, separators=(",", ":")
+        )
+        if len(payload.encode()) > 16 * 1024 * 1024:
+            raise ValueError("A2A_CHECKPOINT_TOO_LARGE")
+        await self._connection.execute(
+            text(
+                "INSERT INTO ksadk_a2a_invocation_checkpoints(owner_scope_key, checkpoint_json) "
+                "VALUES (:key, :payload) ON CONFLICT(owner_scope_key) DO UPDATE SET "
+                "checkpoint_json=EXCLUDED.checkpoint_json"
+            ),
+            {"key": self._key, "payload": payload},
+        )
+        # Commit each send receipt before the wire side effect. An invocation
+        # crash must not roll back its already-recorded possible-send fact.
+        await self._connection.commit()
+
+
+class PostgresA2AInvocationCheckpointStore:
+    """One owner-scoped, process-safe invocation lease with durable receipts.
+
+    PostgreSQL session advisory locks survive receipt commits and are released
+    when a crashed worker's connection closes. Ordinary competing calls fail
+    promptly; recovery may wait within the saved request deadline. No timer
+    grants ownership or forcibly releases another connection's lock.
+    """
+
+    def __init__(
+        self, engine, *, account_id: str, tenant_id: str, runtime_id: str, lease_engine=None
+    ):
+        if not all(isinstance(v, str) and v.strip() for v in (account_id, tenant_id, runtime_id)):
+            raise ValueError("A2A_CHECKPOINT_OWNER_REQUIRED")
+        self._engine = engine
+        self._lease_engine = lease_engine or engine
+        self._owner = (account_id, tenant_id, runtime_id)
+
+    async def initialize(self):
+        async with self._engine.begin() as connection:
+            await connection.execute(
+                text(
+                    "CREATE TABLE IF NOT EXISTS ksadk_a2a_invocation_checkpoints ("
+                    "owner_scope_key TEXT PRIMARY KEY, checkpoint_json TEXT NOT NULL)"
+                )
+            )
+            await connection.execute(
+                text(
+                    "ALTER TABLE ksadk_a2a_invocation_checkpoints ADD COLUMN IF NOT EXISTS "
+                    "cancel_requested BOOLEAN NOT NULL DEFAULT FALSE"
+                )
+            )
+
+    def _key(self, scope_id):
+        return hashlib.sha256(
+            json.dumps([*self._owner, scope_id], separators=(",", ":")).encode()
+        ).hexdigest()
+
+    async def read(self, scope_id):
+        """Read one owner-scoped locator without acquiring its mutation lease."""
+        async with self._engine.connect() as connection:
+            payload = await connection.scalar(
+                text("SELECT checkpoint_json FROM ksadk_a2a_invocation_checkpoints "
+                     "WHERE owner_scope_key=:key"), {"key": self._key(scope_id)},
+            )
+        if payload is None:
+            return None
+        saved = A2AInvocationCheckpoint.model_validate_json(payload)
+        if saved.handle.scope_id != scope_id:
+            raise ValueError("A2A_CHECKPOINT_IDENTITY_CONFLICT")
+        return saved
+
+    async def request_cancel(self, handle):
+        """Record intent without stealing the active invocation's lease."""
+        key = self._key(handle.scope_id)
+        async with self._engine.begin() as connection:
+            payload = await connection.scalar(
+                text(
+                    "SELECT checkpoint_json FROM ksadk_a2a_invocation_checkpoints "
+                    "WHERE owner_scope_key=:key FOR UPDATE"
+                ),
+                {"key": key},
+            )
+            if payload is None:
+                raise KeyError("A2A_INVOCATION_NOT_FOUND")
+            saved = A2AInvocationCheckpoint.model_validate_json(payload)
+            excluded = {"platform_task_id", "remote_task_id", "remote_context_id"}
+            if saved.handle.model_dump(exclude=excluded) != handle.model_dump(exclude=excluded):
+                raise ValueError("A2A_HANDLE_CONFLICT")
+            for locator in excluded:
+                supplied = getattr(handle, locator)
+                if supplied is not None and supplied != getattr(saved.handle, locator):
+                    raise ValueError("A2A_HANDLE_CONFLICT")
+            await connection.execute(
+                text(
+                    "UPDATE ksadk_a2a_invocation_checkpoints SET cancel_requested=TRUE "
+                    "WHERE owner_scope_key=:key"
+                ),
+                {"key": key},
+            )
+            return saved
+
+    @asynccontextmanager
+    async def lease(self, scope_id: str, *, wait_for_recovery: bool = False):
+        if not isinstance(scope_id, str) or not scope_id.strip():
+            raise ValueError("A2A_CHECKPOINT_SCOPE_REQUIRED")
+        key = self._key(scope_id)
+        lock_id = int.from_bytes(bytes.fromhex(key)[:8], "big", signed=True)
+        async with self._lease_engine.connect() as connection:
+            wait_deadline = None
+            while True:
+                if (wait_deadline is not None
+                        and asyncio.get_running_loop().time() >= wait_deadline):
+                    raise TimeoutError("A2A_CHECKPOINT_RECOVERY_TIMEOUT")
+                try:
+                    acquired = await connection.scalar(
+                        text("SELECT pg_try_advisory_lock(:key)"), {"key": lock_id}
+                    )
+                    await connection.commit()
+                except BaseException:
+                    # Acquisition may have succeeded on the server even when its
+                    # response was lost. Never return that session to the pool.
+                    await connection.invalidate()
+                    raise
+                if acquired:
+                    break
+                if not wait_for_recovery:
+                    raise InvocationCheckpointBusy("A2A_INVOCATION_BUSY")
+                if wait_deadline is None:
+                    # This unlocked read supplies only a waiting budget. No
+                    # recovery state or permission to send is taken from it.
+                    payload = await connection.scalar(
+                        text("SELECT checkpoint_json FROM ksadk_a2a_invocation_checkpoints "
+                             "WHERE owner_scope_key=:key"), {"key": key},
+                    )
+                    await connection.commit()
+                    if payload is None:
+                        raise InvocationCheckpointBusy("A2A_INVOCATION_BUSY")
+                    saved = A2AInvocationCheckpoint.model_validate_json(payload)
+                    if saved.handle.scope_id != scope_id:
+                        raise ValueError("A2A_CHECKPOINT_IDENTITY_CONFLICT")
+                    remaining = saved.request.deadline_unix_seconds - time.time()
+                    wait_deadline = asyncio.get_running_loop().time() + remaining
+                remaining = wait_deadline - asyncio.get_running_loop().time()
+                if remaining <= 0:
+                    raise TimeoutError("A2A_CHECKPOINT_RECOVERY_TIMEOUT")
+                await asyncio.sleep(min(0.1, remaining))
+            lease = _CheckpointLease(connection, key, scope_id)
+            try:
+                yield lease
+            finally:
+                lease._active = False
+                # A failed statement may leave a transaction aborted. Roll it
+                # back before releasing the session lock into the pool.
+                try:
+                    if not connection.closed and not connection.invalidated:
+                        await connection.rollback()
+                        await connection.execute(
+                            text("SELECT pg_advisory_unlock(:key)"), {"key": lock_id}
+                        )
+                        await connection.commit()
+                except BaseException:
+                    await connection.invalidate()
+                    raise

ksadk/a2a/invocation_policy.py

@@ -0,0 +1,87 @@
+"""One fail-closed approval and audit boundary for native and tool invocation."""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Literal
+
+from ksadk.runtime_context import get_current_invocation_context
+from ksadk.tools.a2a_policy import A2A_REMOTE_TOOL_POLICY
+from ksadk.tools.gateway import tool_policy_requires_approval
+from ksadk.a2a.policy_refs import resolve_policy_ref
+
+
+class A2AApprovalRequired(PermissionError):
+    """Typed local approval pause; never infer it from a remote error string."""
+
+
+@dataclass(frozen=True)
+class A2AInvocationPolicyDecision:
+    action: Literal["allow", "deny", "approval_required"]
+    reason_code: str | None = None
+
+
+class A2AInvocationPolicy:
+    def __init__(self, *, evaluator=None, audit=None):
+        self._evaluator, self._audit = evaluator, audit
+
+    async def evaluate(self, request, *, route_kind=None):
+        if route_kind not in {None, "hosted_gateway"}:
+            decision = A2AInvocationPolicyDecision("deny", "A2A_HOSTED_ROUTE_REQUIRED")
+        elif getattr(request, "approval_decision", None) == "denied":
+            decision = A2AInvocationPolicyDecision("deny", "A2A_APPROVAL_DENIED")
+        elif self._evaluator is None:
+            decision = A2AInvocationPolicyDecision("deny", "A2A_POLICY_NOT_CONFIGURED")
+        else:
+            decision = self._evaluator(request, route_kind)
+            if inspect.isawaitable(decision):
+                decision = await decision
+            if not isinstance(decision, A2AInvocationPolicyDecision):
+                raise TypeError("A2A_POLICY_INVALID_DECISION")
+            if decision.action not in {"allow", "deny", "approval_required"}:
+                raise ValueError("A2A_POLICY_INVALID_DECISION")
+        if self._audit:
+            receipt = self._audit(request, decision, route_kind)
+            if inspect.isawaitable(receipt):
+                await receipt
+        return decision
+
+
+def runtime_a2a_invocation_policy(*, audit, approval_resolver=None):
+    """Share the runtime tool profile, using only trusted approval resolution.
+
+    The resolver is injected by the runtime interaction owner and must check
+    the exact invocation against its durable approval ledger. Public request
+    correlation fields and a caller-supplied 'approved' value are not proof.
+    Audit is required and awaited before any control-plane or remote request.
+    """
+    if not callable(audit):
+        raise TypeError("A2A_AUDIT_REQUIRED")
+    if approval_resolver is not None and not callable(approval_resolver):
+        raise TypeError("A2A_APPROVAL_RESOLVER_INVALID")
+
+    async def evaluate(request, _route):
+        try:
+            profile = resolve_policy_ref(getattr(getattr(request, "spec", None), "policy_ref", None))
+        except ValueError:
+            return A2AInvocationPolicyDecision("deny", "A2A_POLICY_REF_UNSUPPORTED")
+        context = get_current_invocation_context()
+        if context is None:
+            return A2AInvocationPolicyDecision("deny", "A2A_RUNTIME_CONTEXT_REQUIRED")
+        approval = None
+        if approval_resolver is not None:
+            approval = approval_resolver(request)
+            if inspect.isawaitable(approval):
+                approval = await approval
+            if approval not in {None, "approved", "denied"}:
+                raise ValueError("A2A_APPROVAL_RESOLUTION_INVALID")
+        if approval == "denied":
+            return A2AInvocationPolicyDecision("deny", "A2A_APPROVAL_DENIED")
+        if approval == "approved" or (profile == "runtime" and not tool_policy_requires_approval(
+            A2A_REMOTE_TOOL_POLICY, approval_mode=context.tool_approval_mode
+        )):
+            return A2AInvocationPolicyDecision("allow")
+        return A2AInvocationPolicyDecision("approval_required", "A2A_APPROVAL_REQUIRED")
+
+    return A2AInvocationPolicy(evaluator=evaluate, audit=audit)

ksadk/a2a/invocation_publication.py

@@ -0,0 +1,210 @@
+"""Publish A2A invocation events, descriptors, and task snapshots."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+from a2a.types import Task
+
+from ksadk.a2a.invocation_types import A2AInvocationOutputRef, A2AInvocationSnapshot
+from ksadk.a2a.presentation import (
+    ExecutionScopeDescriptor,
+    ScopeOutputRef,
+    validate_scope_transition,
+)
+from ksadk.events.adapters.a2a import A2AAdapterContext
+from ksadk.events.canonical import (
+    ContinuationCreated,
+    ContinuationResumed,
+    InteractionRequested,
+    InteractionResolved,
+    ItemCompleted,
+    ItemSnapshotReplaced,
+    ItemStarted,
+    SourceRef,
+)
+from ksadk.events.content import ContentSnapshot, DataContent
+from ksadk.events.identity import stable_event_id, stable_item_id, stable_part_id
+
+
+class _InvocationPublicationMixin:
+    async def _publish_one(self, state, event):
+        persisted, _ = await self._event_store.persist_one(state.handle.session_id, event)
+        state.cursor = max(state.cursor, persisted.seq)
+        state.last_event_id = persisted.event_id
+        if self._dispatcher is not None and state.handle.platform_task_id:
+            await self._dispatcher.enqueue(
+                platform_task_id=state.handle.platform_task_id,
+                events=[persisted.model_dump(mode="json")],
+            )
+        return persisted
+
+    async def _flush_pending(self, state):
+        published = []
+        while state.pending:
+            published.append(await self._publish_one(state, state.pending[0]))
+            state.pending.pop(0)
+            await self._save_state(state)
+        return published
+
+    async def _publish_batch(self, state, events):
+        if state.pending:
+            await self._flush_pending(state)
+        state.pending = list(events)
+        # Persist a bounded delivery spool alongside the advanced adapter
+        # state. Recovery drains it idempotently into the sole canonical log.
+        await self._save_state(state)
+        return await self._flush_pending(state)
+
+    async def _publish(self, state, event):
+        return (await self._publish_batch(state, [event]))[0]
+
+    async def _descriptor_event(self, state, status, *, initial=False, cancel=None):
+        values = state.descriptor.model_dump()
+        values.update(
+            status=status,
+            output_refs=tuple(ScopeOutputRef(**ref.model_dump()) for ref in state.refs.values()),
+        )
+        if status in {"completed", "failed", "cancelled"}:
+            values["ended_at"] = float(self._clock())
+        if cancel is not None:
+            values["cancel"] = cancel
+        current = ExecutionScopeDescriptor.model_validate(values)
+        if not initial:
+            validate_scope_transition(state.descriptor, current)
+        state.descriptor = current
+        item = stable_item_id("a2a", state.handle.scope_id, "execution.scope")
+        snapshot = ContentSnapshot(
+            parts=(
+                DataContent(
+                    part_id=stable_part_id("a2a", item, "descriptor"),
+                    data=current.model_dump(mode="json"),
+                ),
+            )
+        )
+        event_type = (
+            "item.started"
+            if initial
+            else "item.completed"
+            if status in {"completed", "failed", "cancelled"}
+            else "item.snapshot_replaced"
+        )
+        common = dict(
+            schema_version=2,
+            seq=state.cursor + 1,
+            timestamp=float(self._clock()),
+            run_id=state.handle.run_id,
+            scope_id=state.handle.scope_id,
+            parent_scope_id=state.handle.parent_scope_id,
+            source=SourceRef(framework="a2a", metadata={"schema": "execution.scope/v1"}),
+            event_id=stable_event_id(
+                "a2a",
+                state.handle.scope_id,
+                item,
+                event_type,
+                "descriptor",
+                current.model_dump_json(),
+                state.cursor,
+            ),
+            item_id=item,
+            item_kind="status",
+        )
+        event = (
+            ItemStarted(**common, initial=snapshot)
+            if initial
+            else ItemCompleted(**common, snapshot=snapshot)
+            if event_type == "item.completed"
+            else ItemSnapshotReplaced(**common, snapshot=snapshot)
+        )
+        return await self._publish(state, event)
+
+    def _snapshot(self, state, status=None, error=None):
+        selected = status or state.descriptor.status
+        state.snapshot = A2AInvocationSnapshot(
+            handle=state.handle,
+            status=selected,
+            source=state.source,
+            observed_at_unix_seconds=float(self._clock()),
+            event_cursor=state.cursor,
+            last_event_id=state.last_event_id,
+            output_refs=tuple(state.refs.values()),
+            continuation_id=state.continuation_id
+            if selected in {"input_required", "auth_required"}
+            else None,
+            error=error,
+        )
+        return state.snapshot
+
+    async def _mapped(self, state, events):
+        mapped = []
+        for event in events:
+            # A remote child never takes ownership of the public root run.
+            if event.event_type.startswith("run."):
+                continue
+            if isinstance(event, ContinuationCreated):
+                state.continuation_id = event.continuation_id
+            if isinstance(event, InteractionRequested) and event.interaction_kind in {
+                "structured_input",
+                "approval",
+            }:
+                state.pending_input = event
+            if (
+                isinstance(event, InteractionResolved)
+                and state.pending_input is not None
+                and event.interaction_id == state.pending_input.interaction_id
+            ):
+                state.pending_input = None
+            if isinstance(event, ContinuationResumed):
+                state.pending_input = None
+            if isinstance(event, ItemCompleted) and event.item_kind in {"message", "artifact"}:
+                for part in event.snapshot.parts:
+                    if getattr(part, "content_type", None) == "text":
+                        ref = A2AInvocationOutputRef(
+                            scope_id=state.handle.scope_id,
+                            item_id=event.item_id,
+                            part_id=part.part_id,
+                        )
+                        state.refs[(event.item_id, part.part_id)] = ref
+            mapped.append(event)
+        for persisted in await self._publish_batch(state, mapped):
+            yield persisted
+
+    async def _bind(self, state, payload, deadline):
+        if isinstance(payload, Task):
+            task_id, context_id = payload.id, payload.context_id
+        else:
+            task_id, context_id = (
+                getattr(payload, "task_id", ""),
+                getattr(payload, "context_id", ""),
+            )
+        if task_id:
+            if state.handle.remote_task_id is None:
+                await self._timed(
+                    self._control_plane.bind_remote_task(
+                        platform_task_id=state.handle.platform_task_id,
+                        remote_task_id=task_id,
+                        remote_context_id=context_id or None,
+                        observed_at=datetime.now(timezone.utc).isoformat(),
+                    ),
+                    deadline,
+                    "prepare",
+                )
+                state.handle = state.handle.bind_remote_task(task_id, context_id or None)
+            elif (task_id, context_id or None) != (
+                state.handle.remote_task_id,
+                state.handle.remote_context_id,
+            ):
+                raise ValueError("A2A_REMOTE_BINDING_CONFLICT")
+        if state.context is None:
+            state.context = A2AAdapterContext(
+                run_id=state.handle.run_id,
+                context_id=context_id or state.handle.scope_id,
+                task_id=task_id or None,
+                fixed_scope_id=state.handle.scope_id,
+                parent_scope_id=state.handle.parent_scope_id,
+            )
+        elif task_id:
+            state.context.task_id, state.context.context_id = (
+                task_id,
+                context_id or state.context.context_id,
+            )