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

ksadk/a2a/codex_runtime.py

@@ -0,0 +1,198 @@
+"""Merge native Codex facts and fixed-expert facts under Kernel ownership."""
+import asyncio
+from contextlib import aclosing
+from dataclasses import replace
+
+from ksadk.a2a.codex import CodexInvocationContextProvider, CodexToolCall
+from ksadk.codex.a2a_tools import denied_tool_result, failed_tool_result
+from ksadk.events.run_authority import current_run_event_authority
+
+
+class CodexRecoveryCancelled(RuntimeError):
+    pass
+
+
+async def managed_codex_events(client, bootstrap, native_events, *, request, owner=None, recovery=None):
+    authority = current_run_event_authority()
+    if authority is None:
+        raise RuntimeError('A2A_KERNEL_AUTHORITY_REQUIRED')
+    if not bootstrap.outbound_enabled:
+        raise RuntimeError('A2A_OUTBOUND_DISABLED')
+    bridge = client._a2a_expert_bridge
+    services = {}
+    authority.register_managed_a2a(bootstrap)
+    policy = authority.managed_invocation_policy()
+
+    def service_for_space(space):
+        if space not in services:
+            recovery_options = {"recover_checkpoint_leases": True} if recovery is not None else {}
+            services[space] = bootstrap.invocation_service_for_space(
+                space, event_store=authority, policy=policy, **recovery_options)
+        return services[space]
+
+    provider = CodexInvocationContextProvider(authority=authority,
+        service_for_space=service_for_space, caller_agent_id=bootstrap.runtime_metadata.agent_id,
+        require_durable=recovery is not None)
+    output = asyncio.Queue(maxsize=1)
+    persisted = asyncio.Condition()
+    triggers = set()
+    calls = {}
+    from ksadk.a2a.codex_checkpoint import CodexCheckpointUnavailable, CodexCheckpointWriter
+
+    checkpoints = CodexCheckpointWriter(
+        client=client, authority=authority, request=request, experts=bridge.experts, recovery=recovery,
+    )
+
+    async def cancel_experts():
+        pending = [task for task in calls.values() if not task.done()]
+        for call_task in pending:
+            call_task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+
+    if owner is not None:
+        owner.cancel_experts = cancel_experts
+
+    async def emit(event):
+        ack = asyncio.Event()
+        await output.put((event, ack))
+        await ack.wait()
+
+    async def invoke(call):
+        try:
+            # A server request may win the race with its earlier native item
+            # notification. Continue draining native events until Kernel has
+            # consumed that trigger; provider then verifies the persisted fact.
+            async with persisted:
+                await persisted.wait_for(lambda: bridge.key(call) in triggers)
+            response = await bridge.experts[call.tool].invoke(
+                call, provider=provider, emit=emit, before_invoke=checkpoints.before_invoke,
+            )
+        except asyncio.CancelledError:
+            client.resolve_a2a_tool(call, failed_tool_result('A2A_CODEX_TOOL_CANCELED'))
+            raise
+        except TimeoutError:
+            # Settle the parent as failed instead of returning a retryable tool
+            # result to the model, which could create another remote task.
+            await output.put((RuntimeError('A2A_CODEX_EXPERT_TIMEOUT'), None))
+            return
+        except CodexCheckpointUnavailable as exc:
+            # Never turn an uncertain private write into a model-retryable tool error.
+            await output.put((exc, None))
+            return
+        except PermissionError as exc:
+            if str(exc) in {'A2A_APPROVAL_DENIED', 'A2A_INPUT_REJECTED'}:
+                bridge.mark_denied(call)
+                response = denied_tool_result()
+            else:
+                response = failed_tool_result('A2A_CODEX_EXPERT_FAILED')
+        except Exception:
+            response = failed_tool_result('A2A_CODEX_EXPERT_FAILED')
+        client.resolve_a2a_tool(call, response)
+
+    async def recover_call(call):
+        # The provider finds the original request/attempt from the same durable
+        # canonical trigger; invoke_with_input rejoins that service checkpoint.
+        try:
+            response = await bridge.experts[call.tool].invoke(
+                call, provider=provider, emit=emit, before_invoke=checkpoints.before_invoke,
+            )
+        except TimeoutError:
+            raise RuntimeError("A2A_CODEX_EXPERT_TIMEOUT") from None
+        items = response.get("contentItems", [])
+        if (len(items) != 1 or items[0].get("type") != "inputText"
+                or not isinstance(items[0].get("text"), str)):
+            raise ValueError("A2A_CODEX_RECOVERY_OUTPUT_INVALID")
+        # Native dynamic tools persist the text content, not their RPC envelope.
+        # Match that representation for calls already resolved before the crash.
+        try:
+            rollout = await client.ensure_managed_a2a_output(call, items[0]["text"])
+        except TimeoutError:
+            # This run's uncertain injection must not close the app-owned
+            # transport and interrupt unrelated threads.
+            raise RuntimeError("A2A_CODEX_RECOVERY_OUTPUT_UNCERTAIN") from None
+        # Persist the flushed output before any new native turn may start.
+        await checkpoints.before_invoke(call)
+        from ksadk.a2a.codex_recovery_events import recovered_expert_terminal
+
+        terminal = await recovered_expert_terminal(authority, call, rollout, response)
+        if terminal is not None:
+            await emit(terminal)
+
+    async def pump():
+        try:
+            if recovery is not None:
+                for call in recovery.calls:
+                    calls[bridge.key(call)] = asyncio.create_task(recover_call(call))
+                await asyncio.gather(*calls.values())
+                await checkpoints.mark_continuation_started()
+            async with aclosing(native_events):
+                async for event in native_events:
+                    if isinstance(event, dict):
+                        call = CodexToolCall.model_validate(event['params'])
+                        key = bridge.key(call)
+                        if event['method'] == 'ksadk/a2a/toolCall':
+                            if key not in calls:
+                                calls[key] = asyncio.create_task(invoke(call))
+                        elif event['method'] == 'ksadk/a2a/toolCancel':
+                            await cancel_experts()
+                            raise RuntimeError('A2A_CODEX_EXPERT_TIMEOUT')
+                        else:
+                            raise ValueError('A2A_CODEX_INTERNAL_EVENT_INVALID')
+                        continue
+                    if event.event_type == 'run.completed' and event.source.framework == 'codex':
+                        await checkpoints.completed_turn(event)
+                    await emit(event)
+                    source = event.source
+                    if getattr(event, 'item_kind', None) == 'tool_call' and source.framework == 'codex':
+                        async with persisted:
+                            triggers.add((source.metadata.get('thread_id'), source.native_run_id, source.native_item_id))
+                            persisted.notify_all()
+            await output.put((None, None))
+        except asyncio.CancelledError:
+            # A canceled recovery child can finish before any native turn
+            # exists. Wake the consumer; ordinary generator teardown must not
+            # enqueue into a queue whose consumer has already gone away.
+            if recovery is not None and not asyncio.current_task().cancelling():
+                await output.put((CodexRecoveryCancelled(), None))
+            else:
+                raise
+        except Exception as exc:
+            await output.put((exc, None))
+
+    from ksadk.runtime_context import (
+        TRUSTED_IDENTITY_METADATA_KEY, get_current_invocation_context,
+        platform_invocation_scope, session_invocation_context,
+    )
+    context = get_current_invocation_context() or session_invocation_context(
+        request.metadata.get("session_context"), agent_id=bootstrap.runtime_metadata.agent_id,
+        user_id=request.user_id, session_id=request.session_id, runner_type="codex",
+        identity=request.metadata.get(TRUSTED_IDENTITY_METADATA_KEY),
+    )
+    profile = request.metadata.get("tool_approval_mode")
+    if isinstance(profile, str) and profile in {"ask", "risk", "full"}:
+        # Kernel's per-run choice must win over an environment default or a
+        # surrounding invocation. Keep the caller's context object unchanged.
+        context = replace(context, tool_approval_mode=profile)
+    with platform_invocation_scope(context):
+        task = asyncio.create_task(pump())
+    try:
+        while True:
+            event, ack = await output.get()
+            if event is None:
+                break
+            if isinstance(event, Exception):
+                if isinstance(event, CodexCheckpointUnavailable):
+                    await cancel_experts()
+                    await client.interrupt_active_turn(event.thread_id)
+                raise event
+            try:
+                yield event
+            finally:
+                ack.set()
+    finally:
+        if owner is not None:
+            owner.cancel_experts = None
+        task.cancel()
+        for call_task in calls.values():
+            call_task.cancel()
+        await asyncio.gather(task, *calls.values(), return_exceptions=True)

ksadk/a2a/codex_session_checkpoint.py

@@ -0,0 +1,162 @@
+"""Private completed-turn handoff, distinct from interrupted invocation recovery."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from ksadk.a2a.invocation_types import RemoteAgentSpec
+from ksadk.codex.rollout_checkpoint import NativeRolloutCheckpoint
+from ksadk.events.canonical import RunCompleted
+from ksadk.kernel.state import RunState
+from ksadk.runtime.adapter import RunHandle, StartRequest
+from ksadk.runtime.executor import handle_digest
+
+
+class CodexSessionCheckpoint(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, hide_input_in_errors=True)
+
+    schema_version: Literal["codex.a2a.session/v1"] = Field(alias="schema")
+    run_id: str = Field(min_length=1)
+    session_id: str = Field(min_length=1)
+    tenant_id: str = Field(min_length=1)
+    agent_instance_id: str = Field(min_length=1)
+    thread_id: str = Field(min_length=1)
+    turn_id: str = Field(min_length=1)
+    start_request: StartRequest
+    bindings: dict[str, RemoteAgentSpec]
+    rollout: NativeRolloutCheckpoint
+    terminal_event: RunCompleted
+
+    @model_validator(mode="after")
+    def validate_completed_scope(self):
+        if (
+            self.rollout.thread_id != self.thread_id
+            or self.start_request.session_id != self.session_id
+        ):
+            raise ValueError("A2A_CODEX_SESSION_SCOPE_CONFLICT")
+        if not self.bindings or any(
+            name != spec.local_name for name, spec in self.bindings.items()
+        ):
+            raise ValueError("A2A_CODEX_SESSION_BINDING_CONFLICT")
+        self.rollout.require_completed_turn(self.turn_id)
+        event = self.terminal_event
+        if (
+            event.run_id != self.run_id
+            or event.source.framework != "codex"
+            or event.source.native_run_id != self.turn_id
+            or event.source.metadata.get("thread_id") != self.thread_id
+        ):
+            raise ValueError("A2A_CODEX_SESSION_TERMINAL_CONFLICT")
+        return self
+
+
+def validated_terminal_recovery_checkpoint(run, handle: RunHandle, *, experts):
+    """Authorize only replay of an already captured terminal fact for this Run."""
+    try:
+        checkpoint = CodexSessionCheckpoint.model_validate(run.metadata.get("runtime_checkpoint"))
+        saved = RunHandle.model_validate(run.metadata.get("handle"))
+        request = checkpoint.start_request
+        if (
+            run.state != RunState.RUNNING
+            or run.run_id != checkpoint.run_id
+            or run.session_id != checkpoint.session_id
+            or run.agent_instance_id != checkpoint.agent_instance_id
+            or run.metadata.get("tenant_id") != checkpoint.tenant_id
+            or saved != handle
+            or handle_digest(saved) != run.metadata.get("handle_digest")
+            or handle.runtime_type != "codex"
+            or handle.session_id != run.session_id
+            or handle.run_id != checkpoint.thread_id
+            or handle.native_ref.get("thread_id") != checkpoint.thread_id
+            or handle.native_ref.get("user_id") != request.user_id
+            or request.metadata.get("run_id") != run.run_id
+            or not request.metadata.get("command_id")
+            or request.metadata.get("command_id") != run.metadata.get("command_id")
+            or (
+                request.metadata.get("invocation_id")
+                and request.metadata["invocation_id"] != run.run_id
+            )
+            or checkpoint.bindings != {name: expert.spec for name, expert in experts.items()}
+        ):
+            raise ValueError("scope mismatch")
+        return checkpoint
+    except (ValueError, TypeError, AttributeError):
+        raise ValueError("A2A_CODEX_TERMINAL_CHECKPOINT_INVALID") from None
+
+
+async def terminal_recovery_event(checkpoint):
+    from ksadk.events.canonical import ItemCompleted
+    from ksadk.events.run_authority import current_run_event_authority
+
+    authority = current_run_event_authority()
+    if authority is None:
+        raise RuntimeError("A2A_KERNEL_AUTHORITY_REQUIRED")
+    identity = authority.runtime_checkpoint_identity()
+    if (
+        authority.run_id != checkpoint.run_id
+        or authority.session_id != checkpoint.session_id
+        or identity["tenant_id"] != checkpoint.tenant_id
+        or identity["agent_instance_id"] != checkpoint.agent_instance_id
+    ):
+        raise ValueError("A2A_CODEX_TERMINAL_AUTHORITY_CONFLICT")
+    prior = await authority.list(checkpoint.session_id)
+    items = {(event.scope_id, event.item_id) for event in prior if isinstance(event, ItemCompleted)}
+    if any(
+        (ref.scope_id, ref.item_id) not in items for ref in checkpoint.terminal_event.output_refs
+    ):
+        raise ValueError("A2A_CODEX_TERMINAL_OUTPUT_NOT_DURABLE")
+    return checkpoint.terminal_event
+
+
+def validated_session_checkpoint(
+    run,
+    handle: RunHandle,
+    request: StartRequest,
+    *,
+    agent_instance_id: str,
+    tenant_id: str,
+    experts,
+):
+    """Validate a Kernel-loaded predecessor against the current trusted owner.
+
+    The caller supplies the canonical predecessor, never a caller-uploaded run.
+    This performs no I/O and never replays an outbound invocation.
+    """
+    try:
+        checkpoint = CodexSessionCheckpoint.model_validate(run.metadata.get("runtime_checkpoint"))
+        saved = RunHandle.model_validate(run.metadata.get("handle"))
+        previous = checkpoint.start_request
+        if (
+            run.state != RunState.COMPLETED
+            or run.run_id != checkpoint.run_id
+            or run.agent_instance_id != checkpoint.agent_instance_id
+            or run.agent_instance_id != agent_instance_id
+            or run.metadata.get("tenant_id") != checkpoint.tenant_id
+            or checkpoint.tenant_id != tenant_id
+            or run.session_id != checkpoint.session_id
+            or request.session_id != checkpoint.session_id
+            or not previous.user_id
+            or request.user_id != previous.user_id
+            or saved != handle
+            or handle_digest(saved) != run.metadata.get("handle_digest")
+            or handle.runtime_type != "codex"
+            or handle.session_id != run.session_id
+            or handle.run_id != checkpoint.thread_id
+            or handle.native_ref.get("thread_id") != checkpoint.thread_id
+            or handle.native_ref.get("user_id") != previous.user_id
+            or request.metadata.get("thread_id") != checkpoint.thread_id
+            or previous.metadata.get("run_id") != run.run_id
+            or not previous.metadata.get("command_id")
+            or previous.metadata.get("command_id") != run.metadata.get("command_id")
+            or (
+                previous.metadata.get("invocation_id")
+                and previous.metadata["invocation_id"] != run.run_id
+            )
+            or not request.metadata.get("run_id")
+            or request.metadata["run_id"] == run.run_id
+            or checkpoint.bindings != {name: expert.spec for name, expert in experts.items()}
+        ):
+            raise ValueError("scope mismatch")
+        return checkpoint
+    except (ValueError, TypeError, AttributeError):
+        raise ValueError("A2A_CODEX_SESSION_CHECKPOINT_INVALID") from None

ksadk/a2a/composition.py

@@ -0,0 +1,88 @@
+"""Load immutable Studio compositions without modifying code-first projects."""
+from __future__ import annotations
+
+import hashlib
+import importlib
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def load_source_module(root: Path, entry_point: str):
+    root = root.resolve()
+    relative = entry_point[:-3] if entry_point.endswith(".py") else entry_point
+    parts = relative.replace("\\", "/").replace("/", ".").split(".")
+    if not parts or any(not part.isidentifier() for part in parts):
+        raise ValueError("A2A_SOURCE_ENTRY_INVALID")
+    source = root.joinpath(*parts).with_suffix(".py").resolve()
+    if not source.is_relative_to(root) or not source.is_file():
+        raise ValueError("A2A_SOURCE_ENTRY_INVALID")
+    # Each bundle gets its own package namespace, including relative imports.
+    namespace = "_ksadk_a2a_bundle_" + hashlib.sha256(str(root).encode()).hexdigest()[:24]
+    if namespace not in sys.modules:
+        package = ModuleType(namespace)
+        package.__path__ = [str(root)]
+        package.__package__ = namespace
+        sys.modules[namespace] = package
+    return importlib.import_module(namespace + "." + ".".join(parts))
+
+
+def compose_project(source, compose, *, framework, local_ids, handoff_sources,
+                    agent_variable, checkpointer=None, resource_tools=()):
+    provider = getattr(source, "ksadk_a2a_local_nodes", None)
+    if provider is not None:
+        if not callable(provider):
+            raise ValueError("A2A_LOCAL_NODE_PROVIDER_INVALID")
+        nodes = provider()
+        if not isinstance(nodes, dict) or set(nodes) != set(local_ids):
+            raise ValueError("A2A_LOCAL_NODE_MAPPING_INVALID")
+        nodes = dict(nodes)
+    elif len(local_ids) == 1:
+        nodes = {local_ids[0]: getattr(source, agent_variable)}
+    elif local_ids:
+        raise ValueError("A2A_LOCAL_NODE_PROVIDER_REQUIRED")
+    else:
+        nodes = {}
+    factories = getattr(source, "ksadk_a2a_local_factories", {})
+    if not isinstance(factories, dict):
+        raise ValueError("A2A_LOCAL_TOOL_FACTORY_REQUIRED")
+    factories = dict(factories)
+    if framework == "langgraph":
+        factory = getattr(source, "ksadk_graph_factory", None)
+        if len(local_ids) == 1 and callable(factory) and provider is None:
+            local_id = local_ids[0]
+            # Nested graphs inherit the outer saver; never keep a private
+            # MemorySaver inside a graph that claims durable recovery.
+            if local_id in handoff_sources and local_id not in factories:
+                import inspect
+                if "resource_tools" not in inspect.signature(factory).parameters:
+                    raise ValueError("A2A_LOCAL_TOOL_FACTORY_REQUIRED")
+                factories[local_id] = lambda tools: factory(
+                    checkpointer=None, resource_tools=tuple(resource_tools) + tuple(tools))
+            elif local_id not in handoff_sources:
+                if resource_tools:
+                    nodes[local_id] = factory(checkpointer=None, resource_tools=tuple(resource_tools))
+                else:
+                    nodes[local_id] = factory(checkpointer=None)
+        elif resource_tools:
+            raise ValueError("A2A_RESOURCE_TOOL_FACTORY_REQUIRED")
+        for local_id, node in nodes.items():
+            if local_id not in handoff_sources and getattr(node, "checkpointer", None) is not None:
+                raise ValueError("A2A_LOCAL_CHECKPOINTER_MUST_INHERIT")
+        return compose(nodes, local_factories=factories, checkpointer=checkpointer)
+    if resource_tools:
+        raise ValueError("A2A_RESOURCE_TOOL_FACTORY_REQUIRED")
+    return compose(nodes)
+
+
+def prepare_composed_state(source, payload, session_context):
+    hook = getattr(source, "ksadk_prepare_state", None)
+    if callable(hook):
+        return hook(payload, session_context)
+    # The default MessagesState contract admits the current textual turn only.
+    # Non-text/custom state requires the project's explicit projection hook.
+    from langchain_core.messages import HumanMessage
+    current = payload.get("input")
+    if not isinstance(current, str):
+        raise ValueError("A2A_PREPARE_STATE_HOOK_REQUIRED")
+    return {"messages": [HumanMessage(content=current)]}

ksadk/a2a/context_store.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import hashlib
+import json
 import os
 import sqlite3
 import uuid
@@ -11,6 +12,9 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
 
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncEngine
+
 
 @dataclass(frozen=True)
 class A2AContextIdentity:
@@ -20,6 +24,7 @@ class A2AContextIdentity:
     tenant_id: str
     caller_principal_type: str
     caller_principal_id: str
+    sub_account_id: str | None = None
 
     def canonical_owner(self) -> str:
         values = (
@@ -29,9 +34,15 @@ class A2AContextIdentity:
             self.caller_principal_id.strip(),
         )
         if not values[0] or not values[2] or not values[3]:
-            raise ValueError(
-                "A2A context identity requires verified account, type, and principal"
-            )
+            raise ValueError("A2A context identity requires verified account, type, and principal")
+        if self.caller_principal_type == "user":
+            # Product identifiers are opaque. Preserve exact values and field
+            # boundaries instead of trimming or joining with an escapable delimiter.
+            owner = [self.account_id, self.tenant_id, self.caller_principal_type,
+                     self.caller_principal_id]
+            if self.sub_account_id is not None:
+                owner.append(self.sub_account_id)
+            return json.dumps(owner, ensure_ascii=False, separators=(",", ":"))
         return "\x1f".join(values)
 
 
@@ -43,6 +54,19 @@ class A2AContextStore(ABC):
         """Verify durable storage is ready before inbound A2A starts serving."""
         raise NotImplementedError
 
+
+def _mapping_key(
+    identity: A2AContextIdentity,
+    external_context_id: str,
+    isolation_scope: str | None,
+) -> tuple[str, str, str]:
+    context = external_context_id.strip()
+    if not context:
+        raise ValueError("external_context_id must be non-empty")
+    owner_hash = hashlib.sha256(identity.canonical_owner().encode("utf-8")).hexdigest()
+    context_hash = hashlib.sha256(context.encode("utf-8")).hexdigest()
+    return owner_hash, context_hash, str(isolation_scope or "")
+
     @abstractmethod
     async def resolve_or_create(
         self,
@@ -97,7 +121,7 @@ class SQLiteA2AContextStore(A2AContextStore):
         isolation_scope: str | None = None,
     ) -> str:
         await self.initialize()
-        key = self._mapping_key(identity, external_context_id, isolation_scope)
+        key = _mapping_key(identity, external_context_id, isolation_scope)
         return await asyncio.to_thread(self._resolve_or_create_sync, key)
 
     async def get(
@@ -108,7 +132,7 @@ class SQLiteA2AContextStore(A2AContextStore):
         isolation_scope: str | None = None,
     ) -> str | None:
         await self.initialize()
-        key = self._mapping_key(identity, external_context_id, isolation_scope)
+        key = _mapping_key(identity, external_context_id, isolation_scope)
         return await asyncio.to_thread(self._get_sync, key)
 
     def _initialize_sync(self) -> None:
@@ -136,19 +160,6 @@ class SQLiteA2AContextStore(A2AContextStore):
             connection.close()
         os.chmod(self._path, 0o600)
 
-    @staticmethod
-    def _mapping_key(
-        identity: A2AContextIdentity,
-        external_context_id: str,
-        isolation_scope: str | None,
-    ) -> tuple[str, str, str]:
-        context = external_context_id.strip()
-        if not context:
-            raise ValueError("external_context_id must be non-empty")
-        owner_hash = hashlib.sha256(identity.canonical_owner().encode("utf-8")).hexdigest()
-        context_hash = hashlib.sha256(context.encode("utf-8")).hexdigest()
-        return owner_hash, context_hash, str(isolation_scope or "")
-
     def _resolve_or_create_sync(self, key: tuple[str, str, str]) -> str:
         owner_hash, context_hash, scope = key
         connection = sqlite3.connect(self._path, timeout=5, isolation_level=None)
@@ -200,4 +211,92 @@ class SQLiteA2AContextStore(A2AContextStore):
             connection.close()
 
 
-__all__ = ["A2AContextIdentity", "A2AContextStore", "SQLiteA2AContextStore"]
+class PostgresA2AContextStore(A2AContextStore):
+    """Shared context mapping for multi-replica managed Runtimes."""
+
+    def __init__(self, engine: AsyncEngine) -> None:
+        self._engine = engine
+
+    async def initialize(self) -> None:
+        async with self._engine.begin() as connection:
+            await connection.execute(
+                text(
+                    """
+                    CREATE TABLE IF NOT EXISTS ksadk_a2a_context_mappings (
+                        schema_version INTEGER NOT NULL,
+                        owner_key_sha256 TEXT NOT NULL,
+                        external_context_sha256 TEXT NOT NULL,
+                        isolation_scope TEXT NOT NULL,
+                        internal_session_id TEXT NOT NULL,
+                        created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                        updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                        PRIMARY KEY(owner_key_sha256, external_context_sha256, isolation_scope)
+                    )
+                    """
+                )
+            )
+
+    async def resolve_or_create(
+        self,
+        identity: A2AContextIdentity,
+        external_context_id: str,
+        *,
+        isolation_scope: str | None = None,
+    ) -> str:
+        owner_hash, context_hash, scope = _mapping_key(
+            identity, external_context_id, isolation_scope
+        )
+        session_id = f"a2a-session-{uuid.uuid4().hex}"
+        async with self._engine.begin() as connection:
+            result = await connection.execute(
+                text(
+                    """
+                    INSERT INTO ksadk_a2a_context_mappings(
+                        schema_version, owner_key_sha256, external_context_sha256,
+                        isolation_scope, internal_session_id
+                    ) VALUES (1, :owner_hash, :context_hash, :scope, :session_id)
+                    ON CONFLICT(owner_key_sha256, external_context_sha256, isolation_scope)
+                    DO UPDATE SET updated_at = CURRENT_TIMESTAMP
+                    RETURNING internal_session_id
+                    """
+                ),
+                {
+                    "owner_hash": owner_hash,
+                    "context_hash": context_hash,
+                    "scope": scope,
+                    "session_id": session_id,
+                },
+            )
+            return str(result.scalar_one())
+
+    async def get(
+        self,
+        identity: A2AContextIdentity,
+        external_context_id: str,
+        *,
+        isolation_scope: str | None = None,
+    ) -> str | None:
+        owner_hash, context_hash, scope = _mapping_key(
+            identity, external_context_id, isolation_scope
+        )
+        async with self._engine.connect() as connection:
+            result = await connection.execute(
+                text(
+                    """
+                    SELECT internal_session_id FROM ksadk_a2a_context_mappings
+                    WHERE owner_key_sha256 = :owner_hash
+                      AND external_context_sha256 = :context_hash
+                      AND isolation_scope = :scope
+                    """
+                ),
+                {"owner_hash": owner_hash, "context_hash": context_hash, "scope": scope},
+            )
+            value = result.scalar_one_or_none()
+            return str(value) if value is not None else None
+
+__all__ = [
+    "A2AContextIdentity",
+    "A2AContextStore",
+    "PostgresA2AContextStore",
+    "SQLiteA2AContextStore",
+]