SaaS fixes (#6578)

Co-authored-by: Claude Opus 4.8 <[email protected]>
Co-authored-by: James Brunton <[email protected]>
Co-authored-by: Reece Browne <[email protected]>
Co-authored-by: ConnorYoh <[email protected]>
Co-authored-by: Reece <[email protected]>
Co-authored-by: EthanHealy01 <[email protected]>
Co-authored-by: Ludy <[email protected]>

engine/src/stirling/api/app.py

@@ -18,6 +18,7 @@ from stirling.agents import (
 )
 from stirling.agents.ledger import MathAuditorAgent
 from stirling.agents.pdf_comment import PdfCommentAgent
+from stirling.api.dependencies import enforce_required_user_id
 from stirling.api.engine_auth import EngineSharedSecretMiddleware
 from stirling.api.middleware import UserIdMiddleware
 from stirling.api.routes import (
@@ -118,15 +119,18 @@ async def lifespan(fast_api: FastAPI):
 app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
 app.add_middleware(UserIdMiddleware)
 app.add_middleware(EngineSharedSecretMiddleware)
-app.include_router(orchestrator_router)
-app.include_router(pdf_edit_router)
-app.include_router(pdf_question_router)
-app.include_router(agent_draft_router)
-app.include_router(execution_router)
-app.include_router(document_router)
-app.include_router(ledger_router)
-app.include_router(pdf_comments_router)
-app.include_router(agent_capabilities_router)
+# Every router gets the same configurable identity gate; /health stays open
+# for liveness probes. See enforce_required_user_id for the policy.
+_user_gate = [Depends(enforce_required_user_id)]
+app.include_router(orchestrator_router, dependencies=_user_gate)
+app.include_router(pdf_edit_router, dependencies=_user_gate)
+app.include_router(pdf_question_router, dependencies=_user_gate)
+app.include_router(agent_draft_router, dependencies=_user_gate)
+app.include_router(execution_router, dependencies=_user_gate)
+app.include_router(document_router, dependencies=_user_gate)
+app.include_router(ledger_router, dependencies=_user_gate)
+app.include_router(pdf_comments_router, dependencies=_user_gate)
+app.include_router(agent_capabilities_router, dependencies=_user_gate)
 
 
 @app.get("/health", response_model=HealthResponse)

engine/src/stirling/api/dependencies.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from fastapi import HTTPException, Request, status
+from typing import Annotated
+
+from fastapi import Depends, HTTPException, Request, status
 
 from stirling.agents import (
     ExecutionPlanningAgent,
@@ -11,6 +13,7 @@ from stirling.agents import (
 )
 from stirling.agents.ledger import MathAuditorAgent
 from stirling.agents.pdf_comment import PdfCommentAgent
+from stirling.config import AppSettings, load_settings
 from stirling.documents import DocumentService
 from stirling.models import UserId
 from stirling.services import AppRuntime, current_user_id
@@ -67,3 +70,16 @@ def require_user_id() -> UserId:
             detail="X-User-Id header is required",
         )
     return user_id
+
+
+def enforce_required_user_id(
+    settings: Annotated[AppSettings, Depends(load_settings)],
+) -> None:
+    """Router-level boundary gate, applied uniformly to every router."""
+    if not settings.require_user_id:
+        return
+    if current_user_id.get() is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="X-User-Id header is required",
+        )

engine/src/stirling/config/__init__.py

@@ -1,10 +1,10 @@
 """Configuration models and loaders for the Stirling AI service."""
 
-from .settings import ENGINE_ROOT, AppSettings, RagBackend, load_settings
+from .settings import ENGINE_ROOT, AppSettings, DocumentsBackend, load_settings
 
 __all__ = [
     "ENGINE_ROOT",
     "AppSettings",
-    "RagBackend",
+    "DocumentsBackend",
     "load_settings",
 ]

engine/src/stirling/config/settings.py

@@ -15,7 +15,7 @@ ENV_FILE = ENGINE_ROOT / ".env"
 ENV_LOCAL_FILE = ENGINE_ROOT / ".env.local"
 
 
-class RagBackend(StrEnum):
+class DocumentsBackend(StrEnum):
     SQLITE = "sqlite"
     PGVECTOR = "pgvector"
 
@@ -27,12 +27,22 @@ class AppSettings(BaseSettings):
     fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL")
     smart_model_max_tokens: int = Field(validation_alias="STIRLING_SMART_MODEL_MAX_TOKENS")
     fast_model_max_tokens: int = Field(validation_alias="STIRLING_FAST_MODEL_MAX_TOKENS")
+    # Process-wide ceiling on concurrent model API calls, shared by both model
+    # tiers. Per-request fan-outs (chunked reasoner, contradiction detection)
+    # carry their own per-request caps, but those multiply under concurrent
+    # traffic; this is the global backstop.
+    model_max_concurrency: int = Field(validation_alias="STIRLING_MODEL_MAX_CONCURRENCY")
 
-    # RAG settings — always on; the backend picks between embedded sqlite-vec and external pgvector.
-    rag_backend: RagBackend = Field(validation_alias="STIRLING_RAG_BACKEND")
+    # Document store: the one database holding vector chunks, ordered page
+    # text, and ACL rows - embedded sqlite-vec or external pgvector.
+    documents_backend: DocumentsBackend = Field(validation_alias="STIRLING_DOCUMENTS_BACKEND")
+    documents_sqlite_path: Path = Field(validation_alias="STIRLING_DOCUMENTS_SQLITE_PATH")
+    documents_pgvector_dsn: str = Field(validation_alias="STIRLING_DOCUMENTS_PGVECTOR_DSN")
+    documents_pgvector_pool_min_size: int = Field(validation_alias="STIRLING_DOCUMENTS_PGVECTOR_POOL_MIN_SIZE")
+    documents_pgvector_pool_max_size: int = Field(validation_alias="STIRLING_DOCUMENTS_PGVECTOR_POOL_MAX_SIZE")
+
+    # RAG settings - always on.
     rag_embedding_model: str = Field(validation_alias="STIRLING_RAG_EMBEDDING_MODEL")
-    rag_store_path: Path = Field(validation_alias="STIRLING_RAG_STORE_PATH")
-    rag_pgvector_dsn: str = Field(validation_alias="STIRLING_RAG_PGVECTOR_DSN")
     rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE")
     rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
     rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
@@ -88,6 +98,12 @@ class AppSettings(BaseSettings):
     max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES")
     max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS")
 
+    # When true, API routes reject requests that lack an X-User-Id header at
+    # the boundary. Self-hosted deployments with security disabled have no
+    # user identity and leave this off; multi-tenant deployments turn it on so
+    # user-scoped work is never processed without a tenant attached.
+    require_user_id: bool = Field(validation_alias="STIRLING_REQUIRE_USER_ID")
+
     log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL")
     log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE")
     # When true, raises httpx + httpcore logger levels so every outgoing

engine/src/stirling/documents/README.md

@@ -54,10 +54,10 @@ everything = RagCapability(runtime.documents)
 Non-secret defaults live in the committed `engine/.env`:
 
 ```
-STIRLING_RAG_BACKEND=sqlite              # or "pgvector"
+STIRLING_DOCUMENTS_BACKEND=sqlite              # or "pgvector"
 STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4
-STIRLING_RAG_STORE_PATH=data/rag.db      # used when backend=sqlite
-STIRLING_RAG_PGVECTOR_DSN=               # used when backend=pgvector
+STIRLING_DOCUMENTS_SQLITE_PATH=data/rag.db      # used when backend=sqlite
+STIRLING_DOCUMENTS_PGVECTOR_DSN=               # used when backend=pgvector
 STIRLING_RAG_CHUNK_SIZE=512
 STIRLING_RAG_CHUNK_OVERLAP=64
 STIRLING_RAG_TOP_K=5
@@ -76,7 +76,7 @@ VOYAGE_API_KEY=your-key
 and self-hosted deployments.
 
 **`pgvector`** - External PostgreSQL with the `vector` extension. Point
-`STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
+`STIRLING_DOCUMENTS_PGVECTOR_DSN` at your Postgres instance.
 
 Both backends implement the same `DocumentStore` interface, so agents and the
 service work identically regardless of which you pick.

engine/src/stirling/documents/pgvector_store.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import asyncio
 import json
 from datetime import datetime
 
 import psycopg
 from pgvector import Vector
 from pgvector.psycopg import register_vector_async
+from psycopg_pool import AsyncConnectionPool
 
 from stirling.contracts.documents import Page, PageRange
 from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
@@ -14,12 +16,20 @@ from stirling.models import OwnerId, PrincipalId
 _READ_PERMISSION = "read"
 
 
+async def _register_vector(conn: psycopg.AsyncConnection) -> None:
+    await register_vector_async(conn)
+
+
 class PgVectorStore(DocumentStore):
     """PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads.
 
     Connects to an external Postgres instance (DSN provided via config) and uses the
     `vector` extension for similarity search. The schema is created on first use.
 
+    Queries run on a shared async connection pool. The pool is opened lazily on
+    first use, after the schema bootstrap, because each pooled connection's
+    configure hook registers the ``vector`` type, which must already exist.
+
     Owned tables:
 
     * ``documents_meta`` - parent row, one per ``(collection, owner_id)``.
@@ -32,28 +42,39 @@ class PgVectorStore(DocumentStore):
     Writes are owner-scoped; reads are ACL-scoped.
     """
 
-    def __init__(self, dsn: str) -> None:
+    def __init__(self, dsn: str, pool_min_size: int, pool_max_size: int) -> None:
         if not dsn:
-            raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_RAG_PGVECTOR_DSN)")
+            raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_DOCUMENTS_PGVECTOR_DSN)")
         self._dsn = dsn
+        self._pool = AsyncConnectionPool(
+            dsn,
+            min_size=pool_min_size,
+            max_size=pool_max_size,
+            open=False,
+            configure=_register_vector,
+            check=AsyncConnectionPool.check_connection,
+        )
         self._initialized = False
+        self._init_lock = asyncio.Lock()
 
-    async def _connect(self) -> psycopg.AsyncConnection:
-        conn = await psycopg.AsyncConnection.connect(self._dsn)
-        await register_vector_async(conn)
-        return conn
-
-    async def _ensure_schema(self) -> None:
+    async def _ensure_ready(self) -> None:
+        """Create the schema and open the pool, exactly once across concurrent callers."""
         if self._initialized:
             return
-        # The `vector` type must exist before register_vector_async (called inside
-        # _connect) can resolve it. On a fresh database it doesn't yet, so create the
-        # extension first on a raw connection that hasn't registered the type.
-        async with await psycopg.AsyncConnection.connect(self._dsn) as bootstrap:
-            async with bootstrap.cursor() as cur:
-                await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
-            await bootstrap.commit()
-        async with await self._connect() as conn:
+        async with self._init_lock:
+            if self._initialized:
+                return
+            await self._bootstrap_schema()
+            await self._pool.open()
+            self._initialized = True
+
+    async def _bootstrap_schema(self) -> None:
+        # The `vector` type must exist before the pool's configure hook
+        # (register_vector_async) can resolve it, and on a fresh database it doesn't
+        # yet. Create the extension + schema on a raw, non-pool connection that hasn't
+        # registered the type; _ensure_ready then opens the pool, whose configure hook
+        # registers the now-existing type per connection.
+        async with await psycopg.AsyncConnection.connect(self._dsn) as conn:
             async with conn.cursor() as cur:
                 await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
                 await cur.execute(
@@ -126,7 +147,6 @@ class PgVectorStore(DocumentStore):
                     "CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
                 )
                 await conn.commit()
-        self._initialized = True
 
     # ── lifecycle of the (collection, owner_id) row ────────────────────────
 
@@ -137,8 +157,8 @@ class PgVectorStore(DocumentStore):
         owner_id: OwnerId,
         expires_at: datetime | None,
     ) -> None:
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.execute(
                     """
@@ -153,8 +173,8 @@ class PgVectorStore(DocumentStore):
                 await conn.commit()
 
     async def purge_owner(self, owner_id: OwnerId) -> int:
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.execute(
                     "DELETE FROM documents_meta WHERE owner_id = %s",
@@ -165,8 +185,8 @@ class PgVectorStore(DocumentStore):
         return deleted
 
     async def reap_expired(self) -> int:
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()")
                 deleted = cur.rowcount
@@ -174,8 +194,8 @@ class PgVectorStore(DocumentStore):
         return deleted
 
     async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 # Cascade FKs handle rag_documents, document_pages, document_acl.
                 await cur.execute(
@@ -200,8 +220,8 @@ class PgVectorStore(DocumentStore):
         if not documents:
             return
 
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 for doc, emb in zip(documents, embeddings):
                     await cur.execute(
@@ -219,8 +239,8 @@ class PgVectorStore(DocumentStore):
                 await conn.commit()
 
     async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.execute(
                     "DELETE FROM document_pages WHERE collection = %s AND owner_id = %s",
@@ -246,8 +266,8 @@ class PgVectorStore(DocumentStore):
     ) -> None:
         if not principals:
             return
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.executemany(
                     """
@@ -265,8 +285,8 @@ class PgVectorStore(DocumentStore):
         owner_id: OwnerId,
         principal: PrincipalId,
     ) -> None:
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.execute(
                     "DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s",
@@ -285,8 +305,8 @@ class PgVectorStore(DocumentStore):
     ) -> list[SearchResult]:
         if not principals:
             return []
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 owner_id = await self._readable_owner_for(cur, collection, principals)
                 if owner_id is None:
@@ -341,8 +361,8 @@ class PgVectorStore(DocumentStore):
     ) -> list[Page]:
         if not principals:
             return []
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 owner_id = await self._readable_owner_for(cur, collection, principals)
                 if owner_id is None:
@@ -367,8 +387,8 @@ class PgVectorStore(DocumentStore):
     async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
         if not principals:
             return False
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 owner_id = await self._readable_owner_for(cur, collection, principals)
                 return owner_id is not None
@@ -376,8 +396,8 @@ class PgVectorStore(DocumentStore):
     async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
         if not principals:
             return []
-        await self._ensure_schema()
-        async with await self._connect() as conn:
+        await self._ensure_ready()
+        async with self._pool.connection() as conn:
             async with conn.cursor() as cur:
                 await cur.execute(
                     """
@@ -392,5 +412,4 @@ class PgVectorStore(DocumentStore):
         return [r[0] for r in rows]
 
     async def close(self) -> None:
-        # Connections are opened and closed per call, so nothing persistent to release.
-        return None
+        await self._pool.close()

engine/src/stirling/services/runtime.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
+import asyncio
 import logging
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import assert_never
+from typing import Any, assert_never
 
 import httpx
-from pydantic_ai.models import Model, infer_model
+from pydantic_ai import RunContext
+from pydantic_ai.messages import ModelMessage, ModelResponse
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, infer_model
 from pydantic_ai.models.anthropic import AnthropicModel
+from pydantic_ai.models.wrapper import WrapperModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
 from pydantic_ai.settings import ModelSettings
 
-from stirling.config import ENGINE_ROOT, AppSettings, RagBackend
+from stirling.config import ENGINE_ROOT, AppSettings, DocumentsBackend
 from stirling.documents import (
     DocumentService,
     DocumentStore,
@@ -43,6 +49,37 @@ def _build_anthropic_http_client() -> httpx.AsyncClient:
     )
 
 
+class ConcurrencyLimitedModel(WrapperModel):
+    """Caps in-flight model API calls with a semaphore shared across the process."""
+
+    def __init__(self, wrapped: Model, semaphore: asyncio.Semaphore) -> None:
+        super().__init__(wrapped)
+        self._semaphore = semaphore
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        async with self._semaphore:
+            return await super().request(messages, model_settings, model_request_parameters)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        async with self._semaphore:
+            async with super().request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as response_stream:
+                yield response_stream
+
+
 @dataclass(frozen=True)
 class AppRuntime:
     settings: AppSettings
@@ -77,17 +114,21 @@ def validate_structured_output_support(model: Model, model_name: str) -> None:
 
 def _build_document_store(settings: AppSettings) -> DocumentStore:
     """Build the configured document store backend."""
-    if settings.rag_backend == RagBackend.SQLITE:
-        store_path = settings.rag_store_path
+    if settings.documents_backend == DocumentsBackend.SQLITE:
+        store_path = settings.documents_sqlite_path
         # Treat ":memory:" as a special in-process token; otherwise resolve against the engine root.
         if str(store_path) != ":memory:" and not store_path.is_absolute():
             store_path = ENGINE_ROOT / store_path
         logger.info("Document store backend=sqlite, db_path=%s", store_path)
         return SqliteVecStore(db_path=store_path)
-    if settings.rag_backend == RagBackend.PGVECTOR:
+    if settings.documents_backend == DocumentsBackend.PGVECTOR:
         logger.info("Document store backend=pgvector, dsn=<configured>")
-        return PgVectorStore(dsn=settings.rag_pgvector_dsn)
-    assert_never(settings.rag_backend)
+        return PgVectorStore(
+            dsn=settings.documents_pgvector_dsn,
+            pool_min_size=settings.documents_pgvector_pool_min_size,
+            pool_max_size=settings.documents_pgvector_pool_max_size,
+        )
+    assert_never(settings.documents_backend)
 
 
 def _build_documents(settings: AppSettings) -> DocumentService:
@@ -108,10 +149,13 @@ def build_runtime(settings: AppSettings) -> AppRuntime:
     validate_structured_output_support(fast_model, settings.fast_model_name)
     validate_structured_output_support(smart_model, settings.smart_model_name)
 
+    # One semaphore across both tiers: the cap protects the provider account
+    # and process resources, which the tiers share.
+    model_semaphore = asyncio.Semaphore(settings.model_max_concurrency)
     return AppRuntime(
         settings=settings,
-        fast_model=fast_model,
-        smart_model=smart_model,
+        fast_model=ConcurrencyLimitedModel(fast_model, model_semaphore),
+        smart_model=ConcurrencyLimitedModel(smart_model, model_semaphore),
         documents=_build_documents(settings),
     )