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

ksadk/a2a/adk.py

@@ -0,0 +1,194 @@
+"""Code-first managed ADK child with explicit invocation identity injection."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from contextlib import aclosing, contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from google.adk.agents import BaseAgent
+from google.adk.events import Event
+from google.genai import types
+from pydantic import PrivateAttr
+
+from ksadk.a2a.exposure import compile_exposure_plan
+from ksadk.a2a.invocation_types import A2AInvocationRequest, RemoteAgentSpec
+from ksadk.a2a.presentation import ScopeTriggerRef
+from ksadk.events.canonical import InteractionRequested, ItemCompleted, ItemSnapshotReplaced, ItemUpdated
+from ksadk.events.content import TextContent
+
+
+@dataclass(frozen=True)
+class ManagedA2ABinding:
+    """Run-local dependencies; never place this object in session state.
+
+    provider(ctx, spec) returns (A2AInvocationRequest, ScopeTriggerRef), optionally
+    awaitable. It supplies the actual public root run/session ID, actual parent
+    scope and triggering canonical item/call ID, binding/attempt IDs, caller and
+    audit identities, stable wire message/idempotency IDs and absolute deadline.
+    """
+
+    invocation_service: Any
+    invocation_context_provider: Callable
+
+
+_binding: ContextVar[ManagedA2ABinding | None] = ContextVar("managed_a2a_binding", default=None)
+
+
+@contextmanager
+def bind_managed_a2a(invocation_service, invocation_context_provider):
+    """Bind dependencies to one async execution context, including child tasks."""
+    token = _binding.set(ManagedA2ABinding(invocation_service, invocation_context_provider))
+    try:
+        yield
+    finally:
+        _binding.reset(token)
+
+
+class ManagedA2AEvent(Event):
+    """Native projection carrying a private reference to the already-stored fact."""
+
+    _canonical: Any = PrivateAttr(default=None)
+
+
+def canonical_event_from_adk(event):
+    return event._canonical if isinstance(event, ManagedA2AEvent) else None
+
+
+class A2ARemoteAgent(BaseAgent):
+    _remote_spec: RemoteAgentSpec = PrivateAttr()
+    _invocation_service: Any = PrivateAttr(default=None)
+    _invocation_context_provider: Any = PrivateAttr(default=None)
+
+    def __init__(self, *, spec, invocation_service=None, invocation_context_provider=None):
+        super().__init__(name=spec.local_name, description=spec.description or "")
+        self._remote_spec = spec
+        self._invocation_service = invocation_service
+        self._invocation_context_provider = invocation_context_provider
+
+    @property
+    def remote_spec(self):
+        return self._remote_spec
+
+    async def _run_async_impl(self, ctx):
+        binding = _binding.get()
+        service = self._invocation_service or (binding.invocation_service if binding else None)
+        provider = self._invocation_context_provider or (
+            binding.invocation_context_provider if binding else None
+        )
+        if provider is None:
+            raise RuntimeError("A2A_MANAGED_INVOCATION_CONTEXT_REQUIRED")
+        resolved = provider(ctx, self._remote_spec)
+        if inspect.isawaitable(resolved):
+            resolved = await resolved
+        from ksadk.a2a.adk_invocation_context import ManagedADKInvocationContext
+
+        prefix_events = ()
+        if isinstance(resolved, ManagedADKInvocationContext):
+            request, trigger, service = resolved.request, resolved.trigger, resolved.service
+            prefix_events = resolved.prefix_events
+        else:
+            request, trigger = resolved
+        if service is None:
+            raise RuntimeError("A2A_MANAGED_INVOCATION_CONTEXT_REQUIRED")
+        if not isinstance(request, A2AInvocationRequest) or not isinstance(
+            trigger, ScopeTriggerRef
+        ):
+            raise TypeError("A2A_MANAGED_INVOCATION_CONTEXT_INVALID")
+        if request.spec != self._remote_spec:
+            raise ValueError("A2A_BINDING_MISMATCH")
+        for event in prefix_events:
+            if event.run_id != request.handle.run_id or event.scope_id != trigger.scope_id:
+                raise ValueError("A2A_TRIGGER_REF_CONFLICT")
+            yield _project(event, self.name, ctx)
+        try:
+            from ksadk.a2a.runtime_input import invoke_with_input
+
+            async with aclosing(
+                invoke_with_input(service, request, trigger_ref=trigger)
+            ) as events:
+                async for canonical in events:
+                    yield _project(canonical, self.name, ctx)
+        except (asyncio.CancelledError, TimeoutError):
+            # Persist an honest cancel request; unsupported/unknown is not a
+            # cancelled child. The runtime owns durable follow-up scheduling.
+            try:
+                handle = service.handle(request.scope_id)
+                await asyncio.shield(service.cancel(handle))
+            finally:
+                raise
+
+
+def _project(canonical, author, ctx):
+    # ADK excludes partial deliveries from session history. Interaction events
+    # carry live waiters/Futures for the Kernel bridge, not native model content;
+    # persisting them makes the next session deepcopy fail. Their durable facts
+    # remain in the Kernel interaction log and the original object still reaches
+    # the canonical adapter for live approval/input delivery.
+    partial = isinstance(canonical, (ItemUpdated, InteractionRequested))
+    parts = []
+    if isinstance(canonical, ItemUpdated) and isinstance(canonical.update, TextContent):
+        parts = [types.Part(text=canonical.update.text)]
+    elif isinstance(canonical, (ItemCompleted, ItemSnapshotReplaced)):
+        parts = [
+            types.Part(text=p.text) for p in canonical.snapshot.parts if isinstance(p, TextContent)
+        ]
+    native = ManagedA2AEvent(
+        # ADK persists every delivery, including recovery projections, with
+        # a unique native row ID. The stable public fact identity lives in
+        # _canonical/custom_metadata; reusing it here breaks native resume.
+        invocation_id=ctx.invocation_id,
+        author=author,
+        branch=ctx.branch,
+        partial=partial,
+        content=types.Content(role="model", parts=parts) if parts else None,
+        custom_metadata={
+            "ksadk_a2a_canonical_ref": {
+                "event_id": canonical.event_id,
+                "run_id": canonical.run_id,
+                "scope_id": canonical.scope_id,
+                "item_id": getattr(canonical, "item_id", None),
+            }
+        },
+    )
+    native._canonical = canonical
+    return native
+
+
+def remote_agent(
+    spec: RemoteAgentSpec | None = None,
+    *,
+    invocation_service=None,
+    invocation_context_provider=None,
+    exposure_plan=None,
+    **fields,
+):
+    """Construct an ADK BaseAgent without discovery or AgentCard network I/O."""
+    if spec is not None and fields:
+        raise ValueError("provide spec or named declaration fields")
+    spec = spec or RemoteAgentSpec(**fields)
+    plan = exposure_plan or compile_exposure_plan(
+        [spec], framework="adk", consumer_id=spec.local_name
+    )
+    matches = [entry for entry in plan.entries if entry.spec.binding_id == spec.binding_id]
+    if (
+        plan.framework != "adk"
+        or len(matches) != 1
+        or matches[0].spec != spec
+        or matches[0].exposure != "native"
+    ):
+        raise ValueError("A2A_EXPOSURE_CONFLICT")
+    return A2ARemoteAgent(
+        spec=spec,
+        invocation_service=invocation_service,
+        invocation_context_provider=invocation_context_provider,
+    )
+
+
+def handoff_tool(spec: RemoteAgentSpec):
+    """Create a fixed-target ADK tool backed by the shared invocation service."""
+    from ksadk.a2a.adk_tools import ManagedA2AHandoffTool
+    return ManagedA2AHandoffTool(spec)

ksadk/a2a/adk_invocation_context.py

@@ -0,0 +1,242 @@
+"""Native ADK trigger identities used by the managed invocation provider."""
+
+import inspect
+import time
+from dataclasses import dataclass
+
+from ksadk.a2a.invocation_types import (
+    A2AInvocationContext,
+    A2AInvocationHandle,
+    A2AInvocationMessage,
+    A2AInvocationRequest,
+)
+from ksadk.a2a.presentation import ScopeTriggerRef
+from ksadk.events.identity import stable_item_id, stable_scope_id
+
+
+class ADKWorkflowTrigger:
+    """Persist a trigger for a static native child slot (sequential/parallel).
+
+    Loops require explicit occurrence injection and are outside this selector.
+    Native handoffs use their real function-call event instead.
+    """
+
+    def __init__(self, authority, *, clock=time.time):
+        self._authority, self._clock = authority, clock
+
+    async def __call__(self, ctx, spec):
+        from google.adk.agents import LoopAgent
+
+        from ksadk.events.canonical import ItemCompleted, ItemStarted, SourceRef
+        from ksadk.events.content import ContentSnapshot, ToolCallContent
+        from ksadk.events.identity import stable_event_id, stable_part_id
+
+        authority = self._authority
+        authority.require_session(authority.session_id)
+        node = ctx.agent
+        if getattr(node, "remote_spec", None) != spec or not ctx.invocation_id:
+            raise ValueError("A2A_NATIVE_EXECUTION_REQUIRED")
+        path, seen = [], set()
+        while node is not None:
+            if id(node) in seen or len(seen) >= 64:
+                raise ValueError("A2A_NATIVE_EXECUTION_PATH_INVALID")
+            seen.add(id(node))
+            if isinstance(node, LoopAgent):
+                raise ValueError("A2A_WORKFLOW_OCCURRENCE_REQUIRED")
+            parent = node.parent_agent
+            slot = (
+                0
+                if parent is None
+                else next(
+                    (index for index, child in enumerate(parent.sub_agents) if child is node), -1
+                )
+            )
+            if slot < 0:
+                raise ValueError("A2A_NATIVE_EXECUTION_PATH_INVALID")
+            path.append(f"{slot}:{node.name}")
+            node = parent
+        native_slot = "/".join(reversed(path))
+        scope = stable_scope_id("adk", authority.run_id)
+        call = stable_item_id(
+            "adk", ctx.invocation_id, ctx.branch or "$root", native_slot, "native-child-call"
+        )
+        item = stable_item_id("adk", scope, call, "tool_call")
+        trigger = ScopeTriggerRef(scope_id=scope, item_id=item, call_id=call)
+        part = stable_part_id("adk", item, "call")
+        snapshot = ContentSnapshot(
+            parts=(ToolCallContent(part_id=part, call_id=call, name=spec.local_name, arguments={}),)
+        )
+        stored = {event.event_id: event for event in await authority.list(authority.session_id)}
+        prefix = []
+        for event_type in ("item.started", "item.completed"):
+            event_id = stable_event_id("adk", scope, item, event_type, part, call, 0)
+            existing = stored.get(event_id)
+            common = dict(
+                schema_version=2,
+                seq=0,
+                timestamp=(existing.timestamp if existing else self._clock()),
+                event_id=event_id,
+                run_id=authority.run_id,
+                scope_id=scope,
+                item_id=item,
+                item_kind="tool_call",
+                source=SourceRef(
+                    framework="adk",
+                    native_run_id=ctx.invocation_id,
+                    metadata={"execution_kind": "remote_agent"},
+                ),
+            )
+            event = (
+                ItemStarted(**common, initial=snapshot, phase="commentary")
+                if event_type == "item.started"
+                else ItemCompleted(**common, snapshot=snapshot)
+            )
+            persisted, _ = await authority.persist_one(authority.session_id, event)
+            prefix.append(persisted)
+        return trigger, tuple(prefix)
+
+
+@dataclass(frozen=True)
+class ManagedADKInvocationContext:
+    request: A2AInvocationRequest
+    trigger: ScopeTriggerRef
+    service: object
+    prefix_events: tuple = ()
+
+
+class ADKInvocationContextProvider:
+    """Run-local composition of identity, message selection, and durable request.
+
+    Selectors are runtime-owned functions, not fields read from model input.
+    A workflow selector returns an explicit trigger and persisted prefix facts.
+    """
+
+    def __init__(
+        self,
+        *,
+        authority,
+        service_for_space,
+        caller_agent_id,
+        message_selector,
+        workflow_trigger=None,
+        clock=time.time,
+    ):
+        if not caller_agent_id:
+            raise ValueError("A2A_CALLER_REQUIRED")
+        self._authority, self._services = authority, service_for_space
+        self._caller, self._message = caller_agent_id, message_selector
+        self._workflow = workflow_trigger or ADKWorkflowTrigger(authority, clock=clock)
+        self._clock = clock
+
+    async def __call__(self, ctx, spec):
+        authority = self._authority
+        authority.require_session(authority.session_id)
+        trigger = native_handoff_trigger(ctx, spec)
+        prefix = ()
+        if trigger is None:
+            if self._workflow is None:
+                raise ValueError("A2A_NATIVE_TRIGGER_REQUIRED")
+            result = self._workflow(ctx, spec)
+            if inspect.isawaitable(result):
+                result = await result
+            trigger, prefix = result
+        if not isinstance(trigger, ScopeTriggerRef):
+            raise TypeError("A2A_TRIGGER_REF_REQUIRED")
+        handle = A2AInvocationHandle(
+            run_id=authority.run_id,
+            session_id=authority.session_id,
+            parent_scope_id=trigger.scope_id,
+            call_id=trigger.call_id,
+            binding_id=spec.binding_id,
+            target_agent_id=spec.target_agent_id,
+            attempt_id=stable_item_id("a2a", trigger.scope_id, trigger.item_id, "attempt-1"),
+        )
+        service = self._services(spec.space_id)
+        if inspect.isawaitable(service):
+            service = await service
+        original = await service.original_request(handle)
+        if original is not None:
+            if original.spec != spec or original.context.caller_agent_id != self._caller:
+                raise ValueError("A2A_RECOVERED_BINDING_CONFLICT")
+            return ManagedADKInvocationContext(original, trigger, service, tuple(prefix))
+        parts = self._message(ctx, spec)
+        if inspect.isawaitable(parts):
+            parts = await parts
+        request = A2AInvocationRequest(
+            spec=spec,
+            handle=handle,
+            message=A2AInvocationMessage(
+                message_id=stable_item_id("a2a", handle.scope_id, "initial-message"),
+                text_parts=parts,
+            ),
+            context=A2AInvocationContext(
+                caller_agent_id=self._caller,
+                audit_id=stable_item_id("a2a", handle.scope_id, "audit"),
+            ),
+            deadline_unix_seconds=float(self._clock()) + (spec.timeout_seconds or 300.0),
+            idempotency_token=stable_item_id("a2a", handle.scope_id, "initial-send"),
+        )
+        return ManagedADKInvocationContext(request, trigger, service, tuple(prefix))
+
+
+def _path(event):
+    return (
+        getattr(getattr(event, "node_info", None), "path", None)
+        or getattr(event, "branch", None)
+        or "$root"
+    )
+
+
+def native_handoff_trigger(ctx, spec):
+    """Resolve a real transfer response back to its exact canonical call item.
+
+    No text/argument hash or target-name-derived call ID is used. Non-handoff
+    execution returns None so its workflow trigger can be allocated explicitly.
+    """
+    events = list(ctx.session.events)
+    for index in range(len(events) - 1, -1, -1):
+        response = events[index]
+        if response.invocation_id != ctx.invocation_id or getattr(
+            response, "branch", None
+        ) != getattr(ctx, "branch", None):
+            continue
+        if getattr(response.actions, "transfer_to_agent", None) != spec.local_name:
+            return None
+        calls = [
+            part.function_response
+            for part in (getattr(response.content, "parts", None) or ())
+            if part.function_response is not None
+            and part.function_response.name == "transfer_to_agent"
+        ]
+        if len(calls) != 1 or not calls[0].id:
+            raise ValueError("A2A_NATIVE_TRIGGER_AMBIGUOUS")
+        call_id = calls[0].id
+        for event in reversed(events[:index]):
+            if (
+                event.invocation_id != ctx.invocation_id
+                or _path(event) != _path(response)
+                or event.author != response.author
+            ):
+                continue
+            matches = [
+                part.function_call
+                for part in (getattr(event.content, "parts", None) or ())
+                if part.function_call is not None
+                and part.function_call.id == call_id
+                and part.function_call.name == "transfer_to_agent"
+            ]
+            if len(matches) > 1:
+                raise ValueError("A2A_NATIVE_TRIGGER_AMBIGUOUS")
+            if matches:
+                if not event.id:
+                    raise ValueError("A2A_NATIVE_TRIGGER_ID_REQUIRED")
+                path = _path(event)
+                return ScopeTriggerRef(
+                    scope_id=stable_scope_id("adk", event.invocation_id, path),
+                    item_id=stable_item_id(
+                        "adk", event.invocation_id, path, event.id, call_id, "tool_call"
+                    ),
+                    call_id=call_id,
+                )
+        raise ValueError("A2A_NATIVE_TRIGGER_CALL_MISSING")
+    return None

ksadk/a2a/adk_runtime.py

@@ -0,0 +1,126 @@
+"""Run-local managed A2A dependencies for the production ADK canonical stream."""
+
+import asyncio
+from contextlib import aclosing, suppress
+
+from ksadk.a2a.adk import bind_managed_a2a
+from ksadk.a2a.adk_invocation_context import ADKInvocationContextProvider
+from ksadk.events.run_authority import current_run_event_authority
+
+
+def managed_message_parts(ctx, _spec):
+    """Forward visible text context, excluding thoughts, tool payloads and state."""
+    parts = []
+    branch = getattr(ctx, "branch", None) or ""
+    for event in ctx.session.events:
+        event_branch = getattr(event, "branch", None) or ""
+        if event_branch and branch != event_branch and not branch.startswith(event_branch + "."):
+            continue
+        if getattr(event, "partial", False):
+            continue
+        text = "".join(
+            part.text
+            for part in (getattr(event.content, "parts", None) or ())
+            if part.text and not getattr(part, "thought", False)
+        )
+        if text.strip():
+            parts.append(f"[{event.author}]\n{text}")
+    if not parts:
+        parts = [
+            part.text
+            for part in (getattr(getattr(ctx, "user_content", None), "parts", None) or ())
+            if part.text and not getattr(part, "thought", False)
+        ]
+    return tuple(parts)
+
+
+async def managed_native_events(runner, native_events):
+    bootstrap = getattr(runner, "_managed_a2a_bootstrap", None)
+    authority = current_run_event_authority()
+    from ksadk.sessions import bind_session_service, resolve_session_service
+
+    session_service = resolve_session_service()
+    async with aclosing(native_events):
+        if bootstrap is None or not bootstrap.outbound_enabled or authority is None:
+            while True:
+                try:
+                    # Pull the native generator while the Kernel-owned service
+                    # is explicitly active.  Async-generator resumption may
+                    # otherwise occur in a context that lost the worker binding.
+                    with bind_session_service(session_service):
+                        event = await anext(native_events)
+                except StopAsyncIteration:
+                    break
+                yield event
+            return
+        authority.register_managed_a2a(bootstrap)
+        policy = authority.managed_invocation_policy()
+        services = {}
+
+        def service_for_space(space_id):
+            if space_id not in services:
+                services[space_id] = bootstrap.invocation_service_for_space(
+                    space_id, event_store=authority, policy=policy
+                )
+            return services[space_id]
+
+        provider = ADKInvocationContextProvider(
+            authority=authority,
+            service_for_space=service_for_space,
+            caller_agent_id=bootstrap.runtime_metadata.agent_id,
+            message_selector=managed_message_parts,
+        )
+        from ksadk.a2a.adk_tools import (
+            ManagedADKTools,
+            bind_managed_adk_tools,
+            install_managed_tools,
+        )
+
+        queue = asyncio.Queue(maxsize=1)
+        # The native ADK iterator is consumed in a child task.  Capture the
+        # Kernel-bound service before that task is created and bind it explicitly
+        # inside the pump; relying on an implicit ContextVar copy is unsafe when
+        # async-generator finalization moves across task contexts.
+        async def emit(event):
+            acknowledged = asyncio.Event()
+            await queue.put((event, acknowledged))
+            await acknowledged.wait()
+
+        async def pump():
+            with bind_session_service(session_service):
+                try:
+                    async for event in native_events:
+                        await emit(event)
+                except asyncio.CancelledError as exc:
+                    if asyncio.current_task().cancelling():
+                        raise
+                    await queue.put((exc, None))
+                except Exception as exc:
+                    await queue.put((exc, None))
+                else:
+                    await queue.put((None, None))
+
+        install_managed_tools(runner._agent)
+        runtime = ManagedADKTools(
+            bootstrap=bootstrap,
+            authority=authority,
+            services=service_for_space,
+            emit=emit,
+        )
+        with bind_managed_a2a(None, provider), bind_managed_adk_tools(runtime):
+            task = asyncio.create_task(pump())
+            try:
+                while True:
+                    event, acknowledged = await queue.get()
+                    if event is None:
+                        break
+                    if isinstance(event, BaseException):
+                        raise event
+                    try:
+                        yield event
+                    finally:
+                        acknowledged.set()
+            finally:
+                task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await task