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

ksadk/sessions/in_memory.py

@@ -1,21 +1,35 @@
 from __future__ import annotations
 
 import asyncio
+import bisect
 import copy
 import time
-from typing import Optional
+from typing import AsyncIterator, Optional
 
 from ksadk.ids import new_session_id
 from ksadk.sessions.base import (
     CANONICAL_EVENT_STORAGE_CAPABILITIES,
     BaseSessionService,
+    CheckpointEventQuery,
     Session,
     SessionEvent,
+    SessionEventQuery,
     SessionState,
     generate_id,
 )
 
 
+class _CheckpointSnapshotView:
+    """Compatibility view over task-scoped checkpoint snapshots."""
+
+    def __init__(self, snapshots: dict[asyncio.Task[object], int]) -> None:
+        self._snapshots = snapshots
+
+    def get(self) -> int | None:
+        task = asyncio.current_task()
+        return self._snapshots.get(task) if task is not None else None
+
+
 class InMemorySessionService(BaseSessionService):
     storage_capabilities = CANONICAL_EVENT_STORAGE_CAPABILITIES
 
@@ -24,7 +38,14 @@ class InMemorySessionService(BaseSessionService):
         self._events_by_id: dict[str, SessionEvent] = {}
         self._events_by_invocation: dict[tuple[str, str], list[SessionEvent]] = {}
         self._states: dict[tuple[str, str, str, str], SessionState] = {}
+        self._event_order: list[tuple[float, str, int, str, int, SessionEvent]] = []
+        self._event_generation = 0
+        self._checkpoint_snapshot_by_task: dict[asyncio.Task[object], int] = {}
+        self._checkpoint_snapshot_generation = _CheckpointSnapshotView(
+            self._checkpoint_snapshot_by_task
+        )
         self._lock = asyncio.Lock()
+        self._checkpoint_scan_lock = asyncio.Lock()
 
     async def create_session(
         self,
@@ -97,25 +118,37 @@ class InMemorySessionService(BaseSessionService):
                 if session.agent_id == agent_id and (user_id is None or session.user_id == user_id)
             )
 
-    async def delete_session(self, session_id: str) -> bool:
+    async def list_session_metadata(
+        self, agent_id: Optional[str] = None, user_id: Optional[str] = None
+    ) -> list[Session]:
         async with self._lock:
-            session = self._sessions.pop(session_id, None)
-            if not session:
-                return False
-            for event in session.events:
-                self._events_by_id.pop(event.id, None)
-                if event.invocation_id is not None:
-                    self._events_by_invocation.pop((session_id, event.invocation_id), None)
-            self._states.pop(
-                self._state_key(
-                    "session",
-                    session.agent_id,
-                    session.user_id,
-                    session_id,
-                ),
-                None,
-            )
-            return True
+            sessions = [
+                self._session_metadata(session)
+                for session in self._sessions.values()
+                if (agent_id is None or session.agent_id == agent_id)
+                and (user_id is None or session.user_id == user_id)
+            ]
+            sessions.sort(key=lambda item: (item.updated_at, item.created_at), reverse=True)
+            return sessions
+
+    async def delete_session(self, session_id: str) -> bool:
+        async with self._checkpoint_scan_lock:
+            async with self._lock:
+                session = self._sessions.pop(session_id, None)
+                if not session:
+                    return False
+                for event in session.events:
+                    self._events_by_id.pop(event.id, None)
+                    if event.invocation_id is not None:
+                        self._events_by_invocation.pop((session_id, event.invocation_id), None)
+                self._event_order = [item for item in self._event_order if item[1] != session_id]
+                self._states.pop(
+                    self._state_key(
+                        "session", session.agent_id, session.user_id, session_id
+                    ),
+                    None,
+                )
+                return True
 
     async def update_session_metadata(
         self,
@@ -164,6 +197,18 @@ class InMemorySessionService(BaseSessionService):
             if not stored.id:
                 stored.id = generate_id()
             session.events.append(stored)
+            self._event_generation += 1
+            bisect.insort(
+                self._event_order,
+                (
+                    stored.timestamp,
+                    stored.session_id,
+                    stored.seq_id,
+                    stored.id,
+                    self._event_generation,
+                    stored,
+                ),
+            )
             self._events_by_id[stored.id] = stored
             if stored.invocation_id is not None:
                 self._events_by_invocation.setdefault(
@@ -255,6 +300,251 @@ class InMemorySessionService(BaseSessionService):
                 events = [event for event in events if event.seq_id < before_seq_id]
             return len(events)
 
+    async def get_sessions_by_ids(self, session_ids: list[str]) -> list[Session]:
+        async with self._lock:
+            return [
+                self._session_metadata(self._sessions[session_id])
+                for session_id in session_ids
+                if session_id in self._sessions
+            ]
+
+    @staticmethod
+    def _session_metadata(session: Session) -> Session:
+        metadata = copy.deepcopy(session)
+        metadata.events = []
+        return metadata
+
+    async def query_events(self, query: SessionEventQuery) -> list[SessionEvent]:
+        return await self._query_events(query, count_only=False)
+
+    async def count_event_query(self, query: SessionEventQuery) -> int:
+        return int(await self._query_events(query, count_only=True))
+
+    async def _query_events(
+        self, query: SessionEventQuery, *, count_only: bool
+    ) -> list[SessionEvent] | int:
+        async with self._lock:
+            selected_ids = (
+                list(dict.fromkeys(query.session_ids))
+                if query.session_ids is not None
+                else list(self._sessions)
+            )
+            allowed_types = set(query.event_types or [])
+            allowed_checkpoint_ids = set(query.checkpoint_ids or [])
+            events = [
+                event
+                for session_id in selected_ids
+                for session in [self._sessions.get(session_id)]
+                if session is not None
+                and (query.agent_id is None or session.agent_id == query.agent_id)
+                for event in session.events
+                if (query.after_seq_id is None or event.seq_id > query.after_seq_id)
+                and (query.before_seq_id is None or event.seq_id < query.before_seq_id)
+                and (not allowed_types or event.event_type in allowed_types)
+                and (query.invocation_id is None or event.invocation_id == query.invocation_id)
+                and (
+                    query.run_id is None
+                    or str((event.metadata or {}).get("run_id") or "") == query.run_id
+                    or (
+                        event.event_type == "continuation.created"
+                        and str(
+                            ((event.content or {}).get("runtime_event") or {}).get(
+                                "run_id", ""
+                            )
+                            or ""
+                        )
+                        == query.run_id
+                    )
+                )
+                and (
+                    query.checkpoint_id is None
+                    or str((event.metadata or {}).get("checkpoint_id") or "")
+                    == query.checkpoint_id
+                    or (
+                        event.event_type == "continuation.created"
+                        and str(
+                            ((event.content or {}).get("runtime_event") or {}).get(
+                                "continuation_id", ""
+                            )
+                            or ""
+                        )
+                        == query.checkpoint_id
+                    )
+                )
+                and (
+                    not allowed_checkpoint_ids
+                    or str((event.metadata or {}).get("checkpoint_id") or "")
+                    in allowed_checkpoint_ids
+                    or (
+                        event.event_type == "continuation.created"
+                        and str(
+                            ((event.content or {}).get("runtime_event") or {}).get(
+                                "continuation_id", ""
+                            )
+                            or ""
+                        )
+                        in allowed_checkpoint_ids
+                    )
+                )
+            ]
+            if count_only:
+                return len(events)
+            if query.order_by_seq:
+                events.sort(key=lambda event: (event.session_id, event.seq_id, event.id))
+            else:
+                events.sort(
+                    key=lambda event: (
+                        event.timestamp,
+                        event.session_id,
+                        event.seq_id,
+                        event.id,
+                    )
+                )
+            if query.from_start:
+                return copy.deepcopy(events[query.offset : query.offset + query.limit])
+            end = max(len(events) - query.offset, 0)
+            return copy.deepcopy(events[max(end - query.limit, 0) : end])
+
+    async def get_checkpoint_lookup_stats(
+        self, session_id: str, run_id: str, checkpoint_id: str
+    ) -> dict[str, object]:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            candidate = None
+            max_seq_id = 0
+            resume_count = 0
+            last_resumed_at = None
+            for event in session.events if session else []:
+                metadata = event.metadata or {}
+                if str(metadata.get("run_id") or "") != run_id:
+                    continue
+                if event.event_type == "run_checkpoint":
+                    max_seq_id = max(max_seq_id, int(event.seq_id or 0))
+                    if str(metadata.get("checkpoint_id") or "") == checkpoint_id and (
+                        candidate is None or event.seq_id > candidate.seq_id
+                    ):
+                        candidate = copy.deepcopy(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 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")
+        async with self._lock:
+            return self._scan_checkpoint_events_locked(query)
+
+    def _scan_checkpoint_events_locked(
+        self, query: CheckpointEventQuery, snapshot_generation: int | None = None
+    ) -> list[SessionEvent]:
+        selected_ids = None if query.session_ids is None else set(query.session_ids)
+        checkpoint_ids = set(query.checkpoint_ids or [])
+        framework = str(query.framework or "").lower()
+        matches: list[SessionEvent] = []
+        for _, session_id, _, _, generation, event in self._event_order:
+            session = self._sessions.get(session_id)
+            metadata = event.metadata or {}
+            if (
+                session is None
+                or (snapshot_generation is not None and generation > snapshot_generation)
+                or event.event_type not in {"run_checkpoint", "continuation.created"}
+                or (selected_ids is not None and session_id not in selected_ids)
+                or (query.agent_id is not None and session.agent_id != query.agent_id)
+                or (
+                    event.event_type == "run_checkpoint"
+                    and checkpoint_ids
+                    and str(metadata.get("checkpoint_id") or "") not in checkpoint_ids
+                )
+                or (
+                    event.event_type == "run_checkpoint"
+                    and query.run_id is not None
+                    and str(metadata.get("run_id") or "") != query.run_id
+                )
+                or (
+                    event.event_type == "run_checkpoint"
+                    and framework
+                    and str(metadata.get("framework") or "").lower() != framework
+                )
+            ):
+                continue
+            matches.append(event)
+        return copy.deepcopy(matches[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")
+        async with self._checkpoint_scan_lock:
+            async with self._lock:
+                snapshot_generation = self._event_generation
+            task = asyncio.current_task()
+            if task is not None:
+                self._checkpoint_snapshot_by_task[task] = snapshot_generation
+            try:
+                offset = query.offset
+                while True:
+                    async with self._lock:
+                        batch = self._scan_checkpoint_events_locked(
+                            CheckpointEventQuery(**{**query.__dict__, "offset": offset}),
+                            snapshot_generation,
+                        )
+                    if not batch:
+                        break
+                    yield batch
+                    offset += len(batch)
+                    if len(batch) < query.limit:
+                        break
+            finally:
+                if task is not None:
+                    self._checkpoint_snapshot_by_task.pop(task, None)
+
+    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 = {(session_id, run_id) for session_id, run_id, _ in unique_keys}
+        latest_seq_ids = {key: 0 for key in run_keys}
+        async with self._lock:
+            task = asyncio.current_task()
+            snapshot_generation = (
+                self._checkpoint_snapshot_by_task.get(task) if task is not None else None
+            )
+            for _, session_id, _, _, generation, event in self._event_order:
+                if snapshot_generation is not None and generation > snapshot_generation:
+                    continue
+                metadata = event.metadata or {}
+                run_id = str(metadata.get("run_id") or "")
+                run_key = (session_id, run_id)
+                if run_key not in run_keys:
+                    continue
+                if event.event_type == "run_checkpoint":
+                    latest_seq_ids[run_key] = max(
+                        latest_seq_ids[run_key], int(event.seq_id or 0)
+                    )
+                elif event.event_type == "run_resume":
+                    key = (session_id, run_id, str(metadata.get("checkpoint_id") or ""))
+                    if key in audits:
+                        audit = audits[key]
+                        audit["resume_count"] = int(audit["resume_count"]) + 1
+                        audit["last_resumed_at"] = max(
+                            audit["last_resumed_at"] or event.timestamp, event.timestamp
+                        )
+        return {"audits": audits, "latest_seq_ids": latest_seq_ids}
+
     async def get_events_for_agent(
         self,
         agent_id: str,

ksadk/sessions/invocation_identity.py

@@ -0,0 +1,235 @@
+"""Trusted business-identity binding for canonical KsADK sessions.
+
+The public ``UserId`` field remains a compatibility hint.  Hosted callers are
+authorized by a verified four-part business identity, which is converted to an
+opaque native user key and durably bound to the session before transcript
+access.  This prevents a known ``SessionId`` from becoming an authorization
+credential while still allowing an authorized legacy session to be adopted.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Mapping
+from typing import Any
+
+from fastapi import HTTPException
+
+from ksadk.runtime_context import PlatformIdentityContext
+from ksadk.sessions.base import BaseSessionService, Session
+
+SESSION_IDENTITY_SCOPE = "platform_identity"
+SESSION_IDENTITY_BINDING_VERSION = 1
+DEFAULT_IAM_IDENTITY_NAMESPACE = "kscloud-iam"
+
+
+def coerce_platform_identity(value: Any) -> PlatformIdentityContext:
+    if isinstance(value, PlatformIdentityContext):
+        return value
+    return PlatformIdentityContext.from_payload(value if isinstance(value, Mapping) else None)
+
+
+def require_complete_platform_identity(value: Any) -> PlatformIdentityContext:
+    identity = coerce_platform_identity(value)
+    if identity.is_empty:
+        return identity
+    if not identity.is_complete:
+        raise HTTPException(status_code=401, detail="Incomplete trusted business identity context")
+    return identity
+
+
+def identity_owner_payload(value: Any) -> dict[str, str]:
+    identity = require_complete_platform_identity(value)
+    if identity.is_empty:
+        return {}
+    return {
+        "identity_namespace": identity.identity_namespace,
+        "tenant_id": identity.tenant_id,
+        "subject_type": identity.subject_type,
+        "subject_id": identity.subject_id,
+    }
+
+
+def identity_digest(value: Any) -> str:
+    owner = identity_owner_payload(value)
+    if not owner:
+        return ""
+    canonical = json.dumps(owner, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+
+
+def identity_native_user_id(value: Any) -> str:
+    """Return an opaque collision-resistant native storage user key."""
+
+    digest = identity_digest(value)
+    return f"aeu_{digest}" if digest else ""
+
+
+def identity_scope_ref(value: Any) -> str:
+    digest = identity_digest(value)
+    return f"aei_{digest}" if digest else ""
+
+
+def identity_session_is_legacy(value: Any, session_user_id: str) -> bool:
+    """Return whether an authorized canonical row still uses its legacy user key."""
+
+    identity = require_complete_platform_identity(value)
+    return bool(
+        not identity.is_empty
+        and session_user_id
+        and session_user_id != identity_native_user_id(identity)
+    )
+
+
+def identity_can_adopt_any_legacy_user(value: Any) -> bool:
+    """Whether this verified principal owns the whole legacy Agent runtime."""
+
+    identity = require_complete_platform_identity(value)
+    return bool(
+        not identity.is_empty
+        and identity.identity_namespace == DEFAULT_IAM_IDENTITY_NAMESPACE
+        and identity.subject_type == "account"
+    )
+
+
+def identity_can_adopt_legacy_user(value: Any, legacy_user_id: str) -> bool:
+    """Whether a verified IAM principal deterministically owns a legacy user."""
+
+    identity = require_complete_platform_identity(value)
+    # ``aeu_`` is reserved for identity-aware rows. If a concurrent creator
+    # has written the opaque key but has not persisted its binding yet, no
+    # other principal may mistake that row for legacy data and claim it.
+    if str(legacy_user_id or "").startswith("aeu_"):
+        return False
+    if identity_can_adopt_any_legacy_user(identity):
+        return True
+    return bool(
+        not identity.is_empty
+        and identity.identity_namespace == DEFAULT_IAM_IDENTITY_NAMESPACE
+        and identity.subject_type == "user"
+        and str(legacy_user_id or "") == identity.subject_id
+    )
+
+
+def _binding_payload(identity: PlatformIdentityContext, native_user_id: str) -> dict[str, Any]:
+    return {
+        "version": SESSION_IDENTITY_BINDING_VERSION,
+        "owner": identity_owner_payload(identity),
+        "scope_ref": identity_scope_ref(identity),
+        "native_user_id": native_user_id,
+    }
+
+
+async def _read_binding(
+    service: BaseSessionService,
+    session: Session,
+) -> dict[str, Any]:
+    state = await service.get_state(
+        agent_id=session.agent_id,
+        user_id=session.user_id,
+        session_id=session.id,
+        scope=SESSION_IDENTITY_SCOPE,
+    )
+    return dict(state.state) if state is not None else {}
+
+
+async def session_identity_binding_matches(
+    *,
+    service: BaseSessionService,
+    session: Session,
+    identity: Any,
+) -> bool:
+    resolved = require_complete_platform_identity(identity)
+    if resolved.is_empty:
+        return False
+    binding = await _read_binding(service, session)
+    return bool(binding and _binding_matches(binding, resolved, session_user_id=session.user_id))
+
+
+def _binding_matches(
+    binding: Mapping[str, Any],
+    identity: PlatformIdentityContext,
+    *,
+    session_user_id: str,
+) -> bool:
+    return (
+        dict(binding.get("owner") or {}) == identity_owner_payload(identity)
+        and str(binding.get("scope_ref") or "") == identity_scope_ref(identity)
+        and str(binding.get("native_user_id") or "") == session_user_id
+    )
+
+
+def _ownership_conflict() -> HTTPException:
+    # Do not reveal which principal owns a guessed session id.
+    return HTTPException(status_code=404, detail="Session not found")
+
+
+async def bind_or_validate_session_identity(
+    *,
+    service: BaseSessionService,
+    session: Session,
+    identity: Any,
+    requested_user_id: str,
+) -> Session:
+    """Validate an existing session or durably adopt an authorized legacy row."""
+
+    resolved = require_complete_platform_identity(identity)
+    if resolved.is_empty:
+        # Legacy/local callers historically resume with SessionId alone. Keep
+        # that behavior when no trusted identity exists, while an explicitly
+        # supplied UserId must still match. Hosted traffic receives a verified
+        # default IAM identity and therefore never relies on this fallback.
+        if requested_user_id and session.user_id != requested_user_id:
+            raise HTTPException(
+                status_code=409,
+                detail="Session id belongs to a different agent or user",
+            )
+        return session
+
+    binding = await _read_binding(service, session)
+    if binding:
+        if not _binding_matches(binding, resolved, session_user_id=session.user_id):
+            raise _ownership_conflict()
+        return session
+
+    expected_native_user_id = identity_native_user_id(resolved)
+    # A newly-created identity-aware row already carries the opaque user key.
+    # Legacy adoption is intentionally limited to the built-in IAM mapping:
+    # an account principal owns the historical Agent runtime, while an IAM
+    # user may only recover the row whose old UserId equals its verified
+    # subject. Customer-defined tenants must use a controlled migration; a
+    # caller-supplied UserId can never make the first visitor the owner.
+    if session.user_id != expected_native_user_id:
+        if not identity_can_adopt_legacy_user(resolved, session.user_id):
+            raise _ownership_conflict()
+
+    await service.update_state(
+        agent_id=session.agent_id,
+        user_id=session.user_id,
+        session_id=session.id,
+        scope=SESSION_IDENTITY_SCOPE,
+        state_delta=_binding_payload(resolved, session.user_id),
+    )
+    # Re-read after the write.  Backends serialize individual state updates;
+    # this also fails closed if a competing claimant replaced the binding.
+    persisted = await _read_binding(service, session)
+    if not _binding_matches(persisted, resolved, session_user_id=session.user_id):
+        raise _ownership_conflict()
+    return session
+
+
+__all__ = [
+    "SESSION_IDENTITY_SCOPE",
+    "bind_or_validate_session_identity",
+    "coerce_platform_identity",
+    "identity_digest",
+    "identity_can_adopt_any_legacy_user",
+    "identity_can_adopt_legacy_user",
+    "identity_native_user_id",
+    "identity_owner_payload",
+    "identity_session_is_legacy",
+    "identity_scope_ref",
+    "require_complete_platform_identity",
+    "session_identity_binding_matches",
+]