agentengine-sdk-python 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

ksadk/a2a/identity.py

@@ -0,0 +1,304 @@
+"""Trusted Gateway identity boundary for inbound A2A protocol requests."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from ksadk.a2a.context_store import A2AContextIdentity
+from ksadk.a2a.ids import require_a2a_resource_id
+
+_CALLER_PRINCIPAL_TYPES = frozenset({"user", "runtime", "service"})
+_AUTHN_MODE_BY_CALLER_TYPE = {
+    "user": "api_key",
+    "runtime": "workload_permit",
+    "service": "gateway_service",
+}
+
+
+@dataclass(frozen=True)
+class A2AIngressTargetBinding:
+    """The Gateway-verified destination to which one Runtime is bound."""
+
+    account_id: str
+    tenant_id: str
+    agent_id: str
+    runtime_id: str
+    a2a_agent_id: str
+
+
+@dataclass(frozen=True)
+class A2AIngressIdentity:
+    """Caller identity verified by Gateway mTLS and target-bound forwarding."""
+
+    account_id: str
+    tenant_id: str
+    caller_principal_type: str
+    caller_principal_id: str
+    target_agent_id: str
+    target_runtime_id: str
+    target_a2a_agent_id: str
+    authn_mode: str
+    caller_runtime_id: str | None = None
+
+    def validate(self) -> None:
+        values = {
+            "account_id": self.account_id,
+            "caller_principal_type": self.caller_principal_type,
+            "caller_principal_id": self.caller_principal_id,
+            "target_agent_id": self.target_agent_id,
+            "target_runtime_id": self.target_runtime_id,
+            "target_a2a_agent_id": self.target_a2a_agent_id,
+            "authn_mode": self.authn_mode,
+        }
+        missing = [name for name, value in values.items() if not str(value).strip()]
+        if missing:
+            raise PermissionError(
+                "verified Gateway identity is missing " + ", ".join(missing)
+            )
+        if self.caller_principal_type not in _CALLER_PRINCIPAL_TYPES:
+            raise PermissionError("verified Gateway identity has an unsupported caller type")
+        if self.authn_mode != _AUTHN_MODE_BY_CALLER_TYPE[self.caller_principal_type]:
+            raise PermissionError("verified Gateway identity has an invalid authentication mode")
+        if self.caller_principal_type == "runtime" and not str(
+            self.caller_runtime_id or ""
+        ).strip():
+            raise PermissionError("verified runtime caller identity is missing caller runtime")
+        if self.caller_principal_type != "runtime" and self.caller_runtime_id is not None:
+            raise PermissionError("verified non-runtime caller identity includes a caller runtime")
+        for field_name, limit in (
+            ("account_id", 64),
+            ("tenant_id", 64),
+            ("caller_principal_type", 16),
+            ("caller_principal_id", 128),
+            ("target_agent_id", 64),
+            ("target_runtime_id", 64),
+            ("target_a2a_agent_id", 64),
+            ("authn_mode", 32),
+            ("caller_runtime_id", 64),
+        ):
+            value = getattr(self, field_name)
+            if value is not None and len(str(value)) > limit:
+                raise PermissionError(
+                    f"verified Gateway identity field {field_name} exceeds {limit} characters"
+                )
+        if not self.target_agent_id.startswith("ar-"):
+            raise PermissionError("verified Gateway identity has an invalid target Agent")
+        try:
+            require_a2a_resource_id(
+                self.target_a2a_agent_id,
+                "a2a-agent-",
+                field_name="verified Gateway identity target A2A Agent",
+            )
+        except ValueError as exc:
+            raise PermissionError(
+                "verified Gateway identity has an invalid target A2A Agent"
+            ) from exc
+
+    def context_identity(self) -> A2AContextIdentity:
+        return A2AContextIdentity(
+            account_id=self.account_id,
+            tenant_id=self.tenant_id,
+            caller_principal_type=self.caller_principal_type,
+            caller_principal_id=self.caller_principal_id,
+        )
+
+    def owner_key(self) -> str:
+        return "/".join(
+            (
+                self.account_id.strip(),
+                self.tenant_id.strip(),
+                self.caller_principal_type.strip(),
+                self.caller_principal_id.strip(),
+            )
+        )
+
+
+class GatewayIdentityVerifier(ABC):
+    """Platform adapter that verifies Gateway mTLS and target-bound forwarding."""
+
+    @abstractmethod
+    async def verify(self, request: Request) -> A2AIngressIdentity:
+        raise NotImplementedError
+
+
+class GatewayProbeVerifier(ABC):
+    """Verifies a trusted Gateway/server probe for the Runtime-local AgentCard."""
+
+    @abstractmethod
+    async def verify_probe(self, request: Request) -> None:
+        raise NotImplementedError
+
+
+class CallableGatewayIdentityVerifier(GatewayIdentityVerifier):
+    """Small adapter for product injectors and tests."""
+
+    def __init__(
+        self,
+        verify: Callable[[Request], Awaitable[A2AIngressIdentity]],
+    ) -> None:
+        self._verify = verify
+
+    async def verify(self, request: Request) -> A2AIngressIdentity:
+        return await self._verify(request)
+
+
+class CallableGatewayProbeVerifier(GatewayProbeVerifier):
+    """Small probe-verifier adapter for product injectors and tests."""
+
+    def __init__(self, verify_probe: Callable[[Request], Awaitable[None]]) -> None:
+        self._verify_probe = verify_probe
+
+    async def verify_probe(self, request: Request) -> None:
+        await self._verify_probe(request)
+
+
+class A2ATrustedIdentityResolver:
+    """Reads only identity written by verified ingress middleware."""
+
+    state_key = "a2a_identity"
+
+    def __init__(
+        self,
+        *,
+        expected_account_id: str | None = None,
+        expected_tenant_id: str | None = None,
+        expected_target_agent_id: str | None = None,
+        expected_target_runtime_id: str | None = None,
+        expected_target_a2a_agent_id: str | None = None,
+        target_binding: A2AIngressTargetBinding | None = None,
+    ) -> None:
+        if target_binding is not None:
+            if any(
+                value is not None
+                for value in (
+                    expected_account_id,
+                    expected_tenant_id,
+                    expected_target_agent_id,
+                    expected_target_runtime_id,
+                    expected_target_a2a_agent_id,
+                )
+            ):
+                raise ValueError(
+                    "pass target_binding or individual expected target fields, not both"
+                )
+            expected_account_id = target_binding.account_id
+            expected_tenant_id = target_binding.tenant_id
+            expected_target_agent_id = target_binding.agent_id
+            expected_target_runtime_id = target_binding.runtime_id
+            expected_target_a2a_agent_id = target_binding.a2a_agent_id
+        self._expected_account_id = str(expected_account_id or "").strip() or None
+        self._expected_tenant_id = expected_tenant_id
+        self._expected_target_agent_id = str(expected_target_agent_id or "").strip() or None
+        self._expected_target_runtime_id = str(expected_target_runtime_id or "").strip() or None
+        self._expected_target_a2a_agent_id = str(expected_target_a2a_agent_id or "").strip() or None
+
+    def resolve(self, request: Request) -> A2AIngressIdentity:
+        raw_state = request.scope.get("state")
+        state = raw_state if isinstance(raw_state, dict) else {}
+        identity = state.get(self.state_key)
+        if not isinstance(identity, A2AIngressIdentity):
+            raise PermissionError("verified Gateway identity is required for inbound A2A")
+        self._validate(identity)
+        return identity
+
+    def _validate(self, identity: A2AIngressIdentity) -> None:
+        identity.validate()
+        identity.context_identity().canonical_owner()
+        if (
+            self._expected_account_id is not None
+            and identity.account_id != self._expected_account_id
+        ):
+            raise PermissionError("verified Gateway identity targets a different account")
+        if (
+            self._expected_tenant_id is not None
+            and identity.tenant_id != self._expected_tenant_id
+        ):
+            raise PermissionError("verified Gateway identity targets a different tenant")
+        if (
+            self._expected_target_agent_id is not None
+            and identity.target_agent_id != self._expected_target_agent_id
+        ):
+            raise PermissionError("verified Gateway identity targets a different Agent")
+        if (
+            self._expected_target_runtime_id is not None
+            and identity.target_runtime_id != self._expected_target_runtime_id
+        ):
+            raise PermissionError("verified Gateway identity targets a different Runtime")
+        if (
+            self._expected_target_a2a_agent_id is not None
+            and identity.target_a2a_agent_id != self._expected_target_a2a_agent_id
+        ):
+            raise PermissionError("verified Gateway identity targets a different A2A Agent")
+
+
+class A2AGatewayIdentityMiddleware(BaseHTTPMiddleware):
+    """Verify Gateway identity before JSON-RPC or HTTP+JSON A2A handlers run."""
+
+    def __init__(
+        self,
+        app: Any,
+        *,
+        verifier: GatewayIdentityVerifier,
+        probe_verifier: GatewayProbeVerifier | None = None,
+        expected_account_id: str | None = None,
+        expected_tenant_id: str | None = None,
+        expected_target_agent_id: str | None = None,
+        expected_target_runtime_id: str | None = None,
+        expected_target_a2a_agent_id: str | None = None,
+        target_binding: A2AIngressTargetBinding | None = None,
+    ) -> None:
+        super().__init__(app)
+        self._verifier = verifier
+        self._probe_verifier = probe_verifier
+        self._identity_resolver = A2ATrustedIdentityResolver(
+            expected_account_id=expected_account_id,
+            expected_tenant_id=expected_tenant_id,
+            expected_target_agent_id=expected_target_agent_id,
+            expected_target_runtime_id=expected_target_runtime_id,
+            expected_target_a2a_agent_id=expected_target_a2a_agent_id,
+            target_binding=target_binding,
+        )
+
+    async def dispatch(self, request: Request, call_next: Any):
+        if request.url.path == "/.well-known/agent-card.json":
+            try:
+                if self._probe_verifier is None:
+                    raise PermissionError("A2A AgentCard probe verifier is not configured")
+                await self._probe_verifier.verify_probe(request)
+            except Exception:
+                return JSONResponse(
+                    status_code=401,
+                    content={"detail": "trusted Gateway identity is required for A2A AgentCard"},
+                )
+        elif request.url.path == "/a2a/jsonrpc" or request.url.path.startswith("/a2a/v1/"):
+            try:
+                identity = await self._verifier.verify(request)
+                self._identity_resolver._validate(identity)
+            except Exception:
+                return JSONResponse(
+                    status_code=401,
+                    content={"detail": "verified Gateway identity is required for inbound A2A"},
+                )
+            raw_state = request.scope.get("state")
+            state = raw_state if isinstance(raw_state, dict) else {}
+            state[A2ATrustedIdentityResolver.state_key] = identity
+            request.scope["state"] = state
+        return await call_next(request)
+
+
+__all__ = [
+    "A2AGatewayIdentityMiddleware",
+    "A2AIngressTargetBinding",
+    "A2AIngressIdentity",
+    "A2ATrustedIdentityResolver",
+    "CallableGatewayIdentityVerifier",
+    "CallableGatewayProbeVerifier",
+    "GatewayIdentityVerifier",
+    "GatewayProbeVerifier",
+]

ksadk/a2a/ids.py

@@ -0,0 +1,22 @@
+"""AgentEngine A2A resource ID validation."""
+
+from __future__ import annotations
+
+import re
+
+_UUID4_HEX = re.compile(r"^[0-9a-f]{12}4[0-9a-f]{3}[89ab][0-9a-f]{15}$")
+
+
+def is_a2a_resource_id(value: str, prefix: str) -> bool:
+    """Return whether value is `<prefix><uuid4_hex>` with a lowercase compact UUID."""
+    return value.startswith(prefix) and bool(_UUID4_HEX.fullmatch(value.removeprefix(prefix)))
+
+
+def require_a2a_resource_id(value: str, prefix: str, *, field_name: str) -> str:
+    """Validate an AgentEngine A2A resource ID and return it unchanged."""
+    if not is_a2a_resource_id(value, prefix):
+        raise ValueError(f"{field_name} must use the {prefix}<uuid4_hex> format")
+    return value
+
+
+__all__ = ["is_a2a_resource_id", "require_a2a_resource_id"]

ksadk/a2a/langgraph.py

@@ -0,0 +1,285 @@
+"""LangGraph 编排方调远端 A2A agent 的流式 helper。
+
+ADK 有原生 ``RemoteA2aAgent``;LangGraph 没有等价封装。这里提供结构化流式
+helper,保留远端 ``thinking`` / ``text`` 类型,并可直接写入 LangGraph custom
+stream,无需 task store、无需 Space 绑定——直接走 a2a-sdk 原生 client。
+"""
+
+from __future__ import annotations
+
+import inspect
+import uuid
+from collections.abc import AsyncIterator
+from typing import Any, Callable, Literal, TypedDict
+
+import httpx
+from a2a.client import ClientCallContext, ClientConfig, create_client
+from a2a.client.card_resolver import A2ACardResolver
+from a2a.types import Message, Part, Role, SendMessageRequest
+
+
+class A2AStreamEvent(TypedDict):
+    """可直接交给 LangGraph custom writer 的远端流式事件。"""
+
+    type: Literal["thinking", "text"]
+    delta: str
+    replace: bool
+
+
+def _base_url(agent_card_url: str) -> str:
+    """从 agent_card_url 取 base_url(剥掉 well-known 路径后缀)。"""
+    url = agent_card_url.strip().rstrip("/")
+    for suffix in ("/.well-known/agent-card.json", "/.well-known/agent.json"):
+        if url.endswith(suffix):
+            url = url[: -len(suffix)]
+    return url
+
+
+async def stream_a2a_agent_events(
+    agent_card_url: str,
+    message: str,
+    *,
+    httpx_client: httpx.AsyncClient | None = None,
+    timeout: float = 60.0,
+) -> AsyncIterator[A2AStreamEvent]:
+    """流式调用远端 A2A agent,保留 reasoning 与正文的结构化类型。
+
+    在 langgraph node 里配合 custom stream 透传::
+
+        from langgraph.config import get_stream_writer
+        from ksadk.a2a.langgraph import stream_a2a_agent_events
+
+        async def call_demo(state, config):
+            writer = get_stream_writer()
+            collected = ""
+            async for event in stream_a2a_agent_events(
+                "http://127.0.0.1:8094", state["query"]
+            ):
+                writer(event)
+                if event["type"] == "text":
+                    collected += event["delta"]
+            return {"result": collected}
+
+    Args:
+        agent_card_url: 远端 agent 的 base URL(如 ``http://127.0.0.1:8094``)
+            或完整 agent-card.json URL。
+        message: 发给远端 agent 的用户消息。
+        httpx_client: 可选复用的 ``httpx.AsyncClient``;不传则内部创建。
+        timeout: httpx 请求超时秒数。
+
+    Yields:
+        ``{"type": "thinking"|"text", "delta": str, "replace": bool}``。
+        reasoning 由 ``adk_thought`` Part metadata 识别,status_update 不伪装
+        成思考；``replace`` 标识权威快照替换既有同类输出。
+    """
+    own_client = httpx_client is None
+    client_httpx = httpx_client or httpx.AsyncClient(timeout=timeout)
+    try:
+        resolver = A2ACardResolver(httpx_client=client_httpx, base_url=_base_url(agent_card_url))
+        card = await resolver.get_agent_card()
+        client = await create_client(
+            agent=card,
+            client_config=ClientConfig(httpx_client=client_httpx, streaming=True),
+        )
+        request = SendMessageRequest(
+            message=Message(
+                role=Role.ROLE_USER,
+                parts=[Part(text=message)],
+                message_id=f"lg-a2a-{uuid.uuid4().hex}",
+            )
+        )
+        snapshots: dict[tuple[str, str], str] = {}
+        async for event in client.send_message(request, context=ClientCallContext()):
+            for stream_event in _extract_artifact_events(event, snapshots):
+                yield stream_event
+    finally:
+        if own_client:
+            await client_httpx.aclose()
+
+
+async def stream_a2a_agent(
+    agent_card_url: str,
+    message: str,
+    *,
+    httpx_client: httpx.AsyncClient | None = None,
+    timeout: float = 60.0,
+) -> AsyncIterator[str]:
+    """兼容旧接口:只产出正文文本增量,reasoning 不混入最终答案。
+
+    该字符串流无法表达 ``replace`` 语义。需要正确处理远端权威快照的调用方
+    应使用 ``stream_a2a_agent_events`` 或 ``stream_a2a_agent_to_writer``。
+    """
+    async for event in stream_a2a_agent_events(
+        agent_card_url,
+        message,
+        httpx_client=httpx_client,
+        timeout=timeout,
+    ):
+        if event["type"] == "text":
+            yield event["delta"]
+
+
+async def stream_a2a_agent_to_writer(
+    agent_card_url: str,
+    message: str,
+    *,
+    writer: Callable[[A2AStreamEvent], Any] | None = None,
+    httpx_client: httpx.AsyncClient | None = None,
+    timeout: float = 60.0,
+) -> str:
+    """将远端 typed events 写入 LangGraph custom stream,返回纯正文结果。
+
+    ``writer`` 不传时自动使用当前 LangGraph node 的 ``get_stream_writer()``。
+    demo 只需调用本函数,无需了解 A2A wire artifact、append 或快照去重。
+    """
+    if writer is None:
+        from langgraph.config import get_stream_writer
+
+        writer = get_stream_writer()
+
+    output: list[str] = []
+    async for event in stream_a2a_agent_events(
+        agent_card_url,
+        message,
+        httpx_client=httpx_client,
+        timeout=timeout,
+    ):
+        write_result = writer(event)
+        if inspect.isawaitable(write_result):
+            await write_result
+        if event["type"] == "text":
+            if event["replace"]:
+                output = [event["delta"]]
+            else:
+                output.append(event["delta"])
+    return "".join(output)
+
+
+def _extract_artifact_events(
+    event: Any,
+    snapshots: dict[tuple[str, str], str],
+) -> list[A2AStreamEvent]:
+    """从 A2A artifact_update 提取去重后的 typed events。
+
+    快照按 ``(artifact_id, type)`` 隔离,避免 reasoning 与 response 互相覆盖。
+    ``append=False`` 按累计快照 diff,``append=True`` 按增量直接透传。
+    """
+    which = None
+    try:
+        which = event.WhichOneof("payload")
+    except Exception:  # 非 protobuf 或结构不符时静默跳过
+        which = None
+    if which == "artifact_update":
+        artifact_update = getattr(event, "artifact_update", None)
+        artifact = getattr(artifact_update, "artifact", None)
+        if artifact is None:
+            return []
+        return _events_from_artifact(
+            artifact,
+            append=bool(getattr(artifact_update, "append", False)),
+            snapshots=snapshots,
+        )
+    if which == "task":
+        result: list[A2AStreamEvent] = []
+        task = getattr(event, "task", None)
+        for artifact in getattr(task, "artifacts", []):
+            result.extend(_events_from_artifact(artifact, append=False, snapshots=snapshots))
+        return result
+    if which == "message":
+        message = getattr(event, "message", None)
+        if message is None:
+            return []
+        artifact = _MessageArtifact(message)
+        return _events_from_artifact(artifact, append=False, snapshots=snapshots)
+    return []
+
+
+class _MessageArtifact:
+    """把直接 Message 响应投影成 artifact-like 对象复用解析逻辑。"""
+
+    def __init__(self, message: Any) -> None:
+        self.artifact_id = f"message:{getattr(message, 'message_id', '') or 'anonymous'}"
+        self.name = "response"
+        self.parts = getattr(message, "parts", [])
+        self.metadata = getattr(message, "metadata", None)
+
+
+def _events_from_artifact(
+    artifact: Any,
+    *,
+    append: bool,
+    snapshots: dict[tuple[str, str], str],
+) -> list[A2AStreamEvent]:
+    """解析一个 artifact-like 对象，支持增量、累计快照和权威替换。"""
+
+    artifact_id = str(
+        getattr(artifact, "artifact_id", "")
+        or getattr(artifact, "name", "")
+        or "anonymous"
+    )
+    artifact_is_reasoning = (
+        str(getattr(artifact, "name", "")).lower() == "reasoning"
+        or _metadata_flag(getattr(artifact, "metadata", None), "adk_thought")
+    )
+    grouped: list[tuple[str, str, bool]] = []
+    for part in getattr(artifact, "parts", []):
+        text = str(getattr(part, "text", "") or "")
+        if not text:
+            continue
+        event_type = (
+            "thinking"
+            if artifact_is_reasoning
+            or _metadata_flag(getattr(part, "metadata", None), "adk_thought")
+            else "text"
+        )
+        replaces_output = _metadata_flag(
+            getattr(part, "metadata", None),
+            "ksadk_output_snapshot",
+        )
+        if grouped and grouped[-1][0] == event_type:
+            grouped[-1] = (
+                event_type,
+                grouped[-1][1] + text,
+                grouped[-1][2] or replaces_output,
+            )
+        else:
+            grouped.append((event_type, text, replaces_output))
+
+    result: list[A2AStreamEvent] = []
+    for event_type, text, replaces_output in grouped:
+        key = (artifact_id, event_type)
+        snapshot = snapshots.get(key, "")
+        if append:
+            delta = text
+            snapshots[key] = snapshot + text
+            replace = False
+        else:
+            delta = text[len(snapshot) :] if text.startswith(snapshot) else text
+            snapshots[key] = text
+            replace = replaces_output or bool(snapshot and not text.startswith(snapshot))
+        if delta:
+            result.append(
+                {
+                    "type": event_type,  # type: ignore[typeddict-item]
+                    "delta": delta,
+                    "replace": replace,
+                }
+            )
+    return result
+
+
+def _metadata_flag(metadata: Any, key: str) -> bool:
+    if metadata is None:
+        return False
+    try:
+        return bool(metadata[key])
+    except (KeyError, TypeError, ValueError):
+        return False
+
+
+__all__ = [
+    "A2AStreamEvent",
+    "stream_a2a_agent",
+    "stream_a2a_agent_events",
+    "stream_a2a_agent_to_writer",
+]