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

ksadk/server/persistence_capability.py

@@ -0,0 +1,516 @@
+"""App-owned persistence and checkpoint capability coordination."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+import os
+import time
+from collections.abc import Awaitable, Callable, Mapping
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+from ksadk.sessions.persistence import gate_runtime_capabilities
+
+logger = logging.getLogger(__name__)
+
+StatusProvider = Callable[..., Awaitable[Mapping[str, Any]]]
+SessionServiceProvider = Callable[[], Any]
+
+
+@dataclass(frozen=True)
+class PersistenceCapabilitySnapshot:
+    """One internally consistent generation of persistence capability state."""
+
+    state: str
+    generation: int
+    session_persistence: Mapping[str, Any]
+    checkpoint_persistence: Mapping[str, Any]
+    native_runtime_capabilities: Mapping[str, Any]
+    runtime_capabilities: Mapping[str, Any]
+    refreshed_at: float
+    next_refresh_at: float
+
+    @property
+    def resume_supported(self) -> bool:
+        resume = self.runtime_capabilities.get("ResumeRun") or {}
+        return bool(isinstance(resume, Mapping) and resume.get("Supported") is True)
+
+    @property
+    def persistence_gate(self) -> Mapping[str, Any]:
+        checkpoint = self.runtime_capabilities.get("Checkpoint") or {}
+        if not isinstance(checkpoint, Mapping):
+            return {}
+        gate = checkpoint.get("PersistenceGate") or {}
+        return gate if isinstance(gate, Mapping) else {}
+
+
+class PersistenceCapabilityCoordinator:
+    """Single-flight, bidirectional refresh for one runtime app."""
+
+    def __init__(
+        self,
+        *,
+        clock: Callable[[], float] = time.monotonic,
+        refresh_interval: float | None = None,
+        initial_retry_delay: float = 5.0,
+        background: bool = True,
+    ) -> None:
+        self._clock = clock
+        self._refresh_interval = max(
+            0.1,
+            float(
+                refresh_interval
+                if refresh_interval is not None
+                else os.getenv("KSADK_PERSISTENCE_PROBE_CACHE_TTL") or "30"
+            ),
+        )
+        self._initial_retry_delay = max(0.1, float(initial_retry_delay))
+        self._background_enabled = background
+        self._snapshot: PersistenceCapabilitySnapshot | None = None
+        self._runner: Any | None = None
+        self._framework = ""
+        self._status_provider: StatusProvider | None = None
+        self._session_service_provider: SessionServiceProvider | None = None
+        self._refresh_task: asyncio.Task[PersistenceCapabilitySnapshot] | None = None
+        self._background_task: asyncio.Task[None] | None = None
+        self._consecutive_failures = 0
+
+    @property
+    def snapshot(self) -> PersistenceCapabilitySnapshot | None:
+        return self._snapshot
+
+    def start(
+        self,
+        *,
+        runner: Any | None,
+        framework: str,
+        status_provider: StatusProvider,
+        session_service_provider: SessionServiceProvider | None = None,
+    ) -> None:
+        """Bind the app and eagerly begin its first capability refresh."""
+        self._bind(
+            runner=runner,
+            framework=framework,
+            status_provider=status_provider,
+            session_service_provider=session_service_provider,
+        )
+        if self._refresh_task is None or self._refresh_task.done():
+            self._refresh_task = asyncio.create_task(self._refresh())
+        self._ensure_background()
+
+    def invalidate(self, runner: Any | None = None) -> None:
+        """Discard state when the app swaps its active runner."""
+        for task in (self._background_task, self._refresh_task):
+            if task is not None and not task.done():
+                task.cancel()
+        self._background_task = None
+        self._refresh_task = None
+        self._snapshot = None
+        self._consecutive_failures = 0
+        if runner is not None:
+            self._runner = runner
+
+    async def get_snapshot(
+        self,
+        *,
+        runner: Any | None,
+        framework: str,
+        status_provider: StatusProvider,
+        session_service_provider: SessionServiceProvider | None = None,
+        force: bool = False,
+        wait_timeout: float | None = None,
+    ) -> PersistenceCapabilitySnapshot:
+        self._bind(
+            runner=runner,
+            framework=framework,
+            status_provider=status_provider,
+            session_service_provider=session_service_provider,
+        )
+        now = self._clock()
+        if not force and self._snapshot is not None and now < self._snapshot.next_refresh_at:
+            self._ensure_background()
+            return self._snapshot
+
+        task = self._refresh_task
+        if task is None or task.done():
+            task = asyncio.create_task(self._refresh())
+            self._refresh_task = task
+        self._ensure_background()
+
+        try:
+            if wait_timeout is None:
+                return await asyncio.shield(task)
+            return await asyncio.wait_for(asyncio.shield(task), timeout=max(0.0, wait_timeout))
+        except asyncio.TimeoutError:
+            if self._snapshot is not None and not self._snapshot.resume_supported:
+                return self._snapshot
+            return self._checking_snapshot(runner)
+
+    def _bind(
+        self,
+        *,
+        runner: Any | None,
+        framework: str,
+        status_provider: StatusProvider,
+        session_service_provider: SessionServiceProvider | None,
+    ) -> None:
+        if self._runner is not None and self._runner is not runner:
+            self.invalidate(runner)
+        self._runner = runner
+        self._framework = str(framework or "").strip().lower()
+        self._status_provider = status_provider
+        self._session_service_provider = session_service_provider
+
+    async def _refresh(self) -> PersistenceCapabilitySnapshot:
+        runner = self._runner
+        provider = self._status_provider
+        if provider is None:
+            raise RuntimeError("Persistence capability coordinator is not bound")
+
+        started = self._clock()
+        refresh_runner = getattr(runner, "refresh_runtime_capabilities", None)
+        if not callable(refresh_runner):
+            refresh_runner = getattr(runner, "prepare_runtime_capabilities", None)
+        runner_refresh_failed = False
+
+        async def prepare_runner() -> None:
+            nonlocal runner_refresh_failed
+            if callable(refresh_runner):
+                try:
+                    result = refresh_runner()
+                    if inspect.isawaitable(result):
+                        await result
+                except Exception:
+                    runner_refresh_failed = True
+
+        async def probe() -> Mapping[str, Any]:
+            try:
+                try:
+                    return await provider(framework=self._framework, use_cache=False)
+                except TypeError as exc:
+                    if "use_cache" not in str(exc):
+                        raise
+                    return await provider(framework=self._framework)
+            except Exception:
+                unavailable = self._refresh_failed_status()
+                return {"Session": unavailable, "Checkpoint": dict(unavailable)}
+
+        _, persistence_status = await asyncio.gather(prepare_runner(), probe())
+        session, checkpoint = self._split_persistence_status(persistence_status)
+        service_provider = self._session_service_provider
+        service = service_provider() if service_provider is not None else None
+        refresh_service = getattr(service, "refresh_persistence_capability", None)
+        if getattr(service, "degraded", False) is True and callable(refresh_service):
+            try:
+                result = refresh_service()
+                if inspect.isawaitable(result):
+                    await result
+            except Exception:
+                pass
+        if getattr(service, "degraded", False) is True:
+            session.update(
+                {
+                    "Status": "degraded",
+                    "Ready": False,
+                    "SharedAcrossPods": False,
+                    "ReasonCode": "SESSION_STORE_DEGRADED",
+                    "Reason": "Session persistence is temporarily degraded",
+                }
+            )
+        native = self._runner_capabilities(runner)
+        if runner_refresh_failed:
+            native = self._runtime_refresh_failed_capabilities(native)
+        effective = self._gate(native, session, checkpoint)
+        effective = self._gate_runner_failure(native, effective)
+        state = "ready" if self._is_ready(native, session, checkpoint, effective) else "degraded"
+        if state == "ready":
+            self._consecutive_failures = 0
+            delay = self._refresh_interval
+        else:
+            self._consecutive_failures += 1
+            delay = min(
+                self._refresh_interval,
+                self._initial_retry_delay * (2 ** (self._consecutive_failures - 1)),
+            )
+        refreshed_at = self._clock()
+        generation = (self._snapshot.generation if self._snapshot else 0) + 1
+        snapshot = PersistenceCapabilitySnapshot(
+            state=state,
+            generation=generation,
+            session_persistence=deepcopy(session),
+            checkpoint_persistence=deepcopy(checkpoint),
+            native_runtime_capabilities=deepcopy(native),
+            runtime_capabilities=deepcopy(effective),
+            refreshed_at=refreshed_at,
+            next_refresh_at=refreshed_at + delay,
+        )
+        previous_state = self._snapshot.state if self._snapshot else "checking"
+        self._snapshot = snapshot
+        if previous_state != state:
+            checkpoint_capability = effective.get("Checkpoint") or {}
+            logger.info(
+                "KSADK persistence capability state changed: %s -> %s",
+                previous_state,
+                state,
+                extra={
+                    "persistence_capability_state": state,
+                    "persistence_capability_generation": generation,
+                    "persistence_reason_code": str(
+                        checkpoint_capability.get("ReasonCode") or "READY"
+                    ),
+                    "persistence_backend": str(
+                        checkpoint_capability.get("Backend") or "none"
+                    ),
+                    "persistence_refresh_ms": int((refreshed_at - started) * 1000),
+                },
+            )
+        return snapshot
+
+    @staticmethod
+    def _split_persistence_status(
+        status: Mapping[str, Any],
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        if isinstance(status.get("Session"), Mapping):
+            return dict(status["Session"]), dict(status.get("Checkpoint") or {})
+        legacy = dict(status)
+        return legacy, dict(legacy)
+
+    @staticmethod
+    def _runner_capabilities(runner: Any) -> dict[str, Any]:
+        describe = getattr(runner, "get_runtime_capabilities", None)
+        return dict(describe() if callable(describe) else {})
+
+    @staticmethod
+    def _requires_persistence(capabilities: Mapping[str, Any]) -> bool:
+        checkpoint = capabilities.get("Checkpoint") or {}
+        if not isinstance(checkpoint, Mapping):
+            return False
+        backend = str(checkpoint.get("Backend") or "").lower()
+        return bool(
+            checkpoint.get("Durable") is True
+            and checkpoint.get("SharedAcrossPods") is True
+        ) or "postgres" in backend
+
+    @classmethod
+    def _gate(
+        cls,
+        capabilities: Mapping[str, Any],
+        session: Mapping[str, Any],
+        checkpoint: Mapping[str, Any],
+    ) -> dict[str, Any]:
+        if not cls._requires_persistence(capabilities):
+            return deepcopy(dict(capabilities))
+        return gate_runtime_capabilities(capabilities, session, checkpoint)
+
+    @classmethod
+    def _gate_runner_failure(
+        cls,
+        native: Mapping[str, Any],
+        effective: Mapping[str, Any],
+    ) -> dict[str, Any]:
+        if not cls._requires_persistence(native):
+            return deepcopy(dict(effective))
+        resume = effective.get("ResumeRun") or {}
+        if isinstance(resume, Mapping) and resume.get("Supported") is True:
+            return deepcopy(dict(effective))
+
+        gated = deepcopy(dict(effective))
+        checkpoint = dict(gated.get("Checkpoint") or {})
+        existing_gate = checkpoint.get("PersistenceGate") or {}
+        if isinstance(existing_gate, Mapping) and existing_gate:
+            return gated
+        native_checkpoint = native.get("Checkpoint") or {}
+        reason_code = str(
+            checkpoint.get("ReasonCode")
+            or (resume.get("ReasonCode") if isinstance(resume, Mapping) else "")
+            or "RUNTIME_CAPABILITY_UNAVAILABLE"
+        )
+        reason = str(
+            checkpoint.get("Reason")
+            or (resume.get("Reason") if isinstance(resume, Mapping) else "")
+            or "Runtime checkpoint capability is unavailable"
+        )
+        native_backend = str(
+            checkpoint.get("NativeBackend")
+            or (
+                native_checkpoint.get("Backend")
+                if isinstance(native_checkpoint, Mapping)
+                else ""
+            )
+            or checkpoint.get("Backend")
+            or "none"
+        )
+        checkpoint.update(
+            {
+                "Supported": False,
+                "Backend": "none",
+                "NativeBackend": native_backend,
+                "Durable": False,
+                "SharedAcrossPods": False,
+                "ResumeMode": "none",
+                "ReasonCode": reason_code,
+                "Reason": reason,
+                "PersistenceGate": {
+                    "BlockedStore": "checkpoint",
+                    "Source": "runtime",
+                    "ReasonCode": reason_code,
+                    "Reason": reason,
+                },
+            }
+        )
+        gated["Checkpoint"] = checkpoint
+        gated["ResumeRun"] = {
+            "Supported": False,
+            "ResumeMode": "none",
+            "ReasonCode": reason_code,
+            "Reason": reason,
+        }
+        return gated
+
+    @staticmethod
+    def _refresh_failed_status() -> dict[str, Any]:
+        return {
+            "Configured": True,
+            "Status": "error",
+            "Ready": False,
+            "Backend": "postgres",
+            "SharedAcrossPods": False,
+            "EffectiveFor": "new_runs_only",
+            "Source": "none",
+            "ReasonCode": "PERSISTENCE_REFRESH_FAILED",
+            "Reason": "Persistence capability refresh failed",
+        }
+
+    @staticmethod
+    def _runtime_refresh_failed_capabilities(
+        capabilities: Mapping[str, Any],
+    ) -> dict[str, Any]:
+        failed = deepcopy(dict(capabilities))
+        checkpoint = dict(failed.get("Checkpoint") or {})
+        checkpoint.update(
+            {
+                "Supported": False,
+                "Durable": False,
+                "ReasonCode": "RUNTIME_CAPABILITY_REFRESH_FAILED",
+                "Reason": "Runtime capability refresh failed",
+            }
+        )
+        failed["Checkpoint"] = checkpoint
+        failed["ResumeRun"] = {
+            "Supported": False,
+            "ResumeMode": "none",
+            "ReasonCode": "RUNTIME_CAPABILITY_REFRESH_FAILED",
+            "Reason": "Runtime capability refresh failed",
+        }
+        return failed
+
+    @classmethod
+    def _is_ready(
+        cls,
+        native: Mapping[str, Any],
+        session: Mapping[str, Any],
+        checkpoint: Mapping[str, Any],
+        effective: Mapping[str, Any],
+    ) -> bool:
+        if not cls._requires_persistence(native):
+            return True
+        resume = effective.get("ResumeRun") or {}
+        return bool(
+            session.get("Ready") is True
+            and checkpoint.get("Ready") is True
+            and isinstance(resume, Mapping)
+            and resume.get("Supported") is True
+        )
+
+    def _checking_snapshot(self, runner: Any) -> PersistenceCapabilitySnapshot:
+        now = self._clock()
+        unavailable = {
+            "Configured": True,
+            "Status": "checking",
+            "Ready": False,
+            "Backend": "postgres",
+            "SharedAcrossPods": False,
+            "EffectiveFor": "new_runs_only",
+            "Source": "none",
+            "ReasonCode": "PERSISTENCE_CHECK_TIMEOUT",
+            "Reason": "Persistence capability check is still in progress",
+        }
+        native = self._runner_capabilities(runner)
+        effective = gate_runtime_capabilities(native, unavailable, unavailable)
+        return PersistenceCapabilitySnapshot(
+            state="checking",
+            generation=0,
+            session_persistence=dict(unavailable),
+            checkpoint_persistence=dict(unavailable),
+            native_runtime_capabilities=native,
+            runtime_capabilities=effective,
+            refreshed_at=now,
+            next_refresh_at=now,
+        )
+
+    def _ensure_background(self) -> None:
+        if not self._background_enabled:
+            return
+        if self._background_task is None or self._background_task.done():
+            self._background_task = asyncio.create_task(self._background_loop())
+
+    async def _background_loop(self) -> None:
+        while True:
+            refresh_task = self._refresh_task
+            if refresh_task is not None and not refresh_task.done():
+                try:
+                    await asyncio.shield(refresh_task)
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    logger.warning(
+                        "KSADK persistence capability refresh failed",
+                        extra={
+                            "persistence_capability_state": "degraded",
+                            "persistence_reason_code": "CAPABILITY_REFRESH_FAILED",
+                        },
+                    )
+            snapshot = self._snapshot
+            delay = (
+                max(0.0, snapshot.next_refresh_at - self._clock())
+                if snapshot is not None
+                else self._initial_retry_delay
+            )
+            await asyncio.sleep(delay)
+            try:
+                await self.get_snapshot(
+                    runner=self._runner,
+                    framework=self._framework,
+                    status_provider=self._status_provider,
+                    session_service_provider=self._session_service_provider,
+                    force=True,
+                )
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.warning(
+                    "KSADK persistence capability background refresh failed",
+                    extra={
+                        "persistence_capability_state": "degraded",
+                        "persistence_reason_code": "CAPABILITY_REFRESH_FAILED",
+                    },
+                )
+
+    async def aclose(self) -> None:
+        tasks = [
+            task
+            for task in (self._background_task, self._refresh_task)
+            if task is not None and not task.done()
+        ]
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        self._background_task = None
+        self._refresh_task = None
+
+
+__all__ = ["PersistenceCapabilityCoordinator", "PersistenceCapabilitySnapshot"]

ksadk/server/routes/common.py

@@ -16,6 +16,7 @@ from fastapi.responses import FileResponse, Response
 
 from ksadk.conversations.attachments import compact_attachment_result_for_session
 from ksadk.conversations.model_context import normalize_model_metadata
+from ksadk.runtime_context import PlatformIdentityContext
 from ksadk.runtime_state import load_state as load_runtime_state
 from ksadk.server.factory import (
     RuntimeAppState,
@@ -28,6 +29,7 @@ from ksadk.server.terminal_sessions import (
     register_terminal_routes,
 )
 from ksadk.sessions import Session
+from ksadk.sessions.invocation_identity import identity_scope_ref
 from ksadk.sessions.local_service import resolve_local_session_dir
 from ksadk.ui_config import UI_PROFILE_CUSTOM, resolve_ui_config
 from ksadk_runtime_common.workspace_files import (
@@ -35,6 +37,10 @@ from ksadk_runtime_common.workspace_files import (
     workspace_files_enabled,
 )
 
+from ..invocation_identity import (
+    trusted_invocation_identity_from_headers,
+    trusted_invocation_identity_headers,
+)
 from . import dependencies as deps
 from .routers import health_meta_router
 
@@ -92,8 +98,17 @@ _MAX_REFERENCE_TEXT_BYTES = 3_000_000
 _UPLOAD_URI_SCHEME = "ksadk-upload://"
 
 
-def _workspace_root_dir() -> Path:
-    return Path(resolve_local_session_dir()) / "workspace"
+def _workspace_root_dir(
+    invocation_identity: PlatformIdentityContext | None = None,
+) -> Path:
+    root = Path(resolve_local_session_dir()) / "workspace"
+    if invocation_identity is not None and not invocation_identity.is_empty:
+        root = root / "identities" / identity_scope_ref(invocation_identity)
+    return root
+
+
+def _workspace_request_root(request: Request) -> Path:
+    return _workspace_root_dir(trusted_invocation_identity_from_headers(request.headers))
 
 
 _NATIVE_TUI_FRAMEWORKS = {"hermes", "openclaw"}
@@ -104,10 +119,18 @@ def _current_framework() -> str:
     return str(context.runtime_type if context is not None else "").strip().lower()
 
 
-def _build_native_terminal_capability(framework: str) -> dict[str, Any]:
+def _build_native_terminal_capability(
+    framework: str,
+    invocation_identity: PlatformIdentityContext | None = None,
+) -> dict[str, Any]:
+    isolated_identity = (
+        isinstance(invocation_identity, PlatformIdentityContext)
+        and not invocation_identity.is_empty
+    )
     enabled = (
         native_terminal_supported()
         and str(framework or "").strip().lower() in _NATIVE_TUI_FRAMEWORKS
+        and not isolated_identity
     )
     return {
         "Enabled": enabled,
@@ -127,6 +150,7 @@ def _register_integrated_routers(app: FastAPI, state: RuntimeAppState) -> None:
         create_workspace_files_router(
             root_getter=_workspace_root_dir,
             enabled_getter=lambda: workspace_files_enabled(default=True),
+            request_root_getter=_workspace_request_root,
         )
     )
     state.terminal_manager = TerminalSessionManager(
@@ -137,6 +161,9 @@ def _register_integrated_routers(app: FastAPI, state: RuntimeAppState) -> None:
         app,
         state.terminal_manager,
         bind_context=lambda: bind_runtime_state(state),
+        identity_access_allowed=lambda headers: (
+            trusted_invocation_identity_from_headers(headers).is_empty
+        ),
     )
 
 
@@ -631,6 +658,7 @@ async def _workspace_runtime_request(
     *,
     params: Optional[Dict[str, Any]] = None,
     files: Optional[Dict[str, Any]] = None,
+    invocation_identity: PlatformIdentityContext | None = None,
 ) -> httpx.Response:
     runtime_app = get_state().app
     if runtime_app is None:
@@ -640,6 +668,7 @@ async def _workspace_runtime_request(
         response = await client.request(
             method,
             runtime_path,
+            headers=trusted_invocation_identity_headers(invocation_identity),
             params=params,
             files=files,
         )