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

ksadk/sessions/resilient.py

@@ -2,14 +2,19 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from collections.abc import AsyncIterator
+from contextvars import ContextVar
 from typing import Any, Optional, cast
 
 from ksadk.sessions.base import (
     BaseSessionService,
+    CheckpointEventQuery,
     Session,
     SessionEvent,
+    SessionEventQuery,
     SessionState,
 )
+from ksadk.sessions.errors import CheckpointScanRestartRequired
 from ksadk.sessions.in_memory import InMemorySessionService
 from ksadk.sessions.resilience import is_session_backend_failure
 
@@ -39,17 +44,30 @@ class ResilientSessionService(BaseSessionService):
         self._hydrate_lock = asyncio.Lock()
         self._primary_session_lock = asyncio.Lock()
         self._primary_session_ids: set[str] = set()
+        self._dirty_session_ids: set[str] = set()
+        self._checkpoint_partition: ContextVar[
+            tuple[frozenset[str], frozenset[str]] | None
+        ] = ContextVar(f"checkpoint_partition_{id(self)}", default=None)
         self._probe_task: asyncio.Task[None] | None = None
 
+    @property
+    def canonical_event_service(self) -> BaseSessionService:
+        """Return the single durable authority for canonical runtime facts.
+
+        Legacy session state remains live-first. Canonical event sequences cannot
+        fail over to the independently sequenced fallback without splitting one
+        logical log, so RuntimeEventStore binds directly to the primary service.
+        """
+
+        return self.primary
+
     @property
     def degraded(self) -> bool:
         return not self._primary_enabled
 
     # This service is intentionally live-first and writes two independently
-    # sequenced stores.  Even when both children can atomically bind a local
-    # seq, the wrapper cannot guarantee one shared physical seq/fact across
-    # both writes, so it inherits BaseSessionService's empty canonical storage
-    # capabilities and RuntimeEventStore fails closed before either write.
+    # sequenced stores, so the wrapper itself does not claim canonical storage
+    # capabilities.
 
     async def _call_primary(self, method_name: str, *args: Any, **kwargs: Any) -> tuple[bool, Any]:
         if not self._primary_enabled:
@@ -326,6 +344,346 @@ class ResilientSessionService(BaseSessionService):
             await self.fallback.count_events(session_id, after_seq_id, before_seq_id),
         )
 
+    async def get_sessions_by_ids(self, session_ids: list[str]) -> list[Session]:
+        result: dict[str, Session] = {}
+        ok, durable = await self._call_primary("get_sessions_by_ids", session_ids)
+        if ok:
+            result.update({session.id: session for session in durable or []})
+        live = await self.fallback.get_sessions_by_ids(session_ids)
+        result.update({session.id: session for session in live})
+        return [result[session_id] for session_id in session_ids if session_id in result]
+
+    async def list_session_metadata(
+        self, agent_id: Optional[str] = None, user_id: Optional[str] = None
+    ) -> list[Session]:
+        result: dict[str, Session] = {}
+        ok, durable = await self._call_primary("list_session_metadata", agent_id, user_id)
+        if ok:
+            result.update({session.id: session for session in durable or []})
+        live = await self.fallback.list_session_metadata(agent_id, user_id)
+        result.update({session.id: session for session in live})
+        sessions = list(result.values())
+        sessions.sort(key=lambda item: (item.updated_at, item.created_at, item.id), reverse=True)
+        return sessions
+
+    async def _partition_query_session_ids(
+        self, session_ids: list[str] | None, agent_id: str | None
+    ) -> tuple[list[str], list[str]]:
+        if session_ids is None:
+            sessions = await self.list_session_metadata(agent_id)
+            ids = [session.id for session in sessions]
+        else:
+            ids = list(dict.fromkeys(session_ids))
+        dirty = set(self._dirty_session_ids)
+        durable_ids: set[str] = set()
+        if ids:
+            ok, durable_metadata = await self._call_primary("get_sessions_by_ids", ids)
+            if ok:
+                durable_ids = {session.id for session in durable_metadata or []}
+        live_metadata = await self.fallback.get_sessions_by_ids(ids)
+        live_ids = {session.id for session in live_metadata}
+        clean_ids = [
+            session_id
+            for session_id in ids
+            if session_id not in dirty and session_id in durable_ids
+        ]
+        fallback_ids = [
+            session_id for session_id in ids if session_id in dirty or session_id in live_ids
+        ]
+        return clean_ids, fallback_ids
+
+    async def query_events(self, query: SessionEventQuery) -> list[SessionEvent]:
+        clean_ids, live_ids = await self._partition_query_session_ids(
+            query.session_ids, query.agent_id
+        )
+        window = query.offset + query.limit
+        durable_events: list[SessionEvent] = []
+        if clean_ids:
+            ok, result = await self._call_primary(
+                "query_events",
+                SessionEventQuery(
+                    **{**query.__dict__, "session_ids": clean_ids, "offset": 0, "limit": window}
+                ),
+            )
+            if ok:
+                durable_events = cast(list[SessionEvent], result or [])
+            else:
+                return await self.fallback.query_events(query)
+        live_events = (
+            await self.fallback.query_events(
+                SessionEventQuery(
+                    **{**query.__dict__, "session_ids": live_ids, "offset": 0, "limit": window}
+                )
+            )
+            if live_ids
+            else []
+        )
+        merged_by_id = {event.id: event for event in [*durable_events, *live_events]}
+        merged = list(merged_by_id.values())
+        if query.order_by_seq:
+            merged.sort(key=lambda event: (event.session_id, event.seq_id, event.id))
+        else:
+            merged.sort(
+                key=lambda event: (event.timestamp, event.session_id, event.seq_id, event.id)
+            )
+        if query.from_start:
+            return merged[query.offset : query.offset + query.limit]
+        end = max(len(merged) - query.offset, 0)
+        return merged[max(end - query.limit, 0) : end]
+
+    async def count_event_query(self, query: SessionEventQuery) -> int:
+        # Deduplicate by physical event id because a hydrated clean session may
+        # be present in both children during recovery.
+        events = await self.query_events(
+            SessionEventQuery(
+                **{
+                    **query.__dict__,
+                    "offset": 0,
+                    "limit": 2**31 - 1,
+                    "from_start": True,
+                }
+            )
+        )
+        return len(events)
+
+    async def get_checkpoint_lookup_stats(
+        self, session_id: str, run_id: str, checkpoint_id: str
+    ) -> dict[str, object]:
+        if session_id not in self._dirty_session_ids:
+            ok, stats = await self._call_primary(
+                "get_checkpoint_lookup_stats", session_id, run_id, checkpoint_id
+            )
+            if ok:
+                return cast(dict[str, object], stats)
+        return await self.fallback.get_checkpoint_lookup_stats(
+            session_id, run_id, checkpoint_id
+        )
+
+    async def scan_checkpoint_events(
+        self, query: CheckpointEventQuery
+    ) -> list[SessionEvent]:
+        if query.limit < 1 or query.limit > 50:
+            raise ValueError("checkpoint scan limit must be between 1 and 50")
+        clean_ids, live_ids = await self._partition_query_session_ids(
+            query.session_ids, query.agent_id
+        )
+        if clean_ids and not live_ids:
+            ok, result = await self._call_primary(
+                "scan_checkpoint_events",
+                CheckpointEventQuery(**{**query.__dict__, "session_ids": clean_ids}),
+            )
+            if ok:
+                return cast(list[SessionEvent], result or [])
+            return await self.fallback.scan_checkpoint_events(query)
+        if live_ids and not clean_ids:
+            return await self.fallback.scan_checkpoint_events(
+                CheckpointEventQuery(**{**query.__dict__, "session_ids": live_ids})
+            )
+        window = query.offset + query.limit
+        durable: list[SessionEvent] = []
+        if clean_ids:
+            ok, result = await self._call_primary(
+                "scan_checkpoint_events",
+                CheckpointEventQuery(
+                    **{
+                        **query.__dict__,
+                        "session_ids": clean_ids,
+                        "offset": 0,
+                        "limit": min(window, 50),
+                    }
+                ),
+            )
+            if ok:
+                durable = cast(list[SessionEvent], result or [])
+            else:
+                return await self.fallback.scan_checkpoint_events(query)
+        live = (
+            await self.fallback.scan_checkpoint_events(
+                CheckpointEventQuery(
+                    **{
+                        **query.__dict__,
+                        "session_ids": live_ids,
+                        "offset": 0,
+                        "limit": min(window, 50),
+                    }
+                )
+            )
+            if live_ids
+            else []
+        )
+        merged = list({event.id: event for event in [*durable, *live]}.values())
+        merged.sort(key=lambda event: (event.timestamp, event.session_id, event.seq_id, event.id))
+        return merged[query.offset : query.offset + query.limit]
+
+    async def iter_checkpoint_event_chunks(
+        self, query: CheckpointEventQuery
+    ) -> AsyncIterator[list[SessionEvent]]:
+        if query.limit < 1 or query.limit > 50:
+            raise ValueError("checkpoint scan limit must be between 1 and 50")
+        clean_ids, live_ids = await self._partition_query_session_ids(
+            query.session_ids, query.agent_id
+        )
+        token = self._checkpoint_partition.set(
+            (frozenset(clean_ids), frozenset(live_ids))
+        )
+        batches = self._iter_checkpoint_event_chunks_partitioned(
+            query, clean_ids, live_ids
+        ).__aiter__()
+        try:
+            async for batch in batches:
+                yield batch
+        finally:
+            close = getattr(batches, "aclose", None)
+            if callable(close):
+                await close()
+            self._checkpoint_partition.reset(token)
+
+    async def _iter_checkpoint_event_chunks_partitioned(
+        self,
+        query: CheckpointEventQuery,
+        clean_ids: list[str],
+        live_ids: list[str],
+    ) -> AsyncIterator[list[SessionEvent]]:
+        if query.limit < 1 or query.limit > 50:
+            raise ValueError("checkpoint scan limit must be between 1 and 50")
+        if not clean_ids:
+            batches = self.fallback.iter_checkpoint_event_chunks(
+                CheckpointEventQuery(**{**query.__dict__, "session_ids": live_ids})
+            ).__aiter__()
+            try:
+                async for batch in batches:
+                    yield batch
+            finally:
+                close = getattr(batches, "aclose", None)
+                if callable(close):
+                    await close()
+            return
+        if not live_ids:
+            batches = self.primary.iter_checkpoint_event_chunks(
+                CheckpointEventQuery(**{**query.__dict__, "session_ids": clean_ids})
+            ).__aiter__()
+            try:
+                async for batch in batches:
+                    yield batch
+            except Exception as exc:
+                if not is_session_backend_failure(exc):
+                    raise
+                self._disable_primary(exc)
+                raise CheckpointScanRestartRequired(
+                    "primary failed during clean checkpoint scan"
+                ) from exc
+            finally:
+                close = getattr(batches, "aclose", None)
+                if callable(close):
+                    await close()
+            return
+
+        iterators = {
+            "clean": self.primary.iter_checkpoint_event_chunks(
+                CheckpointEventQuery(
+                    **{**query.__dict__, "session_ids": clean_ids, "offset": 0, "limit": 50}
+                )
+            ).__aiter__(),
+            "live": self.fallback.iter_checkpoint_event_chunks(
+                CheckpointEventQuery(
+                    **{**query.__dict__, "session_ids": live_ids, "offset": 0, "limit": 50}
+                )
+            ).__aiter__(),
+        }
+        pages: dict[str, list[SessionEvent]] = {"clean": [], "live": []}
+        indexes = {"clean": 0, "live": 0}
+        exhausted = {"clean": False, "live": False}
+        skipped = 0
+
+        async def load(kind: str) -> None:
+            try:
+                pages[kind] = await anext(iterators[kind])
+            except StopAsyncIteration:
+                pages[kind] = []
+                exhausted[kind] = True
+            except Exception as exc:
+                if kind != "clean" or not is_session_backend_failure(exc):
+                    raise
+                self._disable_primary(exc)
+                raise CheckpointScanRestartRequired(
+                    "primary failed during mixed checkpoint scan"
+                ) from exc
+            indexes[kind] = 0
+            if len(pages[kind]) < 50:
+                exhausted[kind] = True
+
+        try:
+            await load("clean")
+            await load("live")
+            output: list[SessionEvent] = []
+            while True:
+                for kind in ("clean", "live"):
+                    if indexes[kind] >= len(pages[kind]) and not exhausted[kind]:
+                        await load(kind)
+                candidates = [
+                    (event.timestamp, event.session_id, event.seq_id, event.id, kind, event)
+                    for kind in ("clean", "live")
+                    for event in pages[kind][indexes[kind] : indexes[kind] + 1]
+                ]
+                if not candidates:
+                    if output:
+                        yield output
+                    break
+                *_, kind, event = min(candidates)
+                indexes[kind] += 1
+                if skipped < query.offset:
+                    skipped += 1
+                    continue
+                output.append(event)
+                if len(output) == query.limit:
+                    yield output
+                    output = []
+        finally:
+            for iterator in iterators.values():
+                close = getattr(iterator, "aclose", None)
+                if callable(close):
+                    await close()
+
+    async def get_checkpoint_stats(
+        self, keys: list[tuple[str, str, str]]
+    ) -> dict[str, object]:
+        if len(keys) > 50:
+            raise ValueError("checkpoint stats batch cannot exceed 50 keys")
+        unique_keys = list(dict.fromkeys(keys))
+        active = self._checkpoint_partition.get()
+        if active is None:
+            clean_ids, live_ids = await self._partition_query_session_ids(
+                list(dict.fromkeys(key[0] for key in unique_keys)), None
+            )
+            clean_set, live_set = set(clean_ids), set(live_ids)
+        else:
+            clean_set, live_set = map(set, active)
+        clean_keys = [key for key in unique_keys if key[0] in clean_set]
+        live_keys = [key for key in unique_keys if key[0] in live_set]
+        durable: dict[str, object] = {"audits": {}, "latest_seq_ids": {}}
+        if clean_keys:
+            ok, value = await self._call_primary("get_checkpoint_stats", clean_keys)
+            if not ok:
+                raise CheckpointScanRestartRequired(
+                    "primary failed while reading checkpoint stats"
+                )
+            durable = cast(dict[str, object], value or durable)
+        live = (
+            await self.fallback.get_checkpoint_stats(live_keys)
+            if live_keys
+            else {"audits": {}, "latest_seq_ids": {}}
+        )
+        return {
+            "audits": {
+                **cast(dict, durable["audits"]),
+                **cast(dict, live["audits"]),
+            },
+            "latest_seq_ids": {
+                **cast(dict, durable["latest_seq_ids"]),
+                **cast(dict, live["latest_seq_ids"]),
+            },
+        }
+
     async def get_events_for_agent(
         self,
         agent_id: str,

ksadk/sessions/topology.py

@@ -0,0 +1,135 @@
+"""Resolve the non-secret Session and framework checkpoint storage topology."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Literal
+
+StorageSource = Literal[
+    "none",
+    "explicit",
+    "session_fallback",
+    "checkpoint_fallback",
+    "framework_override",
+]
+
+
+@dataclass(frozen=True)
+class StorageTarget:
+    """A logical persistent-storage target without connection diagnostics."""
+
+    backend: str
+    dsn: str = ""
+    source: StorageSource = "none"
+    explicit_local: bool = False
+
+
+@dataclass(frozen=True)
+class PersistenceTopology:
+    """Effective Session and framework-native checkpoint storage targets."""
+
+    session: StorageTarget
+    checkpoint: StorageTarget
+    framework: str
+
+
+def _first_configured(*names: str) -> str:
+    for name in names:
+        value = os.getenv(name, "").strip()
+        if value:
+            return value
+    return ""
+
+
+def _session_backend(*, override: str | None = None) -> tuple[str, bool, bool]:
+    raw = (
+        override
+        if override is not None
+        else _first_configured(
+            "KSADK_SESSION_BACKEND",
+            "AGENTENGINE_SESSION_BACKEND",
+            "KSADK_STM_BACKEND",
+        )
+    )
+    normalized = raw.strip().lower()
+    if normalized == "sqlite":
+        normalized = "local"
+    return (
+        normalized or "local",
+        normalized in {"local", "memory"},
+        bool(normalized),
+    )
+
+
+def _framework_checkpoint_dsn(framework: str) -> str:
+    if framework == "adk":
+        return _first_configured("KSADK_ADK_SESSION_URL")
+    if framework in {"langgraph", "langchain", "deepagents"}:
+        return _first_configured("KSADK_LANGGRAPH_CHECKPOINT_DSN")
+    return ""
+
+
+def resolve_persistence_topology(
+    *,
+    framework: str | None = None,
+    session_backend: str | None = None,
+) -> PersistenceTopology:
+    """Resolve dual persistence with documented Session/Checkpoint fallbacks.
+
+    ``session_backend`` is an internal compatibility input for the legacy
+    session resolver.  A caller explicitly selecting ``local``, ``sqlite`` or
+    ``memory`` keeps that local choice even when a checkpoint DSN is present.
+    """
+    normalized_framework = (framework or "").strip().lower()
+    backend, explicit_local, explicit_backend = _session_backend(override=session_backend)
+    session_dsn = _first_configured(
+        "KSADK_SESSION_DSN",
+        "KSADK_STM_URL",
+        "KSADK_STM_DB_URL",
+    )
+    checkpoint_dsn = _first_configured("KSADK_CHECKPOINT_DSN")
+    framework_dsn = _framework_checkpoint_dsn(normalized_framework)
+
+    if explicit_local:
+        session = StorageTarget(backend=backend, explicit_local=True)
+    elif session_dsn:
+        session = StorageTarget(
+            backend=backend if explicit_backend else "postgres",
+            dsn=session_dsn,
+            source="explicit",
+        )
+    elif checkpoint_dsn:
+        session = StorageTarget(
+            backend="postgres",
+            dsn=checkpoint_dsn,
+            source="checkpoint_fallback",
+        )
+    else:
+        session = StorageTarget(backend=backend)
+
+    if framework_dsn:
+        checkpoint = StorageTarget(
+            backend="postgres",
+            dsn=framework_dsn,
+            source="framework_override",
+        )
+    elif checkpoint_dsn:
+        checkpoint = StorageTarget(backend="postgres", dsn=checkpoint_dsn, source="explicit")
+    elif session.backend == "postgres" and session.dsn:
+        checkpoint = StorageTarget(
+            backend="postgres",
+            dsn=session.dsn,
+            source="session_fallback",
+        )
+    else:
+        checkpoint = StorageTarget(backend="none")
+
+    return PersistenceTopology(
+        session=session,
+        checkpoint=checkpoint,
+        framework=normalized_framework,
+    )
+
+
+__all__ = ["PersistenceTopology", "StorageTarget", "resolve_persistence_topology"]

ksadk/skills/models.py

@@ -85,6 +85,10 @@ class SkillListResponse:
     skills: list[SkillRef]
     code: int = 0
     message: str = ""
+    total_count: int | None = None
+    truncated: bool = False
+    next_page: int | None = None
+    pagination_warning: str = ""
 
     @classmethod
     def from_payload(
@@ -94,10 +98,24 @@ class SkillListResponse:
         space_id: str = "",
         space_name: str = "",
     ) -> "SkillListResponse":
-        data = payload.get("Data") or payload.get("data") or {}
-        skills_payload = (
-            data.get("Skills") or data.get("skills") or data.get("Items") or data.get("items") or []
+        data = payload.get("Data", payload.get("data", {}))
+        if not isinstance(data, dict):
+            raise ValueError("Skill directory response must contain an object")
+        skills_payload = next(
+            (data[key] for key in ("Skills", "skills", "Items", "items") if key in data), []
         )
+        if skills_payload is None:
+            skills_payload = []
+        if not isinstance(skills_payload, list) or any(
+            not isinstance(item, dict) for item in skills_payload
+        ):
+            raise ValueError("Skill directory items must be objects")
+        total = data.get("TotalCount", data.get("total_count", payload.get("TotalCount")))
+        if total is not None:
+            if isinstance(total, str) and total.isdigit():
+                total = int(total)
+            if type(total) is not int or total < 0:
+                raise ValueError("Skill directory total count is invalid")
         return cls(
             code=int(payload.get("Code") or payload.get("code") or 0),
             message=str(payload.get("Message") or payload.get("message") or ""),
@@ -109,9 +127,13 @@ class SkillListResponse:
             skills=[
                 SkillRef.from_payload(item) for item in skills_payload if isinstance(item, dict)
             ],
+            total_count=total,
         )
 
-    def active_skills(self) -> list[SkillRef]:
+    def active_skills(self, *, allow_partial: bool = False) -> list[SkillRef]:
+        """Require callers consuming a partial directory to acknowledge it."""
+        if self.truncated and not allow_partial:
+            raise ValueError("Skill directory is incomplete; refresh or continue pagination")
         return [skill for skill in self.skills if skill.is_active]