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

ksadk/runners/langgraph_runner.py

@@ -28,9 +28,12 @@ from ksadk.runtime.skill_eval_result import skill_eval_response_fields
 from ksadk.runtime.timing import extract_timing
 from ksadk.sessions import resolve_persistence_topology
 from ksadk.sessions.continuity import LangGraphSessionAdapter
+from ksadk.sessions.db_budget import session_db_budget
 
+from ksadk.runners._langgraph_tool_projection import _LangGraphToolProjectionMixin
 
-class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, BaseRunner):
+
+class LangGraphRunner(_LangGraphToolProjectionMixin, LangGraphSessionIdentityMixin, _LangGraphStreamMixin, BaseRunner):
     """LangGraph 框架运行时
 
     透传原生 LangGraph 功能，支持任意 State 格式
@@ -334,30 +337,63 @@ class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, Base
         }
 
     async def _create_managed_postgres_saver(self, dsn: str) -> tuple[Any, Any]:
-        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
         from psycopg.rows import dict_row
         from psycopg_pool import AsyncConnectionPool
 
-        timeout = max(0.1, float(os.getenv("KSADK_SESSION_CONNECT_TIMEOUT") or "5"))
+        from ksadk.harness.engine.postgres_checkpointer import bounded_postgres_saver_class
+
+        budget = session_db_budget()
+        timeout = budget.connect_seconds
+        # Reserve checkout time for replacing one silent idle socket. Spending
+        # the entire checkout budget on its probe would fail a healthy run
+        # even when a fresh database connection is immediately available.
+        probe_timeout = min(budget.read_acquire_seconds / 2, timeout / 2)
+
+        async def check_connection(connection):
+            # Pool checkout's deadline does not bound its check callback.
+            # A silent idle socket must not hold the saver lock indefinitely.
+            probe = asyncio.create_task(AsyncConnectionPool.check_connection(connection))
+            try:
+                done, _ = await asyncio.wait((probe,), timeout=probe_timeout)
+                if not done:
+                    # Close before cancellation: psycopg's query cancellation
+                    # also needs a responding server and can itself stall.
+                    await connection.close()
+                    raise TimeoutError("checkpoint connection health check timed out")
+                await probe
+            except asyncio.CancelledError:
+                await connection.close()
+                raise
+            finally:
+                if not probe.done():
+                    probe.cancel()
+                await asyncio.gather(probe, return_exceptions=True)
+
         pool = AsyncConnectionPool(
             conninfo=dsn,
             min_size=1,
             max_size=10,
             open=False,
-            timeout=timeout,
+            timeout=budget.write_acquire_seconds,
+            # Discard dead idle sockets before a checkpoint operation starts.
+            # This retries connection checkout, never graph execution or writes.
+            check=check_connection,
             kwargs={
                 "autocommit": True,
                 "prepare_threshold": 0,
                 "row_factory": dict_row,
+                "connect_timeout": max(1, int(budget.connect_seconds + 0.999)),
+                "options": f"-c statement_timeout={max(1, int(budget.write_sql_seconds * 1000))} "
+                           f"-c lock_timeout={max(1, int(budget.lock_seconds * 1000))}",
             },
         )
         try:
             await pool.open(wait=True, timeout=timeout)
-            saver = AsyncPostgresSaver(pool)
+            saver = bounded_postgres_saver_class()(pool, budget=budget)
             await saver.setup()
             return saver, pool
         except Exception:
-            await pool.close()
+            await pool.close(timeout=budget.cleanup_seconds)
             raise
 
     async def prepare_runtime_capabilities(self) -> None:
@@ -605,6 +641,34 @@ class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, Base
             await pool.close()
         await super().close()
 
+    def _resolve_recursion_limit(self) -> int:
+        """解析本次运行的 recursion_limit。
+
+        LangChain 的 ``astream_events`` 入口会用库默认值 25 填充调用 config，
+        随后 ``Pregel.astream`` 的合并里这个显式 25 会覆盖图实例
+        ``with_config`` 设置的图级默认值——业务侧的 recursion_limit 在流式
+        路径上因此静默失效（生产 GraphRecursionError 的根因）。所以 runner
+        必须在调用 config 里显式携带非默认值：优先图自身 config（业务
+        ``with_config`` 的设置），其次环境变量，最后 runner 默认 50。
+        """
+        agent_config = getattr(self._agent, "config", None)
+        if isinstance(agent_config, Mapping):
+            try:
+                value = int(agent_config.get("recursion_limit") or 0)
+            except (TypeError, ValueError):
+                value = 0
+            if value > 0:
+                return value
+        raw = str(os.getenv("KSADK_LANGGRAPH_RECURSION_LIMIT") or "").strip()
+        if raw:
+            try:
+                value = int(raw)
+            except ValueError:
+                value = 0
+            if value > 0:
+                return value
+        return 50
+
     def _get_config(self, session_id: str) -> dict:
         """获取运行配置
 
@@ -612,8 +676,14 @@ class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, Base
         隔离字段。把租户/agent scope 塞进 configurable.checkpoint_ns 会让
         Pregel 的 aget_state 走子图重定向并抛 "Subgraph <scope> not found"，
         静默打掉流式路径的待审批探测。租户隔离由 thread_id 承担。
+
+        recursion_limit 见 ``_resolve_recursion_limit``：astream_events 会用
+        默认 25 覆盖图级设置，必须在这里显式携带非默认值。
         """
-        return {"configurable": {"thread_id": session_id}}
+        return {
+            "configurable": {"thread_id": session_id},
+            "recursion_limit": self._resolve_recursion_limit(),
+        }
 
     @staticmethod
     def _extract_langgraph_checkpoint_ref(payload: Dict[str, Any]) -> dict[str, Any]:
@@ -1202,7 +1272,11 @@ class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, Base
             last_usage = self._extract_last_usage(result)
             if last_usage:
                 output.setdefault("metadata", {})["last_usage"] = last_usage
-            metadata = await self._latest_checkpoint_metadata(config)
+            # A resume config names the historical input checkpoint. Read the
+            # newly committed state in the same thread/namespace after execution.
+            metadata = await self._latest_checkpoint_metadata(
+                self._latest_checkpoint_config(config)
+            )
             if metadata:
                 output["metadata"] = {**(output.get("metadata") or {}), **metadata}
             return output
@@ -1333,81 +1407,3 @@ class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, Base
         metadata = await self._latest_checkpoint_metadata(latest_config)
         if metadata:
             yield {"type": "checkpoint", "metadata": metadata}
-
-    @staticmethod
-    def _tool_events_from_graph_update(
-        node_output: Any,
-        *,
-        tool_names: dict[str, str],
-    ) -> list[dict[str, Any]]:
-        """Preserve LangGraph update-mode tool messages as structured runner events."""
-
-        if isinstance(node_output, Mapping):
-            messages = node_output.get("messages")
-        else:
-            messages = getattr(node_output, "messages", None)
-        if not isinstance(messages, (list, tuple)):
-            return []
-
-        events: list[dict[str, Any]] = []
-        for message in messages:
-            if isinstance(message, Mapping):
-                tool_calls = message.get("tool_calls")
-                tool_call_id = message.get("tool_call_id")
-                message_name = message.get("name")
-                message_content = message.get("content")
-                message_status = message.get("status")
-            else:
-                tool_calls = getattr(message, "tool_calls", None)
-                tool_call_id = getattr(message, "tool_call_id", None)
-                message_name = getattr(message, "name", None)
-                message_content = getattr(message, "content", None)
-                message_status = getattr(message, "status", None)
-
-            if isinstance(tool_calls, (list, tuple)):
-                for tool_call in tool_calls:
-                    if isinstance(tool_call, Mapping):
-                        call_id = str(tool_call.get("id") or tool_call.get("tool_call_id") or "")
-                        name = str(tool_call.get("name") or "tool")
-                        args = tool_call.get("args")
-                    else:
-                        call_id = str(
-                            getattr(tool_call, "id", None)
-                            or getattr(tool_call, "tool_call_id", None)
-                            or ""
-                        )
-                        name = str(getattr(tool_call, "name", None) or "tool")
-                        args = getattr(tool_call, "args", None)
-                    if call_id:
-                        tool_names[call_id] = name
-                    events.append(
-                        {
-                            "type": "tool_call",
-                            "tool_call_id": call_id or name,
-                            "tool_name": name,
-                            "tool_args": args,
-                        }
-                    )
-
-            if tool_call_id:
-                call_id = str(tool_call_id)
-                name = str(message_name or tool_names.get(call_id) or "tool")
-                status = str(message_status or "").lower()
-                events.append(
-                    {
-                        "type": "tool_result",
-                        "tool_call_id": call_id,
-                        "tool_name": name,
-                        "tool_output": message_content,
-                        "error": message_content if status in {"error", "failed"} else None,
-                    }
-                )
-        return events
-
-    def _filter_tool_tags(self, content: str) -> str:
-        """过滤 <tool_call> 标签"""
-        if not isinstance(content, str):
-            return content
-        content = re.sub(r"<tool_call>.*?</tool_call>", "", content, flags=re.DOTALL)
-        content = re.sub(r"</?(?:tool_call|arg_key|arg_value)>", "", content)
-        return content

ksadk/runners/utils/loader.py

@@ -7,8 +7,11 @@
 import importlib
 import sys
 from pathlib import Path
+from threading import RLock
 from typing import Any
 
+_LOAD_LOCK = RLock()
+
 
 def load_agent_module(
     project_dir: str,
@@ -56,20 +59,31 @@ def load_agent_module(
         # ``agent.py`` 为入口；仅靠 ``import_module("agent")`` 会复用上一个
         # Bundle 的 sys.modules 缓存，进而加载错误的 agentVariable。只复用与
         # 当前 entry 文件一致的缓存，其余情况让 import 从当前 project_path 重载。
-        cached = sys.modules.get(module_name)
-        cached_file = getattr(cached, "__file__", None)
-        cached_matches_entry = bool(
-            cached_file is not None
-            and entry_file is not None
-            and Path(str(cached_file)).resolve() == entry_file
-        )
-        if cached is not None and not cached_matches_entry:
-            sys.modules.pop(module_name, None)
-            importlib.invalidate_caches()
-        if force_reload and module_name in sys.modules:
-            module = importlib.reload(sys.modules[module_name])
-        else:
-            module = importlib.import_module(module_name)
+        # Studio loads immutable Build sources in-process.  A normal import
+        # writes ``__pycache__`` beside the entrypoint, mutating the admitted
+        # Bundle and making the next integrity check fail.  Serialize the
+        # global import/cache mutation and suppress bytecode only for this
+        # immutable load; restore the host interpreter setting afterwards.
+        with _LOAD_LOCK:
+            previous_dont_write_bytecode = sys.dont_write_bytecode
+            sys.dont_write_bytecode = True
+            try:
+                cached = sys.modules.get(module_name)
+                cached_file = getattr(cached, "__file__", None)
+                cached_matches_entry = bool(
+                    cached_file is not None
+                    and entry_file is not None
+                    and Path(str(cached_file)).resolve() == entry_file
+                )
+                if cached is not None and not cached_matches_entry:
+                    sys.modules.pop(module_name, None)
+                    importlib.invalidate_caches()
+                if force_reload and module_name in sys.modules:
+                    module = importlib.reload(sys.modules[module_name])
+                else:
+                    module = importlib.import_module(module_name)
+            finally:
+                sys.dont_write_bytecode = previous_dont_write_bytecode
         agent = getattr(module, agent_variable)
         return agent, module
     except ImportError as e:

ksadk/runtime/_runner_adapter/a2ui_surface.py

@@ -0,0 +1,145 @@
+"""Translate validated A2UI tool envelopes into canonical data items."""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from collections.abc import Mapping
+from typing import Any
+
+from ksadk.events.canonical import (
+    ContentSnapshot,
+    ItemCompleted,
+    ItemStarted,
+    RuntimeEvent,
+    SourceRef,
+)
+from ksadk.events.content import DataContent
+from ksadk.events.identity import stable_event_id, stable_item_id, stable_scope_id
+from ksadk.runtime.adapter import RunHandle
+
+logger = logging.getLogger("ksadk.runtime._runner_adapter.stream_mapping")
+
+
+def _a2ui_surface_event(
+    self: Any,
+    handle: RunHandle,
+    chunk: Any,
+) -> tuple[RuntimeEvent, ...]:
+    """Recognize a validated A2UI tool envelope and emit a canonical data-item event.
+
+    The dynamic ``generate_a2ui`` tool returns official v0.9 operations as a
+    JSON tool result. Tool results are otherwise opaque to the runtime, which
+    would leave AG-UI with nothing to project until a page reload reconstructs
+    history. Convert exactly that envelope at the runtime boundary so it is
+    streamed, persisted, and replayed like every other A2UI surface.
+
+    In the canonical schema, A2UI surfaces are modeled as ``item_kind="data"``
+    items with ``source.protocol="a2ui"``.
+    """
+
+    if not isinstance(chunk, dict):
+        return ()
+    value = chunk.get("tool_output", chunk.get("output"))
+    if value is not None and hasattr(value, "content"):
+        value = value.content
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except (TypeError, ValueError):
+            return ()
+    if not isinstance(value, Mapping):
+        return ()
+    operations_raw = value.get("a2ui_operations")
+    if not isinstance(operations_raw, list) or not operations_raw:
+        return ()
+    operations = [dict(operation) for operation in operations_raw if isinstance(operation, Mapping)]
+    if not operations:
+        return ()
+
+    known: list[tuple[str, str]] = []  # (surface_id, lifecycle)
+    for operation in operations:
+        for key, lifecycle in (
+            ("createSurface", "begin"),
+            ("updateComponents", "update"),
+            ("updateDataModel", "update"),
+            ("deleteSurface", "end"),
+        ):
+            detail = operation.get(key)
+            if isinstance(detail, Mapping) and isinstance(detail.get("surfaceId"), str):
+                surface_id = detail["surfaceId"].strip()
+                if surface_id:
+                    known.append((surface_id, lifecycle))
+                    break
+    if not known:
+        return ()
+    surface_ids = {surface_id for surface_id, _lifecycle in known}
+    if len(surface_ids) != 1:
+        logger.warning("ignoring A2UI tool result with multiple surfaces")
+        return ()
+    surface_id = known[0][0]
+    lifecycle = "begin" if any(lc == "begin" for _, lc in known) else known[0][1]
+
+    framework = self._runtime_type
+    run_id = handle.run_id
+    scope_id = stable_scope_id(framework, run_id)
+    # A dynamic A2UI tool result is a complete, immutable operation batch.  It
+    # must therefore be represented by a closed canonical item.  Reusing one
+    # long-lived item per surface leaves a create/update-only batch open at the
+    # run terminal and makes strict reducers reject the stream.
+    start_seq = self._next_seq()
+    batch_ref = str(chunk.get("call_id") or chunk.get("run_id") or chunk.get("id") or "")
+    if not batch_ref:
+        batch_ref = json.dumps(
+            operations,
+            ensure_ascii=False,
+            sort_keys=True,
+            separators=(",", ":"),
+        )
+    item_id = stable_item_id(framework, run_id, "a2ui-batch", surface_id, batch_ref)
+    source = SourceRef(
+        framework=framework,
+        protocol="a2ui",
+        native_run_id=run_id,
+        metadata={
+            "surface_id": surface_id,
+            "operation_batch": True,
+            "surface_lifecycle": lifecycle,
+        },
+    )
+    # TODO(runtime-event-v2): dict chunk 退化路径,chunk_ordinal 用 seq counter;
+    # LangGraph/Codex 切 stream_canonical_events 后清理
+    timestamp = time.time()
+    snapshot = ContentSnapshot(parts=(DataContent(part_id="a2ui-ops", data=operations),))
+    completed_seq = self._next_seq()
+    return (
+        ItemStarted(
+            schema_version=2,
+            event_id=stable_event_id(
+                framework, scope_id, item_id, "item.started", "a2ui", run_id, start_seq
+            ),
+            seq=start_seq,
+            timestamp=timestamp,
+            run_id=run_id,
+            scope_id=scope_id,
+            source=source,
+            item_id=item_id,
+            item_kind="data",
+            initial=snapshot,
+        ),
+        ItemCompleted(
+            schema_version=2,
+            event_id=stable_event_id(
+                framework, scope_id, item_id, "item.completed", "a2ui", run_id, completed_seq
+            ),
+            seq=completed_seq,
+            timestamp=timestamp,
+            run_id=run_id,
+            scope_id=scope_id,
+            source=source,
+            item_id=item_id,
+            item_kind="data",
+            snapshot=snapshot,
+        ),
+    )

ksadk/runtime/_runner_adapter/gateway_approvals.py

@@ -0,0 +1,93 @@
+"""Recognize fresh gateway approvals without adopting historical tool results."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Optional
+
+from ksadk.events.canonical import ItemCompleted, RuntimeEvent
+from ksadk.events.content import DataContent, ToolCallContent, ToolResultContent
+from ksadk.runtime.adapter import RunHandle
+from ksadk.tools.gateway import approval_interrupt_info_from_result
+
+if TYPE_CHECKING:
+    from ksadk.runtime.runner_adapter import _ActiveRun
+
+
+def _canonical_gateway_approvals(
+    self, handle: RunHandle, run: Optional[_ActiveRun], event: ItemCompleted
+) -> list[RuntimeEvent]:
+    """Recognize local results without treating historical state as new calls."""
+    calls = handle.native_ref.setdefault("_gateway_canonical_calls", {})
+    historical_calls = handle.native_ref.setdefault("_gateway_canonical_history_calls", {})
+
+    def key(call_id):
+        return json.dumps([event.run_id, event.scope_id, call_id])
+
+    results = []
+    for part in event.snapshot.parts:
+        if isinstance(part, ToolCallContent):
+            calls[key(part.call_id)] = {"name": part.name, "arguments": part.arguments}
+        elif isinstance(part, ToolResultContent):
+            results.append((part.call_id, part.result))
+        elif (
+            isinstance(part, DataContent)
+            and event.source.metadata.get("channel") in {"updates", "values"}
+            and isinstance(part.data, dict)
+        ):
+            # Some ToolNode versions only emit values. Remember already
+            # present ToolMessage IDs before admitting new native calls,
+            # so loading a checkpoint cannot reopen historical approvals.
+            channel = event.source.metadata.get("channel")
+            updates = [part.data] if channel == "values" else part.data.values()
+            seen = handle.native_ref.setdefault("_gateway_canonical_results", {})
+            for update in updates:
+                if not isinstance(update, dict):
+                    continue
+                messages = update.get("messages", [])
+                if not isinstance(messages, list):
+                    messages = [messages]
+                for message in messages:
+                    if not isinstance(message, dict) or message.get("type") != "tool":
+                        continue
+                    message_id = message.get("id")
+                    call_id = message.get("tool_call_id")
+                    if channel == "values" and call_id and key(call_id) not in calls:
+                        # A restored snapshot can precede replayed typed calls
+                        # and results. Fence their shared call identity as well
+                        # as the snapshot's distinct ToolMessage identity.
+                        historical_calls[key(call_id)] = True
+                    if not message_id:
+                        continue
+                    result_key = key(message_id)
+                    if result_key in seen:
+                        continue
+                    seen[result_key] = True
+                    results.append((call_id, message.get("content")))
+    approvals = []
+    for call_id, result in results:
+        if key(call_id) in historical_calls:
+            continue
+        call = calls.get(key(call_id))
+        if call is None:
+            continue
+        detail = approval_interrupt_info_from_result(
+            result,
+            fallback_tool_name=call["name"],
+            tool_args=call["arguments"],
+            run_id=call_id,
+        )
+        if detail is None:
+            continue
+        approval_id = str(detail.get("approval_request_id") or detail.get("id") or call_id)
+        if approval_id in handle.native_ref.get("pending_approval_ids", []):
+            continue
+        approvals.extend(
+            self._interaction_requested_from_approval(
+                handle,
+                run,
+                detail=detail,
+                call_id=call_id,
+            )
+        )
+    return approvals