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

ksadk/sessions/persistence.py

@@ -0,0 +1,275 @@
+"""Runtime-safe PostgreSQL persistence readiness diagnostics."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import os
+import time
+from copy import deepcopy
+from typing import Any, Awaitable, Callable, Mapping
+
+from ksadk.sessions import resolve_persistence_topology
+from ksadk.sessions.topology import StorageTarget
+
+ConnectCallable = Callable[..., Awaitable[Any]]
+
+_STATUS_CACHE: dict[str, tuple[float, dict[str, Any]]] = {}
+_CACHE_LOCK = asyncio.Lock()
+
+
+class _NoopAsyncLock:
+    async def __aenter__(self) -> None:
+        return None
+
+    async def __aexit__(self, *_args: Any) -> None:
+        return None
+
+
+def _base_status(*, target: StorageTarget, configured: bool) -> dict[str, Any]:
+    return {
+        "Configured": configured,
+        "Status": "checking" if configured else "not_configured",
+        "Ready": False,
+        "Backend": target.backend,
+        "SharedAcrossPods": False,
+        "EffectiveFor": "new_runs_only",
+        "Source": target.source,
+        "ReasonCode": "",
+        "Reason": "",
+    }
+
+
+def _not_configured_target_status(target: StorageTarget, *, store: str) -> dict[str, Any]:
+    status = _base_status(target=target, configured=False)
+    prefix = f"{store}_" if store else ""
+    label = f"{store.title()} storage" if store else "Storage"
+    status.update(
+        {
+            "ReasonCode": f"{prefix}STORE_NOT_CONFIGURED",
+            "Reason": f"{label} is not configured",
+        }
+    )
+    return status
+
+
+def _error_status(target: StorageTarget, *, code: str, reason: str) -> dict[str, Any]:
+    status = _base_status(target=target, configured=True)
+    status.update({"Status": "error", "ReasonCode": code, "Reason": reason})
+    return status
+
+
+def _classify_probe_error(exc: BaseException) -> tuple[str, str]:
+    name = type(exc).__name__.lower()
+    if isinstance(exc, (ModuleNotFoundError, ImportError)):
+        return "DEPENDENCY_MISSING", "asyncpg is required for PostgreSQL persistence"
+    if "password" in name or "authorization" in name or "authentication" in name:
+        return "AUTH_FAILED", "PostgreSQL authentication failed"
+    if "privilege" in name or "permission" in name:
+        return "SCHEMA_PERMISSION_DENIED", "PostgreSQL schema permissions are insufficient"
+    return "STORE_UNREACHABLE", "PostgreSQL persistence is unreachable"
+
+
+def _cache_key(backend: str, dsn: str) -> str:
+    digest = hashlib.sha256(dsn.encode("utf-8")).hexdigest()
+    return f"{backend}:{digest}"
+
+
+def _asyncpg_dsn(dsn: str) -> str:
+    """Translate ADK's SQLAlchemy asyncpg URL to asyncpg's native DSN."""
+    prefix = "postgresql+asyncpg://"
+    if dsn.lower().startswith(prefix):
+        return "postgresql://" + dsn[len(prefix) :]
+    return dsn
+
+
+async def _default_connect(**kwargs: Any) -> Any:
+    import asyncpg
+
+    return await asyncpg.connect(**kwargs)
+
+
+def _status_for_target(
+    status: Mapping[str, Any], target: StorageTarget, *, store: str = ""
+) -> dict[str, Any]:
+    result = deepcopy(dict(status))
+    result.update({"Backend": target.backend, "Source": target.source})
+    if result.get("ReasonCode") == "STORE_UNREACHABLE" and store:
+        result["ReasonCode"] = f"{store}_STORE_UNREACHABLE"
+    return result
+
+
+async def probe_storage_target(
+    target: StorageTarget,
+    *,
+    connect: ConnectCallable | None = None,
+    use_cache: bool = True,
+    store: str = "",
+) -> dict[str, Any]:
+    """Probe one storage target without exposing its DSN in the result."""
+    if target.backend != "postgres" or not target.dsn:
+        return _not_configured_target_status(target, store=store)
+
+    result = await _probe_postgres_target(target, connect=connect, use_cache=use_cache)
+    return _status_for_target(result, target, store=store)
+
+
+async def _probe_postgres_target(
+    target: StorageTarget,
+    *,
+    connect: ConnectCallable | None,
+    use_cache: bool,
+) -> dict[str, Any]:
+    dsn = _asyncpg_dsn(target.dsn)
+
+    timeout = max(0.1, float(os.getenv("KSADK_PERSISTENCE_PROBE_TIMEOUT") or "2"))
+    ttl = max(0.0, float(os.getenv("KSADK_PERSISTENCE_PROBE_CACHE_TTL") or "30"))
+    key = _cache_key(target.backend, dsn)
+    now = time.monotonic()
+    if use_cache:
+        cached = _STATUS_CACHE.get(key)
+        if cached and now - cached[0] < ttl:
+            return deepcopy(cached[1])
+
+    connector = connect or _default_connect
+    connection = None
+    try:
+        lock = _CACHE_LOCK if use_cache else _NoopAsyncLock()
+        async with lock:
+            if use_cache:
+                cached = _STATUS_CACHE.get(key)
+                if cached and time.monotonic() - cached[0] < ttl:
+                    return deepcopy(cached[1])
+            connection = await asyncio.wait_for(
+                connector(dsn=dsn, timeout=timeout, command_timeout=timeout),
+                timeout=timeout,
+            )
+            await asyncio.wait_for(connection.fetchval("SELECT 1"), timeout=timeout)
+            can_create = await asyncio.wait_for(
+                connection.fetchval(
+                    "SELECT has_schema_privilege(current_user, current_schema(), 'CREATE')"
+                ),
+                timeout=timeout,
+            )
+            if can_create is not True:
+                result = _error_status(
+                    target,
+                    code="SCHEMA_PERMISSION_DENIED",
+                    reason=(
+                        "PostgreSQL account cannot create persistence tables "
+                        "in the current schema"
+                    ),
+                )
+            else:
+                result = _base_status(target=target, configured=True)
+                result.update(
+                    {
+                        "Status": "ready",
+                        "Ready": True,
+                        "SharedAcrossPods": True,
+                        "ReasonCode": "READY",
+                        "Reason": "",
+                    }
+                )
+            if use_cache:
+                _STATUS_CACHE[key] = (time.monotonic(), deepcopy(result))
+            return deepcopy(result)
+    except Exception as exc:
+        code, reason = _classify_probe_error(exc)
+        result = _error_status(target, code=code, reason=reason)
+        if use_cache:
+            _STATUS_CACHE[key] = (time.monotonic(), deepcopy(result))
+        return deepcopy(result)
+    finally:
+        if connection is not None:
+            try:
+                await asyncio.wait_for(connection.close(), timeout=timeout)
+            except Exception:
+                pass
+
+
+async def get_persistence_status(
+    *,
+    framework: str | None = None,
+    connect: ConnectCallable | None = None,
+    use_cache: bool = True,
+) -> dict[str, Any]:
+    """Return independent credential-free Session and Checkpoint readiness."""
+    topology = resolve_persistence_topology(framework=framework)
+    probe_tasks: dict[str, asyncio.Task[dict[str, Any]]] = {}
+
+    async def probe(target: StorageTarget, *, store: str) -> dict[str, Any]:
+        if target.backend != "postgres" or not target.dsn:
+            return _not_configured_target_status(target, store=store)
+        key = _cache_key(target.backend, _asyncpg_dsn(target.dsn))
+        if key not in probe_tasks:
+            probe_tasks[key] = asyncio.create_task(
+                _probe_postgres_target(
+                    target,
+                    connect=connect,
+                    use_cache=use_cache,
+                )
+            )
+        return _status_for_target(await probe_tasks[key], target, store=store)
+
+    session_status, checkpoint_status = await asyncio.gather(
+        probe(topology.session, store="SESSION"),
+        probe(topology.checkpoint, store="CHECKPOINT"),
+    )
+    return {"Session": session_status, "Checkpoint": checkpoint_status}
+
+
+def gate_runtime_capabilities(
+    capabilities: Mapping[str, Any] | None,
+    session_persistence: Mapping[str, Any],
+    checkpoint_persistence: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Fail closed unless Session and Checkpoint persistence are both ready."""
+
+    gated = deepcopy(dict(capabilities or {}))
+    checkpoint_persistence = checkpoint_persistence or session_persistence
+    blocked_store = (
+        "session" if session_persistence.get("Ready") is not True else "checkpoint"
+    )
+    unavailable = (
+        session_persistence if blocked_store == "session" else checkpoint_persistence
+    )
+    if unavailable.get("Ready") is True:
+        return gated
+
+    reason = str(unavailable.get("Reason") or "PostgreSQL persistence is unavailable")
+    reason_code = str(
+        unavailable.get("ReasonCode") or "RUNTIME_CAPABILITY_UNAVAILABLE"
+    )
+    checkpoint = dict(gated.get("Checkpoint") or {})
+    native_backend = str(checkpoint.get("Backend") or "").strip()
+    checkpoint.update(
+        {
+            "Supported": False,
+            "Backend": "none",
+            "Durable": False,
+            "SharedAcrossPods": False,
+            "ResumeMode": "none",
+            "ReasonCode": reason_code,
+            "Reason": reason,
+            "PersistenceGate": {
+                "BlockedStore": blocked_store,
+                "Source": str(unavailable.get("Source") or "none"),
+                "ReasonCode": reason_code,
+                "Reason": reason,
+            },
+        }
+    )
+    if native_backend:
+        checkpoint["NativeBackend"] = native_backend
+    gated["Checkpoint"] = checkpoint
+    gated["ResumeRun"] = {
+        "Supported": False,
+        "ResumeMode": "none",
+        "ReasonCode": reason_code,
+        "Reason": reason,
+    }
+    return gated
+
+
+__all__ = ["gate_runtime_capabilities", "get_persistence_status", "probe_storage_target"]

ksadk/sessions/postgres_service.py

@@ -6,6 +6,8 @@ import asyncio
 import json
 import logging
 import time
+from collections.abc import AsyncIterator
+from contextvars import ContextVar
 from typing import Any, Optional
 from urllib.parse import urlsplit, urlunsplit
 
@@ -19,8 +21,10 @@ from ksadk.sessions._postgres_tables import (
 from ksadk.sessions.base import (
     CANONICAL_EVENT_STORAGE_CAPABILITIES,
     BaseSessionService,
+    CheckpointEventQuery,
     Session,
     SessionEvent,
+    SessionEventQuery,
     SessionState,
     generate_id,
 )
@@ -56,6 +60,9 @@ class PostgresSessionService(_PostgresSchemaMixin, BaseSessionService):
         self._pool_lock = asyncio.Lock()
         self._schema_ready = False
         self._schema_lock = asyncio.Lock()
+        self._checkpoint_snapshot_connection: ContextVar[Any | None] = ContextVar(
+            f"checkpoint_snapshot_connection_{id(self)}", default=None
+        )
 
     async def create_session(
         self,
@@ -489,6 +496,357 @@ class PostgresSessionService(_PostgresSchemaMixin, BaseSessionService):
             """
             return int(await connection.fetchval(query, *params) or 0)
 
+    async def get_sessions_by_ids(self, session_ids: list[str]) -> list[Session]:
+        if not session_ids:
+            return []
+        await self._ensure_schema()
+        async with self._pool.acquire() as connection:
+            rows = await connection.fetch(
+                f"SELECT id,agent_id,user_id,title,title_source,summary,first_prompt,"
+                f"last_prompt,state_json,created_at,updated_at,version "
+                f"FROM {KSADK_PG_SESSIONS_TABLE} WHERE namespace=$1 "
+                "AND id=ANY($2::text[])",
+                self.namespace,
+                session_ids,
+            )
+        by_id = {row["id"]: self._session_from_row(row, events=[]) for row in rows}
+        return [by_id[item] for item in session_ids if item in by_id]
+
+    async def list_session_metadata(
+        self, agent_id: Optional[str] = None, user_id: Optional[str] = None
+    ) -> list[Session]:
+        await self._ensure_schema()
+        clauses = ["namespace=$1"]
+        params: list[Any] = [self.namespace]
+        if agent_id is not None:
+            params.append(agent_id)
+            clauses.append(f"agent_id=${len(params)}")
+        if user_id is not None:
+            params.append(user_id)
+            clauses.append(f"user_id=${len(params)}")
+        async with self._pool.acquire() as connection:
+            rows = await connection.fetch(
+                f"SELECT id,agent_id,user_id,title,title_source,summary,first_prompt,"
+                f"last_prompt,state_json,created_at,updated_at,version "
+                f"FROM {KSADK_PG_SESSIONS_TABLE} WHERE {' AND '.join(clauses)} "
+                "ORDER BY updated_at DESC,created_at DESC,id DESC",
+                *params,
+            )
+        return [self._session_from_row(row, events=[]) for row in rows]
+
+    def _batch_event_where(self, query: SessionEventQuery) -> tuple[list[str], list[Any]]:
+        clauses = ["event_row.namespace=$1"]
+        params: list[Any] = [self.namespace]
+        if query.session_ids is not None:
+            if not query.session_ids:
+                return ["FALSE"], []
+            params.append(query.session_ids)
+            clauses.append(f"event_row.session_id=ANY(${len(params)}::text[])")
+        for value, clause in (
+            (query.agent_id, "session_row.agent_id"),
+            (query.after_seq_id, "event_row.seq_id >"),
+            (query.before_seq_id, "event_row.seq_id <"),
+            (query.invocation_id, "event_row.invocation_id"),
+        ):
+            if value is not None:
+                params.append(value)
+                operator = "" if clause.endswith((">", "<")) else "="
+                clauses.append(f"{clause}{operator}${len(params)}")
+        if query.event_types:
+            params.append(query.event_types)
+            clauses.append(f"event_row.event_type=ANY(${len(params)}::text[])")
+        if query.run_id is not None:
+            params.append(query.run_id)
+            clauses.append(
+                "(event_row.metadata_json->>'run_id'="
+                f"${len(params)}"
+                " OR (event_row.event_type='continuation.created'"
+                f" AND event_row.content_json->'runtime_event'->>'run_id'=${len(params)}))"
+            )
+        if query.checkpoint_id is not None:
+            params.append(query.checkpoint_id)
+            clauses.append(
+                "(event_row.metadata_json->>'checkpoint_id'="
+                f"${len(params)}"
+                " OR (event_row.event_type='continuation.created'"
+                f" AND event_row.content_json->'runtime_event'->>'continuation_id'=${len(params)}))"
+            )
+        if query.checkpoint_ids:
+            params.append(query.checkpoint_ids)
+            clauses.append(
+                "(event_row.metadata_json->>'checkpoint_id'=ANY("
+                f"${len(params)}::text[])"
+                " OR (event_row.event_type='continuation.created'"
+                " AND event_row.content_json->'runtime_event'->>'continuation_id'=ANY("
+                f"${len(params)}::text[])))"
+            )
+        return clauses, params
+
+    async def query_events(self, query: SessionEventQuery) -> list[SessionEvent]:
+        await self._ensure_schema()
+        clauses, params = self._batch_event_where(query)
+        direction = "ASC" if query.from_start else "DESC"
+        order = (
+            f"event_row.session_id {direction},event_row.seq_id {direction},"
+            f"event_row.id {direction}"
+            if query.order_by_seq
+            else f"event_row.timestamp {direction},event_row.session_id {direction},"
+            f"event_row.seq_id {direction},event_row.id {direction}"
+        )
+        params.extend([query.limit, query.offset])
+        async with self._pool.acquire() as connection:
+            rows = await connection.fetch(
+                f"SELECT event_row.id,event_row.session_id,event_row.author,event_row.event_type,"
+                f"event_row.content_json,event_row.timestamp,event_row.state_delta_json,"
+                f"event_row.seq_id,event_row.invocation_id,event_row.metadata_json "
+                f"FROM {KSADK_PG_EVENTS_TABLE} event_row "
+                f"JOIN {KSADK_PG_SESSIONS_TABLE} session_row "
+                "ON session_row.namespace=event_row.namespace "
+                "AND session_row.id=event_row.session_id "
+                f"WHERE {' AND '.join(clauses)} ORDER BY {order} "
+                f"LIMIT ${len(params)-1} OFFSET ${len(params)}",
+                *params,
+            )
+        events = [self._event_from_row(row) for row in rows]
+        if not query.from_start:
+            events.reverse()
+        return events
+
+    async def count_event_query(self, query: SessionEventQuery) -> int:
+        await self._ensure_schema()
+        clauses, params = self._batch_event_where(query)
+        async with self._pool.acquire() as connection:
+            return int(
+                await connection.fetchval(
+                    f"SELECT COUNT(*) FROM {KSADK_PG_EVENTS_TABLE} event_row "
+                    f"JOIN {KSADK_PG_SESSIONS_TABLE} session_row "
+                    "ON session_row.namespace=event_row.namespace "
+                    "AND session_row.id=event_row.session_id "
+                    f"WHERE {' AND '.join(clauses)}",
+                    *params,
+                )
+                or 0
+            )
+
+    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")
+        return await self._scan_checkpoint_events_with_connection(None, query)
+
+    async def _scan_checkpoint_events_with_connection(
+        self,
+        connection: Any | None,
+        query: CheckpointEventQuery,
+        cursor: tuple[float, str, int, str] | None = None,
+    ) -> list[SessionEvent]:
+        event_query = SessionEventQuery(
+            session_ids=query.session_ids,
+            agent_id=query.agent_id,
+            event_types=["run_checkpoint", "continuation.created"],
+            offset=query.offset,
+            limit=query.limit,
+            from_start=True,
+        )
+        if query.framework is None and connection is None:
+            return await self.query_events(event_query)
+        await self._ensure_schema()
+        clauses, params = self._batch_event_where(event_query)
+        if query.checkpoint_ids:
+            params.append(query.checkpoint_ids)
+            clauses.append(
+                "(event_row.metadata_json->>'checkpoint_id'=ANY("
+                f"${len(params)}::text[])"
+                " OR (event_row.event_type='continuation.created'"
+                " AND event_row.content_json->'runtime_event'->>'continuation_id'=ANY("
+                f"${len(params)}::text[])))"
+            )
+        if query.run_id is not None:
+            params.append(query.run_id)
+            clauses.append(
+                "(event_row.metadata_json->>'run_id'="
+                f"${len(params)}"
+                " OR (event_row.event_type='continuation.created'"
+                f" AND event_row.content_json->'runtime_event'->>'run_id'=${len(params)}))"
+            )
+        if query.framework is not None:
+            params.append(query.framework.lower())
+            clauses.append(
+                "(lower(event_row.metadata_json->>'framework')="
+                f"${len(params)}"
+                " OR (event_row.event_type='continuation.created'"
+                " AND lower(event_row.content_json->'runtime_event'->'source'->>'framework')="
+                f"${len(params)}))"
+            )
+        if cursor is not None:
+            placeholders = []
+            for value in cursor:
+                params.append(value)
+                placeholders.append(f"${len(params)}")
+            clauses.append(
+                "(event_row.timestamp,event_row.session_id,event_row.seq_id,event_row.id) > "
+                f"({','.join(placeholders)})"
+            )
+        params.extend([query.limit, query.offset])
+        owns_connection = connection is None
+        if owns_connection:
+            connection = await self._pool.acquire()
+        try:
+            rows = await connection.fetch(
+                f"SELECT event_row.id,event_row.session_id,event_row.author,event_row.event_type,"
+                f"event_row.content_json,event_row.timestamp,event_row.state_delta_json,"
+                f"event_row.seq_id,event_row.invocation_id,event_row.metadata_json "
+                f"FROM {KSADK_PG_EVENTS_TABLE} event_row "
+                f"JOIN {KSADK_PG_SESSIONS_TABLE} session_row "
+                "ON session_row.namespace=event_row.namespace "
+                "AND session_row.id=event_row.session_id "
+                f"WHERE {' AND '.join(clauses)} ORDER BY event_row.timestamp,event_row.session_id,"
+                f"event_row.seq_id,event_row.id LIMIT ${len(params)-1} OFFSET ${len(params)}",
+                *params,
+            )
+            return [self._event_from_row(row) for row in rows]
+        finally:
+            if owns_connection:
+                await self._pool.release(connection)
+
+    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")
+        await self._ensure_schema()
+        async with self._pool.acquire() as connection:
+            async with connection.transaction(isolation="repeatable_read", readonly=True):
+                token = self._checkpoint_snapshot_connection.set(connection)
+                try:
+                    cursor: tuple[float, str, int, str] | None = None
+                    first_page = True
+                    while True:
+                        batch = await self._scan_checkpoint_events_with_connection(
+                            connection,
+                            CheckpointEventQuery(
+                                **{
+                                    **query.__dict__,
+                                    "offset": query.offset if first_page else 0,
+                                }
+                            ),
+                            cursor,
+                        )
+                        if not batch:
+                            break
+                        yield batch
+                        if len(batch) < query.limit:
+                            break
+                        last = batch[-1]
+                        cursor = (
+                            last.timestamp,
+                            last.session_id,
+                            last.seq_id,
+                            last.id,
+                        )
+                        first_page = False
+                finally:
+                    self._checkpoint_snapshot_connection.reset(token)
+
+    async def get_checkpoint_lookup_stats(
+        self, session_id: str, run_id: str, checkpoint_id: str
+    ) -> dict[str, object]:
+        events = await self.query_events(
+            SessionEventQuery(
+                session_ids=[session_id], run_id=run_id, limit=2**31 - 1, from_start=True
+            )
+        )
+        candidate = None
+        max_seq_id = 0
+        resume_count = 0
+        last_resumed_at = None
+        for event in events:
+            metadata = event.metadata or {}
+            if event.event_type == "run_checkpoint":
+                max_seq_id = max(max_seq_id, event.seq_id)
+                if str(metadata.get("checkpoint_id") or "") == checkpoint_id:
+                    candidate = event
+            elif event.event_type == "run_resume" and str(
+                metadata.get("checkpoint_id") or ""
+            ) == checkpoint_id:
+                resume_count += 1
+                last_resumed_at = max(last_resumed_at or event.timestamp, event.timestamp)
+        return {"candidate": candidate, "max_seq_id": max_seq_id, "resume_count": resume_count,
+                "last_resumed_at": last_resumed_at}
+
+    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))
+        audits = {
+            key: {"resume_count": 0, "last_resumed_at": None} for key in unique_keys
+        }
+        run_keys = list(
+            dict.fromkeys((session_id, run_id) for session_id, run_id, _ in unique_keys)
+        )
+        latest = {key: 0 for key in run_keys}
+        if not unique_keys:
+            return {"audits": audits, "latest_seq_ids": latest}
+        await self._ensure_schema()
+        connection = self._checkpoint_snapshot_connection.get()
+        owns_connection = connection is None
+        if owns_connection:
+            connection = await self._pool.acquire()
+        try:
+            audit_rows = await connection.fetch(
+                f"""WITH requested(session_id,run_id,checkpoint_id) AS (
+                    SELECT * FROM unnest($2::text[],$3::text[],$4::text[])
+                )
+                SELECT requested.session_id,requested.run_id,requested.checkpoint_id,
+                       COUNT(event_row.id) AS resume_count,
+                       MAX(event_row.timestamp) AS last_resumed_at
+                FROM requested
+                LEFT JOIN {KSADK_PG_EVENTS_TABLE} event_row
+                  ON event_row.namespace=$1
+                 AND event_row.session_id=requested.session_id
+                 AND event_row.event_type='run_resume'
+                 AND event_row.metadata_json->>'run_id'=requested.run_id
+                 AND event_row.metadata_json->>'checkpoint_id'=requested.checkpoint_id
+                GROUP BY requested.session_id,requested.run_id,requested.checkpoint_id""",
+                self.namespace,
+                [key[0] for key in unique_keys],
+                [key[1] for key in unique_keys],
+                [key[2] for key in unique_keys],
+            )
+            latest_rows = await connection.fetch(
+                f"""WITH requested(session_id,run_id) AS (
+                    SELECT * FROM unnest($2::text[],$3::text[])
+                )
+                SELECT requested.session_id,requested.run_id,
+                       COALESCE(MAX(event_row.seq_id),0) AS latest_seq_id
+                FROM requested
+                LEFT JOIN {KSADK_PG_EVENTS_TABLE} event_row
+                  ON event_row.namespace=$1
+                 AND event_row.session_id=requested.session_id
+                 AND event_row.event_type='run_checkpoint'
+                 AND event_row.metadata_json->>'run_id'=requested.run_id
+                GROUP BY requested.session_id,requested.run_id""",
+                self.namespace,
+                [key[0] for key in run_keys],
+                [key[1] for key in run_keys],
+            )
+        finally:
+            if owns_connection:
+                await self._pool.release(connection)
+        for row in audit_rows:
+            audits[(row["session_id"], row["run_id"], row["checkpoint_id"])] = {
+                "resume_count": int(row["resume_count"] or 0),
+                "last_resumed_at": row["last_resumed_at"],
+            }
+        for row in latest_rows:
+            latest[(row["session_id"], row["run_id"])] = int(
+                row["latest_seq_id"] or 0
+            )
+        return {"audits": audits, "latest_seq_ids": latest}
+
     def _agent_events_query_parts(
         self,
         agent_id: str,

ksadk/sessions/resilience.py

@@ -45,8 +45,16 @@ def is_session_backend_failure(exc: BaseException) -> bool:
             current,
             (SessionBackendUnavailable, TimeoutError, OSError, ConnectionError),
         ):
+            logger.warning(
+                "Session backend error; assuming backend failure: %s",
+                current
+            )
             return True
         module = type(current).__module__
         if module.startswith("asyncpg.") or module.startswith("sqlalchemy."):
+            logger.warning(
+                "SQLAlchemy error; assuming backend failure: %s",
+                current
+            )
             return True
     return False