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

ksadk/sessions/_local_tables.py

@@ -0,0 +1,21 @@
+"""Local SQLite session 存储的表名常量（纯移动自 local_service，行为不变）。"""
+
+KSADK_SESSIONS_TABLE = "ksadk_sessions"
+KSADK_EVENTS_TABLE = "ksadk_events"
+KSADK_STATES_TABLE = "ksadk_states"
+
+LEGACY_SESSIONS_TABLE = "sessions"
+LEGACY_EVENTS_TABLE = "events"
+LEGACY_STATES_TABLE = "states"
+
+DEFAULT_SESSION_DB_NAME = "sessions.sqlite"
+
+__all__ = [
+    "DEFAULT_SESSION_DB_NAME",
+    "KSADK_EVENTS_TABLE",
+    "KSADK_SESSIONS_TABLE",
+    "KSADK_STATES_TABLE",
+    "LEGACY_EVENTS_TABLE",
+    "LEGACY_SESSIONS_TABLE",
+    "LEGACY_STATES_TABLE",
+]

ksadk/sessions/_postgres_schema.py

@@ -0,0 +1,305 @@
+"""PostgresSessionService 的 schema 初始化/迁移 DDL（纯移动自 postgres_service，行为不变）。
+
+以 mixin 形式被 :class:`PostgresSessionService` 继承，依赖宿主提供
+``_schema_ready`` / ``_schema_lock`` / ``_ensure_pool`` / ``_pool``。
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ksadk.sessions._postgres_tables import (
+    _PG_SCHEMA_ADVISORY_LOCK_KEY,
+    KSADK_PG_EVENTS_TABLE,
+    KSADK_PG_SESSIONS_TABLE,
+    KSADK_PG_STATES_TABLE,
+    PG_READABLE_EVENTS_VIEW,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class _PostgresSchemaMixin:
+    async def _ensure_schema(self) -> None:
+        if self._schema_ready:
+            return
+        async with self._schema_lock:
+            if self._schema_ready:
+                return
+            await self._ensure_pool()
+            async with self._pool.acquire() as connection:
+                # A new service instance normally points at an already-current
+                # shared schema.  Avoid taking any DDL lock in that hot path;
+                # otherwise a cold initializer can deadlock with a ready pod's
+                # concurrent event INSERT.
+                if await self._core_schema_is_current(connection):
+                    self._schema_ready = True
+                    return
+
+                migrated = False
+                async with connection.transaction():
+                    # Instance-local asyncio locks cannot coordinate pods.  The
+                    # transaction-scoped database lock serializes true schema
+                    # creation/migration, then the second shape check lets a
+                    # waiting initializer skip duplicate DDL.
+                    await connection.execute(
+                        "SELECT pg_advisory_xact_lock($1)",
+                        _PG_SCHEMA_ADVISORY_LOCK_KEY,
+                    )
+                    if not await self._core_schema_is_current(connection):
+                        await self._create_core_schema(connection)
+                        if not await self._core_schema_is_current(connection):
+                            raise RuntimeError(
+                                "Postgres session schema migration did not produce "
+                                "the required core shape"
+                            )
+                        migrated = True
+
+                # The readable view is optional.  Create it only for the one
+                # initializer that changed core schema, after the core DDL has
+                # committed so a view permission error cannot roll it back.
+                if migrated:
+                    try:
+                        await connection.execute(f"""
+                        CREATE OR REPLACE VIEW {PG_READABLE_EVENTS_VIEW} AS
+                        SELECT
+                            event_row.namespace,
+                            event_row.tenant_id,
+                            event_row.workspace_id,
+                            session_row.agent_id,
+                            session_row.user_id,
+                            session_row.title AS session_title,
+                            event_row.session_id,
+                            event_row.seq_id,
+                            event_row.id AS event_id,
+                            event_row.invocation_id,
+                            event_row.author,
+                            event_row.event_type,
+                            CASE
+                                WHEN event_row.event_type = 'user_message' THEN 'user'
+                                WHEN event_row.event_type IN (
+                                    'assistant_message', 'reasoning', 'tool_call'
+                                ) THEN 'assistant'
+                                WHEN event_row.event_type = 'tool_result' THEN 'tool'
+                                ELSE NULL
+                            END AS message_role,
+                            COALESCE(
+                                NULLIF(event_row.content_json #>> '{{parts,0,text}}', ''),
+                                NULLIF(event_row.content_json ->> 'text', ''),
+                                NULLIF(event_row.metadata_json ->> 'reasoning', ''),
+                                NULLIF(event_row.metadata_json ->> 'tool_output', '')
+                            ) AS message_text,
+                            event_row.metadata_json ->> 'tool_name' AS tool_name,
+                            CASE
+                                WHEN event_row.event_type = 'run_status' THEN COALESCE(
+                                    event_row.content_json ->> 'status',
+                                    event_row.metadata_json ->> 'status'
+                                )
+                                ELSE NULL
+                            END AS lifecycle_status,
+                            to_timestamp(event_row.timestamp) AS created_at,
+                            event_row.content_json,
+                            event_row.state_delta_json,
+                            event_row.metadata_json
+                        FROM {KSADK_PG_EVENTS_TABLE} AS event_row
+                        JOIN {KSADK_PG_SESSIONS_TABLE} AS session_row
+                          ON session_row.namespace = event_row.namespace
+                         AND session_row.id = event_row.session_id;
+                            """)
+                    except Exception as exc:
+                        logger.warning("Postgres readable session view unavailable: %s", exc)
+            self._schema_ready = True
+
+    @staticmethod
+    async def _core_schema_is_current(connection: Any) -> bool:
+        return bool(await connection.fetchval(f"""
+                    SELECT
+                        to_regclass('{KSADK_PG_SESSIONS_TABLE}') IS NOT NULL
+                        AND to_regclass('{KSADK_PG_EVENTS_TABLE}') IS NOT NULL
+                        AND to_regclass('{KSADK_PG_STATES_TABLE}') IS NOT NULL
+                        AND to_regclass('idx_ksadk_pg_events_session_seq') IS NOT NULL
+                        AND to_regclass('idx_ksadk_pg_events_session_invocation_seq') IS NOT NULL
+                        AND to_regclass('idx_ksadk_pg_events_session_ts') IS NOT NULL
+                        AND to_regclass('idx_ksadk_pg_sessions_agent_updated') IS NOT NULL
+                        AND NOT EXISTS (
+                            SELECT 1
+                            FROM (
+                                VALUES
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'namespace'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'tenant_id'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'workspace_id'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'id'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'agent_id'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'user_id'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'title'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'title_source'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'summary'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'first_prompt'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'last_prompt'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'state_json'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'created_at'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'updated_at'),
+                                    ('{KSADK_PG_SESSIONS_TABLE}', 'version'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'namespace'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'tenant_id'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'workspace_id'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'id'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'session_id'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'author'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'event_type'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'content_json'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'timestamp'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'state_delta_json'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'seq_id'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'invocation_id'),
+                                    ('{KSADK_PG_EVENTS_TABLE}', 'metadata_json'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'namespace'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'tenant_id'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'workspace_id'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'scope'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'agent_id'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'user_id'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'session_id'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'state_json'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'version'),
+                                    ('{KSADK_PG_STATES_TABLE}', 'updated_at')
+                            ) AS required(table_name, column_name)
+                            WHERE NOT EXISTS (
+                                SELECT 1
+                                FROM pg_attribute AS attribute_row
+                                WHERE attribute_row.attrelid = to_regclass(required.table_name)
+                                  AND attribute_row.attname = required.column_name
+                                  AND NOT attribute_row.attisdropped
+                            )
+                        )
+                        AND EXISTS (
+                            SELECT 1
+                            FROM pg_constraint AS constraint_row
+                            WHERE constraint_row.conrelid = to_regclass('{KSADK_PG_SESSIONS_TABLE}')
+                              AND constraint_row.contype = 'p'
+                              AND pg_get_constraintdef(constraint_row.oid)
+                                  = 'PRIMARY KEY (namespace, id)'
+                        )
+                        AND EXISTS (
+                            SELECT 1
+                            FROM pg_constraint AS constraint_row
+                            WHERE constraint_row.conrelid = to_regclass('{KSADK_PG_EVENTS_TABLE}')
+                              AND constraint_row.contype = 'p'
+                              AND pg_get_constraintdef(constraint_row.oid)
+                                  = 'PRIMARY KEY (namespace, id)'
+                        )
+                        AND EXISTS (
+                            SELECT 1
+                            FROM pg_constraint AS constraint_row
+                            WHERE constraint_row.conrelid = to_regclass('{KSADK_PG_STATES_TABLE}')
+                              AND constraint_row.contype = 'p'
+                              AND pg_get_constraintdef(constraint_row.oid)
+                                  = 'PRIMARY KEY (namespace, scope, agent_id, user_id, session_id)'
+                        )
+                        AND EXISTS (
+                            SELECT 1
+                            FROM pg_constraint AS constraint_row
+                            WHERE constraint_row.conrelid = to_regclass('{KSADK_PG_EVENTS_TABLE}')
+                              AND constraint_row.contype = 'u'
+                              AND pg_get_constraintdef(constraint_row.oid)
+                                  = 'UNIQUE (namespace, session_id, seq_id)'
+                        )
+                        AND EXISTS (
+                            SELECT 1
+                            FROM pg_constraint AS constraint_row
+                            WHERE constraint_row.conrelid = to_regclass('{KSADK_PG_EVENTS_TABLE}')
+                              AND constraint_row.contype = 'f'
+                              AND pg_get_constraintdef(constraint_row.oid)
+                                  = concat(
+                                      'FOREIGN KEY (namespace, session_id) REFERENCES ',
+                                      '{KSADK_PG_SESSIONS_TABLE}(namespace, id) ON DELETE CASCADE'
+                                  )
+                        )
+                    """))
+
+    @staticmethod
+    async def _create_core_schema(connection: Any) -> None:
+        await connection.execute(f"""
+                CREATE TABLE IF NOT EXISTS {KSADK_PG_SESSIONS_TABLE} (
+                    namespace TEXT NOT NULL,
+                    tenant_id TEXT NOT NULL DEFAULT 'default',
+                    workspace_id TEXT NOT NULL DEFAULT 'default',
+                    id TEXT NOT NULL,
+                    agent_id TEXT NOT NULL,
+                    user_id TEXT NOT NULL,
+                    title TEXT NOT NULL DEFAULT '',
+                    title_source TEXT NOT NULL DEFAULT '',
+                    summary TEXT NOT NULL DEFAULT '',
+                    first_prompt TEXT NOT NULL DEFAULT '',
+                    last_prompt TEXT NOT NULL DEFAULT '',
+                    state_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+                    created_at DOUBLE PRECISION NOT NULL,
+                    updated_at DOUBLE PRECISION NOT NULL,
+                    version INTEGER NOT NULL DEFAULT 0,
+                    PRIMARY KEY (namespace, id)
+                );
+
+                CREATE TABLE IF NOT EXISTS {KSADK_PG_EVENTS_TABLE} (
+                    namespace TEXT NOT NULL,
+                    tenant_id TEXT NOT NULL DEFAULT 'default',
+                    workspace_id TEXT NOT NULL DEFAULT 'default',
+                    id TEXT NOT NULL,
+                    session_id TEXT NOT NULL,
+                    author TEXT NOT NULL,
+                    event_type TEXT NOT NULL,
+                    content_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+                    timestamp DOUBLE PRECISION NOT NULL,
+                    state_delta_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+                    seq_id INTEGER NOT NULL,
+                    invocation_id TEXT,
+                    metadata_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+                    PRIMARY KEY (namespace, id),
+                    UNIQUE (namespace, session_id, seq_id),
+                    FOREIGN KEY (namespace, session_id)
+                        REFERENCES {KSADK_PG_SESSIONS_TABLE}(namespace, id)
+                        ON DELETE CASCADE
+                );
+
+                CREATE INDEX IF NOT EXISTS idx_ksadk_pg_events_session_seq
+                ON {KSADK_PG_EVENTS_TABLE} (namespace, session_id, seq_id);
+
+                CREATE INDEX IF NOT EXISTS idx_ksadk_pg_events_session_invocation_seq
+                ON {KSADK_PG_EVENTS_TABLE} (namespace, session_id, invocation_id, seq_id);
+
+                CREATE INDEX IF NOT EXISTS idx_ksadk_pg_events_session_ts
+                ON {KSADK_PG_EVENTS_TABLE} (namespace, session_id, timestamp, id);
+
+                CREATE INDEX IF NOT EXISTS idx_ksadk_pg_sessions_agent_updated
+                ON {KSADK_PG_SESSIONS_TABLE} (namespace, agent_id, updated_at DESC, id);
+
+                CREATE TABLE IF NOT EXISTS {KSADK_PG_STATES_TABLE} (
+                    namespace TEXT NOT NULL,
+                    tenant_id TEXT NOT NULL DEFAULT 'default',
+                    workspace_id TEXT NOT NULL DEFAULT 'default',
+                    scope TEXT NOT NULL,
+                    agent_id TEXT NOT NULL,
+                    user_id TEXT NOT NULL DEFAULT '',
+                    session_id TEXT NOT NULL DEFAULT '',
+                    state_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
+                    version INTEGER NOT NULL DEFAULT 0,
+                    updated_at DOUBLE PRECISION NOT NULL,
+                    PRIMARY KEY (namespace, scope, agent_id, user_id, session_id)
+                );
+
+                ALTER TABLE {KSADK_PG_SESSIONS_TABLE}
+                ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
+                ALTER TABLE {KSADK_PG_SESSIONS_TABLE}
+                ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT 'default';
+                ALTER TABLE {KSADK_PG_EVENTS_TABLE}
+                ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
+                ALTER TABLE {KSADK_PG_EVENTS_TABLE}
+                ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT 'default';
+                ALTER TABLE {KSADK_PG_STATES_TABLE}
+                ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
+                ALTER TABLE {KSADK_PG_STATES_TABLE}
+                ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT 'default';
+                """)
+
+
+__all__ = ["_PostgresSchemaMixin"]

ksadk/sessions/_postgres_tables.py

@@ -0,0 +1,15 @@
+"""Postgres session 存储的表/视图常量（纯移动自 postgres_service，行为不变）。"""
+
+KSADK_PG_SESSIONS_TABLE = "ksadk_sessions"
+KSADK_PG_EVENTS_TABLE = "ksadk_events"
+KSADK_PG_STATES_TABLE = "ksadk_states"
+PG_READABLE_EVENTS_VIEW = "ksadk_session_events_readable"
+_PG_SCHEMA_ADVISORY_LOCK_KEY = 0x4B5341444B53444B
+
+__all__ = [
+    "KSADK_PG_EVENTS_TABLE",
+    "KSADK_PG_SESSIONS_TABLE",
+    "KSADK_PG_STATES_TABLE",
+    "PG_READABLE_EVENTS_VIEW",
+    "_PG_SCHEMA_ADVISORY_LOCK_KEY",
+]

ksadk/sessions/base.py

@@ -5,7 +5,25 @@ import time
 import uuid
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any, Literal, Optional, TypeAlias
+
+SessionEventSeqBinding: TypeAlias = Literal["runtime_event.seq", "session_event.seq"]
+
+
+@dataclass(frozen=True)
+class SessionServiceStorageCapabilities:
+    """Typed storage guarantees required by canonical event persistence."""
+
+    atomic_seq_bindings: frozenset[SessionEventSeqBinding] = frozenset()
+    indexed_event_lookup: bool = False
+    indexed_invocation_lookup: bool = False
+
+
+CANONICAL_EVENT_STORAGE_CAPABILITIES = SessionServiceStorageCapabilities(
+    atomic_seq_bindings=frozenset({"runtime_event.seq", "session_event.seq"}),
+    indexed_event_lookup=True,
+    indexed_invocation_lookup=True,
+)
 
 
 def generate_id() -> str:
@@ -47,6 +65,47 @@ class SessionEvent:
     seq_id: int = 0
     invocation_id: Optional[str] = None
     metadata: dict[str, Any] = field(default_factory=dict)
+    seq_binding: SessionEventSeqBinding | None = None
+
+    def bind_seq_id(self, seq_id: int) -> None:
+        """Bind a store-assigned cursor into an explicitly declared content field.
+
+        Most session events only need the physical ``seq_id`` column.  A typed
+        event envelope may additionally declare ``runtime_event.seq`` through
+        the transient ``seq_binding`` capability so the JSON fact and carrier
+        are written atomically with the same cursor.  The binding is consumed
+        before persistence and never appears in public metadata or content.
+        """
+
+        self.seq_id = int(seq_id)
+        binding = self.seq_binding
+        self.seq_binding = None
+        if binding is None:
+            return
+        if binding == "session_event.seq":
+            envelope = self.content.get("session_event")
+            if not isinstance(envelope, dict):
+                raise ValueError("session_event.seq binding requires session_event content")
+            envelope = dict(envelope)
+            envelope["seq"] = self.seq_id
+            self.content = {**self.content, "session_event": envelope}
+            return
+        if binding != "runtime_event.seq":
+            raise ValueError(f"unsupported SessionEvent seq binding {binding!r}")
+        runtime_event = self.content.get("runtime_event")
+        if not isinstance(runtime_event, dict):
+            raise ValueError("runtime_event.seq binding requires runtime_event content")
+        runtime_event = dict(runtime_event)
+        runtime_event["seq"] = self.seq_id
+        content = {**self.content, "runtime_event": runtime_event}
+        # Keep the embedded generic envelope dump consistent with the same
+        # cursor when both carriers are present.
+        envelope = content.get("session_event")
+        if isinstance(envelope, dict):
+            envelope = dict(envelope)
+            envelope["seq"] = self.seq_id
+            content["session_event"] = envelope
+        self.content = content
 
     @classmethod
     def from_dict(
@@ -251,6 +310,8 @@ def _infer_event_type(payload: dict[str, Any]) -> str:
 
 
 class BaseSessionService(abc.ABC):
+    storage_capabilities = SessionServiceStorageCapabilities()
+
     @abc.abstractmethod
     async def create_session(
         self,
@@ -309,8 +370,26 @@ class BaseSessionService(abc.ABC):
 
     @abc.abstractmethod
     async def append_event(self, session_id: str, event: SessionEvent) -> SessionEvent:
+        """Append an event with a backend-unique ID and session-unique cursor."""
         raise NotImplementedError
 
+    async def get_event_by_id(self, session_id: str, event_id: str) -> Optional[SessionEvent]:
+        """Indexed physical-id lookup for idempotent event insertion."""
+
+        raise NotImplementedError("session backend does not support indexed event lookup")
+
+    async def get_events_by_invocation_id(
+        self,
+        session_id: str,
+        invocation_id: str,
+        *,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        """Indexed invocation read used by canonical run replay/recovery."""
+
+        raise NotImplementedError("session backend does not support indexed invocation lookup")
+
     @abc.abstractmethod
     async def get_events(
         self,

ksadk/sessions/in_memory.py

@@ -6,12 +6,23 @@ import time
 from typing import Optional
 
 from ksadk.ids import new_session_id
-from ksadk.sessions.base import BaseSessionService, Session, SessionEvent, SessionState, generate_id
+from ksadk.sessions.base import (
+    CANONICAL_EVENT_STORAGE_CAPABILITIES,
+    BaseSessionService,
+    Session,
+    SessionEvent,
+    SessionState,
+    generate_id,
+)
 
 
 class InMemorySessionService(BaseSessionService):
+    storage_capabilities = CANONICAL_EVENT_STORAGE_CAPABILITIES
+
     def __init__(self):
         self._sessions: dict[str, Session] = {}
+        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._lock = asyncio.Lock()
 
@@ -91,6 +102,10 @@ class InMemorySessionService(BaseSessionService):
             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",
@@ -135,12 +150,25 @@ class InMemorySessionService(BaseSessionService):
             if not session:
                 raise ValueError(f"Session {session_id} not found")
 
+            # Match the durable Local/Postgres physical primary-key contract.
+            # Canonical RuntimeEvent storage relies on a deterministic
+            # session+event storage id so concurrent insert losers cannot
+            # allocate another seq.  Auto-generated ids retain their existing
+            # behavior because SessionEvent always supplies a fresh id.
+            if event.id in self._events_by_id:
+                raise ValueError(f"SessionEvent id {event.id!r} already exists")
+
             stored = copy.deepcopy(event)
             stored.session_id = session_id
-            stored.seq_id = len(session.events) + 1
+            stored.bind_seq_id(len(session.events) + 1)
             if not stored.id:
                 stored.id = generate_id()
             session.events.append(stored)
+            self._events_by_id[stored.id] = stored
+            if stored.invocation_id is not None:
+                self._events_by_invocation.setdefault(
+                    (session_id, stored.invocation_id), []
+                ).append(stored)
             session.updated_at = time.time()
 
             if stored.state_delta:
@@ -165,6 +193,29 @@ class InMemorySessionService(BaseSessionService):
 
             return copy.deepcopy(stored)
 
+    async def get_event_by_id(self, session_id: str, event_id: str) -> Optional[SessionEvent]:
+        async with self._lock:
+            event = self._events_by_id.get(event_id)
+            if event is None or event.session_id != session_id:
+                return None
+            return copy.deepcopy(event)
+
+    async def get_events_by_invocation_id(
+        self,
+        session_id: str,
+        invocation_id: str,
+        *,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        async with self._lock:
+            events = list(self._events_by_invocation.get((session_id, invocation_id), ()))
+            if after_seq_id is not None:
+                events = [event for event in events if event.seq_id > after_seq_id]
+            if before_seq_id is not None:
+                events = [event for event in events if event.seq_id < before_seq_id]
+            return copy.deepcopy(events)
+
     async def get_events(
         self,
         session_id: str,