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

ksadk/kernel/worker_interactions.py

@@ -0,0 +1,197 @@
+"""Record worker interactions and deliver committed managed requests."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ksadk.kernel.worker import ActiveExecution
+
+from datetime import datetime, timezone
+
+from ksadk.kernel.approval_presentation import approval_presentation
+from ksadk.kernel.contracts import ActivationWriteGuard
+from ksadk.kernel.errors import InvalidCommandError
+from ksadk.kernel.state import RunState
+from ksadk.kernel.store import RunRecord
+
+
+class _WorkerInteractionMixin:
+    async def _record_interaction_request(
+        self,
+        execution: ActiveExecution,
+        run: RunRecord,
+        event: object,
+        guard: ActivationWriteGuard,
+    ) -> RunRecord:
+        """Persist one framework interaction as the durable ledger authority."""
+
+        from ksadk.events.canonical import ApprovalRequest, InteractionRequested
+        from ksadk.interaction.contracts import InteractionRecord
+
+        assert isinstance(event, InteractionRequested)
+        presentation = None
+        if isinstance(event.request, ApprovalRequest):
+            request_schema = {
+                "type": "object",
+                "properties": {
+                    "decision": {
+                        "type": "string",
+                        "enum": ["approve", "reject"],
+                    }
+                },
+                "required": ["decision"],
+            }
+            native_target = {"call_id": event.request.call_id or event.interaction_id}
+            presentation = approval_presentation(event.request)
+        else:
+            request_schema = dict(event.request.schema_)
+            # Codex maps the native JSON-RPC request id to a stable canonical
+            # interaction id for replay.  The live client, however, indexes its
+            # pending callback by the original request id.  Preserve that id as
+            # the provider target or SubmitInteraction can find the durable
+            # record but cannot wake the blocked Codex callback.
+            native_call_id = (
+                event.source.native_event_id or event.source.native_item_id or event.interaction_id
+            )
+            native_target = {"call_id": native_call_id}
+        for key in ("checkpoint_id", "thread_id", "checkpoint_ns"):
+            value = execution.handle.native_ref.get(key)
+            if value is not None:
+                native_target[key] = str(value)
+        provider_id = execution.interaction_provider.provider_id or execution.handle.runtime_type
+        managed_request = None
+        if getattr(event, "_invocation", None) is not None:
+            from ksadk.a2a.runtime_approval import approval_target, bound_approval_request
+
+            managed_request = bound_approval_request(
+                event, run_id=run.run_id, session_id=run.session_id
+            )
+        if managed_request is not None:
+            provider_id = "managed-a2a"
+            native_target = approval_target(managed_request)
+            if event._waiter is not None:
+                execution.managed_interaction_waiters[event.interaction_id] = event._waiter
+        if getattr(event, "_managed_input_waiter", None) is not None:
+            from ksadk.a2a.runtime_input import bound_input_waiter
+
+            waiter = bound_input_waiter(event, run_id=run.run_id, session_id=run.session_id)
+            if waiter is not None:
+                provider_id, native_target = "managed-a2a", waiter.target
+                execution.managed_interaction_waiters[event.interaction_id] = waiter
+        record = InteractionRecord(
+            interaction_id=event.interaction_id,
+            tenant_id=str(run.metadata.get("tenant_id") or ""),
+            agent_instance_id=run.agent_instance_id,
+            session_id=run.session_id,
+            run_id=run.run_id,
+            kind=event.interaction_kind,
+            request_schema=request_schema,
+            created_at=datetime.fromtimestamp(event.timestamp, timezone.utc).isoformat(),
+            expires_at=(
+                datetime.fromtimestamp(
+                    managed_request.deadline_unix_seconds, timezone.utc
+                ).isoformat()
+                if managed_request is not None
+                else None
+            ),
+            presentation=presentation,
+            provider_id=provider_id,
+            native_target=native_target,
+            continuation_metadata={"runtime_run_id": execution.runtime_run_id},
+        )
+        persisted = await self._store.request(record, guard=guard)  # type: ignore[attr-defined]
+        delivered = False
+        if provider_id == "managed-a2a" and persisted.status not in {"pending", "resolving"}:
+            waiter = execution.managed_interaction_waiters.get(record.interaction_id)
+            if waiter is None or not waiter.accepts(persisted):
+                raise InvalidCommandError("managed interaction recovery has no matching waiter")
+            if persisted.status != "resolved":
+                raise InvalidCommandError("managed interaction was cancelled or expired")
+            await self._deliver_committed_managed_interaction(persisted, waiter)
+            execution.managed_interaction_waiters.pop(record.interaction_id, None)
+            delivered = True
+        # A preceding reply can have resumed the durable run while this
+        # stream still holds its earlier WAITING snapshot. Read it again for
+        # each new interaction, including input following an approval.
+        latest = await self._store.load_run(run.run_id)
+        if latest is None:
+            raise InvalidCommandError("interaction run no longer exists")
+        if delivered:
+            if latest.state is RunState.RUNNING:
+                return latest
+            return await self._store.save_run_transition(
+                latest.model_copy(update={"state": RunState.RUNNING}),
+                expected_fence=guard.fencing_token,
+            )
+        if latest.state is RunState.WAITING:
+            return latest
+        update = {"state": RunState.WAITING}
+        if (
+            provider_id == "managed-a2a" and execution.handle.native_ref.get("managed_a2a_takeover")
+        ) or (
+            provider_id == "langgraph"
+            and native_target.get("checkpoint_id")
+            and native_target.get("thread_id")
+        ):
+            # Native input updates the live handle after start. Persist that
+            # checkpoint with WAITING so a new process can attach the same run.
+            from ksadk.runtime.executor import handle_digest
+
+            update["metadata"] = {
+                **latest.metadata,
+                "handle": execution.handle.model_dump(mode="json"),
+                "handle_digest": handle_digest(execution.handle),
+            }
+        return await self._store.save_run_transition(
+            latest.model_copy(update=update),
+            expected_fence=guard.fencing_token,
+        )
+
+    async def _deliver_committed_managed_interaction(self, record, waiter) -> None:
+        """Bridge a crash after ledger commit but before the live waiter woke."""
+        from ksadk.interaction.contracts import InteractionReceipt
+
+        if self._session_events is None:
+            raise InvalidCommandError("managed recovery requires the interaction event log")
+        cursor = 0
+        while True:
+            events = await self._session_events.read(record.session_id, cursor, 256)
+            if not events:
+                raise InvalidCommandError("managed interaction terminal fact is missing")
+            for event in events:
+                payload = event.payload
+                if (
+                    event.family == "interaction"
+                    and event.family_version == 1
+                    and event.event_type == "interaction.resolved"
+                    and event.run_id == record.run_id
+                    and all(
+                        payload.get(key) == getattr(record, key)
+                        for key in (
+                            "interaction_id",
+                            "tenant_id",
+                            "agent_instance_id",
+                            "session_id",
+                            "run_id",
+                            "kind",
+                            "revision",
+                        )
+                    )
+                ):
+                    waiter.deliver(
+                        InteractionReceipt(
+                            interaction_id=record.interaction_id,
+                            revision=record.revision,
+                            status="resolved",
+                            outcome=payload.get("outcome"),
+                            event_id=event.event_id,
+                            accepted_seq=event.seq,
+                        ),
+                        payload.get("response"),
+                    )
+                    return
+            next_cursor = max(event.seq for event in events)
+            if next_cursor <= cursor:
+                raise InvalidCommandError("managed interaction recovery cursor is invalid")
+            cursor = next_cursor

ksadk/knowledge_base/client.py

@@ -33,9 +33,10 @@ import logging
 import os
 from typing import Any, List, Optional
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from ksadk.common.aicp_env import resolve_aicp_connection
+from ksadk.common.cloud_headers import build_cloud_headers, resolve_traffic_source
 
 logger = logging.getLogger(__name__)
 
@@ -74,6 +75,7 @@ class KnowledgeBaseClient(BaseModel):
     access_key: str = ""
     secret_key: str = ""
     session_token: str = ""
+    traffic_source: str = Field(default_factory=resolve_traffic_source)
     region: str = "cn-beijing-6"
     endpoint: str = "aicp.api.ksyun.com"
     scheme: str = "https"
@@ -93,6 +95,11 @@ class KnowledgeBaseClient(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    @field_validator("traffic_source")
+    @classmethod
+    def validate_traffic_source(cls, value):
+        return resolve_traffic_source(value)
+
     def model_post_init(self, __context: Any) -> None:
         if not self.dataset_id:
             raise ValueError("dataset_id is required for KnowledgeBaseClient")
@@ -156,6 +163,13 @@ class KnowledgeBaseClient(BaseModel):
         owner = self
 
         class ObservedAicpClient(aicp_module.AicpClient):
+            def call(self, action, params, options=None):
+                options = dict(options or {})
+                options["CustomHeaders"] = build_cloud_headers(
+                    options.get("CustomHeaders"), traffic_source=owner.traffic_source
+                )
+                return super().call(action, params, options)
+
             def _check_status(self, response):
                 # The upstream SDK discards HTTP status when raising its generic
                 # ServerNetworkError. Capture only the status before delegation;

ksadk/managed_runtime.py

@@ -112,13 +112,33 @@ def installed_runtime_version(runtime_name: str) -> str:
         return ""
 
 
+def native_package_version(catalog_version: str) -> str:
+    """Keep a qualified image identity while pinning its native public version.
+
+    PEP 440 local labels distinguish platform image builds in the catalog.
+    Manifest, lock and control-plane comparisons still use the full identity.
+    Unqualified versions retain their existing exact-string semantics.
+    """
+    if "+" not in catalog_version:
+        return catalog_version
+    from packaging.version import InvalidVersion, Version
+
+    try:
+        parsed = Version(catalog_version)
+    except InvalidVersion as exc:
+        raise ManagedRuntimeError("Invalid qualified runtime version") from exc
+    if not parsed.local:
+        raise ManagedRuntimeError("Qualified runtime version requires a local label")
+    return parsed.public
+
+
 def validate_installed_runtime(resolved: ResolvedRuntime) -> str:
     installed = installed_runtime_version(resolved.name)
     if not installed:
         raise ManagedRuntimeError(
             f"本地缺少 {resolved.name} runtime；请安装 `pip install 'ksadk[{resolved.name}]'`"
         )
-    if resolved.version and installed != resolved.version:
+    if resolved.version and installed != native_package_version(resolved.version):
         raise ManagedRuntimeError(
             f"本地 {resolved.name} runtime 版本为 {installed}，"
             f"配置要求 {resolved.version}；请安装匹配版本后重试"
@@ -149,7 +169,14 @@ def validate_runtime_binary(resolved: ResolvedRuntime) -> str:
             "'ksadk[codex]'`"
         ) from exc
     output = f"{completed.stdout}\n{completed.stderr}".strip()
-    if resolved.version and resolved.version not in output:
+    expected = native_package_version(resolved.version)
+    if "+" in resolved.version:
+        import re
+        match = re.search(r"^codex-cli ([^\s]+)(?:\s|$)", output)
+        binary_matches = match is not None and match.group(1) == expected
+    else:
+        binary_matches = expected in output
+    if resolved.version and not binary_matches:
         raise ManagedRuntimeError(
             f"Codex CLI 版本输出为 {output or '(empty)'}，配置要求 {resolved.version}"
         )

ksadk/memory/adk/backends/http_ltm_backend.py

@@ -14,6 +14,7 @@ from typing import List, Optional
 import httpx
 from pydantic import ConfigDict
 
+from ksadk.common.cloud_headers import build_cloud_headers
 from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
 
 logger = logging.getLogger(__name__)
@@ -68,7 +69,7 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
 
             self._client = httpx.Client(
                 base_url=self.base_url,
-                headers=headers,
+                headers=build_cloud_headers(headers),
                 timeout=self.timeout,
             )
         return self._client

ksadk/memory/adk/backends/sdk_ltm_backend.py

@@ -26,8 +26,9 @@ import time
 import uuid
 from typing import Any, List, Set
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict, Field, field_validator
 
+from ksadk.common.cloud_headers import build_cloud_headers, resolve_traffic_source
 from ksadk.memory.adk.backends.base_ltm_backend import (
     CAP_ADD,
     CAP_DELETE,
@@ -91,6 +92,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
     access_key: str = ""
     secret_key: str = ""
     session_token: str = ""
+    traffic_source: str = Field(default_factory=resolve_traffic_source)
     region: str = "cn-beijing-6"
     endpoint: str = "aicp.api.ksyun.com"
     scheme: str = "https"
@@ -107,6 +109,11 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    @field_validator("traffic_source")
+    @classmethod
+    def validate_traffic_source(cls, value):
+        return resolve_traffic_source(value)
+
     def model_post_init(self, __context) -> None:
         if not self.access_key or not self.secret_key:
             logger.warning(
@@ -181,6 +188,13 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
         owner = self
 
         class ObservedAicpClient(aicp_module.AicpClient):
+            def call(self, action, params, options=None):
+                options = dict(options or {})
+                options["CustomHeaders"] = build_cloud_headers(
+                    options.get("CustomHeaders"), traffic_source=owner.traffic_source
+                )
+                return super().call(action, params, options)
+
             def _check_status(self, response):
                 # Preserve the status that the upstream SDK's generic exception
                 # otherwise discards. No response content is used as authority.

ksadk/memory/adk/bounded_database_session_service.py

@@ -0,0 +1,232 @@
+"""PostgreSQL phase budgets around ADK's native transactional session service."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator
+
+from sqlalchemy import event, text
+
+from ksadk.compat.adk_compat import DatabaseSessionService
+from ksadk.sessions.db_budget import DatabaseBudget, session_db_budget
+
+logger = logging.getLogger(__name__)
+
+
+class ADKSessionLockTimeout(RuntimeError):
+    """Local session serialization is busy; PostgreSQL availability is unknown."""
+
+    sqlstate = "55P03"
+
+
+class BoundedADKDatabaseSessionService(DatabaseSessionService):
+    """Keep ADK's transaction logic while bounding its database phases.
+
+    A cancelled call is never replayed.  In particular, a timed-out commit may
+    have succeeded; callers must reconcile its stable session/event ID.
+    """
+
+    def __init__(self, *, db_url: str, budget: DatabaseBudget | None = None) -> None:
+        self.db_budget = budget or session_db_budget()
+        self._active_connections: dict[asyncio.Task[Any], Any] = {}
+        self._checked_out_drivers: dict[asyncio.Task[Any], set[Any]] = {}
+        self._driver_owners: dict[int, tuple[asyncio.Task[Any], Any]] = {}
+        phase = self.db_budget
+        super().__init__(
+            db_url=db_url,
+            pool_timeout=phase.write_acquire_seconds,
+            connect_args={
+                "timeout": phase.connect_seconds,
+                "command_timeout": phase.write_sql_seconds,
+                "server_settings": {
+                    "statement_timeout": str(max(1, int(phase.write_sql_seconds * 1000))),
+                    "lock_timeout": str(max(1, int(phase.lock_seconds * 1000))),
+                },
+            },
+        )
+        event.listen(self.db_engine.sync_engine, "checkout", self._on_checkout)
+        event.listen(self.db_engine.sync_engine, "checkin", self._on_checkin)
+
+    def _on_checkout(self, dbapi_connection: Any, _record: Any, _proxy: Any) -> None:
+        task = asyncio.current_task()
+        driver = getattr(dbapi_connection, "driver_connection", None)
+        if task is not None and driver is not None:
+            self._driver_owners[id(_record)] = (task, driver)
+            self._checked_out_drivers.setdefault(task, set()).add(driver)
+
+    def _on_checkin(self, dbapi_connection: Any, _record: Any) -> None:
+        owned = self._driver_owners.pop(id(_record), None)
+        if owned is not None:
+            task, driver = owned
+            drivers = self._checked_out_drivers.get(task)
+            if drivers is not None:
+                drivers.discard(driver)
+                if not drivers:
+                    self._checked_out_drivers.pop(task, None)
+
+    async def _bounded(self, operation: Any, seconds: float) -> Any:
+        task = asyncio.create_task(operation)
+        try:
+            return await asyncio.wait_for(asyncio.shield(task), seconds)
+        except (asyncio.CancelledError, TimeoutError):
+            deadline = asyncio.get_running_loop().time() + self.db_budget.cleanup_seconds
+            task.cancel()
+            connection = self._active_connections.get(task)
+            # asyncpg terminate is synchronous and closes the server-side
+            # transaction, including prepare_tables' direct engine checkout.
+            drivers = tuple(self._checked_out_drivers.get(task, ()))
+            if not drivers and connection is not None:
+                try:
+                    drivers = (connection.sync_connection.connection.driver_connection,)
+                except Exception:
+                    pass
+            terminated = False
+            for driver in drivers:
+                try:
+                    driver.terminate()
+                    terminated = True
+                except Exception:
+                    logger.exception("Could not terminate cancelled ADK DB connection")
+            if connection is not None and not terminated:
+                try:
+                    remaining = max(0.0, deadline - asyncio.get_running_loop().time())
+                    if remaining:
+                        await asyncio.wait_for(connection.invalidate(), remaining)
+                except Exception:
+                    logger.exception("Could not invalidate cancelled ADK DB connection")
+            remaining = max(0.0, deadline - asyncio.get_running_loop().time())
+            done, _ = await asyncio.wait(
+                (task,), timeout=min(remaining, self.db_budget.cleanup_seconds / 2)
+            )
+            if not done:
+                task.cancel()
+                remaining = max(0.0, deadline - asyncio.get_running_loop().time())
+                done, _ = await asyncio.wait((task,), timeout=remaining)
+            if done:
+                await asyncio.gather(task, return_exceptions=True)
+            else:
+                if connection is not None:
+                    try:
+                        remaining = max(0.0, deadline - asyncio.get_running_loop().time())
+                        if remaining:
+                            await asyncio.wait_for(connection.invalidate(), remaining)
+                    except Exception:
+                        logger.exception("Could not invalidate cancelled ADK DB connection")
+                # A pending connection attempt has no checked-out socket to
+                # terminate; retire the pool and consume the late task result.
+                if connection is None:
+                    remaining = max(0.0, deadline - asyncio.get_running_loop().time())
+                    if remaining:
+                        try:
+                            await asyncio.wait_for(self.db_engine.dispose(close=False), remaining)
+                        except Exception:
+                            logger.exception("Could not retire cancelled ADK DB pool")
+                task.add_done_callback(
+                    lambda finished: finished.exception() if not finished.cancelled() else None
+                )
+            raise
+
+    @asynccontextmanager
+    async def _rollback_on_exception_session(
+        self, *, read_only: bool = False
+    ) -> AsyncIterator[Any]:
+        factory = self._get_database_session_factory(read_only=read_only)
+        async with factory() as sql_session:
+            try:
+                acquire = (
+                    self.db_budget.read_acquire_seconds
+                    if read_only
+                    else self.db_budget.write_acquire_seconds
+                )
+                async with asyncio.timeout(acquire):
+                    connection = await sql_session.connection()
+                current = asyncio.current_task()
+                if current is not None:
+                    self._active_connections[current] = connection
+                sql_seconds = (
+                    self.db_budget.read_sql_seconds
+                    if read_only
+                    else self.db_budget.write_sql_seconds
+                )
+                # SET LOCAL shares ADK's transaction, so it cannot split an
+                # event/state update or leak limits to the next pool borrower.
+                await sql_session.execute(
+                    text("SELECT set_config('statement_timeout', :milliseconds, true)"),
+                    {"milliseconds": str(max(1, int(sql_seconds * 1000)))},
+                )
+                await sql_session.execute(
+                    text("SELECT set_config('lock_timeout', :milliseconds, true)"),
+                    {"milliseconds": str(max(1, int(self.db_budget.lock_seconds * 1000)))},
+                )
+                yield sql_session
+            except BaseException:
+                await sql_session.rollback()
+                raise
+            finally:
+                if current := asyncio.current_task():
+                    self._active_connections.pop(current, None)
+
+    @asynccontextmanager
+    async def _with_session_lock(self, **kwargs: Any) -> AsyncIterator[None]:
+        context = super()._with_session_lock(**kwargs)
+        try:
+            await asyncio.wait_for(context.__aenter__(), self.db_budget.lock_seconds)
+        except TimeoutError as exc:
+            raise ADKSessionLockTimeout("ADK session append lock timed out") from exc
+        try:
+            yield
+        except BaseException as exc:
+            await context.__aexit__(type(exc), exc, exc.__traceback__)
+            raise
+        else:
+            await context.__aexit__(None, None, None)
+
+    async def create_session(self, **kwargs: Any) -> Any:
+        return await self._bounded(
+            super().create_session(**kwargs), self.db_budget.write_operation_seconds
+        )
+
+    async def get_session(self, **kwargs: Any) -> Any:
+        budget = (
+            self.db_budget.write_operation_seconds
+            if not self._tables_created
+            else self.db_budget.read_acquire_seconds + self.db_budget.read_sql_seconds
+        )
+        return await self._bounded(
+            super().get_session(**kwargs), budget,
+        )
+
+    async def list_sessions(self, **kwargs: Any) -> Any:
+        budget = (
+            self.db_budget.write_operation_seconds
+            if not self._tables_created
+            else self.db_budget.read_acquire_seconds + self.db_budget.read_sql_seconds
+        )
+        return await self._bounded(
+            super().list_sessions(**kwargs), budget,
+        )
+
+    async def get_user_state(self, **kwargs: Any) -> Any:
+        budget = (
+            self.db_budget.write_operation_seconds
+            if not self._tables_created
+            else self.db_budget.read_acquire_seconds + self.db_budget.read_sql_seconds
+        )
+        return await self._bounded(
+            super().get_user_state(**kwargs), budget,
+        )
+
+    async def delete_session(self, *args: Any, **kwargs: Any) -> None:
+        await self._bounded(
+            super().delete_session(*args, **kwargs), self.db_budget.write_operation_seconds
+        )
+
+    async def append_event(self, *args: Any, **kwargs: Any) -> Any:
+        return await self._bounded(
+            super().append_event(*args, **kwargs), self.db_budget.write_operation_seconds
+        )
+
+    async def close(self) -> None:
+        await self.db_engine.dispose()