fix(api): make the 504 name the layer that ate the request budget (ax-0917-h-01/h-02)

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 ``"<VERB> <route template>"`` — 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 ------------------------------------------------
 

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:

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;

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)

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:

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)