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

ksadk/runners/_langgraph_state_reader.py

@@ -0,0 +1,106 @@
+"""LangGraph 图状态读取的唯一入口。
+
+审批探测（goal-18）、checkpoint 元数据、usage 提取、可恢复性校验共用同一套
+读取与 interrupt 解析逻辑；除本模块外不得直接调用 ``aget_state``/``get_state``。
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_MAX_NESTED_DEPTH = 4
+
+
+def graph_state_reader(agent: Any) -> Any | None:
+    """返回该图的 state 读取函数（async 优先），没有则 None。"""
+    return getattr(agent, "aget_state", None) or getattr(agent, "get_state", None)
+
+
+async def read_graph_state(
+    agent: Any,
+    config: Mapping[str, Any],
+    *,
+    include_subgraphs: bool = False,
+) -> tuple[Any | None, Exception | None]:
+    """统一的图状态读取入口。
+
+    返回 ``(state, error)``；读取失败时 ``error`` 非 None。调用方按语义决定后续：
+    usage / metadata / 可恢复性这类旁路读取可以降级；审批判定必须显性失败。
+    任何情况下都不静默吞掉异常——至少留一条 warning 日志。
+
+    ``include_subgraphs=True`` 会把子图状态一并取回，供 interrupt 探测遍历嵌套暂停点。
+    """
+    reader = graph_state_reader(agent)
+    if reader is None:
+        return None, None
+    try:
+        if include_subgraphs:
+            try:
+                state = reader(config, subgraphs=True)
+            except TypeError:
+                # 老版本 graph 不接受 subgraphs 关键字。
+                state = reader(config)
+        else:
+            state = reader(config)
+        if inspect.isawaitable(state):
+            state = await state
+    except Exception as exc:
+        logger.warning("LangGraph state read failed: %r", exc)
+        return None, exc
+    return state, None
+
+
+def interrupt_info_from_state(state: Any, *, _depth: int = 0) -> dict:
+    """从 state 中获取 interrupt 信息（含子图嵌套任务）"""
+    if state is None or _depth > _MAX_NESTED_DEPTH:
+        return {}
+    tasks = getattr(state, "tasks", None)
+    for task in tasks or ():
+        interrupts = getattr(task, "interrupts", None)
+        for intr in interrupts or ():
+            if hasattr(intr, "value"):
+                return _interrupt_info_from_value(intr)
+        # subgraphs=True 时子图快照挂在 task.state 上，递归下钻。
+        nested = getattr(task, "state", None)
+        if nested is not None:
+            info = interrupt_info_from_state(nested, _depth=_depth + 1)
+            if info:
+                return info
+    return {}
+
+
+def _interrupt_info_from_value(intr: Any) -> dict:
+    value = intr.value
+    info = dict(value) if isinstance(value, Mapping) else {"value": value}
+    interrupt_id = str(getattr(intr, "id", "") or "")
+    if interrupt_id:
+        info.setdefault("approval_request_id", interrupt_id)
+    # HumanInTheLoopMiddleware 的 interrupt value 把 tool_name/arguments 嵌在
+    # action_requests[0] 里；下游期望顶层字段，这里提取并保留原 action_requests。
+    action_requests = info.get("action_requests")
+    if isinstance(action_requests, list) and action_requests:
+        first = action_requests[0]
+        if isinstance(first, Mapping):
+            info.setdefault("tool_name", str(first.get("name") or ""))
+            raw_args = first.get("args") or first.get("arguments")
+            if raw_args is not None:
+                info.setdefault("arguments", raw_args)
+            if first.get("description") is not None:
+                info.setdefault("description", str(first.get("description")))
+    return info
+
+
+def known_subgraph_namespaces(agent: Any) -> set[str]:
+    """当前图里真实存在的子图 namespace 集合（无法探测时返回空集）。"""
+    get_subgraphs = getattr(agent, "get_subgraphs", None)
+    if not callable(get_subgraphs):
+        return set()
+    try:
+        return {str(ns) for ns, _ in get_subgraphs()}
+    except Exception:
+        return set()

ksadk/runners/_session_identity.py

@@ -0,0 +1,285 @@
+"""Framework-native session bindings for verified platform identities."""
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+from collections.abc import Mapping
+from typing import Any
+
+from ksadk.runtime_context import get_current_invocation_context
+from ksadk.sessions import resolve_session_service
+from ksadk.sessions.continuity import ConversationSessionCore
+from ksadk.sessions.invocation_identity import (
+    identity_native_user_id,
+    identity_scope_ref,
+    identity_session_is_legacy,
+)
+
+
+class ADKSessionIdentityMixin:
+    """Bind canonical sessions to ADK's native ``user_id`` and session id."""
+
+    def _invocation_session_owner(self) -> tuple[str, str]:
+        context = get_current_invocation_context()
+        if context is None or context.identity.is_empty:
+            return "ksadk_user", ""
+        if not context.identity.is_complete:
+            raise ValueError("Incomplete trusted business identity context")
+        return (
+            str(context.user_id or identity_native_user_id(context.identity)),
+            identity_scope_ref(context.identity),
+        )
+
+    @staticmethod
+    def _invocation_uses_legacy_session(native_user_id: str) -> bool:
+        context = get_current_invocation_context()
+        if context is None or context.identity.is_empty or not context.identity.is_complete:
+            return False
+        return identity_session_is_legacy(context.identity, native_user_id)
+
+    @staticmethod
+    def _session_cache_key(external_session_id: str, owner_scope_ref: str) -> str:
+        return (
+            f"{owner_scope_ref}:{external_session_id}" if owner_scope_ref else external_session_id
+        )
+
+    async def _get_adk_session(self, *, user_id: str, session_id: str) -> Any | None:
+        if self._session_service is None or not session_id:
+            return None
+        try:
+            return await self._session_service.get_session(
+                app_name=self._agent.name,
+                user_id=user_id,
+                session_id=session_id,
+            )
+        except Exception:
+            return None
+
+    async def _create_adk_session(
+        self,
+        *,
+        user_id: str,
+        session_id: str = "",
+        preserve_requested_id: bool = False,
+    ) -> Any:
+        if self._short_term_memory:
+            session = await self._short_term_memory.create_session(
+                app_name=self._agent.name,
+                user_id=user_id,
+                session_id=session_id,
+            )
+        else:
+            if self._session_service is None:
+                raise RuntimeError("ADK session service is not initialized")
+            create_kwargs: dict[str, Any] = {
+                "app_name": self._agent.name,
+                "user_id": user_id,
+            }
+            if session_id and preserve_requested_id:
+                create_kwargs["session_id"] = session_id
+            try:
+                session = await self._session_service.create_session(**create_kwargs)
+            except TypeError:
+                # Older ADK in-memory services did not accept session_id.
+                create_kwargs.pop("session_id", None)
+                session = await self._session_service.create_session(**create_kwargs)
+        if session is None:
+            raise RuntimeError("Failed to create ADK session")
+        return session
+
+    async def _ensure_session(self, external_session_id: str | None = None) -> str:
+        native_user_id, owner_scope_ref = self._invocation_session_owner()
+        if external_session_id:
+            cache_key = self._session_cache_key(external_session_id, owner_scope_ref)
+            if cache_key in self._session_map:
+                return self._session_map[cache_key]
+
+            core = ConversationSessionCore(resolve_session_service())
+            binding = await core.get_binding_by_session_id(external_session_id, "adk")
+            bound_scope_ref = str(binding.get("owner_scope_ref") or "")
+            if bound_scope_ref and bound_scope_ref != owner_scope_ref:
+                raise ValueError("ADK session belongs to a different business identity")
+
+            internal_session_id = str(binding.get("internal_session_id") or "")
+            bound_user_id = str(binding.get("native_user_id") or "")
+            if binding:
+                native_user_id = bound_user_id or "ksadk_user"
+            elif owner_scope_ref and self._invocation_uses_legacy_session(native_user_id):
+                # Canonical preprocessing has already proved that this exact
+                # platform session is an authorized legacy row.
+                legacy_session = await self._get_adk_session(
+                    user_id="ksadk_user",
+                    session_id=external_session_id,
+                )
+                if legacy_session is not None:
+                    native_user_id = "ksadk_user"
+                    internal_session_id = str(legacy_session.id)
+
+            session = None
+            if internal_session_id:
+                session = await self._get_adk_session(
+                    user_id=native_user_id,
+                    session_id=internal_session_id,
+                )
+            if session is None:
+                session = await self._create_adk_session(
+                    user_id=native_user_id,
+                    session_id=internal_session_id or external_session_id,
+                    preserve_requested_id=bool(owner_scope_ref or internal_session_id),
+                )
+            internal_session_id = str(session.id)
+            self._session_map[cache_key] = internal_session_id
+            self._session_user_map[cache_key] = native_user_id
+            await core.set_binding_by_session_id(
+                external_session_id,
+                "adk",
+                {
+                    "external_session_id": str(external_session_id),
+                    "internal_session_id": internal_session_id,
+                    "native_user_id": native_user_id,
+                    "owner_scope_ref": owner_scope_ref,
+                },
+            )
+            return internal_session_id
+
+        # Identity-aware requests normally receive a canonical external id
+        # before reaching a runner. Keep this fallback isolated as well.
+        if owner_scope_ref:
+            cache_key = self._session_cache_key("__default__", owner_scope_ref)
+            if cache_key not in self._session_map:
+                session = await self._create_adk_session(user_id=native_user_id)
+                self._session_map[cache_key] = str(session.id)
+                self._session_user_map[cache_key] = native_user_id
+            return self._session_map[cache_key]
+
+        if self._default_session_id is None:
+            session = await self._create_adk_session(user_id=native_user_id)
+            self._default_session_id = session.id
+        return str(self._default_session_id)
+
+    def _native_user_for_session(self, external_session_id: str | None) -> str:
+        native_user_id, owner_scope_ref = self._invocation_session_owner()
+        if not external_session_id:
+            return native_user_id
+        cache_key = self._session_cache_key(str(external_session_id), owner_scope_ref)
+        return self._session_user_map.get(cache_key, native_user_id)
+
+
+class LangGraphSessionIdentityMixin:
+    """Bind canonical sessions to LangGraph thread/checkpoint addresses."""
+
+    def _invocation_identity_scope_ref(self) -> str:
+        context = get_current_invocation_context()
+        if context is None or context.identity.is_empty:
+            return ""
+        if not context.identity.is_complete:
+            raise ValueError("Incomplete trusted business identity context")
+        return identity_scope_ref(context.identity)
+
+    @staticmethod
+    def _invocation_uses_legacy_session() -> bool:
+        context = get_current_invocation_context()
+        if context is None or context.identity.is_empty or not context.identity.is_complete:
+            return False
+        return identity_session_is_legacy(context.identity, context.user_id)
+
+    @staticmethod
+    def _invocation_agent_scope_ref() -> str:
+        context = get_current_invocation_context()
+        if context is None:
+            return ""
+        return str(getattr(context, "agent_id", "") or "")
+
+    @classmethod
+    def _identity_thread_id(
+        cls, session_id: str, owner_scope_ref: str, agent_scope_ref: str = ""
+    ) -> str:
+        # 隔离口径：业务身份 + agent + 外部 session 共同决定 thread。
+        # agent 必须参与哈希——共享 checkpoint 库下，同租户的不同 agent 若复用
+        # 同一个 external_session_id，缺 agent 维度会落到同一个 thread 互相覆盖。
+        digest = hashlib.sha256(
+            f"{owner_scope_ref}\0{agent_scope_ref}\0{session_id}".encode()
+        ).hexdigest()
+        return f"lgt_{digest}"
+
+    async def _thread_has_checkpoint(self, thread_id: str) -> bool:
+        get_state = getattr(self._agent, "aget_state", None) or getattr(
+            self._agent, "get_state", None
+        )
+        if not callable(get_state):
+            return False
+        try:
+            state = get_state(self._get_config(thread_id))
+            if inspect.isawaitable(state):
+                state = await state
+        except Exception:
+            return False
+        config = (
+            state.get("config") if isinstance(state, Mapping) else getattr(state, "config", None)
+        )
+        values = (
+            state.get("values") if isinstance(state, Mapping) else getattr(state, "values", None)
+        )
+        metadata = (
+            state.get("metadata")
+            if isinstance(state, Mapping)
+            else getattr(state, "metadata", None)
+        )
+        created_at = (
+            state.get("created_at")
+            if isinstance(state, Mapping)
+            else getattr(state, "created_at", None)
+        )
+        if values or metadata or created_at:
+            return True
+        configurable = config.get("configurable") if isinstance(config, Mapping) else None
+        return bool(
+            isinstance(configurable, Mapping)
+            and str(configurable.get("checkpoint_id") or "").strip()
+        )
+
+    async def _get_session_config(self, session_id: str) -> dict[str, Any]:
+        owner_scope_ref = self._invocation_identity_scope_ref()
+        if not owner_scope_ref:
+            return self._get_config(session_id)
+
+        cache_key = f"{owner_scope_ref}:{session_id}"
+        cached = self._identity_thread_bindings.get(cache_key)
+        if cached:
+            return self._get_config(cached)
+
+        async with self._identity_thread_lock:
+            cached = self._identity_thread_bindings.get(cache_key)
+            if cached:
+                return self._get_config(cached)
+
+            core = ConversationSessionCore(resolve_session_service())
+            binding = await core.get_binding_by_session_id(session_id, "langgraph")
+            bound_scope_ref = str(binding.get("owner_scope_ref") or "")
+            if bound_scope_ref and bound_scope_ref != owner_scope_ref:
+                raise ValueError("LangGraph session belongs to a different business identity")
+
+            thread_id = str(binding.get("thread_id") or "")
+            if not thread_id:
+                thread_id = self._identity_thread_id(
+                    session_id,
+                    owner_scope_ref,
+                    self._invocation_agent_scope_ref(),
+                )
+                if self._invocation_uses_legacy_session() and await self._thread_has_checkpoint(
+                    session_id
+                ):
+                    thread_id = session_id
+
+            await core.set_binding_by_session_id(
+                session_id,
+                "langgraph",
+                {
+                    "external_session_id": session_id,
+                    "thread_id": thread_id,
+                    "owner_scope_ref": owner_scope_ref,
+                },
+            )
+            self._identity_thread_bindings[cache_key] = thread_id
+            return self._get_config(thread_id)

ksadk/runners/adk_runner.py

@@ -21,6 +21,7 @@ from opentelemetry import trace
 from ksadk.compat.adk_compat import genai_types as types
 from ksadk.conversations.attachments import classify_attachment_kind, read_attachment_uri_bytes
 from ksadk.conversations.model_context import supports_native_image_input
+from ksadk.runners._session_identity import ADKSessionIdentityMixin
 from ksadk.runners.base_runner import BaseRunner
 from ksadk.runners.usage_accumulator import accumulate_usage
 from ksadk.runners.utils import load_agent_module
@@ -45,7 +46,7 @@ def _part_metadata_flag(part: Any, key: str) -> bool:
     return False
 
 
-class ADKRunner(BaseRunner):
+class ADKRunner(ADKSessionIdentityMixin, BaseRunner):
     """ADK 框架运行时"""
 
     def __init__(self, detection_result: Any, project_dir: str):
@@ -54,6 +55,7 @@ class ADKRunner(BaseRunner):
         self._session_service: Any = None
         # Map external session_ids (e.g. from run_interactive or web) to ADK internal session IDs
         self._session_map: Dict[str, str] = {}
+        self._session_user_map: Dict[str, str] = {}
         # Fallback default session
         self._default_session_id: Optional[str] = None
         # Memory integration
@@ -304,9 +306,7 @@ class ADKRunner(BaseRunner):
         if self._short_term_memory is None:
             self._short_term_memory = self._init_short_term_memory()
         if self._short_term_memory is not None and self._session_service is None:
-            self._session_service = getattr(
-                self._short_term_memory, "session_service", None
-            )
+            self._session_service = getattr(self._short_term_memory, "session_service", None)
         resumable = self._resolve_resumability()
         self._resume_disabled_reason = None
         self._resume_disabled_reason_code = None
@@ -645,7 +645,7 @@ class ADKRunner(BaseRunner):
 
             agent_name = self._agent.name if self._agent else "default"
             ltm = LongTermMemory.from_env(app_name=agent_name)
-            logger.info(f"LongTermMemory initialized: backend={backend}, " f"app_name={agent_name}")
+            logger.info(f"LongTermMemory initialized: backend={backend}, app_name={agent_name}")
             return ltm
         except Exception as e:
             logger.warning(f"Failed to init LongTermMemory: {e}.")
@@ -669,7 +669,7 @@ class ADKRunner(BaseRunner):
 
             kb = KnowledgeBaseClient.from_env()
             logger.info(
-                f"KnowledgeBase initialized: dataset_id={kb.dataset_id}, " f"region={kb.region}"
+                f"KnowledgeBase initialized: dataset_id={kb.dataset_id}, region={kb.region}"
             )
             return kb
         except ImportError:
@@ -1197,49 +1197,6 @@ class ADKRunner(BaseRunner):
 
         return prepare_trace_metadata(detection_result=getattr(self, "detection_result", None))
 
-    async def _ensure_session(self, external_session_id: Optional[str] = None) -> str:
-        """Get or create ADK session ID based on external ID
-
-        When ShortTermMemory is configured, uses its create_session method
-        which supports session retrieval (if session_id already exists).
-        """
-        # Case 1: External ID provided
-        if external_session_id:
-            if external_session_id in self._session_map:
-                return self._session_map[external_session_id]
-
-            # Create new ADK session and map it
-            if self._short_term_memory:
-                session = await self._short_term_memory.create_session(
-                    app_name=self._agent.name,
-                    user_id="ksadk_user",
-                    session_id=external_session_id,
-                )
-            else:
-                if self._session_service is None:
-                    raise RuntimeError("ADK session service is not initialized")
-                session = await self._session_service.create_session(
-                    app_name=self._agent.name, user_id="ksadk_user"
-                )
-            self._session_map[external_session_id] = session.id
-            return str(session.id)
-
-        # Case 2: No external ID (use default singleton)
-        if self._default_session_id is None:
-            if self._short_term_memory:
-                session = await self._short_term_memory.create_session(
-                    app_name=self._agent.name,
-                    user_id="ksadk_user",
-                )
-            else:
-                if self._session_service is None:
-                    raise RuntimeError("ADK session service is not initialized")
-                session = await self._session_service.create_session(
-                    app_name=self._agent.name, user_id="ksadk_user"
-                )
-            self._default_session_id = session.id
-        return str(self._default_session_id)
-
     async def save_session_to_long_term_memory(
         self, session_id: str, user_id: str = "ksadk_user"
     ) -> bool:
@@ -1482,7 +1439,7 @@ class ADKRunner(BaseRunner):
             return max_seq
         except Exception as exc:
             logger.warning(
-                "ADKRunner: failed to query max checkpoint_seq " "for run_id=%s: %s",
+                "ADKRunner: failed to query max checkpoint_seq for run_id=%s: %s",
                 run_id,
                 exc,
             )
@@ -1589,8 +1546,14 @@ class ADKRunner(BaseRunner):
                     session_id,
                     "adk",
                     {
-                        "external_session_id": str(session_id),
-                        "internal_session_id": str(session_id),
+                        "external_session_id": str(
+                            binding.get("external_session_id") or session_id
+                        ),
+                        "internal_session_id": str(
+                            binding.get("internal_session_id") or session_id
+                        ),
+                        "native_user_id": str(binding.get("native_user_id") or "ksadk_user"),
+                        "owner_scope_ref": str(binding.get("owner_scope_ref") or ""),
                         "invocation_map": invocation_map,
                     },
                 )
@@ -1752,7 +1715,7 @@ class ADKRunner(BaseRunner):
             metadata=metadata,
         )
         logger.debug(
-            "ADKRunner: wrote checkpoint adk-ckpt-%d at boundary " "(session=%s, invocation_id=%s)",
+            "ADKRunner: wrote checkpoint adk-ckpt-%d at boundary (session=%s, invocation_id=%s)",
             checkpoint_seq,
             session_id,
             adk_invocation_id,
@@ -1793,9 +1756,7 @@ class ADKRunner(BaseRunner):
         persistence_degraded = bool(
             shared_across_pods and getattr(session_service, "degraded", False)
         )
-        platform_resumable = (
-            self._resumable and shared_across_pods and not persistence_degraded
-        )
+        platform_resumable = self._resumable and shared_across_pods and not persistence_degraded
         metadata["is_resumable"] = platform_resumable
         metadata["resume_status"] = "resumable" if platform_resumable else "disabled"
         metadata["backend"] = stm_backend or "in_memory"
@@ -1810,8 +1771,7 @@ class ADKRunner(BaseRunner):
                 self._resume_disabled_reason
                 if not self._resumable and self._resume_disabled_reason
                 else (
-                    "ADK database session persistence was degraded when this checkpoint "
-                    "was written"
+                    "ADK database session persistence was degraded when this checkpoint was written"
                 )
                 if persistence_degraded
                 else "ADK checkpoint uses an in-memory or local-only session backend; "
@@ -1852,7 +1812,7 @@ class ADKRunner(BaseRunner):
             )
         if not adk_invocation_id:
             logger.error(
-                "Resume requested but ADK invocation_id not found for " "session=%s ksadk_inv=%s",
+                "Resume requested but ADK invocation_id not found for session=%s ksadk_inv=%s",
                 session_id,
                 ksadk_invocation_id,
             )
@@ -1869,6 +1829,7 @@ class ADKRunner(BaseRunner):
         *,
         input_data: Dict[str, Any],
         session_id: str,
+        native_user_id: str = "ksadk_user",
         user_input: str,
         is_resume: bool,
         run_config: Optional[Any] = None,
@@ -1898,7 +1859,7 @@ class ADKRunner(BaseRunner):
 
         run_kwargs: Dict[str, Any] = {
             "session_id": session_id,
-            "user_id": "ksadk_user",
+            "user_id": native_user_id,
         }
         if run_config is not None:
             run_kwargs["run_config"] = run_config
@@ -1968,6 +1929,7 @@ class ADKRunner(BaseRunner):
             # Use external session ID if provided
             req_session_id = input_data.get("session_id")
             session_id = await self._ensure_session(req_session_id)
+            native_user_id = self._native_user_for_session(req_session_id)
 
             # 准备 Metadata 并设置 Span Attributes
             # Langfuse Exporter 会读取这些 span attributes
@@ -1982,6 +1944,7 @@ class ADKRunner(BaseRunner):
             wrapped_async = await self._prepare_run_events(
                 input_data=input_data,
                 session_id=session_id,
+                native_user_id=native_user_id,
                 user_input=user_input,
                 is_resume=is_resume,
             )
@@ -2073,6 +2036,7 @@ class ADKRunner(BaseRunner):
             # Use external session ID if provided
             req_session_id = input_data.get("session_id")
             session_id = await self._ensure_session(req_session_id)
+            native_user_id = self._native_user_for_session(req_session_id)
 
             # 准备 Metadata 并设置 Span Attributes
             agent_user_id, tags, _, _ = self._prepare_trace_metadata(session_id)
@@ -2087,6 +2051,7 @@ class ADKRunner(BaseRunner):
             wrapped_async = await self._prepare_run_events(
                 input_data=input_data,
                 session_id=session_id,
+                native_user_id=native_user_id,
                 user_input=user_input,
                 is_resume=is_resume,
                 run_config=run_config,