perf(server): collapse the per-event access-control reads into one checkout (#4737)

Every event posted to POST /v1/sessions/{id}/events runs an access-control
check before it does anything else, and that check's reads dominate at scale.
For a session whose permissions/metadata are stable for the turn, each streamed
chunk re-pays: resolve_access (session_permissions + users) plus
get_conversation, which internally reads the conversation row + labels in one
pool checkout and the metadata row in a separate one. That is three distinct
checkouts, each with a pool_pre_ping liveness round-trip, per event — the bulk
of the ~16 queries / ~10 checkouts per event measured in #3004.

shared_read_scope() (db/utils.py): a read-only, per-request scope in which
managed_session() reuses one session per engine instead of opening a fresh
checkout on every store call. Wrapped around the access-control burst, so the
permission + conversation + metadata reads collapse to a single checkout
(single-DB); split-DB deployments keep independent checkouts per engine. Write
makers (immediate=True) never participate, so BEGIN IMMEDIATE isolation is
untouched, and the scope is a strict no-op everywhere it is not activated. The
scope is deliberately kept off any path that spans runner network I/O so it
never pins a pooled connection across a multi-second turn.

No ACL semantics change: the same rows are read, just once per request over one
connection. Verified by a new checkout-counting test that the access-control
burst issues exactly one pool checkout, plus unit tests for the scope's reuse,
nesting, write-maker bypass, per-engine keying, and cleanup.

Closes #3004

Co-authored-by: Isaac


Register the scope's session before its SQLite PRAGMAs run: those executes
force the pool checkout, so a failure there must leave the session tracked by
the scope's cleanup — otherwise the checked-out connection would leak.

Signed-off-by: Edwin He <41037314+Edwinhe03@users.noreply.github.com>

omnigent/db/utils.py

@@ -10,6 +10,7 @@ import time
 import uuid
 from collections.abc import Callable, Iterator
 from contextlib import AbstractContextManager, contextmanager
+from contextvars import ContextVar
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -553,6 +554,54 @@ def clear_engine_cache() -> None:
 # ── Managed session ────────────────────────────────────
 
 
+# Ambient per-engine sessions for a read-only "share one checkout" scope. When
+# active (see :func:`shared_read_scope`), ``managed_session()`` reuses the
+# scope's session for its engine instead of opening a fresh pool checkout,
+# collapsing several back-to-back reads (e.g. the access-control check's
+# permission + conversation lookups) into a single connection round-trip.
+# Keyed by ``id(engine)`` so distinct engines (split-DB) still get independent
+# checkouts. Unset outside a scope, so it is a strict no-op for every ordinary
+# caller.
+_shared_read_sessions: ContextVar[dict[int, Session] | None] = ContextVar(
+    "omnigent_shared_read_sessions", default=None
+)
+
+
+@contextmanager
+def shared_read_scope() -> Iterator[None]:
+    """Collapse back-to-back reads into one pool checkout per engine.
+
+    Within this scope, ``managed_session()`` reuses a single session per
+    engine rather than checking out a fresh pooled connection (plus a
+    ``pool_pre_ping`` round-trip) on every store call. Intended for a short,
+    strictly READ-ONLY burst — an access-control check, a snapshot assembly —
+    where the per-call checkout dominates the actual query time.
+
+    Nesting reuses the outer scope. Write makers (``immediate=True``) never
+    participate, so they keep their own ``BEGIN IMMEDIATE`` isolation even
+    when nested here. Never hold this open across network I/O: it pins a
+    pooled connection for the scope's whole duration.
+    """
+    if _shared_read_sessions.get() is not None:
+        # Already inside a scope — the outer one owns the sessions.
+        yield
+        return
+    sessions: dict[int, Session] = {}
+    token = _shared_read_sessions.set(sessions)
+    try:
+        yield
+        for session in sessions.values():
+            session.commit()
+    except BaseException:
+        for session in sessions.values():
+            session.rollback()
+        raise
+    finally:
+        for session in sessions.values():
+            session.close()
+        _shared_read_sessions.reset(token)
+
+
 def make_managed_session_maker(
     engine: Engine,
     *,
@@ -592,7 +641,27 @@ def make_managed_session_maker(
         Commits on clean exit, rolls back on exception. For SQLite
         backends, enables foreign key enforcement and sets a
         busy timeout before yielding.
+
+        Inside a :func:`shared_read_scope` (and only for read makers), the
+        scope's per-engine session is reused instead of a fresh checkout;
+        the scope — not this block — owns its commit/close.
         """
+        if not immediate:
+            shared = _shared_read_sessions.get()
+            if shared is not None:
+                key = id(engine)
+                session = shared.get(key)
+                if session is None:
+                    session = factory()
+                    # Register before the PRAGMAs: those executes force the pool
+                    # checkout, so if one raises the scope must already track the
+                    # session to close it (otherwise the connection would leak).
+                    shared[key] = session
+                    if is_sqlite:
+                        session.execute(text("PRAGMA foreign_keys = ON"))
+                        session.execute(text("PRAGMA busy_timeout = 20000"))  # 20s
+                yield session
+                return
         with factory() as session:
             try:
                 if is_sqlite:

omnigent/server/routes/_auth_helpers.py

@@ -23,6 +23,7 @@ import dataclasses
 
 from fastapi import Request
 
+from omnigent.db.utils import shared_read_scope
 from omnigent.entities import Conversation
 from omnigent.errors import ErrorCode, OmnigentError
 from omnigent.server.auth import (
@@ -290,52 +291,58 @@ def _require_access_and_level_sync(
             code=ErrorCode.UNAUTHORIZED,
         )
 
-    # Single round-trip: admin flag + the user's and public grants on the
-    # conversation the caller asked about. The displayed level is the direct
-    # grant (no parent walk), matching get_permission_level exactly.
-    access = permission_store.resolve_access(user_id, conversation_id)
-    level = resolved_level(access)
+    # One read-only burst: the permission resolve, the conversation lookup,
+    # and any parent-chain walk all share a single pool checkout instead of
+    # one per store call. On the per-streamed-event path this is re-run for a
+    # session whose data is stable for the turn, so the checkout — plus
+    # ``pool_pre_ping`` — is the cost that matters.
+    with shared_read_scope():
+        # Single round-trip: admin flag + the user's and public grants on the
+        # conversation the caller asked about. The displayed level is the direct
+        # grant (no parent walk), matching get_permission_level exactly.
+        access = permission_store.resolve_access(user_id, conversation_id)
+        level = resolved_level(access)
 
-    # Admins bypass the conversation lookup entirely (mirrors
-    # check_session_access's admin short-circuit, which never reads the
-    # conversation). A missing conversation is left for the snapshot builder
-    # to 404 on, exactly as today.
-    if access.is_admin:
-        return SessionAccess(level=level, conversation=None)
+        # Admins bypass the conversation lookup entirely (mirrors
+        # check_session_access's admin short-circuit, which never reads the
+        # conversation). A missing conversation is left for the snapshot builder
+        # to 404 on, exactly as today.
+        if access.is_admin:
+            return SessionAccess(level=level, conversation=None)
 
-    conv = conversation_store.get_conversation(conversation_id)
-    if conv is None:
-        raise OmnigentError(
-            "Conversation not found",
-            code=ErrorCode.NOT_FOUND,
-        )
+        conv = conversation_store.get_conversation(conversation_id)
+        if conv is None:
+            raise OmnigentError(
+                "Conversation not found",
+                code=ErrorCode.NOT_FOUND,
+            )
 
-    if conv.parent_conversation_id is None:
-        # Top-level session: the access-governing grant lives on this same
-        # conversation, so reuse the rows already fetched — no extra reads.
-        allowed = resolved_allows(access, required_level)
-    else:
-        # Sub-agent: access delegates to the parent chain. Defer to the
-        # canonical recursive checker (its own reads); sub-agents are rare
-        # and the parent's grants are a different conversation's rows.
-        allowed = check_session_access(
-            user_id,
-            conv.parent_conversation_id,
-            required_level,
-            permission_store,
-            conversation_store,
-        )
-    if allowed:
-        return SessionAccess(level=level, conversation=conv)
+        if conv.parent_conversation_id is None:
+            # Top-level session: the access-governing grant lives on this same
+            # conversation, so reuse the rows already fetched — no extra reads.
+            allowed = resolved_allows(access, required_level)
+        else:
+            # Sub-agent: access delegates to the parent chain. Defer to the
+            # canonical recursive checker (its own reads); sub-agents are rare
+            # and the parent's grants are a different conversation's rows.
+            allowed = check_session_access(
+                user_id,
+                conv.parent_conversation_id,
+                required_level,
+                permission_store,
+                conversation_store,
+            )
+        if allowed:
+            return SessionAccess(level=level, conversation=conv)
 
-    # Denied — distinguish "has some access but not enough" (403) from
-    # "no access at all" (404, to avoid leaking session existence).
-    if conv.parent_conversation_id is None:
-        has_any = resolved_allows(access, 1)
-    else:
-        has_any = check_session_access(
-            user_id, conv.parent_conversation_id, 1, permission_store, conversation_store
-        )
+        # Denied — distinguish "has some access but not enough" (403) from
+        # "no access at all" (404, to avoid leaking session existence).
+        if conv.parent_conversation_id is None:
+            has_any = resolved_allows(access, 1)
+        else:
+            has_any = check_session_access(
+                user_id, conv.parent_conversation_id, 1, permission_store, conversation_store
+            )
     if has_any:
         level_name = _LEVEL_NAMES.get(required_level, str(required_level))
         raise OmnigentError(