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

ksadk/a2a/langgraph_runtime.py

@@ -0,0 +1,89 @@
+"""Join LangGraph native events and managed A2A facts in one owned run stream."""
+
+import asyncio
+from contextlib import aclosing, suppress
+
+from ksadk.a2a.langgraph_invocation_context import LangGraphInvocationContextProvider
+from ksadk.a2a.langgraph_managed import bind_managed_langgraph
+from ksadk.events.run_authority import current_run_event_authority
+
+
+async def managed_langgraph_events(runner, native_events):
+    bootstrap = getattr(runner, "_managed_a2a_bootstrap", None)
+    async with aclosing(native_events):
+        if bootstrap is None or not bootstrap.outbound_enabled:
+            async for event in native_events:
+                yield event
+            return
+        authority = current_run_event_authority()
+        if authority is None:
+            raise RuntimeError("A2A_KERNEL_AUTHORITY_REQUIRED")
+        authority.register_managed_a2a(bootstrap)
+        policy = authority.managed_invocation_policy()
+        services = {}
+
+        def service_for_space(space):
+            if space not in services:
+                services[space] = bootstrap.invocation_service_for_space(
+                    space, event_store=authority, policy=policy,
+                )
+            return services[space]
+
+        from ksadk.a2a.langgraph_tool_runtime import ManagedLangGraphTools, graph_exposure_plan
+
+        runtime = ManagedLangGraphTools(
+            bootstrap=bootstrap, authority=authority, plan=graph_exposure_plan(runner._agent),
+        )
+        provider = LangGraphInvocationContextProvider(
+            authority=authority, service_for_space=service_for_space,
+            caller_agent_id=bootstrap.runtime_metadata.agent_id,
+            tool_trigger=runtime.tool_trigger,
+        )
+        from ksadk.sessions import bind_session_service, resolve_session_service
+
+        session_service = resolve_session_service()
+        queue = asyncio.Queue(maxsize=1)
+
+        async def emit(event):
+            from ksadk.a2a.langgraph_recovery import managed_checkpoint
+
+            checkpoint = await managed_checkpoint(runner, authority, event)
+            if checkpoint is not None:
+                checkpoint_ack = asyncio.Event()
+                await queue.put((checkpoint, checkpoint_ack))
+                await checkpoint_ack.wait()
+            acknowledged = asyncio.Event()
+            await queue.put((event, acknowledged))
+            await acknowledged.wait()
+
+        async def pump():
+            with bind_session_service(session_service), bind_managed_langgraph(None, provider, emit, runtime=runtime):
+                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))
+
+        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()
+                    await runtime.notify_event()
+        finally:
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task

ksadk/a2a/langgraph_tool_runtime.py

@@ -0,0 +1,132 @@
+"""Managed discovery and dynamic calls tied to real LangGraph tool execution."""
+
+import asyncio
+
+from ksadk.a2a.exposure import compile_exposure_plan
+from ksadk.a2a.invocation_types import RemoteAgentSpec
+from ksadk.a2a.langgraph_managed import _invoke
+from ksadk.a2a.langgraph_tools import native_tool_call_id
+from ksadk.a2a.presentation import ScopeTriggerRef
+from ksadk.events.content import ToolCallContent
+from ksadk.events.identity import stable_item_id
+
+
+def graph_exposure_plan(graph):
+    pending, seen, specs = [graph], set(), {}
+    while pending:
+        value = pending.pop()
+        if id(value) in seen:
+            continue
+        seen.add(id(value))
+        spec = getattr(value, "remote_spec", None)
+        metadata = getattr(value, "metadata", None)
+        if spec is None and isinstance(metadata, dict) and metadata.get("ksadk_a2a_remote_spec"):
+            spec = RemoteAgentSpec.model_validate(metadata["ksadk_a2a_remote_spec"])
+        if spec is not None:
+            previous = specs.get(spec.binding_id)
+            if previous is not None and previous != spec:
+                raise ValueError("A2A_BINDING_ID_CONFLICT")
+            specs[spec.binding_id] = spec
+        for attr in ("bound", "func", "afunc", "__self__"):
+            child = getattr(value, attr, None)
+            if child is not None:
+                pending.append(child)
+        for attr in ("nodes", "tools_by_name"):
+            children = getattr(value, attr, None)
+            if isinstance(children, dict):
+                pending.extend(children.values())
+        children = getattr(value, "steps", None)
+        if isinstance(children, (list, tuple)):
+            pending.extend(children)
+    return compile_exposure_plan(list(specs.values()), framework="langgraph", consumer_id="graph")
+
+
+class ManagedLangGraphTools:
+    def __init__(self, *, bootstrap, authority, plan):
+        self.bootstrap, self.authority, self.plan = bootstrap, authority, plan
+        self.changed = asyncio.Condition()
+
+    async def notify_event(self):
+        async with self.changed:
+            self.changed.notify_all()
+
+    async def discover(self, agent_ref=None):
+        from ksadk.toolsets.a2a import _a2a_space_id, _normalize_agent_ref
+
+        self.authority.require_session(self.authority.session_id)
+        space = _a2a_space_id()
+        if not space:
+            raise ValueError("A2A_SPACE_REQUIRED")
+        discovered = await self.bootstrap.client_for_space(space).discover()
+        agents = [{"agent_id": a.agent_id, "name": a.agent_card.name,
+                   "description": getattr(a.agent_card, "description", ""),
+                   "version_id": a.version_id} for a in discovered]
+        if agent_ref is None:
+            return {"ok": True, "space_id": space, "agents": agents}
+        target = _normalize_agent_ref(str(agent_ref))
+        matches = [a for a in agents if target and target in {
+            _normalize_agent_ref(a["agent_id"]), _normalize_agent_ref(a["name"])}]
+        if len(matches) != 1:
+            raise ValueError("A2A_AGENT_NOT_FOUND" if not matches else "A2A_AGENT_AMBIGUOUS")
+        return {"ok": True, "space_id": space, "agent": matches[0]}
+
+    async def call(self, args, config):
+        try:
+            native_tool_call_id()
+            message = args.get("message")
+            if not isinstance(message, str) or not message.strip():
+                raise ValueError("A2A_MESSAGE_REQUIRED")
+            if not isinstance(args.get("agent"), str) or not args["agent"].strip():
+                raise ValueError("A2A_AGENT_REQUIRED")
+            found = await self.discover(args.get("agent"))
+            agent, space = found["agent"], found["space_id"]
+            self.plan.assert_generic_target_allowed(agent["agent_id"])
+            spec = RemoteAgentSpec(
+                local_name="call_a2a_agent",
+                binding_id=stable_item_id("a2a", space, agent["agent_id"]),
+                space_id=space, target_agent_id=agent["agent_id"],
+                expected_version_id=agent["version_id"], exposure="tool",
+            )
+            return await _invoke(spec, {"message": message}, config,
+                                 surface="tool", result_envelope=True)
+        except Exception as exc:
+            code = str(exc)
+            if not code.startswith("A2A_") or len(code) >= 128:
+                code = "A2A_INVOCATION_FAILED"
+            return {"ok": False, "error_type": code, "error_message": "Managed A2A invocation failed"}
+
+    async def tool_trigger(self, spec, config):
+        call_id = native_tool_call_id()
+        namespace = config["metadata"]["langgraph_checkpoint_ns"]
+        components = namespace.split("|") if namespace else []
+        async def resolve():
+            async with self.changed:
+                while True:
+                    facts = await self.authority.list(self.authority.session_id)
+                    # v3 message facts carry the graph namespace, while the
+                    # running ToolNode config adds its own node:task segment.
+                    # Some supported ToolNodes emit no tools-channel events;
+                    # the real ainvoke ToolCall envelope plus the persisted
+                    # model call is the proof, not a fabricated tool-start.
+                    calls = {}
+                    for event in facts:
+                        if event.source.framework != "langgraph":
+                            continue
+                        if event.source.metadata.get("namespace") != components[:-1]:
+                            continue
+                        snapshot = getattr(event, "snapshot", None) or getattr(event, "initial", None)
+                        if snapshot and any(isinstance(p, ToolCallContent)
+                                            and p.call_id == call_id and p.name == spec.local_name
+                                            for p in snapshot.parts):
+                            calls[(event.scope_id, event.item_id)] = event
+                    if len(calls) == 1:
+                        event = next(iter(calls.values()))
+                        return ScopeTriggerRef(scope_id=event.scope_id,
+                                               item_id=event.item_id, call_id=call_id)
+                    if len(calls) > 1:
+                        raise ValueError("A2A_NATIVE_TRIGGER_AMBIGUOUS")
+                    await self.changed.wait()
+        try:
+            return await asyncio.wait_for(resolve(), timeout=5)
+        except TimeoutError as exc:
+            raise ValueError("A2A_NATIVE_TRIGGER_NOT_PERSISTED") from exc

ksadk/a2a/langgraph_tools.py

@@ -0,0 +1,63 @@
+"""Preserve native LangChain tool identity while routing managed A2A calls."""
+
+from contextvars import ContextVar
+
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool, StructuredTool
+
+_native_call: ContextVar[str | None] = ContextVar("langgraph_a2a_native_tool_call", default=None)
+
+
+def native_tool_call_id():
+    call_id = _native_call.get()
+    if not call_id:
+        raise ValueError("A2A_NATIVE_TOOL_TRIGGER_REQUIRED")
+    return call_id
+
+
+class ManagedLangGraphTool(StructuredTool):
+    async def ainvoke(self, input, config=None, **kwargs):
+        call_id = input.get("id") if isinstance(input, dict) and input.get("type") == "tool_call" else None
+        token = _native_call.set(call_id)
+        try:
+            # StructuredTool sends sync-only discovery straight to invoke in
+            # an executor, bypassing our managed _arun. BaseTool keeps the
+            # async dispatch; the fallback _arun still offloads sync work.
+            return await BaseTool.ainvoke(self, input, config, **kwargs)
+        finally:
+            _native_call.reset(token)
+
+    async def _arun(self, *args, config: RunnableConfig, run_manager=None, **kwargs):
+        from ksadk.a2a.langgraph_managed import _binding
+        from ksadk.toolsets.a2a import call_a2a_agent, get_a2a_agent_card, list_a2a_agents
+
+        binding = _binding.get()
+        runtime = binding.runtime if binding else None
+        function = self.coroutine or self.func
+        if runtime is not None:
+            if function is call_a2a_agent:
+                return await runtime.call(kwargs, config)
+            if function in (list_a2a_agents, get_a2a_agent_card):
+                try:
+                    if function is get_a2a_agent_card:
+                        agent = kwargs.get("agent")
+                        if not isinstance(agent, str) or not agent.strip():
+                            raise ValueError("A2A_AGENT_REQUIRED")
+                        return await runtime.discover(agent)
+                    return await runtime.discover()
+                except Exception as exc:
+                    code = str(exc)
+                    if not code.startswith("A2A_") or len(code) >= 128:
+                        code = "A2A_DISCOVERY_FAILED"
+                    return {"ok": False, "error_type": code,
+                            "error_message": "Managed A2A discovery failed"}
+        return await super()._arun(*args, config=config, run_manager=run_manager, **kwargs)
+
+
+def managed_tool(func):
+    import inspect
+
+    return ManagedLangGraphTool.from_function(
+        func=None if inspect.iscoroutinefunction(func) else func,
+        coroutine=func if inspect.iscoroutinefunction(func) else None,
+    )

ksadk/a2a/message_digest.py

@@ -0,0 +1,40 @@
+"""Canonical A2A Message digest shared by request and control-plane contracts."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+from google.protobuf.json_format import MessageToDict
+
+MAX_A2A_MESSAGE_BYTES = 1024 * 1024
+MAX_A2A_MESSAGE_PARTS = 64
+MAX_A2A_MESSAGE_ID_LENGTH = 128
+
+
+def canonical_message_bytes(message: Any) -> bytes:
+    """Serialize a pinned A2A protobuf using canonical lowerCamelCase JSON."""
+
+    payload = MessageToDict(message, preserving_proto_field_name=False)
+    return json.dumps(
+        payload,
+        ensure_ascii=False,
+        sort_keys=True,
+        separators=(",", ":"),
+    ).encode("utf-8")
+
+
+def canonical_message_sha256(message: Any) -> str:
+    """Hash the logical Message, not transport or permit metadata."""
+
+    return hashlib.sha256(canonical_message_bytes(message)).hexdigest()
+
+
+__all__ = [
+    "MAX_A2A_MESSAGE_BYTES",
+    "MAX_A2A_MESSAGE_ID_LENGTH",
+    "MAX_A2A_MESSAGE_PARTS",
+    "canonical_message_bytes",
+    "canonical_message_sha256",
+]