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

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,10 +55,12 @@ 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
         self._short_term_memory: Any = None
+        self._checkpoint_storage_source = "none"
         self._long_term_memory: Any = None
         # Knowledge base integration
         self._knowledge_base: Any = None
@@ -66,6 +69,7 @@ class ADKRunner(BaseRunner):
         # ADK resumability state
         self._resumable: bool = False
         self._resume_disabled_reason: Optional[str] = None
+        self._resume_disabled_reason_code: Optional[str] = None
         # P1.1 sub-issue: guard invocation_map read-modify-write so concurrent
         # invocations on the same session don't lose each other's mappings.
         self._invocation_map_lock = asyncio.Lock()
@@ -75,7 +79,21 @@ class ADKRunner(BaseRunner):
         ).strip()
 
     async def close(self) -> None:
-        """Close runtime toolsets owned by this runner."""
+        """Close database sessions and runtime toolsets owned by this runner."""
+        session_service = self._session_service
+        self._session_service = None
+        if session_service is not None:
+            close_session = getattr(session_service, "aclose", None) or getattr(
+                session_service, "close", None
+            )
+            if callable(close_session):
+                try:
+                    result = close_session()
+                    if inspect.isawaitable(result):
+                        await result
+                except Exception as exc:
+                    logger.warning("Failed to close ADK session service: %s", exc)
+
         toolsets = list(self._runtime_toolsets)
         self._runtime_toolsets.clear()
         for toolset in toolsets:
@@ -252,14 +270,23 @@ class ADKRunner(BaseRunner):
             "KSADK_STM_DB_URL",
             "KSADK_SESSION_BACKEND",
             "KSADK_SESSION_DSN",
+            "KSADK_CHECKPOINT_DSN",
         )
         if not any(str(os.environ.get(name, "")).strip() for name in configured_names):
             return None
 
         try:
             from ksadk.memory.adk import ShortTermMemory
+            from ksadk.sessions import resolve_persistence_topology
 
-            stm = ShortTermMemory.from_env()
+            topology = resolve_persistence_topology(framework="adk")
+            self._checkpoint_storage_source = topology.checkpoint.source
+            if topology.checkpoint.backend == "postgres" and topology.checkpoint.dsn:
+                stm = ShortTermMemory.from_persistence_target(topology.checkpoint)
+            else:
+                # Preserve legacy SQLite ADK/STM path handling when no remote
+                # checkpoint target has been selected.
+                stm = ShortTermMemory.from_env()
             logger.info(
                 "ShortTermMemory initialized: backend=%s path=%s",
                 stm.backend,
@@ -273,6 +300,69 @@ class ADKRunner(BaseRunner):
     def get_session_adapter(self):
         return ADKSessionAdapter()
 
+    async def prepare_runtime_capabilities(self) -> None:
+        """Prepare persistence-backed capability state without importing user code."""
+
+        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)
+        resumable = self._resolve_resumability()
+        self._resume_disabled_reason = None
+        self._resume_disabled_reason_code = None
+        if resumable.enabled:
+            compatible, reason = self._check_adk_resume_compatibility()
+            if not compatible:
+                self._resumable = False
+                self._resume_disabled_reason = reason
+                self._resume_disabled_reason_code = "ADK_VERSION_UNSUPPORTED"
+                return
+        self._resumable = resumable.enabled
+
+    async def refresh_runtime_capabilities(self) -> None:
+        await super().refresh_runtime_capabilities()
+        session_service = (
+            getattr(self._short_term_memory, "session_service", None)
+            if self._short_term_memory
+            else None
+        )
+        refresh = getattr(session_service, "refresh_persistence_capability", None)
+        if callable(refresh):
+            result = refresh()
+            if inspect.isawaitable(result):
+                await result
+
+    async def attach_runtime_handle(self, handle: Any) -> bool:
+        """Validate that a persisted ADK invocation can be resumed here.
+
+        Invocation existence is resolved by ADK when ``run_async`` receives the
+        invocation id.  This seam proves the deployment-side prerequisites
+        before Runtime v2 accepts a handle restored from session metadata.
+        """
+
+        if str(getattr(handle, "runtime_type", "") or "").strip().lower() != "adk":
+            return False
+        native_ref = getattr(handle, "native_ref", None)
+        if not isinstance(native_ref, Mapping):
+            return False
+        framework_ref = native_ref.get("framework_ref")
+        adk_ref = framework_ref.get("adk") if isinstance(framework_ref, Mapping) else None
+        invocation_id = str(
+            native_ref.get("invocation_id")
+            or (adk_ref.get("invocation_id") if isinstance(adk_ref, Mapping) else "")
+            or ""
+        ).strip()
+        if not invocation_id:
+            return False
+
+        await self.prepare_runtime_capabilities()
+        capability = self.describe_checkpoint_capability()
+        return bool(
+            capability.get("Supported")
+            and capability.get("Durable")
+            and capability.get("SharedAcrossPods")
+        )
+
     def describe_checkpoint_capability(self) -> dict[str, Any]:
         resumable = getattr(self, "_resumable", False)
         stm_backend = (
@@ -285,8 +375,16 @@ class ADKRunner(BaseRunner):
             elif stm_backend == "database":
                 backend = "adk_invocation+postgres"
             shared_across_pods = stm_backend == "database"
-            return {
-                "Supported": shared_across_pods,
+            session_service = (
+                getattr(self._short_term_memory, "session_service", None)
+                if self._short_term_memory
+                else None
+            )
+            persistence_degraded = bool(
+                shared_across_pods and getattr(session_service, "degraded", False)
+            )
+            capability = {
+                "Supported": shared_across_pods and not persistence_degraded,
                 "Backend": backend,
                 "Scope": "invocation",
                 "Durable": stm_backend is not None and stm_backend != "local",
@@ -296,12 +394,19 @@ class ADKRunner(BaseRunner):
                 "Reason": (
                     "ADK ResumabilityConfig and shared database session backend enabled; "
                     "resume via invocation_id"
-                    if shared_across_pods
+                    if shared_across_pods and not persistence_degraded
+                    else "ADK database session persistence is temporarily degraded"
+                    if persistence_degraded
                     else "ADK ResumabilityConfig enabled, but the session backend is "
                     "process-local or SQLite and cannot support cross-pod recovery"
                 ),
             }
-        return {
+            if persistence_degraded:
+                capability["ReasonCode"] = "CHECKPOINT_STORE_DEGRADED"
+            elif not shared_across_pods:
+                capability["ReasonCode"] = "CHECKPOINTER_NOT_DURABLE"
+            return capability
+        capability = {
             "Supported": False,
             "Backend": "none",
             "Scope": "unknown",
@@ -314,6 +419,9 @@ class ADKRunner(BaseRunner):
                 "KSADK_ADK_RESUMABLE=1 or configure App with resumability_config"
             ),
         }
+        if self._resume_disabled_reason_code:
+            capability["ReasonCode"] = self._resume_disabled_reason_code
+        return capability
 
     def get_runtime_capabilities(self) -> dict[str, Any]:
         capabilities = super().get_runtime_capabilities()
@@ -365,6 +473,9 @@ class ADKRunner(BaseRunner):
                 ),
             }
         capabilities["ResumeRun"]["Reason"] = capabilities["Checkpoint"]["Reason"]
+        reason_code = str(capabilities["Checkpoint"].get("ReasonCode") or "")
+        if reason_code:
+            capabilities["ResumeRun"]["ReasonCode"] = reason_code
         return capabilities
 
     @dataclass
@@ -445,6 +556,8 @@ class ADKRunner(BaseRunner):
 
         resumable = self._resolve_resumability()
         resumability_enabled = resumable.enabled
+        self._resume_disabled_reason = None
+        self._resume_disabled_reason_code = None
 
         # 版本兼容性检查：低于最低版本时强制关闭恢复
         resume_compatible, resume_reason = self._check_adk_resume_compatibility()
@@ -453,6 +566,7 @@ class ADKRunner(BaseRunner):
             resumable = self._ResolvabilityResult(enabled=False, source="version_check", app=None)
             resumability_enabled = False
             self._resume_disabled_reason = resume_reason
+            self._resume_disabled_reason_code = "ADK_VERSION_UNSUPPORTED"
 
         if resumable.app is not None:
             runner_kwargs = dict(
@@ -474,7 +588,7 @@ class ADKRunner(BaseRunner):
                 )
             except ImportError:
                 logger.warning(
-                    "ADK ResumabilityConfig not available (requires google-adk >= 1.14.0); "
+                    "ADK ResumabilityConfig not available (requires google-adk >= 1.16.0); "
                     "falling back to non-resumable Runner"
                 )
                 runner_kwargs = dict(
@@ -484,6 +598,7 @@ class ADKRunner(BaseRunner):
                 )
                 resumability_enabled = False
                 self._resume_disabled_reason = "ADK ResumabilityConfig is unavailable"
+                self._resume_disabled_reason_code = "ADK_VERSION_UNSUPPORTED"
         else:
             runner_kwargs = dict(
                 agent=self._agent,
@@ -530,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}.")
@@ -554,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:
@@ -938,7 +1053,8 @@ class ADKRunner(BaseRunner):
             raise TypeError("加载的对象不是有效的 ADK Agent")
 
         # 初始化记忆体 (从环境变量读取配置)
-        self._short_term_memory = self._init_short_term_memory()
+        if self._short_term_memory is None:
+            self._short_term_memory = self._init_short_term_memory()
         self._long_term_memory = self._init_long_term_memory()
 
         # 初始化知识库 (从环境变量读取配置)
@@ -1081,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:
@@ -1366,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,
             )
@@ -1473,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,
                     },
                 )
@@ -1636,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,
@@ -1669,17 +1748,32 @@ class ADKRunner(BaseRunner):
             getattr(self._short_term_memory, "backend", None) if self._short_term_memory else None
         )
         shared_across_pods = stm_backend == "database"
-        platform_resumable = self._resumable and shared_across_pods
+        session_service = (
+            getattr(self._short_term_memory, "session_service", None)
+            if self._short_term_memory
+            else None
+        )
+        persistence_degraded = bool(
+            shared_across_pods and getattr(session_service, "degraded", False)
+        )
+        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"
         metadata["scope"] = "invocation"
-        metadata["durable"] = stm_backend is not None and stm_backend != "local"
-        metadata["shared_across_pods"] = shared_across_pods
+        metadata["durable"] = (
+            stm_backend is not None and stm_backend != "local" and not persistence_degraded
+        )
+        metadata["shared_across_pods"] = shared_across_pods and not persistence_degraded
+        metadata["source"] = self._checkpoint_storage_source
         if not platform_resumable:
             metadata["resume_disabled_reason"] = (
                 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"
+                )
+                if persistence_degraded
                 else "ADK checkpoint uses an in-memory or local-only session backend; "
                 "cross-pod resume is unavailable"
             )
@@ -1718,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,
             )
@@ -1735,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,
@@ -1764,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
@@ -1834,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
@@ -1848,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,
             )
@@ -1939,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)
@@ -1953,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,
@@ -2132,6 +2231,7 @@ class ADKRunner(BaseRunner):
                                 fc_args = {}
                         yield {
                             "type": "tool_call",
+                            "tool_call_id": str(fc_id),
                             "tool_name": getattr(fc, "name", "unknown"),
                             "tool_args": fc_args,
                         }
@@ -2146,6 +2246,7 @@ class ADKRunner(BaseRunner):
                             emitted_tool_call_ids.add(tc_id)
                             yield {
                                 "type": "tool_call",
+                                "tool_call_id": str(tc_id),
                                 "tool_name": tc_name,
                                 "tool_args": getattr(tool_call, "input", {}),
                             }
@@ -2174,6 +2275,7 @@ class ADKRunner(BaseRunner):
                         fr = getattr(part, "function_response", None)
                         if fr is not None:
                             fr_name = getattr(fr, "name", "unknown")
+                            fr_id = getattr(fr, "id", "") or fr_name
                             fr_output = getattr(fr, "response", None) or {}
                             if not isinstance(fr_output, dict):
                                 try:
@@ -2187,6 +2289,7 @@ class ADKRunner(BaseRunner):
                                     fr_output = {"raw": str(fr_output)}
                             yield {
                                 "type": "tool_result",
+                                "tool_call_id": str(fr_id),
                                 "tool_name": fr_name,
                                 "tool_output": fr_output,
                             }

ksadk/runners/base_runner.py

@@ -151,6 +151,18 @@ class BaseRunner(ABC):
             },
         }
 
+    async def prepare_runtime_capabilities(self) -> None:
+        """Prepare optional async resources required for capability discovery."""
+        return None
+
+    async def refresh_runtime_capabilities(self) -> None:
+        """Refresh transient resources used for capability discovery.
+
+        The default implementation preserves compatibility for runners whose
+        capability preparation is either synchronous or permanently stable.
+        """
+        await self.prepare_runtime_capabilities()
+
     async def close(self) -> None:
         """释放 runner 持有的运行期资源。"""
         return None
@@ -193,7 +205,9 @@ class BaseRunner(ABC):
         return kwargs
 
     @staticmethod
-    def build_native_context(platform_context: Any) -> dict[str, Any] | None:
+    def build_native_context(
+        platform_context: Any, *, context_schema: Any = None
+    ) -> dict[str, Any] | None:
         if not isinstance(platform_context, dict):
             return None
         native_context = {
@@ -201,6 +215,20 @@ class BaseRunner(ABC):
             for key in ("agent_id", "user_id", "session_id")
             if platform_context.get(key) is not None
         }
+        from typing import get_args, get_type_hints
+        from ksadk.session_context import SessionContext
+
+        try:
+            fields = get_type_hints(context_schema) if context_schema is not None else {}
+        except (TypeError, NameError):
+            fields = {}
+        session_type = fields.get("session")
+        if session_type is SessionContext or set(get_args(session_type)) == {
+            SessionContext,
+            type(None),
+        }:
+            native_context = {key: value for key, value in native_context.items() if key in fields}
+            native_context["session"] = SessionContext.from_payload(platform_context.get("session"))
         return native_context or None
 
     @staticmethod