diff --git a/core-api/src/core_api/clients/storage_client.py b/core-api/src/core_api/clients/storage_client.py index 7614266a7..bd8277121 100644 --- a/core-api/src/core_api/clients/storage_client.py +++ b/core-api/src/core_api/clients/storage_client.py @@ -19,6 +19,7 @@ from core_api.clients.identity_token import fetch_auth_header from core_api.config import settings from core_api.constants import STORAGE_CONNECT_TIMEOUT_SECONDS, STORAGE_READ_TIMEOUT_SECONDS +from core_api.request_phase import phase logger = logging.getLogger(__name__) @@ -426,11 +427,17 @@ def _shielded() -> Awaitable[httpx.Response]: return self._cancel_safe(do_request()) observed_gen = self._pool_generation - try: - return await retry(_shielded, label=label) - except httpx.PoolTimeout: - await self._recycle_pools(observed_gen=observed_gen, label=label) - return await retry(_shielded, label=label) + # ``label`` is already ``" "`` — bounded, no path + # params — so it is safe as a phase name and reads as the hop it is. + # Wrapping the whole retry policy, not one attempt: a request that + # burns the budget across three retries spent that time HERE, and + # per-attempt phases would report the last one's few hundred ms. + with phase(f"storage.{label}"): + try: + return await retry(_shielded, label=label) + except httpx.PoolTimeout: + await self._recycle_pools(observed_gen=observed_gen, label=label) + return await retry(_shielded, label=label) # -- internal helpers ------------------------------------------------ diff --git a/core-api/src/core_api/middleware/per_tenant_concurrency.py b/core-api/src/core_api/middleware/per_tenant_concurrency.py index d2c571940..dfb3a25c2 100644 --- a/core-api/src/core_api/middleware/per_tenant_concurrency.py +++ b/core-api/src/core_api/middleware/per_tenant_concurrency.py @@ -38,6 +38,7 @@ from fastapi import HTTPException from core_api.config import settings +from core_api.request_phase import phase logger = logging.getLogger(__name__) @@ -119,7 +120,8 @@ async def per_tenant_slot( sem = _get_semaphore(scope, tenant_id) try: async with asyncio.timeout(settings.per_tenant_acquire_timeout_seconds): - await sem.acquire() + with phase(f"slot_acquire.{scope}"): + await sem.acquire() except TimeoutError: logger.info( "per-tenant concurrency cap reached", @@ -177,7 +179,12 @@ async def per_tenant_storage_slot( "per-tenant storage slot saturated; queuing", extra={"scope": scope, "tenant_id": tenant_id, "cap": _cap_for(scope)}, ) - await sem.acquire() + # DEBUG is off in prod, so the log above is not evidence there. This + # queue is unbounded by design and explicitly relies on the request + # budget as its only cap — which makes it a prime candidate for eating + # that budget, and the one hop that had no signal at all when it did. + with phase(f"slot_acquire.{scope}"): + await sem.acquire() try: yield finally: diff --git a/core-api/src/core_api/middleware/request_observation.py b/core-api/src/core_api/middleware/request_observation.py index fa1bf45b5..fe6e3b330 100644 --- a/core-api/src/core_api/middleware/request_observation.py +++ b/core-api/src/core_api/middleware/request_observation.py @@ -71,6 +71,7 @@ from starlette.types import ASGIApp as ASGIApplication from starlette.types import Receive, Scope, Send +from core_api import request_phase from core_api.constants import PROBE_ROUTES from core_api.services.capability_usage import record_usage @@ -211,11 +212,13 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # ``http.response.start`` we still emit an event, and a crash that # never produced a status line is most accurately reported as 5xx. status_code = 500 + response_started = False async def _send(message: MutableMapping[str, Any]) -> None: - nonlocal status_code + nonlocal status_code, response_started if message["type"] == "http.response.start": status_code = message["status"] + response_started = True await send(message) # Captured BEFORE the downstream call: a Starlette ``Mount`` appends @@ -229,6 +232,19 @@ async def _send(message: MutableMapping[str, Any]) -> None: start = time.monotonic() try: await self.app(scope, receive, _send) + except BaseException: + # This middleware sits INSIDE RequestTimeoutMiddleware, so a + # request the budget kills unwinds through here with no + # ``http.response.start`` ever sent — and the 500 default above + # then filed every 45s timeout in this metric as a crash. That is + # the first dashboard an incident reaches for, and on 2026-09-17 + # it said the wrong thing: /search and /recall read as 500s while + # the callers were holding 504s. The docstring's claim that + # "504s/429s aren't observed" was true of the intent and false of + # the output. + if not response_started and request_phase.past_deadline(): + status_code = 504 + raise finally: duration_ms = (time.monotonic() - start) * 1000.0 # ``scope["route"]`` is set by the router during the call above; diff --git a/core-api/src/core_api/middleware/request_timeout.py b/core-api/src/core_api/middleware/request_timeout.py index c749e5c84..64385832d 100644 --- a/core-api/src/core_api/middleware/request_timeout.py +++ b/core-api/src/core_api/middleware/request_timeout.py @@ -24,6 +24,7 @@ from starlette.types import ASGIApp as ASGIApplication from starlette.types import Receive, Scope, Send +from core_api import request_phase from core_api.constants import is_mcp_path from core_api.errors import REQUEST_BUDGET_EXCEEDED, make_error_payload @@ -83,11 +84,17 @@ async def _send(message: MutableMapping[str, Any]) -> None: await send(message) started_at = time.monotonic() + # Armed HERE and nowhere deeper: this middleware sits OUTSIDE + # SlowAPI's ``BaseHTTPMiddleware``, which runs the rest of the app in + # a separate task. A recorder created below that split would be bound + # in a context copy this frame never sees. See ``request_phase``. + phases, phase_token = request_phase.begin(self.timeout_seconds) try: async with asyncio.timeout(self.timeout_seconds): await self.app(scope, receive, _send) except TimeoutError: elapsed = round(time.monotonic() - started_at, 3) + attribution = phases.snapshot() # Structured, not interpolated. The 2026-09-17 incident was # triaged from a log line that carried the path and nothing # else, so "which route, how long, how often" could not be @@ -99,6 +106,13 @@ async def _send(message: MutableMapping[str, Any]) -> None: "elapsed_seconds": elapsed, "method": scope.get("method", "?"), "path": scope["path"], + # The one field that turns this line from "a request on + # /search was slow" into "the embedding hop was slow". + # Flat and top-level so it is groupable in the log + # backend without unpacking a nested object. + "phase": attribution["phase"], + "phases_cancelled": attribution["phases_cancelled"], + "phases_completed": attribution["phases_completed"], }, ) if response_started: @@ -110,13 +124,20 @@ async def _send(message: MutableMapping[str, Any]) -> None: REQUEST_BUDGET_EXCEEDED, ( f"Request exceeded the {self.timeout_seconds}s server budget and was " - f"cancelled. No upstream reported a failure — this deadline is ours. " - f"Retry; if it recurs on the same route, the handler is the slow part." + f"cancelled" + + (f" while running {attribution['phase']}" if attribution["phase"] else "") + + ". No upstream reported a failure — this deadline is ours. " + "Retry; if it recurs on the same route, the handler is the slow part." ), details={ "budget_seconds": self.timeout_seconds, "elapsed_seconds": elapsed, "path": scope["path"], + # Attribution, not decoration: without it a stalled + # embedding provider and a stalled storage read are the + # same 504, and the only lead an operator has is the + # wall-clock number that every one of them shares. + **attribution, }, ) body = json.dumps(payload).encode() @@ -143,3 +164,5 @@ async def _send(message: MutableMapping[str, Any]) -> None: "more_body": False, } ) + finally: + request_phase.end(phase_token) diff --git a/core-api/src/core_api/pipeline/runner.py b/core-api/src/core_api/pipeline/runner.py index 458e63557..33d7f6963 100644 --- a/core-api/src/core_api/pipeline/runner.py +++ b/core-api/src/core_api/pipeline/runner.py @@ -11,6 +11,7 @@ from core_api.pipeline.context import PipelineContext from core_api.pipeline.step import Step, StepOutcome, StepResult +from core_api.request_phase import phase logger = logging.getLogger(__name__) @@ -39,7 +40,13 @@ async def run(self, ctx: PipelineContext) -> PipelineResult: for step in self._steps: t_step = time.perf_counter() try: - step_result = await step.execute(ctx) + # The per-step timing below only exists for steps that FINISH. + # A request cancelled by the budget dies inside a step, so the + # step that ate the time is precisely the one with no log line + # — the absence an operator has to notice rather than read. + # ``phase`` records the entry, so the unwind reports it. + with phase(f"{self._name}.{step.name}"): + step_result = await step.execute(ctx) if step_result is None: step_result = StepResult(outcome=StepOutcome.SUCCESS) except HTTPException: diff --git a/core-api/src/core_api/request_phase.py b/core-api/src/core_api/request_phase.py new file mode 100644 index 000000000..d1399df9c --- /dev/null +++ b/core-api/src/core_api/request_phase.py @@ -0,0 +1,199 @@ +"""Request-scoped phase attribution for the request budget (ax-0917-h-01/h-02). + +``RequestTimeoutMiddleware`` cancels a handler at +``request_timeout_seconds`` and answers ``REQUEST_BUDGET_EXCEEDED`` with the +budget, the elapsed time and the path. That says the deadline passed; it does +not say WHICH layer consumed it. Two entirely different incidents — a stalled +embedding provider and a stalled storage read — produced byte-identical +evidence, so "the next occurrence will say which layer ate the budget" was not +true of the shipped instrumentation. + +The same principle is already written down one layer lower, beside +``EMBEDDING_GATE_TIMEOUT_SECONDS``: *"a cancellation carries no attribution — +it says the deadline passed, not which layer ate it."* That module solved it +for itself by failing under the caller's budget. A blanket middleware has no +equivalent move: it is the outermost deadline, so nothing below it can fail +first on its behalf. It has to be TOLD what was running. + +So each bounded hop announces itself with :func:`phase`, and the recorder +keeps the stack. On timeout the middleware reports the phases that were still +open — innermost first, which is the layer — plus the ones that had already +completed and how long each took, which is what turns "slow" into "slow HERE". + +Why a mutable object behind a ContextVar rather than the ContextVar carrying +the data: ``BaseHTTPMiddleware`` (SlowAPI) runs the downstream app in a +separate anyio task, and so does every ``asyncio.ensure_future`` on the search +path. A child task inherits a COPY of the context, so a ``set()`` down there +is invisible up here — but the copy binds the same object, so mutations to it +are not. That is why the middleware creates the recorder before the split and +never re-sets the var. + +Cost: one ``perf_counter`` and one list append per phase, on a path that +already times every pipeline step. ``phase()`` is a no-op when no recorder is +bound (MCP transport, opt-out routes, background tasks), so instrumented code +is safe to call from anywhere. +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar, Token + +# Caps on what reaches the 504 body and the log line. A search runs 11 steps +# plus its storage hops; a write pipeline is comparable. 32 completed phases +# covers both with headroom, and the count of what was dropped is reported so +# a truncated list never reads as a complete one. +_MAX_COMPLETED = 32 +_MAX_OPEN = 16 + + +class RequestPhases: + """The phase stack for ONE request. Mutated from any task in its context.""" + + __slots__ = ( + "_cancelled", + "_completed", + "_deadline", + "_dropped", + "_next_id", + "_open", + "_started_at", + ) + + def __init__(self, budget_seconds: float | None = None) -> None: + self._started_at = time.monotonic() + self._deadline = None if budget_seconds is None else self._started_at + budget_seconds + # id -> (name, entered_at). A dict rather than a list because phases + # nest AND overlap: the search step races the embedding and the entity + # boost in separate tasks, so exits are not LIFO. + self._open: dict[int, tuple[str, float]] = {} + self._completed: list[tuple[str, float]] = [] + # Phases that unwound on an exception — under the budget timeout that + # is the cancellation, so this list IS the stack that was in flight, + # innermost first (the deepest ``finally`` runs first). + self._cancelled: list[tuple[str, float]] = [] + self._dropped = 0 + self._next_id = 0 + + def enter(self, name: str) -> int: + token = self._next_id + self._next_id += 1 + if len(self._open) < _MAX_OPEN: + self._open[token] = (name, time.monotonic()) + return token + + def exit(self, token: int, *, failed: bool) -> None: + entry = self._open.pop(token, None) + if entry is None: + return + name, entered_at = entry + record = (name, round(time.monotonic() - entered_at, 3)) + bucket = self._cancelled if failed else self._completed + if len(bucket) >= _MAX_COMPLETED: + self._dropped += 1 + return + bucket.append(record) + + def snapshot(self) -> dict: + """What the budget-exceeded response and log line report. + + ``phase`` is the single-field answer to "which layer": the innermost + thing that was still running. It is read from ``_cancelled`` rather + than ``_open`` because by the time the middleware catches + ``TimeoutError`` the cancellation has already unwound every ``with`` + block below it — the stack is in the unwind record, not in ``_open``. + ``_open`` is still consulted for the case where a phase is held by a + task the cancellation has not reached yet. + """ + in_flight = [ + {"phase": name, "seconds": round(time.monotonic() - at, 3)} + for name, at in sorted(self._open.values(), key=lambda e: e[1], reverse=True) + ] + cancelled = [{"phase": n, "seconds": s} for n, s in self._cancelled] + deepest = None + if cancelled: + deepest = cancelled[0]["phase"] + elif in_flight: + deepest = in_flight[0]["phase"] + out: dict = { + "phase": deepest, + "phases_cancelled": cancelled, + "phases_completed": [{"phase": n, "seconds": s} for n, s in self._completed], + } + if in_flight: + out["phases_open"] = in_flight + if self._dropped: + out["phases_dropped"] = self._dropped + return out + + def past_deadline(self) -> bool: + """Has the request budget already expired? + + Lets a layer INSIDE the timeout middleware tell "cancelled by our own + deadline" from "crashed". It has to be asked rather than told: the + middleware only learns the budget blew after every inner ``finally`` + has already run, so by the time it could set a flag, the layers that + need the answer are gone. The comparison is exact, not approximate — + ``asyncio.timeout`` fires off ``loop.time()``, which is + ``time.monotonic()`` for the default event loop, so on the unwinding + path the clock has provably passed this deadline. + """ + return self._deadline is not None and time.monotonic() >= self._deadline + + +_phases: ContextVar[RequestPhases | None] = ContextVar("request_phases", default=None) + + +def begin(budget_seconds: float | None = None) -> tuple[RequestPhases, Token]: + """Arm a recorder for this request. Caller MUST :func:`end` the token. + + Resetting matters even though production serves each request in its own + task: an in-process ASGI transport (every integration test, and the + storage bridge) calls the app inside the CALLER's context, so a leaked + binding would let one request's phases land in the next one's report. + """ + recorder = RequestPhases(budget_seconds) + return recorder, _phases.set(recorder) + + +def end(token: Token) -> None: + _phases.reset(token) + + +def current() -> RequestPhases | None: + return _phases.get() + + +def past_deadline() -> bool: + """``True`` only inside a budgeted request whose budget has expired.""" + recorder = _phases.get() + return recorder is not None and recorder.past_deadline() + + +@contextmanager +def phase(name: str) -> Iterator[None]: + """Mark ``name`` as the layer running inside this block. + + Synchronous by design: the block it wraps is almost always an ``await``, + and a sync context manager wraps awaited code perfectly well while costing + a fraction of ``@asynccontextmanager``'s per-use machinery. Names are + literals or bounded identifiers (step names, storage route labels, slot + scopes) — never user input, which would make this a cardinality bomb in + the log line. + """ + recorder = _phases.get() + if recorder is None: + yield + return + token = recorder.enter(name) + try: + yield + except BaseException: + # ``BaseException``, not ``Exception``: the case this exists for is + # ``CancelledError``, which is neither. + recorder.exit(token, failed=True) + raise + else: + recorder.exit(token, failed=False) diff --git a/core-api/src/core_api/routes/memories.py b/core-api/src/core_api/routes/memories.py index ca1ebe9bc..a55924c67 100644 --- a/core-api/src/core_api/routes/memories.py +++ b/core-api/src/core_api/routes/memories.py @@ -25,6 +25,7 @@ from common import permanent_failure from common.enrichment.constants import SERVER_RESERVED_MEMORY_TYPES from core_api import openapi_responses as _oar +from core_api import request_phase from core_api.agent_ids import ( ALWAYS_RESERVED_AGENT_IDS, DEFAULT_AGENT_ID, @@ -2423,20 +2424,33 @@ async def _search_inner( # Auth / tenant errors raised downstream are expected outcomes, # not DB/network failures — don't flag them as ``error=True``. raise - except Exception: + except BaseException: + # ``BaseException``, not ``Exception``: the request budget kills this + # handler with ``CancelledError``, which is neither — so a search the + # server gave up on logged ``error=false`` with ``row_count=0`` and + # a duration exactly equal to the budget. In this route's own + # telemetry a timed-out search was indistinguishable from a search + # that legitimately matched nothing, which put every 45s timeout into + # the empty-result rate and none into the error rate. success = False raise finally: if logger.isEnabledFor(logging.INFO): + cancelled = request_phase.past_deadline() logger.info( "search request completed", extra={ "path": "memory-search", "tenant_id": body.tenant_id, "top_k": body.top_k, + # Meaningless on a cancelled request — the pipeline never + # filled it — and reported anyway so the pair + # (row_count=0, cancelled=true) reads as one fact rather + # than as an empty result set. "row_count": len(results), "total_ms": (time.perf_counter() - t_start) * 1000, "error": not success, + "cancelled": cancelled, }, ) recall_tracked = bool(recall_ctx.get("recall_tracked")) diff --git a/core-api/src/core_api/services/memory_service.py b/core-api/src/core_api/services/memory_service.py index e5764bcc1..cdbba3b8d 100644 --- a/core-api/src/core_api/services/memory_service.py +++ b/core-api/src/core_api/services/memory_service.py @@ -16,6 +16,7 @@ from core_api.clients.storage_client import DuplicateMemoryError, get_storage_client from core_api.config import settings from core_api.middleware.per_tenant_concurrency import per_tenant_slot, per_tenant_storage_slot +from core_api.request_phase import phase from core_api.services.agent_identity import ReservedAgentIdError, enforce_reserved_write_id from core_api.tasks import track_task @@ -4732,7 +4733,14 @@ async def _get_or_cache_embedding(query: str, tenant_id: str, tenant_config): # propagates through the ``except`` below (future + joiners) and # the ``finally`` still pops the in-flight entry. async with per_tenant_slot("embed", tenant_id): - embedding = await asyncio.wait_for(get_query_embedding(query, tenant_config), timeout=10.0) + # h-02's shape — both semantic endpoints down while CRUD stayed + # healthy — points at exactly this hop, because it is the one + # /search and /recall share and CRUD never touches. Naming it + # separately from ``slot_acquire.embed`` is the whole point: a + # stalled provider and a queue behind other tenants' embeds are + # different incidents with different owners. + with phase("embed.query"): + embedding = await asyncio.wait_for(get_query_embedding(query, tenant_config), timeout=10.0) if embedding is None: # Two different things arrive as ``None`` and they are not the # same incident. A blank query cannot be embedded by anyone, and diff --git a/tests/test_ax_h01_h02_timeout_attribution.py b/tests/test_ax_h01_h02_timeout_attribution.py new file mode 100644 index 000000000..3fca05d2e --- /dev/null +++ b/tests/test_ax_h01_h02_timeout_attribution.py @@ -0,0 +1,470 @@ +"""ax-0917-h-01 / h-02 — does the 45s 504 say WHICH layer ate the budget? + +#1633 gave ``RequestTimeoutMiddleware`` the canonical envelope +(``REQUEST_BUDGET_EXCEEDED`` plus ``budget_seconds`` / ``elapsed_seconds`` / +``path``), a structured log line and ``Retry-After``. Both rows then closed on +the plan "the next occurrence will say which layer ate the budget" — so that +plan's premise is what these tests check, before an occurrence depends on it. + +They drive the REAL app stack: the real ``/api/v1/search`` route, the real +pipeline, the real middleware order, with only the budget lowered so it fires +in under a second. One hop is stalled at a time, as deep as the hop goes — +the embedding provider call, the per-tenant embed slot, the storage service's +own query method — and the evidence the two produce is then COMPARED. The +question is not "does the middleware emit an envelope" (``test_request_timeout`` +pins that against a synthetic app) but "given the evidence, can an operator +name the layer". + +Before ``core_api.request_phase``, the answer was no: a stalled embedding +provider and a stalled storage read returned byte-identical bodies and +byte-identical log records. ``test_the_504_names_the_layer_that_ate_the_budget`` +is the test that failed. + +No provider calls and no LLM calls — every stalled hop is patched. +""" + +from __future__ import annotations + +import asyncio +import logging + +import pytest + +from core_api.middleware.request_timeout import RequestTimeoutMiddleware +from tests.conftest import get_test_auth + +pytestmark = pytest.mark.integration + +_BUDGET_S = 0.4 +# Far past the budget: a stall that could finish would make the test a race. +_STALL_S = 30.0 +_FAKE_VECTOR = [0.0] * 8 + + +def _timeout_middleware(): + """The live ``RequestTimeoutMiddleware`` instance inside the built stack. + + The budget is captured at ``add_middleware`` time from + ``app_settings.request_timeout_seconds``, so patching the settings object + does nothing once the stack exists. Reaching the instance is what lets + these tests exercise the production path at test speed rather than + re-assembling an app that resembles it. + """ + from core_api.app import app + + if app.middleware_stack is None: + app.middleware_stack = app.build_middleware_stack() + node = app.middleware_stack + while node is not None: + if isinstance(node, RequestTimeoutMiddleware): + return node + node = getattr(node, "app", None) + raise AssertionError( + "RequestTimeoutMiddleware is not in the app's middleware stack" + ) + + +@pytest.fixture +def tight_budget(monkeypatch): + monkeypatch.setattr(_timeout_middleware(), "timeout_seconds", _BUDGET_S) + + +# Released in teardown so a stalled hop does not outlive its test. +# ``storage_client._cancel_safe`` SHIELDS the in-flight request on purpose (a +# cancelled httpx call strands its pooled connection — incident 2026-06-16), so +# cancelling the caller does not stop the stall; under the session-scoped test +# loop it would keep sleeping through later tests and be collected mid-flight. +_release: asyncio.Event | None = None + + +@pytest.fixture(autouse=True) +def stall_gate(): + global _release + _release = asyncio.Event() + yield + _release.set() + _release = None + + +async def _stall_forever(*_a, **_kw): + gate = _release + if gate is None: + await asyncio.sleep(_STALL_S) + else: + # Past the budget by two orders of magnitude, so the request is + # cancelled long before this returns; the wait only ends at teardown. + await asyncio.wait_for(gate.wait(), _STALL_S) + # Whatever awaited this is gone by now — a shielded task draining at + # teardown must unwind quietly, not raise into a done-callback. + return [] + + +async def _ready(value): + return value + + +def _stall_embedding_provider(monkeypatch): + """Hang the SHARED query-embedding hop — the h-02 shape. + + ``memory_service.get_query_embedding`` is the provider call both + ``/search`` and ``/recall`` funnel into and the one CRUD never touches, + which is why "both semantic endpoints down, CRUD healthy" points here. + """ + monkeypatch.setattr( + "core_api.services.memory_service.get_query_embedding", _stall_forever + ) + + +def _stall_storage_query(monkeypatch): + """Embedding resolves instantly; the storage read hangs instead. + + Patched at ``postgres_service.memory_scored_search`` — past the storage + client, past its retry policy, past the storage app's own router — so the + request really does traverse every layer between core-api and the query. + """ + monkeypatch.setattr( + "core_api.pipeline.steps.search.parallel_embed_entity_boost._get_or_cache_embedding", + lambda *_a, **_kw: _ready(_FAKE_VECTOR), + ) + from core_storage_api.services.postgres_service import PostgresService + + monkeypatch.setattr(PostgresService, "memory_scored_search", _stall_forever) + + +class _NeverAcquires: + """A semaphore nobody ever gets into.""" + + def locked(self) -> bool: + return True + + async def acquire(self) -> bool: + await _stall_forever() + return True + + def release(self) -> None: # pragma: no cover - never reached + pass + + +def _stall_storage_bulkhead(monkeypatch): + """Saturate the storage-search bulkhead so the QUEUE is what blocks. + + ``per_tenant_storage_slot`` is the one wait on this path that is + deliberately unbounded — its own docstring names the request budget as the + only thing capping it — so it is the hop most able to eat 45s, and it had + no INFO-level signal at all when it did. A queue with no free slots and a + slow storage backend look identical end to end and are fixed differently + (capacity vs the backend). + """ + monkeypatch.setattr( + "core_api.pipeline.steps.search.parallel_embed_entity_boost._get_or_cache_embedding", + lambda *_a, **_kw: _ready(_FAKE_VECTOR), + ) + import core_api.middleware.per_tenant_concurrency as ptc + + real = ptc._get_semaphore + monkeypatch.setattr( + ptc, + "_get_semaphore", + lambda scope, tenant_id: ( + _NeverAcquires() if scope == "storage_search" else real(scope, tenant_id) + ), + ) + + +async def _timed_out_search(client) -> tuple[int, dict, dict]: + tenant_id, headers = get_test_auth() + resp = await client.post( + "/api/v1/search", + json={"tenant_id": tenant_id, "query": "who runs fleet ops"}, + headers=headers, + ) + return resp.status_code, resp.json(), dict(resp.headers) + + +def _budget_log(caplog) -> logging.LogRecord: + hits = [r for r in caplog.records if r.getMessage() == "request exceeded budget"] + assert hits, ( + f"no budget log emitted; saw {[r.getMessage() for r in caplog.records]}" + ) + assert len(hits) == 1, "one timeout must emit exactly one budget log" + return hits[0] + + +async def _timeout_with(client, caplog, stall) -> tuple[dict, logging.LogRecord]: + caplog.clear() + with caplog.at_level(logging.WARNING): + stall() + status, body, headers = await _timed_out_search(client) + assert status == 504, body + assert headers["retry-after"] == "1" + return body, _budget_log(caplog) + + +# --------------------------------------------------------------------------- +# 1 — the timeout path really fires through the production stack +# --------------------------------------------------------------------------- + + +async def test_the_real_app_returns_the_budget_envelope_on_a_stalled_hop( + client, tight_budget, monkeypatch, caplog +): + caplog.clear() + with caplog.at_level(logging.WARNING): + _stall_embedding_provider(monkeypatch) + status, body, headers = await _timed_out_search(client) + + assert status == 504, body + err = body["error"] + assert err["code"] == "REQUEST_BUDGET_EXCEEDED" + assert err["details"]["budget_seconds"] == _BUDGET_S + assert err["details"]["path"] == "/api/v1/search" + assert err["details"]["elapsed_seconds"] >= _BUDGET_S + assert headers["retry-after"] == "1" + assert headers["content-type"] == "application/json" + + rec = _budget_log(caplog) + assert rec.levelno == logging.WARNING + assert rec.path == "/api/v1/search" + assert rec.method == "POST" + assert rec.budget_seconds == _BUDGET_S + assert rec.elapsed_seconds >= _BUDGET_S + + +async def test_the_declared_content_length_matches_the_body_it_sends( + client, tight_budget, monkeypatch, caplog +): + """The body now carries the phase list, so its length varies with what was + running. A raw ASGI ``http.response.start`` sends exactly the headers it is + given — nothing downstream recomputes this.""" + with caplog.at_level(logging.WARNING): + _stall_embedding_provider(monkeypatch) + tenant_id, headers = get_test_auth() + resp = await client.post( + "/api/v1/search", + json={"tenant_id": tenant_id, "query": "q"}, + headers=headers, + ) + assert resp.status_code == 504 + assert int(resp.headers["content-length"]) == len(resp.content) + + +# --------------------------------------------------------------------------- +# 2 — the question both rows are actually waiting on +# --------------------------------------------------------------------------- + + +def _evidence(body: dict, rec: logging.LogRecord) -> dict: + """Everything an operator holds after ONE 504, minus the wall-clock noise. + + Every duration is dropped, not just ``elapsed_seconds``: on a timeout the + in-flight phases all read within a few ms of the budget by construction, + so letting a duration discriminate would claim an attribution that a real + incident — where every 504 lands at the same ~45s — does not have. What is + left is only the NAMES, which is the claim being tested. + """ + details = dict(body["error"]["details"]) + details.pop("elapsed_seconds", None) + for key in ("phases_cancelled", "phases_completed", "phases_open"): + if key in details: + details[key] = [p["phase"] for p in details[key]] + return { + "code": body["error"]["code"], + "message": body["error"]["message"], + "details": details, + "log_phase": getattr(rec, "phase", None), + "log_cancelled": [p["phase"] for p in getattr(rec, "phases_cancelled", [])], + } + + +async def test_the_504_names_the_layer_that_ate_the_budget( + client, tight_budget, monkeypatch, caplog +): + """Two different incidents must not produce one indistinguishable 504. + + A stalled embedding provider and a stalled storage read have different + owners and different fixes. Before ``request_phase`` the evidence for both + was the same three fields — code, budget, path — and the plan h-01/h-02 + closed on ("the next occurrence will say which layer") was not met by the + code that closed them. + """ + embed_body, embed_log = await _timeout_with( + client, caplog, lambda: _stall_embedding_provider(monkeypatch) + ) + embed_evidence = _evidence(embed_body, embed_log) + + monkeypatch.undo() + monkeypatch.setattr(_timeout_middleware(), "timeout_seconds", _BUDGET_S) + + storage_body, storage_log = await _timeout_with( + client, caplog, lambda: _stall_storage_query(monkeypatch) + ) + storage_evidence = _evidence(storage_body, storage_log) + + assert embed_evidence != storage_evidence, ( + "a stalled embedding hop and a stalled storage hop emit identical " + f"evidence; the 504 cannot attribute the budget to a layer:\n{embed_evidence}" + ) + assert embed_evidence["details"]["phase"] == "embed.query" + assert ( + storage_evidence["details"]["phase"] == "storage.POST /memories/scored-search" + ) + + +async def test_the_log_line_carries_the_phase_for_aggregation( + client, tight_budget, monkeypatch, caplog +): + """Body attribution only helps the one caller who got the 504. "Which layer + is eating budgets, how often" is a log question, so the phase has to be a + flat top-level field on the record, not something to parse out of text.""" + _, rec = await _timeout_with( + client, caplog, lambda: _stall_storage_query(monkeypatch) + ) + assert rec.phase == "storage.POST /memories/scored-search" + assert [p["phase"] for p in rec.phases_cancelled][:2] == [ + "storage.POST /memories/scored-search", + "search.execute_scored_search", + ] + # What already finished is the other half of the answer: it rules the + # completed hops OUT, which is how "the pipeline was fine until here" gets + # said at all. + assert "search.classify_query" in [p["phase"] for p in rec.phases_completed] + + +async def test_a_saturated_bulkhead_is_not_reported_as_a_slow_backend( + client, tight_budget, monkeypatch, caplog +): + """h-02's shape — both semantic endpoints down, CRUD healthy — fits BOTH a + stalled storage backend and a storage bulkhead with no free slots, and + those have opposite fixes. The queue wait is where the time goes, so the + queue is what the 504 has to name — and it is the one hop whose only prior + signal was a DEBUG line that prod never emits.""" + body, rec = await _timeout_with( + client, caplog, lambda: _stall_storage_bulkhead(monkeypatch) + ) + assert body["error"]["details"]["phase"] == "slot_acquire.storage_search" + assert rec.phase == "slot_acquire.storage_search" + assert "storage.POST /memories/scored-search" not in [ + p["phase"] for p in rec.phases_cancelled + ], "a request still queued for a slot never reached the storage call" + + +async def test_the_message_says_the_phase_in_words( + client, tight_budget, monkeypatch, caplog +): + """A caller that logs ``error.message`` and nothing else — the common case + for an agent SDK — still gets the layer.""" + body, _ = await _timeout_with( + client, caplog, lambda: _stall_embedding_provider(monkeypatch) + ) + assert "embed.query" in body["error"]["message"] + + +# --------------------------------------------------------------------------- +# 3 — the recorder must not leak between requests +# --------------------------------------------------------------------------- + + +async def test_phases_do_not_leak_from_one_request_into_the_next( + client, tight_budget, monkeypatch, caplog +): + """The recorder is bound to a ContextVar. Under an in-process ASGI + transport the app runs in the CALLER's context, so a binding left behind + would put one request's phases in the next request's 504 — attribution + that is worse than none.""" + from core_api import request_phase + + await _timeout_with(client, caplog, lambda: _stall_embedding_provider(monkeypatch)) + assert request_phase.current() is None + + monkeypatch.undo() + monkeypatch.setattr(_timeout_middleware(), "timeout_seconds", _BUDGET_S) + + _, rec = await _timeout_with( + client, caplog, lambda: _stall_storage_query(monkeypatch) + ) + assert "embed.query" not in [p["phase"] for p in rec.phases_cancelled] + assert request_phase.current() is None + + +async def test_a_request_that_finishes_leaves_no_recorder_bound(client, monkeypatch): + from core_api import request_phase + + monkeypatch.setattr( + "core_api.routes.memories.search_memories", + lambda *_a, **_kw: _ready([]), + ) + tenant_id, headers = get_test_auth() + resp = await client.post( + "/api/v1/search", + json={"tenant_id": tenant_id, "query": "q"}, + headers=headers, + ) + assert resp.status_code == 200, resp.text + assert request_phase.current() is None + + +# --------------------------------------------------------------------------- +# 4 — the evidence that pointed the WRONG way +# --------------------------------------------------------------------------- + + +async def test_a_budget_timeout_is_not_filed_as_a_500_in_the_access_log( + client, tight_budget, monkeypatch, caplog +): + """``http.request`` backs the per-endpoint dashboard — the first place an + incident looks. ``RequestObservationMiddleware`` sits INSIDE the budget, so + a cancelled request unwinds through it with no status line ever sent and + fell to its 500 default. The caller held a 504 and the dashboard said + crash, which is a different investigation.""" + caplog.clear() + with caplog.at_level(logging.INFO, logger="core_api.access"): + _stall_embedding_provider(monkeypatch) + status, _, _ = await _timed_out_search(client) + assert status == 504 + + events = [r for r in caplog.records if r.getMessage() == "http.request"] + assert len(events) == 1, [r.getMessage() for r in caplog.records] + assert events[0].http_status_code == 504 + # Not asserted here: the route LABEL, which the six pre-existing + # failures in ``test_request_observation.py`` already own. + assert events[0].http_route.endswith("/search") + + +async def test_the_route_does_not_log_a_cancelled_search_as_a_completed_one( + client, tight_budget, monkeypatch, caplog +): + """``search request completed`` reported ``error=false, row_count=0`` on a + request the server itself killed: ``CancelledError`` is not an + ``Exception``, so the arm that sets ``success = False`` never ran. In this + route's own telemetry a timed-out search was a legitimately empty one.""" + caplog.clear() + with caplog.at_level(logging.INFO, logger="core_api.routes.memories"): + _stall_embedding_provider(monkeypatch) + status, _, _ = await _timed_out_search(client) + assert status == 504 + + done = [r for r in caplog.records if r.getMessage() == "search request completed"] + assert len(done) == 1 + assert done[0].error is True + assert done[0].cancelled is True + + +async def test_a_genuinely_empty_search_is_still_not_an_error( + client, monkeypatch, caplog +): + """The other side of the same line: matching nothing is a successful + search, and must not be dragged into the error rate by the fix above.""" + monkeypatch.setattr( + "core_api.routes.memories.search_memories", lambda *_a, **_kw: _ready([]) + ) + caplog.clear() + with caplog.at_level(logging.INFO, logger="core_api.routes.memories"): + tenant_id, headers = get_test_auth() + resp = await client.post( + "/api/v1/search", + json={"tenant_id": tenant_id, "query": "q"}, + headers=headers, + ) + assert resp.status_code == 200, resp.text + done = [r for r in caplog.records if r.getMessage() == "search request completed"] + assert done and done[-1].error is False and done[-1].cancelled is False