diff --git a/.env.example b/.env.example index 0631050d41..412cf4ecc8 100644 --- a/.env.example +++ b/.env.example @@ -112,6 +112,12 @@ HINDSIGHT_API_LOG_LEVEL=info # 'failed' (not 'completed'), surfacing silently-dropped facts. Default false. # HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS=false +# Wall-clock ceiling (seconds) for one retain task in the worker. A retain that +# blocks indefinitely is cancelled and marked 'failed' — and so becomes +# retryable — instead of holding its worker slot until the process restarts. +# Set well above your slowest healthy retain; 0 disables. Default 3600. +# HINDSIGHT_API_RETAIN_WALL_TIMEOUT=3600 + # Dry-run extraction preview endpoint (POST /memories/dry-run-extract). Enabled by default; it makes # a real LLM call but stores nothing. Set to false to remove the endpoint (returns 404). # HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=true diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index 45a4057770..915479c891 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -690,6 +690,7 @@ def _parse_boolean_env(env_name: str, default: bool) -> bool: } ENV_WORKER_CONSOLIDATION_BANK_PRIORITY = "HINDSIGHT_API_WORKER_CONSOLIDATION_BANK_PRIORITY" ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT" +ENV_RETAIN_WALL_TIMEOUT = "HINDSIGHT_API_RETAIN_WALL_TIMEOUT" # Reflect agent settings ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS" @@ -1168,6 +1169,14 @@ def _parse_strategy_boosts(raw: str | None) -> dict[str, str]: DEFAULT_OPERATION_RETENTION_DAYS = 0 DEFAULT_OPERATION_CLEANUP_BATCH_SIZE = 1000 DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention. +# Wall-clock ceiling for one retain task in the worker (0 disables). This is a +# deadlock/wedge backstop, not a latency target: a retain that blocks forever on +# a lock, an LLM permit or a queue put would otherwise hold its worker slot until +# the process restarts, and 'processing' is neither retryable nor cancellable +# through the API. Set well above any healthy retain so it only ever fires on a +# genuine wedge — the per-attempt LLM timeout and the retry budget already bound +# the normal slow path. +DEFAULT_RETAIN_WALL_TIMEOUT = 3600 # seconds (1 hour) # Reflect agent settings DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response @@ -2104,6 +2113,7 @@ class HindsightConfig: operation_retention_days: int operation_cleanup_batch_size: int retain_max_concurrent: int + retain_wall_timeout: int # Reflect agent settings reflect_max_iterations: int @@ -3215,6 +3225,7 @@ def from_env(cls) -> "HindsightConfig": DEFAULT_OPERATION_CLEANUP_BATCH_SIZE, ), retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))), + retain_wall_timeout=int(os.getenv(ENV_RETAIN_WALL_TIMEOUT, str(DEFAULT_RETAIN_WALL_TIMEOUT))), # Reflect agent settings reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))), reflect_prompt_cache_enabled=os.getenv( diff --git a/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py b/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py index e6cdb6908c..9b018fd735 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py +++ b/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py @@ -253,11 +253,20 @@ async def bulk_insert_entities( entity_names: list[str], entity_dates: list, ) -> dict[str, str]: + # ORDER BY LOWER(name) so every concurrent batch inserts in the same order + # as the conflict target (bank_id, LOWER(canonical_name)). ON CONFLICT DO + # NOTHING takes a ShareLock on the inserting transaction of any speculative + # row it collides with, so two batches with overlapping names inserting in + # different orders deadlock. The caller already sorts by Python's + # ``str.lower()``, which agrees with the index for ASCII but not for every + # locale (see the Turkish-İ note in entity_resolver) — ordering in SQL makes + # the database's own collation the single arbiter for all writers. inserted_rows = await conn.fetch( f""" INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count) SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0 FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date) + ORDER BY LOWER(name) ON CONFLICT (bank_id, LOWER(canonical_name)) DO NOTHING RETURNING id, LOWER(canonical_name) AS name_lower diff --git a/hindsight-api-slim/hindsight_api/engine/db/postgresql.py b/hindsight-api-slim/hindsight_api/engine/db/postgresql.py index 701e86bf4e..552fa22687 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/postgresql.py +++ b/hindsight-api-slim/hindsight_api/engine/db/postgresql.py @@ -78,6 +78,7 @@ def run_migrations(self, dsn: str, *, schema: str | None = None) -> None: def __init__(self) -> None: self._pool: asyncpg.Pool | None = None self._acquire_warn_threshold_s: float = 1.0 + self._acquire_timeout_s: float | None = None async def initialize( self, @@ -93,6 +94,13 @@ async def initialize( from ...config import get_config self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0 + # Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a + # *connect* kwarg (how long establishing a new connection may take), and + # ``Pool.acquire()`` defaults to waiting for a free connection forever. + # Passing it here alone made HINDSIGHT_API_DB_ACQUIRE_TIMEOUT a no-op for + # the wait it names: a pool-exhaustion stall never surfaced as an error, + # it just hung (#3002). 0 restores the unbounded behaviour. + self._acquire_timeout_s = acquire_timeout if acquire_timeout > 0 else None self._pool = await asyncpg.create_pool( dsn, min_size=min_size, @@ -138,7 +146,9 @@ def _pool_stats(self) -> PoolStats | None: async def acquire(self) -> AsyncIterator[PostgresConnection]: pool = self._ensure_pool() async with instrument_acquire( - pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s + pool.acquire(timeout=self._acquire_timeout_s), + pool_stats=self._pool_stats, + warn_threshold_s=self._acquire_warn_threshold_s, ) as conn: yield PostgresConnection(conn) @@ -146,7 +156,9 @@ async def acquire(self) -> AsyncIterator[PostgresConnection]: async def transaction(self) -> AsyncIterator[PostgresConnection]: pool = self._ensure_pool() async with instrument_acquire( - pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s + pool.acquire(timeout=self._acquire_timeout_s), + pool_stats=self._pool_stats, + warn_threshold_s=self._acquire_warn_threshold_s, ) as conn: async with conn.transaction(): yield PostgresConnection(conn) diff --git a/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py b/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py index 0f22fb5e18..23e1302109 100644 --- a/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py @@ -888,7 +888,13 @@ async def call( from ..worker.stage import set_stage structured = "+structured" if response_format is not None else "" - set_stage(f"llm.{self.provider}.{scope}{structured}") + # `.queued` until the concurrency permits are in hand — see the acquire + # below. Without it, a call waiting on a saturated semaphore is + # indistinguishable from one the provider is actively running, and the + # label points at the provider (#3002: an operator lost an hour to + # "llm.bedrock.*" for tasks that had never reached Bedrock). + base_stage = f"llm.{self.provider}.{scope}{structured}" + set_stage(f"{base_stage}.queued") # Resolve the retry policy: explicit per-call arg wins, else the provider's # configured per-operation/global default, else this method's own fallback. @@ -948,6 +954,7 @@ async def call( async with AsyncExitStack() as stack: for sem in _semaphores_for_scope(scope): await stack.enter_async_context(sem) + set_stage(base_stage) # cached_prefix is only set for providers that returned a handle # from get_or_create_cached_prefix() (e.g. Gemini); it's None for @@ -1039,7 +1046,9 @@ async def call_with_tools( """ from ..worker.stage import set_stage - set_stage(f"llm.{self.provider}.{scope}+tools") + # `.queued` until the permits are held — see the structured path above. + base_stage = f"llm.{self.provider}.{scope}+tools" + set_stage(f"{base_stage}.queued") # Resolve the retry policy: explicit per-call arg wins, else the provider's # configured per-operation/global default, else this method's own fallback. @@ -1081,6 +1090,7 @@ async def call_with_tools( async with AsyncExitStack() as stack: for sem in _semaphores_for_scope(scope): await stack.enter_async_context(sem) + set_stage(base_stage) # cached_prefix is only set for providers that returned a handle # from get_or_create_cached_prefix() / create_incremental_cache(); diff --git a/hindsight-api-slim/hindsight_api/engine/providers/gemini_llm.py b/hindsight-api-slim/hindsight_api/engine/providers/gemini_llm.py index ba205eeb50..6c91f2fa0f 100644 --- a/hindsight-api-slim/hindsight_api/engine/providers/gemini_llm.py +++ b/hindsight-api-slim/hindsight_api/engine/providers/gemini_llm.py @@ -424,8 +424,7 @@ def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentCon last_exception = None for attempt in range(max_retries + 1): - if attempt > 0: - set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}") try: response = await asyncio.wait_for( self._client.aio.models.generate_content( @@ -769,8 +768,7 @@ def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig": last_exception = None for attempt in range(max_retries + 1): - if attempt > 0: - set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}") try: # With the cache active, send only the un-cached tail (delta); # on the uncached fallback path send the full conversation so the diff --git a/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py b/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py index db57668314..ebc06c9f14 100644 --- a/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py +++ b/hindsight-api-slim/hindsight_api/engine/providers/litellm_llm.py @@ -255,8 +255,7 @@ async def call( last_exception = None for attempt in range(max_retries + 1): - if attempt > 0: - set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}") try: response = await asyncio.wait_for( self._acompletion(**call_kwargs), @@ -447,8 +446,7 @@ async def call_with_tools( last_exception = None for attempt in range(max_retries + 1): - if attempt > 0: - set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}") try: response = await asyncio.wait_for( self._acompletion(**call_kwargs), diff --git a/hindsight-api-slim/hindsight_api/engine/providers/openai_compatible_llm.py b/hindsight-api-slim/hindsight_api/engine/providers/openai_compatible_llm.py index 1b6d1c65e9..297fc035c3 100644 --- a/hindsight-api-slim/hindsight_api/engine/providers/openai_compatible_llm.py +++ b/hindsight-api-slim/hindsight_api/engine/providers/openai_compatible_llm.py @@ -899,8 +899,7 @@ async def call( # Surface attempt count in worker stage so JSON-schema retry loops # are visible from logs (small models on strict structured output # often loop here). Cheap no-op outside worker context. - if attempt > 0: - set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}") try: if response_format is not None: response = await self._client.chat.completions.create(**call_params) @@ -1269,8 +1268,7 @@ async def call_with_tools( last_exception = None for attempt in range(max_retries + 1): - if attempt > 0: - set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}") try: response = await self._client.chat.completions.create(**call_params) @@ -1474,8 +1472,7 @@ async def _call_ollama_native( async with httpx.AsyncClient(timeout=300.0) as client: for attempt in range(max_retries + 1): - if attempt > 0: - set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}") + set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}") try: response = await client.post(native_url, json=payload, headers=headers) response.raise_for_status() diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index 00661485ec..f3bc1e3ce8 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -1412,26 +1412,38 @@ async def _extract_one(global_idx: int, chunk_text: str) -> None: tasks: list[asyncio.Task] = [] skipped_total = 0 - for i, chunk_text in enumerate(all_pre_chunks): - chunk_hash = chunk_storage.compute_chunk_hash(chunk_text) - if chunk_hash in existing_chunk_hashes: - # Memory: skipped chunks aren't needed either. - all_pre_chunks[i] = "" - skipped_total += 1 - continue - tasks.append(asyncio.create_task(_extract_one(i, chunk_text))) - - if skipped_total > 0: - log_buffer.append(f"[streaming] Producer: skipped {skipped_total}/{total_chunks} already-committed chunks") + try: + for i, chunk_text in enumerate(all_pre_chunks): + chunk_hash = chunk_storage.compute_chunk_hash(chunk_text) + if chunk_hash in existing_chunk_hashes: + # Memory: skipped chunks aren't needed either. + all_pre_chunks[i] = "" + skipped_total += 1 + continue + tasks.append(asyncio.create_task(_extract_one(i, chunk_text))) - # Wait for all extractions; collect exceptions - results = await asyncio.gather(*tasks, return_exceptions=True) - for r in results: - if isinstance(r, BaseException): - producer_error.append(r) + if skipped_total > 0: + log_buffer.append( + f"[streaming] Producer: skipped {skipped_total}/{total_chunks} already-committed chunks" + ) - # Signal the consumer that production is done - await chunk_queue.put(None) + # Wait for all extractions; collect exceptions + results = await asyncio.gather(*tasks, return_exceptions=True) + for r in results: + if isinstance(r, BaseException): + producer_error.append(r) + + # Signal the consumer that production is done + await chunk_queue.put(None) + finally: + # Cancellation arriving mid-fan-out (the consumer failed, or the worker's + # wall-clock ceiling fired) must not strand extraction tasks. Cancelling + # the gather above already propagates to them, but tasks created before + # we reach it would otherwise survive and park on `chunk_queue.put()` + # for the life of the process. + for extraction in tasks: + if not extraction.done(): + extraction.cancel() # ---- DB Consumer ---- # Drains enriched chunks from the queue in batches and runs @@ -1869,8 +1881,27 @@ async def _run_mini_batch_db_work() -> None: logger.warning("Failed to check operation recovery state", exc_info=True) if not facts_already_committed: - # Run producer and consumer concurrently - await asyncio.gather(_llm_producer(), _db_consumer()) + # Run producer and consumer concurrently. + # + # Cancellation is explicit because plain gather() leaks: when the consumer + # raises (a deadlock victim, a lock timeout) gather propagates that error + # immediately but leaves the producer — and every extraction task under it + # — running. Those tasks then block forever on `chunk_queue.put()` into a + # queue nobody drains, pinning their chunk payloads and still spending LLM + # permits and tokens on an operation that already failed (#3002). The same + # applies when the worker's wall-clock ceiling cancels us from above. + producer_task = asyncio.create_task(_llm_producer()) + consumer_task = asyncio.create_task(_db_consumer()) + try: + await asyncio.gather(producer_task, consumer_task) + finally: + for pipeline_task in (producer_task, consumer_task): + if not pipeline_task.done(): + pipeline_task.cancel() + # Await the cancellations so neither half outlives this call; the + # results are already accounted for by the gather above (or by the + # exception that is propagating). + await asyncio.gather(producer_task, consumer_task, return_exceptions=True) # Propagate producer errors (e.g. LLM failures) if producer_error: diff --git a/hindsight-api-slim/hindsight_api/worker/poller.py b/hindsight-api-slim/hindsight_api/worker/poller.py index 9829aa9dc8..a97ce0bebd 100644 --- a/hindsight-api-slim/hindsight_api/worker/poller.py +++ b/hindsight-api-slim/hindsight_api/worker/poller.py @@ -19,6 +19,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any +from ..config import get_config from ..engine.schema import fq_table_explicit as fq_table from ..metrics import get_metrics_collector from .exceptions import DeferOperation, RetryTaskAt @@ -37,6 +38,34 @@ def _metric_operation_label(operation_type: str | None) -> str: return operation_type or "unknown" +def _wall_timeout_for(task_type: str) -> float | None: + """Wall-clock ceiling for one task of this type, or None when unbounded. + + A task that wedges (a lock wait with no deadlock cycle to break it, an LLM + permit that never frees, a producer blocked on a queue nobody drains) holds + its worker slot forever: the operation stays 'processing', which the API + refuses to retry *or* cancel, and once every slot is held the worker stops + claiming work entirely (#3002). Per-operation timeouts elsewhere bound one + LLM call or one query, never the whole task — this is the outer backstop + that turns "wedged until restart" into "failed and retryable". + + Only retain is bounded today; reflect self-bounds inside the engine + (``reflect_wall_timeout``) and the remaining types have no reported wedge. + """ + if task_type in _RETAIN_OP_TYPES: + timeout = get_config().retain_wall_timeout + return float(timeout) if timeout > 0 else None + return None + + +class _WallTimeoutExceeded(Exception): + """A task was cancelled because it blew through its wall-clock ceiling.""" + + def __init__(self, timeout: float) -> None: + super().__init__(f"wall-clock timeout after {timeout:.0f}s") + self.timeout = timeout + + def _updated_row_count(result: Any) -> int: """Extract a row count from backend execute() results.""" if isinstance(result, int): @@ -727,6 +756,26 @@ async def _cleanup_task(self, operation_id: str, operation_type: str): if self._in_flight_by_type[operation_type] == 0: del self._in_flight_by_type[operation_type] + async def _run_executor(self, task: ClaimedTask, task_type: str) -> None: + """Run the task executor under its type's wall-clock ceiling, if it has one.""" + wall_timeout = _wall_timeout_for(task_type) + if wall_timeout is None: + await self._executor(task.task_dict) + return + + # asyncio.timeout() rather than wait_for(): `expired()` distinguishes our + # ceiling firing from an inner TimeoutError merely bubbling out (an asyncpg + # command timeout, say), which wait_for would surface as the same exception. + # Reporting a task's own timeout as a wedge would send operators hunting for + # the wrong thing. + try: + async with asyncio.timeout(wall_timeout) as cm: + await self._executor(task.task_dict) + except asyncio.TimeoutError as e: + if cm.expired(): + raise _WallTimeoutExceeded(wall_timeout) from e + raise + async def _execute_task_inner(self, task: ClaimedTask, holder: StageHolder | None = None): """Inner task execution with retry/fail handling. @@ -769,10 +818,25 @@ async def _execute_task_inner(self, task: ClaimedTask, holder: StageHolder | Non logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})") if task.schema: task.task_dict["_schema"] = task.schema - await self._executor(task.task_dict) + await self._run_executor(task, task_type) logger.debug(f"Task {task.operation_id} execution finished") await self._mark_completed(task.operation_id, task.schema) terminal_success = True + except _WallTimeoutExceeded as e: + # The executor has already been cancelled; all that's left is to say so + # clearly. Handled apart from the generic branch below so the operator + # sees the wedge for what it is rather than a bare "TimeoutError", and + # so the stage that was current when the ceiling fired is preserved — + # that breadcrumb is the only pointer to where the task was stuck. + stage = holder.stage if holder is not None else "unknown" + message = ( + f"Task exceeded the {e.timeout:.0f}s wall-clock limit for '{task_type}' " + f"(stage={stage}) and was cancelled. Raise HINDSIGHT_API_RETAIN_WALL_TIMEOUT " + f"if this is a legitimately long operation, or set it to 0 to disable the limit." + ) + logger.error(f"Task {task.operation_id} timed out: {message}") + await self._mark_failed(task.operation_id, message, task.schema) + terminal_success = False except DeferOperation as e: # Deferral is not a terminal outcome — do not record a completion. await self._defer_operation(task.operation_id, e.exec_date, e.reason, task.schema) diff --git a/hindsight-api-slim/tests/test_llm_stage_breadcrumbs.py b/hindsight-api-slim/tests/test_llm_stage_breadcrumbs.py new file mode 100644 index 0000000000..39f8f4245b --- /dev/null +++ b/hindsight-api-slim/tests/test_llm_stage_breadcrumbs.py @@ -0,0 +1,68 @@ +"""The LLM stage breadcrumb must distinguish waiting from calling (#3002). + +`[WORKER_TASK] stage=llm.bedrock.retain_extract_facts+structured` was stamped +before the concurrency permits were acquired, so a task queued behind a +saturated semaphore looked identical to one the provider was actively running. +An operator debugging a wedged worker spent hours on Bedrock for tasks that had +never reached Bedrock. +""" + +import asyncio + +import pytest + +from hindsight_api.engine.llm_wrapper import LLMConfig +from hindsight_api.worker.stage import StageHolder, bind_holder + + +def _mock_llm() -> LLMConfig: + return LLMConfig(provider="mock", api_key="", base_url="", model="m") + + +@pytest.mark.asyncio +async def test_stage_says_queued_while_waiting_for_a_permit(monkeypatch): + llm = _mock_llm() + holder = StageHolder() + gate = asyncio.Semaphore(0) # never free: stands in for a saturated cap + + monkeypatch.setattr("hindsight_api.engine.llm_wrapper._semaphores_for_scope", lambda scope: [gate]) + + async def run_call(): + bind_holder(holder) + await llm.call(messages=[{"role": "user", "content": "hi"}], scope="retain_extract_facts") + + task = asyncio.create_task(run_call()) + for _ in range(5): # let it reach the acquire + await asyncio.sleep(0) + + assert holder.stage == "llm.mock.retain_extract_facts.queued" + + gate.release() + await task + assert not holder.stage.endswith(".queued"), "stage stayed 'queued' after the permit was granted" + + +@pytest.mark.asyncio +async def test_stage_drops_queued_once_the_call_is_in_flight(monkeypatch): + """With permits free the call proceeds, and the stage names the in-flight + call — the state the label always claimed to describe.""" + llm = _mock_llm() + holder = StageHolder() + seen: list[str] = [] + + async def fake_call(**_kwargs): + seen.append(holder.stage) + return "ok" + + monkeypatch.setattr(llm._provider_impl, "call", fake_call) + + async def run_call(): + bind_holder(holder) + await llm.call( + messages=[{"role": "user", "content": "hi"}], + scope="retain_extract_facts", + ) + + await asyncio.create_task(run_call()) + + assert seen == ["llm.mock.retain_extract_facts"] diff --git a/hindsight-api-slim/tests/test_pool_acquire_timeout.py b/hindsight-api-slim/tests/test_pool_acquire_timeout.py new file mode 100644 index 0000000000..db891e6a14 --- /dev/null +++ b/hindsight-api-slim/tests/test_pool_acquire_timeout.py @@ -0,0 +1,88 @@ +"""HINDSIGHT_API_DB_ACQUIRE_TIMEOUT must bound the wait it names (#3002). + +The value was only passed to ``asyncpg.create_pool(timeout=...)``, which is a +*connect* kwarg — how long establishing a new connection may take. The wait the +knob is named for, ``Pool.acquire()`` blocking until a connection frees up, was +left at asyncpg's default of "wait forever", so pool exhaustion never surfaced +as an error: it just hung, and none of the deployment's configured timeouts +applied. + +Deterministic (no DB): a fake pool records the kwargs acquire() is called with. +""" + +from contextlib import asynccontextmanager + +import pytest + +from hindsight_api.engine.db.postgresql import PostgreSQLBackend + + +class _FakeConn: + @asynccontextmanager + async def transaction(self): + yield + + +class _FakePool: + def __init__(self): + self.acquire_kwargs: list[dict] = [] + + def acquire(self, **kwargs): + self.acquire_kwargs.append(kwargs) + + @asynccontextmanager + async def _cm(): + yield _FakeConn() + + return _cm() + + def get_size(self): + return 1 + + def get_idle_size(self): + return 1 + + def get_max_size(self): + return 1 + + +def _backend(acquire_timeout: float | None) -> tuple[PostgreSQLBackend, _FakePool]: + """Build a backend around a fake pool, skipping initialize()'s real connect.""" + backend = PostgreSQLBackend() + pool = _FakePool() + backend._pool = pool + backend._acquire_timeout_s = acquire_timeout + return backend, pool + + +@pytest.mark.asyncio +async def test_acquire_passes_the_configured_timeout(): + backend, pool = _backend(30.0) + + async with backend.acquire(): + pass + + assert pool.acquire_kwargs == [{"timeout": 30.0}] + + +@pytest.mark.asyncio +async def test_transaction_passes_the_configured_timeout(): + """transaction() acquires too — it must not keep the unbounded default.""" + backend, pool = _backend(30.0) + + async with backend.transaction(): + pass + + assert pool.acquire_kwargs == [{"timeout": 30.0}] + + +@pytest.mark.asyncio +async def test_zero_restores_the_unbounded_wait(): + """0 is the documented escape hatch for deployments that would rather queue + than fail; asyncpg reads timeout=None as 'wait forever'.""" + backend, pool = _backend(None) + + async with backend.acquire(): + pass + + assert pool.acquire_kwargs == [{"timeout": None}] diff --git a/hindsight-api-slim/tests/test_retain_pipeline_cancellation.py b/hindsight-api-slim/tests/test_retain_pipeline_cancellation.py new file mode 100644 index 0000000000..accc0f8d47 --- /dev/null +++ b/hindsight-api-slim/tests/test_retain_pipeline_cancellation.py @@ -0,0 +1,98 @@ +"""The streaming retain pipeline must not outlive the call that started it (#3002). + +`_streaming_retain_batch` runs an LLM producer and a DB consumer concurrently. +A plain `asyncio.gather` propagates the consumer's exception immediately but +leaves the producer — and every extraction task it fanned out — running. Those +tasks then park forever on `chunk_queue.put()` into a queue nobody drains: they +pin their chunk payloads for the life of the process and still spend LLM permits +and provider tokens on an operation that already failed. + +The test drives the real pipeline with a DB error injected into the consumer +(standing in for the deadlock victim / lock timeout that triggered this in +production) while extractions are still in flight. +""" + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from hindsight_api.engine.retain import orchestrator +from hindsight_api.engine.response_models import TokenUsage + + +class _ExplodingPool: + """Backend-shaped pool whose acquire() raises, failing the consumer's batch. + + `_wraps_backend` puts `acquire_with_retry` on its backend path; RuntimeError + is not in the retryable set, so it propagates on the first attempt the way a + deadlock victim's error would. + """ + + _wraps_backend = True + ops = None + + def acquire(self): + raise RuntimeError("deadlock detected") + + +@pytest.mark.asyncio +async def test_consumer_failure_cancels_in_flight_extractions(monkeypatch): + hanging_started = asyncio.Event() + cancelled = 0 + calls = 0 + + async def fake_extract_and_embed(*_args, **_kwargs): + """First chunk returns instantly (so the consumer runs and fails); the + rest hang the way a call queued behind a saturated LLM semaphore does.""" + nonlocal calls, cancelled + calls += 1 + if calls == 1: + return [], [], [], TokenUsage() + hanging_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled += 1 + raise + return [], [], [], TokenUsage() + + monkeypatch.setattr(orchestrator, "_extract_and_embed", fake_extract_and_embed) + + chunks = ["chunk one", "chunk two", "chunk three", "chunk four"] + before = {t for t in asyncio.all_tasks()} + + with pytest.raises(RuntimeError, match="deadlock detected"): + await orchestrator._streaming_retain_batch( + pool=_ExplodingPool(), + embeddings_model=MagicMock(), + llm_config=MagicMock(), + entity_resolver=MagicMock(), + format_date_fn=lambda d: str(d), + bank_id="bank-cancel", + contents_dicts=[{"content": "\n".join(chunks)}], + contents=[], + config=MagicMock(), + document_id="doc-cancel", + is_first_batch=True, + fact_type_override=None, + document_tags=None, + agent_name="agent", + log_buffer=[], + start_time=0.0, + all_pre_chunks=list(chunks), + chunk_to_content=[0] * len(chunks), + # One chunk per consumer batch, so the consumer reaches the DB (and + # fails) while the remaining extractions are still in flight. + chunk_batch_size=1, + ) + + assert hanging_started.is_set(), "test never reached the state it means to cover" + + # Let the cancellations land. + for _ in range(5): + await asyncio.sleep(0) + + leaked = [t for t in asyncio.all_tasks() if t not in before and not t.done()] + assert not leaked, f"pipeline tasks outlived the failed operation: {leaked}" + assert cancelled >= 1, "in-flight extractions were not cancelled" diff --git a/hindsight-api-slim/tests/test_worker_wall_timeout.py b/hindsight-api-slim/tests/test_worker_wall_timeout.py new file mode 100644 index 0000000000..a4688aa2f6 --- /dev/null +++ b/hindsight-api-slim/tests/test_worker_wall_timeout.py @@ -0,0 +1,179 @@ +"""Wall-clock ceiling for worker tasks (#3002). + +A retain that blocks forever — a lock wait with no deadlock cycle for Postgres +to break, an LLM permit that never frees, a producer parked on a queue nobody +drains — used to hold its worker slot until the process restarted. The operation +stayed 'processing', which the API refuses to either retry or cancel, so once +every slot was held the worker stopped claiming retains entirely. + +The per-call timeouts that already existed (LLM request, DB statement, DB +acquire) each bound one step; none bounds the task. These tests cover the outer +ceiling that does, and — just as importantly — that it stays out of the way of +everything else. +""" + +import uuid +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from hindsight_api.config import clear_config_cache +from hindsight_api.worker.poller import ClaimedTask, _wall_timeout_for + + +@pytest.fixture(autouse=True) +def _reset_config_cache(): + """Config is cached process-wide; clear it around each test so an env patch + here can't leak into another test.""" + clear_config_cache() + yield + clear_config_cache() + + +def _make_poller(executor): + from hindsight_api.worker import WorkerPoller + + poller = WorkerPoller(backend=MagicMock(), worker_id="w-test", executor=executor) + # Stub the terminal-state handlers so the poller never touches the DB. + poller._mark_completed = AsyncMock() + poller._mark_failed = AsyncMock() + poller._defer_operation = AsyncMock() + poller._schedule_retry = AsyncMock() + return poller + + +async def _run(executor, task_type="batch_retain"): + poller = _make_poller(executor) + task = ClaimedTask( + operation_id=str(uuid.uuid4()), + task_dict={"type": task_type, "operation_type": task_type, "bank_id": "bank-1"}, + schema=None, + ) + with patch("hindsight_api.worker.poller.get_metrics_collector", return_value=MagicMock()): + await poller._execute_task_inner(task) + return poller, task + + +class TestWallTimeoutResolution: + def test_retain_variants_are_bounded(self): + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "1800"}): + clear_config_cache() + assert _wall_timeout_for("retain") == 1800.0 + assert _wall_timeout_for("batch_retain") == 1800.0 + assert _wall_timeout_for("file_convert_retain") == 1800.0 + + def test_zero_disables(self): + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "0"}): + clear_config_cache() + assert _wall_timeout_for("batch_retain") is None + + def test_other_task_types_are_unbounded(self): + """Only retain is bounded: reflect self-bounds inside the engine, and the + rest have no reported wedge. Bounding them here would be a behaviour + change nobody asked for — consolidation on a large bank is legitimately + long-running.""" + assert _wall_timeout_for("consolidation") is None + assert _wall_timeout_for("graph_maintenance") is None + assert _wall_timeout_for("reflect") is None + + +class TestWallTimeoutEnforcement: + @pytest.mark.asyncio + async def test_wedged_retain_is_cancelled_and_marked_failed(self): + """The whole point: the executor is cancelled (freeing the slot) and the + operation lands in 'failed', which the API *will* retry — unlike + 'processing', which it refuses to retry or cancel.""" + import asyncio + + cancelled = asyncio.Event() + + async def wedged(_task_dict): + try: + await asyncio.Event().wait() # never completes + except asyncio.CancelledError: + cancelled.set() + raise + + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "1"}): + clear_config_cache() + poller, task = await _run(wedged) + + assert cancelled.is_set(), "executor was not cancelled — the worker slot would stay held" + poller._mark_completed.assert_not_awaited() + poller._mark_failed.assert_awaited_once() + message = poller._mark_failed.await_args.args[1] + assert "wall-clock limit" in message + assert "HINDSIGHT_API_RETAIN_WALL_TIMEOUT" in message + + @pytest.mark.asyncio + async def test_failure_message_carries_the_stage(self): + """The stage at the moment the ceiling fires is the only breadcrumb + pointing at *where* the task was stuck, so it has to survive into the + error the operator reads.""" + import asyncio + + from hindsight_api.worker.stage import set_stage + + async def wedged(_task_dict): + set_stage("llm.bedrock.retain_extract_facts+structured.queued") + await asyncio.Event().wait() + + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "1"}): + clear_config_cache() + poller = _make_poller(wedged) + task = ClaimedTask( + operation_id=str(uuid.uuid4()), + task_dict={"type": "batch_retain", "operation_type": "batch_retain", "bank_id": "bank-1"}, + schema=None, + ) + from hindsight_api.worker.stage import StageHolder + + holder = StageHolder(stage="queued.batch_retain") + with patch("hindsight_api.worker.poller.get_metrics_collector", return_value=MagicMock()): + await poller._execute_task_inner(task, holder) + + message = poller._mark_failed.await_args.args[1] + assert "llm.bedrock.retain_extract_facts+structured.queued" in message + + @pytest.mark.asyncio + async def test_fast_task_is_untouched(self): + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "60"}): + clear_config_cache() + poller, task = await _run(AsyncMock()) + + poller._mark_completed.assert_awaited_once_with(task.operation_id, task.schema) + poller._mark_failed.assert_not_awaited() + + @pytest.mark.asyncio + async def test_inner_timeout_is_not_reported_as_a_wedge(self): + """A TimeoutError raised *by* the task (an asyncpg command timeout, say) + must not be dressed up as the wall-clock ceiling firing — that would send + an operator hunting for a wedge that never happened.""" + + async def inner_timeout(_task_dict): + raise TimeoutError("statement timeout") + + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "600"}): + clear_config_cache() + poller, task = await _run(inner_timeout) + + poller._mark_failed.assert_awaited_once() + message = poller._mark_failed.await_args.args[1] + assert "wall-clock limit" not in message + assert "statement timeout" in message + + @pytest.mark.asyncio + async def test_unbounded_type_is_not_cancelled(self): + """A slow non-retain task runs to completion even past the retain + ceiling — the timeout is per operation type, not global.""" + import asyncio + + async def slow(_task_dict): + await asyncio.sleep(0.2) + + with patch.dict("os.environ", {"HINDSIGHT_API_RETAIN_WALL_TIMEOUT": "1"}): + clear_config_cache() + poller, task = await _run(slow, task_type="consolidation") + + poller._mark_completed.assert_awaited_once() + poller._mark_failed.assert_not_awaited() diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 40c042e541..7bbf44818f 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -53,7 +53,7 @@ Migrations will automatically create the schema if it doesn't exist and create a | `HINDSIGHT_API_READ_DB_POOL_MIN_SIZE` | Minimum connections in the read-replica pool (only used when `READ_DATABASE_URL` is set) | Falls back to `DB_POOL_MIN_SIZE` | | `HINDSIGHT_API_READ_DB_POOL_MAX_SIZE` | Maximum connections in the read-replica pool (only used when `READ_DATABASE_URL` is set) | Falls back to `DB_POOL_MAX_SIZE` | | `HINDSIGHT_API_DB_COMMAND_TIMEOUT` | PostgreSQL command timeout in seconds (asyncpg client-side) | `60` | -| `HINDSIGHT_API_DB_ACQUIRE_TIMEOUT` | Connection acquisition timeout in seconds | `30` | +| `HINDSIGHT_API_DB_ACQUIRE_TIMEOUT` | Connection acquisition timeout in seconds. Bounds how long a caller waits for a free pool connection before failing (retried by the caller); `0` waits indefinitely. | `30` | | `HINDSIGHT_API_DB_STATEMENT_TIMEOUT` | Postgres `statement_timeout` applied to every pool connection, in seconds. Server-side safety net for runaway queries. Does **not** apply to Alembic migrations (which run on a separate psycopg2 engine). Set to `0` to disable. | `600` | | `HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER` | Optional Postgres `max_parallel_workers_per_gather` applied to every pool connection of this process. Unset leaves the server default. Set to `0` on background-worker processes so bulk maintenance queries (consolidation, graph upkeep) run serially instead of fanning out across CPU cores shared with latency-sensitive traffic. | unset | @@ -1127,6 +1127,7 @@ Controls the retain (memory ingestion) pipeline. | `HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS` | Extract causal relationships between facts | `true` | | `HINDSIGHT_API_RETAIN_BATCH_ENABLED` | Use LLM Batch API for fact extraction (50% cost savings, only with async operations) | `false` | | `HINDSIGHT_API_RETAIN_MAX_CONCURRENT` | Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention during high-concurrency ingestion. | `4` | +| `HINDSIGHT_API_RETAIN_WALL_TIMEOUT` | Wall-clock ceiling in seconds for one retain task in the worker. A retain that blocks indefinitely (lock contention, an unreachable LLM endpoint) is cancelled and marked `failed` instead of holding its worker slot until the process restarts, so it can be retried. Set well above your slowest healthy retain; `0` disables. | `3600` | | `HINDSIGHT_API_RETAIN_BATCH_TOKENS` | Max characters per sub-batch for async retain auto-splitting | `10000` | | `HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE` | Max chunks per streaming batch when retain ingests long documents. Each chunk produces roughly 17 facts, so the default 100 chunks ≈ 1700 facts per batch. Lower to cap memory/LLM pressure on large documents; raise for smaller chunks. Configurable per bank. | `100` | | `HINDSIGHT_API_RETAIN_ENTITY_LOOKUP` | Entity lookup method during retain: `full` (exact match) or `trigram` (fuzzy trigram matching) | `trigram` | diff --git a/hindsight-embed/hindsight_embed/env.example b/hindsight-embed/hindsight_embed/env.example index 0631050d41..412cf4ecc8 100644 --- a/hindsight-embed/hindsight_embed/env.example +++ b/hindsight-embed/hindsight_embed/env.example @@ -112,6 +112,12 @@ HINDSIGHT_API_LOG_LEVEL=info # 'failed' (not 'completed'), surfacing silently-dropped facts. Default false. # HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS=false +# Wall-clock ceiling (seconds) for one retain task in the worker. A retain that +# blocks indefinitely is cancelled and marked 'failed' — and so becomes +# retryable — instead of holding its worker slot until the process restarts. +# Set well above your slowest healthy retain; 0 disables. Default 3600. +# HINDSIGHT_API_RETAIN_WALL_TIMEOUT=3600 + # Dry-run extraction preview endpoint (POST /memories/dry-run-extract). Enabled by default; it makes # a real LLM call but stores nothing. Set to false to remove the endpoint (returns 404). # HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=true diff --git a/skills/hindsight-docs/references/developer/configuration.md b/skills/hindsight-docs/references/developer/configuration.md index 7e65d10973..45ce6c2383 100644 --- a/skills/hindsight-docs/references/developer/configuration.md +++ b/skills/hindsight-docs/references/developer/configuration.md @@ -53,7 +53,7 @@ Migrations will automatically create the schema if it doesn't exist and create a | `HINDSIGHT_API_READ_DB_POOL_MIN_SIZE` | Minimum connections in the read-replica pool (only used when `READ_DATABASE_URL` is set) | Falls back to `DB_POOL_MIN_SIZE` | | `HINDSIGHT_API_READ_DB_POOL_MAX_SIZE` | Maximum connections in the read-replica pool (only used when `READ_DATABASE_URL` is set) | Falls back to `DB_POOL_MAX_SIZE` | | `HINDSIGHT_API_DB_COMMAND_TIMEOUT` | PostgreSQL command timeout in seconds (asyncpg client-side) | `60` | -| `HINDSIGHT_API_DB_ACQUIRE_TIMEOUT` | Connection acquisition timeout in seconds | `30` | +| `HINDSIGHT_API_DB_ACQUIRE_TIMEOUT` | Connection acquisition timeout in seconds. Bounds how long a caller waits for a free pool connection before failing (retried by the caller); `0` waits indefinitely. | `30` | | `HINDSIGHT_API_DB_STATEMENT_TIMEOUT` | Postgres `statement_timeout` applied to every pool connection, in seconds. Server-side safety net for runaway queries. Does **not** apply to Alembic migrations (which run on a separate psycopg2 engine). Set to `0` to disable. | `600` | | `HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER` | Optional Postgres `max_parallel_workers_per_gather` applied to every pool connection of this process. Unset leaves the server default. Set to `0` on background-worker processes so bulk maintenance queries (consolidation, graph upkeep) run serially instead of fanning out across CPU cores shared with latency-sensitive traffic. | unset | @@ -1127,6 +1127,7 @@ Controls the retain (memory ingestion) pipeline. | `HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS` | Extract causal relationships between facts | `true` | | `HINDSIGHT_API_RETAIN_BATCH_ENABLED` | Use LLM Batch API for fact extraction (50% cost savings, only with async operations) | `false` | | `HINDSIGHT_API_RETAIN_MAX_CONCURRENT` | Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention during high-concurrency ingestion. | `4` | +| `HINDSIGHT_API_RETAIN_WALL_TIMEOUT` | Wall-clock ceiling in seconds for one retain task in the worker. A retain that blocks indefinitely (lock contention, an unreachable LLM endpoint) is cancelled and marked `failed` instead of holding its worker slot until the process restarts, so it can be retried. Set well above your slowest healthy retain; `0` disables. | `3600` | | `HINDSIGHT_API_RETAIN_BATCH_TOKENS` | Max characters per sub-batch for async retain auto-splitting | `10000` | | `HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE` | Max chunks per streaming batch when retain ingests long documents. Each chunk produces roughly 17 facts, so the default 100 chunks ≈ 1700 facts per batch. Lower to cap memory/LLM pressure on large documents; raise for smaller chunks. Configurable per bank. | `100` | | `HINDSIGHT_API_RETAIN_ENTITY_LOOKUP` | Entity lookup method during retain: `full` (exact match) or `trigram` (fuzzy trigram matching) | `trigram` |