add read db (#773)

* feat: implement read DB and fix queue stale cleanup

* fix: use read_db in internal methods

* fix: mention read db in the CLAUDE.md

* fix: make TRACING checkout hook autocommit-safe; sample cleanup-gate jitter once

The DB.TRACING checkout hook ran `SELECT set_config(...)` at pool checkout,
before the dialect applies the read engine's AUTOCOMMIT isolation level. That
statement autobegins a transaction, and psycopg then refuses to switch the
connection into AUTOCOMMIT ("can't change 'autocommit' now: connection in
transaction status INTRANS"), so every read_only session 500s under TRACING and
the INTRANS connection leaks back to poison later write checkouts. Run the hook
in autocommit and restore the prior mode so it never leaves an open transaction;
set_config(..., is_local=false) is session-scoped and survives the boundary.
Add a regression test (fails without the fix) covering read_only + TRACING.

Also sample the stale-cleanup gate's jittered interval once per attempt instead
of re-rolling it every poll, so the spacing is a fixed deadline per cycle rather
than a random walk (and is testable at non-zero jitter ratios).

* fix: reset request_context in TRACING checkout-hook test

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja 2026-06-10 13:28:36 -04:00 committed by GitHub
parent f75b336a3c
commit bf494257b8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 469 additions and 60 deletions

View File

@ -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

View File

@ -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

View File

@ -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[

View File

@ -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,

View File

@ -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:

View File

@ -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,

View File

@ -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)

View File

@ -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)

View File

@ -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()

View File

@ -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))

View File

@ -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()

View File

@ -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.

View File

@ -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(

View File

@ -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,

View File

@ -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.

View File

@ -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,

View File

@ -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"

View File

@ -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)

View File

@ -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

View File

@ -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")

View File

@ -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

View File

@ -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",
}

View File

@ -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

View File

@ -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)