Harden Engram runtime and verification

Author: macayu17

.env.example

@@ -23,6 +23,6 @@ LOG_LEVEL=info
 MCP_SERVICE_KEY=
 ENGRAM_SERVICE_KEY=
 API_PUBLIC_URL=http://localhost:8000
-CORS_ORIGINS=http://localhost:3001
+CORS_ORIGINS=http://localhost:3001,http://localhost:3011
 NEXT_PUBLIC_CLERK_PUBLISHABLE_KEY=
 CLERK_SECRET_KEY=

.github/workflows/ci.yml

@@ -13,6 +13,8 @@ jobs:
         with:
           python-version: "3.11"
       - run: python -m compileall api
+      - run: python -m pip install asyncpg==0.29.0 pydantic==2.7.0 pydantic-settings==2.3.0 httpx==0.27.0 pytest==8.2.2
+      - run: python -m pytest api/test_user_auth_cache.py api/test_provider_and_extraction_parsing.py
 
   mcp:
     runs-on: ubuntu-latest
@@ -27,6 +29,8 @@ jobs:
         working-directory: mcp
       - run: npm run build
         working-directory: mcp
+      - run: npm run verify:defaults
+        working-directory: mcp
 
   dashboard:
     runs-on: ubuntu-latest

api/config.py

@@ -24,6 +24,8 @@ class Settings(BaseSettings):
     log_level: str = "info"
     cors_origins: str = "http://localhost:3001"
     engram_service_key: str = ""
+    proxy_auth_cache_ttl_seconds: int = 300
+    proxy_auth_cache_max_entries: int = 4096
     engram_test_api_url: str = "http://localhost:8000"
     engram_test_provider: str = "openai"
     engram_test_model: str = "gpt-4o-mini"

api/db/schema.sql

@@ -5,10 +5,17 @@ CREATE TABLE IF NOT EXISTS users (
     id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
     external_id TEXT UNIQUE NOT NULL,
     api_key_hash TEXT UNIQUE NOT NULL,
+    max_memories_injected INT NOT NULL DEFAULT 5,
+    retrieval_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.5,
+    dedup_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.95,
     created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
     updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
 );
 
+ALTER TABLE users ADD COLUMN IF NOT EXISTS max_memories_injected INT NOT NULL DEFAULT 5;
+ALTER TABLE users ADD COLUMN IF NOT EXISTS retrieval_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.5;
+ALTER TABLE users ADD COLUMN IF NOT EXISTS dedup_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.95;
+
 CREATE TABLE IF NOT EXISTS memories (
     id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
     user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,

api/main.py

@@ -7,6 +7,7 @@
 
 from api.config import settings
 from api.db.connection import check_database, close_pool, init_pool
+from api.db.schema import apply_schema
 from api.routes import logs, memories, proxy, users
 from api.services.embedding import is_model_loaded, load_model
 
@@ -16,6 +17,7 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+    await apply_schema()
     await init_pool()
     load_model()
     yield

api/models/user.py

@@ -20,3 +20,15 @@ class UserResponse(BaseModel):
 
 class UserCreateResponse(UserResponse):
     api_key: str
+
+
+class UserConfigUpdate(BaseModel):
+    max_memories_injected: int | None = Field(default=None, ge=1, le=20)
+    retrieval_threshold: float | None = Field(default=None, ge=0, le=1)
+    dedup_threshold: float | None = Field(default=None, ge=0, le=1)
+
+
+class UserConfigResponse(BaseModel):
+    max_memories_injected: int
+    retrieval_threshold: float
+    dedup_threshold: float

api/routes/memories.py

@@ -46,7 +46,7 @@ async def create_memory_route(
     user: asyncpg.Record = Depends(get_current_user),
     db: asyncpg.Connection = Depends(get_db),
 ) -> dict[str, object]:
-    return await create_memory(user["id"], payload.content, db)
+    return await create_memory(user["id"], payload.content, db, float(user["dedup_threshold"]))
 
 
 @router.get("/{memory_id}", response_model=MemoryResponse)

api/routes/proxy.py

@@ -1,14 +1,15 @@
 import asyncio
+from collections.abc import Mapping
 import json
 import logging
-from uuid import UUID
 
 import asyncpg
-from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
+from fastapi import APIRouter, Header, HTTPException, Request, Response
 
-from api.dependencies import get_current_user, get_db
+from api.db.connection import get_pool
 from api.services.extraction import run_extraction_task
-from api.services.proxy import build_proxy_result
+from api.services.proxy import ProxyResult, build_proxy_passthrough_result, build_proxy_result
+from api.services.users import get_cached_user_by_api_key, get_user_by_api_key
 
 
 logger = logging.getLogger(__name__)
@@ -18,36 +19,105 @@
 @router.post("/v1/chat")
 async def proxy_chat(
     request: Request,
+    x_engram_key: str = Header(...),
     x_engram_user_id: str = Header(...),
     x_engram_provider: str = Header(default="openai"),
     x_engram_disable_injection: bool = Header(default=False),
     x_engram_disable_extraction: bool = Header(default=False),
-    db: asyncpg.Connection = Depends(get_db),
-    user: asyncpg.Record = Depends(get_current_user),
 ) -> Response:
     body = await parse_request_body(request)
+    result = await build_proxy_response_with_available_auth(
+        x_engram_key,
+        x_engram_user_id,
+        x_engram_provider,
+        x_engram_disable_injection,
+        x_engram_disable_extraction,
+        body,
+        request.headers,
+    )
+    return create_response(result)
+
+
+async def build_proxy_response_with_available_auth(
+    api_key: str,
+    requested_external_id: str,
+    provider: str,
+    disable_injection: bool,
+    disable_extraction: bool,
+    body: dict[str, object],
+    headers: Mapping[str, str],
+) -> ProxyResult:
+    try:
+        pool = get_pool()
+    except RuntimeError as error:
+        logger.warning("Database pool unavailable, trying cached proxy auth: %s", error)
+        return await build_cached_proxy_response(api_key, requested_external_id, provider, body, headers)
+    acquire_context = pool.acquire()
     try:
+        db = await acquire_context.__aenter__()
+    except (asyncpg.PostgresError, OSError, ConnectionError, RuntimeError) as error:
+        logger.warning("Database unavailable, trying cached proxy auth: %s", error)
+        return await build_cached_proxy_response(api_key, requested_external_id, provider, body, headers)
+    try:
+        try:
+            user = await get_user_by_api_key(api_key, db)
+        except (asyncpg.PostgresError, OSError, ConnectionError) as error:
+            logger.warning("Database auth unavailable, trying cached proxy auth: %s", error)
+            return await build_cached_proxy_response(api_key, requested_external_id, provider, body, headers)
+        if user is None:
+            raise HTTPException(status_code=401, detail="Invalid API key")
         result = await build_proxy_result(
             user["id"],
             user["external_id"],
-            x_engram_user_id,
+            requested_external_id,
             body,
-            x_engram_provider,
-            x_engram_disable_injection,
-            request.headers,
+            provider,
+            disable_injection,
+            headers,
             db,
+            int(user["max_memories_injected"]),
+            float(user["retrieval_threshold"]),
         )
     except PermissionError as error:
         raise HTTPException(status_code=403, detail=str(error)) from error
     except ValueError as error:
         raise HTTPException(status_code=422, detail=str(error)) from error
     except RuntimeError as error:
         raise HTTPException(status_code=502, detail=str(error)) from error
-    if not x_engram_disable_extraction and result.status_code < 400:
+    finally:
+        try:
+            await acquire_context.__aexit__(None, None, None)
+        except (asyncpg.PostgresError, OSError, ConnectionError) as error:
+            logger.warning("Database connection release failed: %s", error)
+    if not disable_extraction and result.status_code < 400:
         try:
-            asyncio.create_task(run_extraction_task(user["id"], result.conversation_id, body, result.content))
+            asyncio.create_task(run_extraction_task(user["id"], result.conversation_id, body, result.content, float(user["dedup_threshold"])))
         except Exception as error:
             logger.warning("Failed to schedule extraction: %s", error)
+    return result
+
+
+async def build_cached_proxy_response(
+    api_key: str,
+    requested_external_id: str,
+    provider: str,
+    body: dict[str, object],
+    headers: Mapping[str, str],
+) -> ProxyResult:
+    user = get_cached_user_by_api_key(api_key)
+    if user is None:
+        raise HTTPException(status_code=503, detail="Database unavailable and API key is not cached")
+    try:
+        return await build_proxy_passthrough_result(str(user["external_id"]), requested_external_id, body, provider, headers)
+    except PermissionError as error:
+        raise HTTPException(status_code=403, detail=str(error)) from error
+    except ValueError as error:
+        raise HTTPException(status_code=422, detail=str(error)) from error
+    except RuntimeError as error:
+        raise HTTPException(status_code=502, detail=str(error)) from error
+
+
+def create_response(result: ProxyResult) -> Response:
     return Response(
         content=result.content,
         status_code=result.status_code,

api/routes/users.py

@@ -2,8 +2,8 @@
 from fastapi import APIRouter, Depends, HTTPException, Response, status
 
 from api.dependencies import get_current_user, get_db, require_service_key
-from api.models.user import ServiceUserKeyCreate, UserCreate, UserCreateResponse, UserResponse
-from api.services.users import create_or_issue_user_key, create_user, delete_user
+from api.models.user import ServiceUserKeyCreate, UserConfigResponse, UserConfigUpdate, UserCreate, UserCreateResponse, UserResponse
+from api.services.users import create_or_issue_user_key, create_user, delete_user, get_user_config, regenerate_user_key, update_user_config
 
 
 router = APIRouter()
@@ -47,6 +47,43 @@ async def get_current_user_route(user: asyncpg.Record = Depends(get_current_user
     }
 
 
+@router.get("/me/config", response_model=UserConfigResponse)
+async def get_current_user_config_route(
+    user: asyncpg.Record = Depends(get_current_user),
+    db: asyncpg.Connection = Depends(get_db),
+) -> dict[str, object]:
+    return await get_user_config(user["id"], db)
+
+
+@router.patch("/me/config", response_model=UserConfigResponse)
+async def update_current_user_config_route(
+    payload: UserConfigUpdate,
+    user: asyncpg.Record = Depends(get_current_user),
+    db: asyncpg.Connection = Depends(get_db),
+) -> dict[str, object]:
+    return await update_user_config(
+        user["id"],
+        payload.max_memories_injected,
+        payload.retrieval_threshold,
+        payload.dedup_threshold,
+        db,
+    )
+
+
+@router.post("/me/api-key", response_model=UserCreateResponse)
+async def regenerate_current_user_key_route(
+    user: asyncpg.Record = Depends(get_current_user),
+    db: asyncpg.Connection = Depends(get_db),
+) -> dict[str, object]:
+    row, api_key = await regenerate_user_key(user, db)
+    return {
+        "id": row["id"],
+        "external_id": row["external_id"],
+        "api_key": api_key,
+        "created_at": row["created_at"],
+    }
+
+
 @router.delete("/me", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_current_user_route(
     user: asyncpg.Record = Depends(get_current_user),

api/services/deduplication.py

@@ -18,7 +18,9 @@ async def store_memory_with_deduplication(
     conversation_id: object | None,
     confidence: float,
     db: asyncpg.Connection,
+    dedup_threshold: float | None = None,
 ) -> StoreMemoryResult:
+    dedup_threshold_value = settings.dedup_threshold if dedup_threshold is None else dedup_threshold
     vector_literal = format_embedding_for_pgvector(embedding)
     nearest = await db.fetchrow(
         """
@@ -32,7 +34,7 @@ async def store_memory_with_deduplication(
         vector_literal,
         user_id,
     )
-    if nearest is not None and nearest["score"] > settings.dedup_threshold:
+    if nearest is not None and nearest["score"] > dedup_threshold_value:
         return {"action": "skipped", "memory": dict(nearest)}
     if nearest is not None and nearest["score"] > settings.memory_refinement_threshold:
         updated = await db.fetchrow(