diff --git a/CLAUDE.md b/CLAUDE.md index fbd05871..87a95c3e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -116,6 +116,7 @@ cd sdks/typescript && bun run tsc --noEmit - Explicit error handling with appropriate exception types - Docstrings: Use Google style docstrings - **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection. +- **Never write through a read-only session** (`tracked_db(..., read_only=True)`, `get_read_db`, `ReadSessionLocal`). These run in AUTOCOMMIT mode with no transaction: writes are NOT blocked by the database — they silently commit immediately, and `begin_nested()` savepoints break. There is no runtime guard; this is enforced by convention only. Use `read_only=True` strictly for SELECT-only windows; anything that mutates (including get-or-create paths) must use a regular write session. ### Runtime Architecture diff --git a/config.toml.example b/config.toml.example index dbaf7e3b..3aacfbb7 100644 --- a/config.toml.example +++ b/config.toml.example @@ -98,6 +98,10 @@ POLLING_BACKOFF_MULTIPLIER = 2.0 POLLING_STARTUP_JITTER_SECONDS = 30.0 POLLING_JITTER_RATIO = 0.5 STALE_SESSION_TIMEOUT_MINUTES = 5 +# Minimum (jittered) spacing between stale-work-unit cleanup runs per instance. +# Staleness is a minutes-timescale condition, so cleanup doesn't need to run on +# every seconds-scale poll (0.0 = run every poll, legacy behavior). +STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS = 60.0 # QUEUE_ERROR_RETENTION_SECONDS = 2592000 # 30 days DEDUPLICATE = true LOG_OBSERVATIONS = false diff --git a/src/config.py b/src/config.py index 0fc9b55e..834982f4 100644 --- a/src/config.py +++ b/src/config.py @@ -769,6 +769,10 @@ class DeriverSettings(HonchoSettings): # to 0.0 to disable. POLLING_JITTER_RATIO: Annotated[float, Field(default=0.5, ge=0.0, le=1.0)] = 0.5 STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5 + # Minimum (jittered) spacing between stale-work-unit cleanup runs + STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS: Annotated[ + float, Field(default=60.0, ge=0.0, le=3600.0) + ] = 60.0 # Retention window (seconds) for keeping errored items in the queue QUEUE_ERROR_RETENTION_SECONDS: Annotated[ diff --git a/src/crud/document.py b/src/crud/document.py index fd068d6a..f8560629 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -369,7 +369,7 @@ async def query_documents( max_distance, top_k, ) - async with tracked_db("query_documents.pgvector") as managed_db: + async with tracked_db("query_documents.pgvector", read_only=True) as managed_db: docs = await _query_documents_pgvector( managed_db, workspace_name, @@ -407,7 +407,7 @@ async def query_documents( document_ids=document_ids, filters=filters, ) - async with tracked_db("query_documents.fetch") as managed_db: + async with tracked_db("query_documents.fetch", read_only=True) as managed_db: docs = await fetch_documents_by_ids( db=managed_db, workspace_name=workspace_name, diff --git a/src/crud/message.py b/src/crud/message.py index ec673bb0..42f18f9a 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -817,7 +817,7 @@ async def _semantic_search_messages( # Pre-fetch peer session scope if needed (short-lived DB session) allowed_session_names: list[str] | None = None if observer and not session_name: - async with tracked_db(f"{operation_name}.peer_scope") as db: + async with tracked_db(f"{operation_name}.peer_scope", read_only=True) as db: allowed_session_names = await get_peer_session_names( db, workspace_name, observer ) @@ -837,7 +837,7 @@ async def _semantic_search_messages( if not message_ids: return [] - async with tracked_db(operation_name) as db: + async with tracked_db(operation_name, read_only=True) as db: matched_messages = ( await _fetch_messages_by_ids( db, @@ -853,7 +853,7 @@ async def _semantic_search_messages( _expunge_snippets(db, snippets) return snippets - async with tracked_db(operation_name) as db: + async with tracked_db(operation_name, read_only=True) as db: snippets = await _search_messages_pgvector( db, workspace_name, @@ -985,7 +985,7 @@ async def grep_messages( List of tuples: (matched_messages, context_messages) Each snippet may contain multiple matches if they were close together. """ - async with tracked_db("message.grep_messages") as db: + async with tracked_db("message.grep_messages", read_only=True) as db: # Pre-fetch peer session scope if needed allowed_session_names = None if observer and not session_name: diff --git a/src/crud/representation.py b/src/crud/representation.py index bb969120..4ade79a4 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -276,7 +276,7 @@ class RepresentationManager: ) async with tracked_db( - "representation_manager.get_working_representation" + "representation_manager.get_working_representation", read_only=True ) as new_db: return await self._get_working_representation_internal( new_db, diff --git a/src/db.py b/src/db.py index dc9d7c67..005d52eb 100644 --- a/src/db.py +++ b/src/db.py @@ -66,6 +66,24 @@ SessionLocal = async_sessionmaker( class_=AsyncSession, ) +# Read-only engine: shares `engine`'s pool, but checks connections out in DBAPI +# AUTOCOMMIT mode, so psycopg emits NO BEGIN — a SELECT never autobegins a +# transaction. The backend therefore returns to state 'idle' (not 'idle in +# transaction') the moment a statement completes. +read_engine = engine.execution_options(isolation_level="AUTOCOMMIT") + +# Sessions for SELECT-only work (same lazy-checkout semantics as SessionLocal). +# MUST NOT be used for writes: with no enclosing transaction, begin_nested() +# savepoints (see the crud get-or-create paths) break, and every flush would +# commit immediately. Use SessionLocal for anything that mutates. +ReadSessionLocal = async_sessionmaker( + autocommit=False, + autoflush=False, + expire_on_commit=False, + bind=read_engine, + class_=AsyncSession, +) + def _set_application_name_on_checkout( dbapi_connection: Any, _connection_record: Any, _connection_proxy: Any @@ -76,16 +94,30 @@ def _set_application_name_on_checkout( reused pooled connection is re-tagged for the new caller), reading the per-task ``request_context`` the request/task scope has already set. Best-effort: a failure here must never break the checkout. + + Runs in autocommit so it never leaves the connection 'idle in transaction' + at checkout: this hook fires BEFORE the dialect applies execution-option + isolation levels, and psycopg refuses to switch a connection into AUTOCOMMIT + (which the read engine does) while a transaction opened by this statement is + still in progress. set_config(..., is_local=false) is session-scoped, so it + persists past the autocommit boundary. """ context = request_context.get() or "unknown" try: - cursor = dbapi_connection.cursor() + previous_autocommit = dbapi_connection.autocommit + if not previous_autocommit: + dbapi_connection.autocommit = True try: - cursor.execute( - "SELECT set_config('application_name', %s, false)", (context,) - ) + cursor = dbapi_connection.cursor() + try: + cursor.execute( + "SELECT set_config('application_name', %s, false)", (context,) + ) + finally: + cursor.close() finally: - cursor.close() + if not previous_autocommit: + dbapi_connection.autocommit = False except Exception: logger.debug("setting application_name on checkout failed", exc_info=True) diff --git a/src/dependencies.py b/src/dependencies.py index 31461496..060186b4 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -4,7 +4,7 @@ from contextlib import asynccontextmanager from fastapi import Depends from sqlalchemy.ext.asyncio import AsyncSession -from src.db import SessionLocal, request_context +from src.db import ReadSessionLocal, SessionLocal, request_context async def get_db(): @@ -32,12 +32,39 @@ async def get_db(): await db.close() +async def get_read_db(): + """FastAPI Dependency Generator for SELECT-only handlers. + + Same lazy-checkout semantics as get_db, but the session is bound to the + AUTOCOMMIT read engine: no BEGIN is ever emitted, so the connection can not + sit 'idle in transaction' between the query and this teardown — a delayed + finally here is harmless (the backend is plain 'idle'). close() is still + required to release the connection itself back to the pool. + + MUST only be used by handlers that never mutate; see ReadSessionLocal. + """ + db: AsyncSession = ReadSessionLocal() + try: + yield db + finally: + # rollback is a wire-level no-op under AUTOCOMMIT; kept to reset any + # Python-side session state before close, mirroring get_db. + await db.rollback() + await db.close() + + @asynccontextmanager -async def tracked_db(operation_name: str | None = None): +async def tracked_db(operation_name: str | None = None, *, read_only: bool = False): """Context manager for tracked database sessions. Sets a task-scoped request_context so the lazy session picks it up for tracing/attribution, then yields a lazy session (see get_db). + + Pass read_only=True for SELECT-only windows: the session is then bound to + the AUTOCOMMIT read engine, so the work inside the block never holds an + open transaction (no idle-in-transaction parking; the pooler can reclaim + the backend between statements). Never use read_only=True on a path that + mutates — see ReadSessionLocal. """ # Get request ID if available, or create operation-specific one context = request_context.get() @@ -47,14 +74,15 @@ async def tracked_db(operation_name: str | None = None): context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}" token = request_context.set(context) - db = SessionLocal() + db = (ReadSessionLocal if read_only else SessionLocal)() try: yield db except Exception: await db.rollback() raise finally: - # Always send ROLLBACK unconditionally — see get_db() comment. + # Always send ROLLBACK unconditionally — see get_db() comment. (Under + # read_only/AUTOCOMMIT it is a wire-level no-op.) await db.rollback() await db.close() if token: # Only reset if we set it @@ -62,3 +90,4 @@ async def tracked_db(operation_name: str | None = None): db: AsyncSession = Depends(get_db) +read_db: AsyncSession = Depends(get_read_db) diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 3d131293..e97e67ed 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -2,6 +2,7 @@ import asyncio import contextlib import random import signal +import time from asyncio import Task from collections.abc import Sequence from dataclasses import dataclass, field @@ -134,6 +135,16 @@ class QueueManager: settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS ) + # Monotonic timestamp of the last stale-work-unit cleanup ATTEMPT. + # None -> the first poll always runs cleanup (recovers rows left stale + # by a crashed predecessor immediately). + self._last_stale_cleanup_attempt: float | None = None + # Jittered gate width (seconds) sampled ONCE per attempt, so the deadline + # for the next run is fixed when the timestamp is set rather than + # re-rolled on every poll (which would make the effective spacing a + # random walk and untestable at non-zero jitter ratios). + self._stale_cleanup_gate_seconds: float = 0.0 + # Initialize from settings self.workers: int = settings.DERIVER.WORKERS self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers) @@ -258,6 +269,35 @@ class QueueManager: # Polling and Scheduling # ########################## + async def _maybe_cleanup_stale_work_units(self) -> None: + """Run stale-work-unit cleanup at most once per (jittered) interval. + + Staleness is a minutes-timescale condition (STALE_SESSION_TIMEOUT_MINUTES), + but the polling loop fires on a seconds timescale on every deriver + instance — running cleanup unconditionally per poll multiplies into + unnecessary write transactions. Gate it locally: + concurrent cleaners on other instances remain safe via FOR UPDATE SKIP + LOCKED, so no cross-instance coordination is required, and the jittered + gate (sampled once per attempt) keeps instances from re-synchronizing + their cleanup runs. The gate tracks the last ATTEMPT (set before + running), so a failing cleanup waits a full interval instead of retrying + every poll against a DB that is already struggling. An interval of 0 + preserves run-every-poll behavior. + """ + interval = settings.DERIVER.STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS + if ( + interval > 0.0 + and self._last_stale_cleanup_attempt is not None + and time.monotonic() - self._last_stale_cleanup_attempt + < self._stale_cleanup_gate_seconds + ): + return + # Record the attempt and fix the next deadline before running, so the + # gate width is stable for this cycle and a failing cleanup still waits. + self._last_stale_cleanup_attempt = time.monotonic() + self._stale_cleanup_gate_seconds = self._jitter(interval) + await self.cleanup_stale_work_units() + async def cleanup_stale_work_units(self) -> None: """Clean up stale work units""" async with tracked_db("cleanup_stale_work_units") as db: @@ -457,7 +497,7 @@ class QueueManager: continue try: - await self.cleanup_stale_work_units() + await self._maybe_cleanup_stale_work_units() claimed_work_units = await self.get_and_claim_work_units() if claimed_work_units: self._reset_poll_interval() diff --git a/src/dialectic/chat.py b/src/dialectic/chat.py index e9659a21..ae6ea290 100644 --- a/src/dialectic/chat.py +++ b/src/dialectic/chat.py @@ -40,7 +40,7 @@ async def agentic_chat( The synthesized answer string """ # Short-lived DB session for validation + config - async with tracked_db("dialectic.preflight") as db: + async with tracked_db("dialectic.preflight", read_only=True) as db: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) if observer != observed: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) @@ -101,7 +101,7 @@ async def agentic_chat_stream( Chunks of the response text as they are generated """ # Short-lived DB session for validation + config - async with tracked_db("dialectic.preflight") as db: + async with tracked_db("dialectic.preflight", read_only=True) as db: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) if observer != observed: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 5a5f690b..dbbfed35 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -121,7 +121,7 @@ class DialecticAgent: token_limit=max_tokens, reverse=False, # chronological order ) - async with tracked_db("dialectic.session_history") as db: + async with tracked_db("dialectic.session_history", read_only=True) as db: result = await db.execute(stmt) messages = result.scalars().all() diff --git a/src/routers/conclusions.py b/src/routers/conclusions.py index 20aadb75..3a25a5d7 100644 --- a/src/routers/conclusions.py +++ b/src/routers/conclusions.py @@ -6,7 +6,7 @@ from fastapi_pagination.ext.sqlalchemy import apaginate from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas -from src.dependencies import db +from src.dependencies import db, read_db from src.exceptions import ResourceNotFoundException, ValidationException from src.security import require_auth from src.telemetry.events import EmbeddingCallPurpose @@ -67,7 +67,7 @@ async def list_conclusions( False, description="Whether to reverse the order of results", ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """ List Conclusions using optional filters, ordered by recency unless `reverse` is true. Results are paginated. @@ -97,7 +97,7 @@ async def query_conclusions( ..., description="Semantic search parameters for Conclusions", ), - db: AsyncSession = db, + db: AsyncSession = read_db, ) -> list[schemas.Conclusion]: """ Query Conclusions using semantic search. Use `top_k` to control the number of results returned. diff --git a/src/routers/messages.py b/src/routers/messages.py index 917ca713..b58a3287 100644 --- a/src/routers/messages.py +++ b/src/routers/messages.py @@ -18,7 +18,7 @@ from sqlalchemy.orm.attributes import flag_modified from src import crud, schemas from src.config import settings -from src.dependencies import db +from src.dependencies import db, read_db from src.deriver import enqueue from src.exceptions import FileTooLargeError, ResourceNotFoundException from src.security import require_auth @@ -260,7 +260,7 @@ async def get_messages( reverse: bool | None = Query( False, description="Whether to reverse the order of results" ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all messages for a Session with optional filters. Results are paginated.""" try: @@ -288,7 +288,7 @@ async def get_message( workspace_id: str = Path(...), session_id: str = Path(...), message_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get a single message by ID from a Session.""" honcho_message = await crud.get_message( diff --git a/src/routers/peers.py b/src/routers/peers.py index efa19a78..e1aa45dd 100644 --- a/src/routers/peers.py +++ b/src/routers/peers.py @@ -14,7 +14,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas from src.config import settings -from src.dependencies import db, tracked_db +from src.dependencies import db, read_db, tracked_db from src.dialectic.chat import agentic_chat, agentic_chat_stream from src.embedding_client import embedding_client from src.exceptions import AuthenticationException, ResourceNotFoundException @@ -43,7 +43,7 @@ async def get_peers( None, description="Filtering options for the peers list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Peers for a Workspace, paginated with optional filters.""" filter_param = None @@ -134,7 +134,7 @@ async def get_sessions_for_peer( None, description="Filtering options for the sessions list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Sessions for a Peer, paginated with optional filters.""" filter_param = None @@ -318,7 +318,7 @@ async def get_peer_card( None, description="Optional target peer to retrieve a card for, from the observer's perspective. If not provided, returns the observer's own card", ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get a peer card for a specific peer relationship. @@ -412,7 +412,6 @@ async def get_peer_context( le=100, description="Maximum number of conclusions to include in the representation", ), - db: AsyncSession = db, ): """ Get context for a peer, including their representation and peer card. @@ -459,10 +458,12 @@ async def get_peer_context( parent_category="api", ) - # Get the peer card - peer_card = await crud.get_peer_card( - db, workspace_id, observer=peer_id, observed=observed - ) + async with tracked_db( + "peers.get_peer_context.peer_card", read_only=True + ) as card_db: + peer_card = await crud.get_peer_card( + card_db, workspace_id, observer=peer_id, observed=observed + ) response = schemas.PeerContext( peer_id=peer_id, diff --git a/src/routers/sessions.py b/src/routers/sessions.py index 98a8714a..eb8794d7 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -12,7 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import config, crud, schemas from src.cache.client import safe_cache_delete from src.crud.session import session_cache_key -from src.dependencies import db +from src.dependencies import db, read_db from src.deriver.enqueue import enqueue_deletion from src.embedding_client import embedding_client from src.exceptions import ( @@ -251,7 +251,7 @@ async def get_sessions( None, description="Filtering and pagination options for the sessions list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Sessions for a Workspace, paginated with optional filters.""" filter_param = None @@ -544,7 +544,7 @@ async def get_peer_config( workspace_id: str = Path(...), session_id: str = Path(...), peer_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get the configuration for a Peer in a Session.""" return await crud.get_peer_config( @@ -599,7 +599,7 @@ async def set_peer_config( async def get_session_peers( workspace_id: str = Path(...), session_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Peers in a Session. Results are paginated.""" try: @@ -622,7 +622,7 @@ async def get_session_peers( async def get_session_context( workspace_id: str = Path(...), session_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, tokens: int | None = Query( None, le=config.settings.GET_CONTEXT_MAX_TOKENS, @@ -814,7 +814,7 @@ async def get_session_context( async def get_session_summaries( workspace_id: str = Path(...), session_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ) -> schemas.SessionSummaries: """ Get available summaries for a Session. diff --git a/src/routers/workspaces.py b/src/routers/workspaces.py index a2298480..8b19fdfd 100644 --- a/src/routers/workspaces.py +++ b/src/routers/workspaces.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas from src.config import settings -from src.dependencies import db +from src.dependencies import db, read_db from src.deriver.enqueue import enqueue_deletion, enqueue_dream from src.exceptions import AuthenticationException from src.security import JWTParams, require_auth @@ -67,7 +67,7 @@ async def get_all_workspaces( None, description="Filtering and pagination options for the workspaces list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Workspaces, paginated with optional filters.""" filter_param = None @@ -169,7 +169,7 @@ async def get_queue_status( session_id: str | None = Query( None, description="Optional session ID to filter by" ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """ Get the processing queue status for a Workspace, optionally scoped to an observer, sender, diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index de5b2b09..dae38577 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -1607,7 +1607,7 @@ async def _handle_get_recent_history( ) -> "str | ToolResult": """Handle get_recent_history tool.""" _ = tool_input - async with tracked_db("tool.get_recent_history") as db: + async with tracked_db("tool.get_recent_history", read_only=True) as db: history: list[models.Message] = await get_recent_history( db, workspace_name=ctx.workspace_name, @@ -1723,7 +1723,7 @@ async def _handle_get_observation_context( ctx: ToolContext, tool_input: dict[str, Any] ) -> "str | ToolResult": """Handle get_observation_context tool.""" - async with tracked_db("tool.get_observation_context") as db: + async with tracked_db("tool.get_observation_context", read_only=True) as db: messages = await get_observation_context( db, workspace_name=ctx.workspace_name, @@ -1862,7 +1862,7 @@ async def _handle_get_messages_by_date_range( if isinstance(before_date, str): return before_date # Error message - async with tracked_db("tool.get_messages_by_date_range") as db: + async with tracked_db("tool.get_messages_by_date_range", read_only=True) as db: messages = await crud.get_messages_by_date_range( db, workspace_name=ctx.workspace_name, @@ -1980,7 +1980,7 @@ async def _handle_get_recent_observations( ) -> str: """Handle get_recent_observations tool.""" session_only = tool_input.get("session_only", False) - async with tracked_db("tool.get_recent_observations") as db: + async with tracked_db("tool.get_recent_observations", read_only=True) as db: documents = await crud.query_documents_recent( db=db, workspace_name=ctx.workspace_name, @@ -2006,7 +2006,7 @@ async def _handle_get_most_derived_observations( ctx: ToolContext, tool_input: dict[str, Any] ) -> str: """Handle get_most_derived_observations tool.""" - async with tracked_db("tool.get_most_derived_observations") as db: + async with tracked_db("tool.get_most_derived_observations", read_only=True) as db: documents = await crud.query_documents_most_derived( db=db, workspace_name=ctx.workspace_name, @@ -2038,7 +2038,7 @@ async def _handle_get_session_summary( if summary_type == "long" else summarizer.SummaryType.SHORT ) - async with tracked_db("tool.get_session_summary") as db: + async with tracked_db("tool.get_session_summary", read_only=True) as db: summary = await summarizer.get_summary( db, ctx.workspace_name, ctx.session_name, st ) @@ -2050,7 +2050,7 @@ async def _handle_get_session_summary( async def _handle_get_peer_card(ctx: ToolContext, tool_input: dict[str, Any]) -> str: """Handle get_peer_card tool.""" _ = tool_input - async with tracked_db("tool.get_peer_card") as db: + async with tracked_db("tool.get_peer_card", read_only=True) as db: peer_card = await crud.get_peer_card( db, workspace_name=ctx.workspace_name, @@ -2209,7 +2209,7 @@ async def _handle_get_reasoning_chain( return f"ERROR: Invalid direction '{direction}'. Must be 'premises', 'conclusions', or 'both'" # Get the observation itself - async with tracked_db("tool.get_reasoning_chain") as db: + async with tracked_db("tool.get_reasoning_chain", read_only=True) as db: docs = await crud.get_documents_by_ids(db, ctx.workspace_name, [observation_id]) if not docs or not docs[0]: return f"ERROR: Observation '{observation_id}' not found" diff --git a/src/utils/search.py b/src/utils/search.py index 15721933..761b63e0 100644 --- a/src/utils/search.py +++ b/src/utils/search.py @@ -448,7 +448,7 @@ async def search( return search_results[0][:limit] return [] - async with tracked_db("search.messages") as managed_db: + async with tracked_db("search.messages", read_only=True) as managed_db: combined_results = await _run_search(managed_db) for message in combined_results: managed_db.expunge(message) diff --git a/tests/conftest.py b/tests/conftest.py index 90fc269b..df2b194d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,7 +32,7 @@ from src import models from src.cache.client import cache from src.config import settings from src.db import Base -from src.dependencies import get_db +from src.dependencies import get_db, get_read_db from src.exceptions import HonchoException from src.main import app from src.models import Peer, Workspace @@ -340,6 +340,10 @@ async def client( yield db_session app.dependency_overrides[get_db] = override_get_db + # Read-only routes use get_read_db (AUTOCOMMIT engine) in production; in + # tests they must see the same per-test database/session as writes, both + # for isolation and so data written by a test is visible to its reads. + app.dependency_overrides[get_read_db] = override_get_db # No-op the startup embedding-schema validator inside the lifespan. The # global `engine` it would inspect points to a DB that isn't migrated in @@ -792,7 +796,10 @@ def mock_tracked_db(request: pytest.FixtureRequest): session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) @asynccontextmanager - async def mock_tracked_db_context(_: str | None = None): + async def mock_tracked_db_context(_: str | None = None, *, read_only: bool = False): + # read_only is accepted (and ignored): in tests both engines resolve to + # the same per-test database session. + del read_only async with session_factory() as session: yield session diff --git a/tests/integration/test_message_embeddings.py b/tests/integration/test_message_embeddings.py index 198f2d36..83ae2ae2 100644 --- a/tests/integration/test_message_embeddings.py +++ b/tests/integration/test_message_embeddings.py @@ -474,7 +474,10 @@ async def test_search_messages_external_lookup_happens_before_tracked_db( return [([message], [message])] @asynccontextmanager - async def fake_tracked_db(_operation_name: str | None = None): + async def fake_tracked_db( + _operation_name: str | None = None, *, read_only: bool = False + ): + del read_only call_order.append("enter") yield fake_db call_order.append("exit") @@ -578,7 +581,10 @@ async def test_search_messages_temporal_external_lookup_happens_before_tracked_d return [([message], [message])] @asynccontextmanager - async def fake_tracked_db(_operation_name: str | None = None): + async def fake_tracked_db( + _operation_name: str | None = None, *, read_only: bool = False + ): + del read_only call_order.append("enter") yield fake_db call_order.append("exit") diff --git a/tests/sdk_typescript/conftest.py b/tests/sdk_typescript/conftest.py index 74fe148b..15505a0e 100644 --- a/tests/sdk_typescript/conftest.py +++ b/tests/sdk_typescript/conftest.py @@ -20,7 +20,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker from uvicorn.config import Config from uvicorn.server import Server -from src.dependencies import get_db +from src.dependencies import get_db, get_read_db from src.main import app @@ -95,6 +95,9 @@ def ts_test_server( yield session app.dependency_overrides[get_db] = override_get_db + # Read-only routes use get_read_db (AUTOCOMMIT engine) in production; in + # tests they must resolve to the same per-test database. + app.dependency_overrides[get_read_db] = override_get_db # No-op the lifespan's startup embedding-schema validator — same # reasoning as the `client` fixture in tests/conftest.py: the module- @@ -133,7 +136,9 @@ def mock_tracked_db(ts_db_session: async_sessionmaker[AsyncSession]): # Create a tracked_db that uses fresh sessions (not shared) @asynccontextmanager - async def ts_tracked_db(_: str | None = None): + async def ts_tracked_db(_: str | None = None, *, read_only: bool = False): + # read_only accepted (and ignored): tests use one per-test database. + del read_only async with ts_db_session() as session: yield session diff --git a/tests/startup/test_embedding_validator.py b/tests/startup/test_embedding_validator.py index 962e07f0..1d257ba5 100644 --- a/tests/startup/test_embedding_validator.py +++ b/tests/startup/test_embedding_validator.py @@ -14,6 +14,7 @@ from sqlalchemy import text from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import AsyncEngine +from src.config import settings from src.startup.embedding_validator import ( StartupValidationError, _assert_pgvector_dims_match, # pyright: ignore[reportPrivateUsage] @@ -120,18 +121,24 @@ async def test_validator_fails_closed_when_introspection_keeps_failing( @pytest.mark.asyncio async def test_validator_passes_against_test_database( db_engine: AsyncEngine, + monkeypatch: pytest.MonkeyPatch, ) -> None: """The test DB is provisioned at the default dim (1536); the validator should accept it without raising.""" + # conftest provisions the test tables in `public`; pin the validator to it + # so a developer's local .env DB_SCHEMA can't point it at another schema. + monkeypatch.setattr(settings.DB, "SCHEMA", "public") await validate_embedding_schema(db_engine) @pytest.mark.asyncio async def test_validator_raises_when_schema_dim_diverges_from_settings( db_engine: AsyncEngine, + monkeypatch: pytest.MonkeyPatch, ) -> None: """ALTER one of the embedding columns to a non-1536 dim and confirm the validator raises with an actionable message.""" + monkeypatch.setattr(settings.DB, "SCHEMA", "public") # see test above async with db_engine.begin() as conn: await conn.execute( text( @@ -181,8 +188,13 @@ def test_non_1536_pgvector_without_migrated_no_longer_raises_at_config_time() -> """The dim-vs-MIGRATED guard has been removed. Constructing AppSettings with non-1536 + default pgvector + MIGRATED=false should now succeed (the runtime schema validator at startup is the safety net).""" + # Minimal env, NOT a copy of os.environ: load_dotenv() in the app mutates + # the parent pytest process's environ, so inheriting it would leak a + # developer's local .env (DB_SCHEMA, VECTOR_STORE_MIGRATED, ...) into the + # child despite PYTHON_DOTENV_DISABLED. The child must see pure defaults + # plus exactly the overrides below. env = { - **os.environ, + "PATH": os.environ.get("PATH", ""), "PYTHON_DOTENV_DISABLED": "1", "EMBEDDING_VECTOR_DIMENSIONS": "768", } diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py index e4bab0b5..faffa89a 100644 --- a/tests/test_db_resilience.py +++ b/tests/test_db_resilience.py @@ -50,6 +50,11 @@ class _FakeCursor: class _FakeDBAPIConn: def __init__(self, recorder: list[Any], raise_exc: Exception | None = None) -> None: self._cursor: _FakeCursor = _FakeCursor(recorder, raise_exc) + # Real pooled connections are checked out in non-autocommit mode; the + # hook flips this to True for its statement then restores it so it never + # leaves an open transaction that would block the read engine's + # AUTOCOMMIT switch. + self.autocommit: bool = False def cursor(self) -> _FakeCursor: return self._cursor @@ -69,6 +74,8 @@ def test_checkout_hook_sets_application_name_from_request_context() -> None: assert "set_config" in sql and "application_name" in sql assert params == ("request:trace-ctx",) assert conn._cursor.closed is True # pyright: ignore[reportPrivateUsage] + # The hook restored the original (non-autocommit) mode after its statement. + assert conn.autocommit is False def test_checkout_hook_defaults_to_unknown_without_context() -> None: @@ -198,3 +205,106 @@ def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None: # gauge negative. tracker.on_error(SimpleNamespace(connection=SimpleNamespace(info={}))) assert value() == start + + +@pytest.mark.asyncio +async def test_stale_cleanup_time_gate(monkeypatch: pytest.MonkeyPatch) -> None: + """cleanup_stale_work_units runs at most once per gate interval per + instance (staleness is a minutes-timescale condition; per-poll cleanup + multiplies into needless fleet-wide write transactions). First poll always + runs it so a crashed predecessor's stale rows are recovered immediately.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + monkeypatch.setattr( + settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0 + ) + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + runs = {"n": 0} + + async def fake_cleanup() -> None: + runs["n"] += 1 + + monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) + + clock = {"now": 1_000.0} + monkeypatch.setattr( + "src.deriver.queue_manager.time.monotonic", lambda: clock["now"] + ) + + # First call runs (no prior attempt recorded). + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 1 + + # Inside the gate window: skipped. + clock["now"] += 10.0 + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 1 + + # Past the gate window: runs again. + clock["now"] += 60.0 + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 2 + + +@pytest.mark.asyncio +async def test_stale_cleanup_gate_failed_attempt_waits_full_interval( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The gate records the ATTEMPT before running, so a failing cleanup is not + retried on every poll against a DB that is already struggling.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + monkeypatch.setattr( + settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0 + ) + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + attempts = {"n": 0} + + async def failing_cleanup() -> None: + attempts["n"] += 1 + raise RuntimeError("db unavailable") + + monkeypatch.setattr(qm, "cleanup_stale_work_units", failing_cleanup) + + clock = {"now": 1_000.0} + monkeypatch.setattr( + "src.deriver.queue_manager.time.monotonic", lambda: clock["now"] + ) + + with pytest.raises(RuntimeError): + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert attempts["n"] == 1 + + # Immediately after the failure: still gated, no hammering. + clock["now"] += 1.0 + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert attempts["n"] == 1 + + +@pytest.mark.asyncio +async def test_stale_cleanup_gate_zero_interval_runs_every_poll( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Interval 0.0 preserves legacy run-on-every-poll behavior.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + monkeypatch.setattr( + settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 0.0 + ) + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + runs = {"n": 0} + + async def fake_cleanup() -> None: + runs["n"] += 1 + + monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) + + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 2 diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index d5ba3ea1..1b5c219b 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -150,3 +150,161 @@ async def test_tracked_db_rolls_back_open_transaction_on_exit( assert fake_db.rollback_calls == 1 assert fake_db.close_calls == 1 + + +@pytest.mark.asyncio +async def test_tracked_db_read_only_uses_read_sessionmaker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # read_only=True must construct the session from ReadSessionLocal (the + # AUTOCOMMIT engine) — and never touch SessionLocal — while keeping the + # same rollback/close teardown. + read_fake = FakeSession() + monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake) + monkeypatch.setattr( + dependencies_module, + "SessionLocal", + lambda: pytest.fail("read_only window constructed a write session"), + ) + + async with real_tracked_db("read_op", read_only=True) as db: + assert db is read_fake + + assert read_fake.rollback_calls == 1 + assert read_fake.close_calls == 1 + + +@pytest.mark.asyncio +async def test_get_read_db_rolls_back_and_closes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + read_fake = FakeSession() + monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake) + + dep_gen = dependencies_module.get_read_db() + try: + db = await anext(dep_gen) + assert db is read_fake + assert read_fake.connection_calls == 0 # still lazy, no eager checkout + finally: + await dep_gen.aclose() + + assert read_fake.rollback_calls == 1 + assert read_fake.close_calls == 1 + + +def test_read_engine_is_autocommit_and_shares_pool() -> None: + # The read engine must differ from the write engine ONLY by isolation + # level: AUTOCOMMIT (so reads never autobegin a transaction) on the same + # underlying pool (no second connection budget). + from src.db import engine, read_engine + + assert ( + read_engine.sync_engine._execution_options.get( # pyright: ignore[reportPrivateUsage] + "isolation_level" + ) + == "AUTOCOMMIT" + ) + assert read_engine.sync_engine.pool is engine.sync_engine.pool + + +@pytest.mark.asyncio +async def test_read_only_session_runs_in_autocommit_on_the_wire() -> None: + # Wire-level guarantee behind the whole read-path fix: a read_only session's + # connection has the DBAPI autocommit flag set, so psycopg emits no BEGIN + # and the backend sits in state 'idle' (not 'idle in transaction') after a + # statement returns. NOTE: get_isolation_level() can NOT verify this — it + # reports the server's transaction_isolation GUC (READ COMMITTED), because + # autocommit is a driver behavior, not a server isolation level. + from sqlalchemy import text + + from src.db import read_engine + + async with real_tracked_db("read_op", read_only=True) as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + conn = await db.connection() + raw = (await conn.get_raw_connection()).driver_connection + assert raw is not None + assert raw.autocommit is True + + # Definitive check, from a second connection: after the SELECT above, + # the session's backend must be plain 'idle' — an open transaction + # would report 'idle in transaction' and be reapable in production. + async with read_engine.connect() as observer: + state = ( + await observer.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + assert state == "idle" + + +@pytest.mark.asyncio +async def test_write_session_holds_idle_in_transaction_after_select() -> None: + # Contrast guard documenting WHY the read engine exists: the default + # (transactional) session autobegins on the first statement and leaves the + # backend 'idle in transaction' until rollback/close — the state that + # Postgres's idle_in_transaction_session_timeout reaps and that pins a + # transaction-mode pooler backend. + from sqlalchemy import text + + from src.db import read_engine + + async with real_tracked_db("write_op") as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + async with read_engine.connect() as observer: + state = ( + await observer.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + assert state == "idle in transaction" + + +@pytest.mark.asyncio +async def test_read_only_session_works_with_tracing_checkout_hook() -> None: + # Regression: the DB.TRACING checkout hook runs set_config() at pool + # checkout, BEFORE the dialect applies the read engine's AUTOCOMMIT + # isolation level. If that statement is allowed to autobegin a transaction, + # psycopg then refuses to switch the connection into AUTOCOMMIT + # ("can't change 'autocommit' now: connection in transaction") and every + # read_only session 500s under TRACING. The hook must run in autocommit so + # it leaves the connection idle. This combination is otherwise untested + # because DB.TRACING defaults to false. + from sqlalchemy import event, text + + from src.db import ( + _set_application_name_on_checkout, # pyright: ignore[reportPrivateUsage] + engine, + read_engine, + ) + + context_token = request_context.set("tracing-regression") + event.listen(engine.sync_engine, "checkout", _set_application_name_on_checkout) + try: + async with real_tracked_db("read_op", read_only=True) as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + app_name = (await db.execute(text("SHOW application_name"))).scalar() + conn = await db.connection() + raw = (await conn.get_raw_connection()).driver_connection + assert raw is not None + # AUTOCOMMIT was applied despite the checkout hook running first. + assert raw.autocommit is True + # The hook still tagged the connection (set_config is session-scoped, + # so it survives the autocommit boundary). + assert app_name == "tracing-regression" + # Backend is idle, not idle-in-transaction: the no-BEGIN guarantee + # holds even with the hook firing. + async with read_engine.connect() as observer: + state = ( + await observer.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + assert state == "idle" + finally: + event.remove(engine.sync_engine, "checkout", _set_application_name_on_checkout) + request_context.reset(context_token)