813 lines
31 KiB
Python
813 lines
31 KiB
Python
"""_cvdiag_backend.py — backend-layer CVDIAG boundary instrumentation.
|
|
|
|
This module wires the spec §3 / §5 **11 backend boundaries** into a Python
|
|
showcase integration, emitting schema-v1 CVDIAG envelopes through the shared
|
|
``_shared.cvdiag_bootstrap.emit_cvdiag`` sink. It is the per-integration
|
|
companion to the header-forwarding shim (``_header_forwarding.py``): that file
|
|
forwards correlation headers onto outbound LLM calls and logs lightweight
|
|
``CVDIAG component=backend-<fw> boundary=...`` breadcrumbs; THIS file emits the
|
|
full structured ``CVDIAG {<json>}`` envelopes the harness/classifier consume.
|
|
|
|
The 11 backend boundaries (spec §5 / §6 tier matrix):
|
|
|
|
1. ``backend.request.ingress`` — HTTP request received (default)
|
|
2. ``backend.agent.enter`` — agent loop entered (default)
|
|
3. ``backend.llm.call.start`` — outbound LLM call dispatched (verbose)
|
|
4. ``backend.llm.call.heartbeat`` — fires ~10s while an LLM call is
|
|
outstanding (verbose)
|
|
5. ``backend.llm.call.response`` — LLM response received (verbose)
|
|
6. ``backend.sse.first_byte`` — first SSE byte written (verbose)
|
|
7. ``backend.sse.event`` — every SSE event written (debug)
|
|
8. ``backend.sse.aborted`` — stream terminated abnormally (default)
|
|
9. ``backend.agent.exit`` — agent loop exited (default)
|
|
10. ``backend.response.complete`` — HTTP response stream closed (default)
|
|
11. ``backend.error.caught`` — exception caught in the agent loop
|
|
(default)
|
|
|
|
Guarding
|
|
--------
|
|
ALL emission is gated behind the ``CVDIAG_BACKEND_EMITTER`` env flag, default
|
|
OFF. With the flag off this module is byte-for-byte inert — no envelope is
|
|
built, no stdout line is written, the middleware passes the request straight
|
|
through. This is the canary-safe default: the flag is flipped ON only after a
|
|
deploy is confirmed healthy.
|
|
|
|
Tier gating
|
|
-----------
|
|
Each boundary carries a tier per the §6 matrix. ``_shared.cvdiag_bootstrap``
|
|
resolves the active tier (default | verbose | debug) once at import; this
|
|
module suppresses a boundary whose tier exceeds the active tier so the
|
|
default-tier production emit stays within the §7 event-count budget.
|
|
|
|
Pure instrumentation
|
|
--------------------
|
|
Nothing here may throw into the request path. ``emit_cvdiag`` already swallows
|
|
its own errors; the helpers below additionally guard envelope construction so a
|
|
malformed metadata bag degrades to a dropped emit, never a 500.
|
|
|
|
Plan unit: L1-C.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import re
|
|
import secrets
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, Optional
|
|
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
|
|
from _shared.cvdiag_bootstrap import _resolve_tier, current_tier, emit_cvdiag
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Framework tag — mirrors ``_header_forwarding._CVDIAG_FRAMEWORK`` so the
|
|
# structured envelopes and the breadcrumb log lines agree on the integration
|
|
# identity. (L1-D: change this single constant when copying to a sibling.)
|
|
_CVDIAG_FRAMEWORK = "langroid"
|
|
|
|
# ── Env gate ─────────────────────────────────────────────────────────────────
|
|
|
|
_BACKEND_EMITTER_ENV = "CVDIAG_BACKEND_EMITTER"
|
|
|
|
|
|
def cvdiag_backend_enabled() -> bool:
|
|
"""True iff the backend emitter is explicitly enabled (default OFF).
|
|
|
|
Read live (not cached) so a test can toggle the env var per-case via
|
|
``monkeypatch.setenv``; the cost is one ``os.environ`` lookup per emit,
|
|
which is negligible against the JSON serialization that follows.
|
|
"""
|
|
return os.environ.get(_BACKEND_EMITTER_ENV) == "1"
|
|
|
|
|
|
# ── Tier ordering (spec §6) ────────────────────────────────────────────────
|
|
|
|
_TIER_RANK = {"default": 0, "verbose": 1, "debug": 2}
|
|
|
|
# Per-boundary minimum tier required to emit (spec §6 matrix, backend rows).
|
|
_BOUNDARY_TIER: Dict[str, str] = {
|
|
"backend.request.ingress": "verbose",
|
|
"backend.agent.enter": "default",
|
|
"backend.llm.call.start": "verbose",
|
|
"backend.llm.call.heartbeat": "verbose",
|
|
"backend.llm.call.response": "verbose",
|
|
"backend.sse.first_byte": "verbose",
|
|
"backend.sse.event": "debug",
|
|
"backend.sse.aborted": "default",
|
|
"backend.agent.exit": "default",
|
|
"backend.response.complete": "default",
|
|
"backend.error.caught": "default",
|
|
}
|
|
|
|
|
|
def _active_tier() -> str:
|
|
"""Resolve the verbosity tier from a LIVE env read.
|
|
|
|
``cvdiag_backend_enabled()`` reads ``CVDIAG_BACKEND_EMITTER`` live, so the
|
|
tier MUST be read from the same live source — otherwise flipping
|
|
``CVDIAG_VERBOSE`` / ``CVDIAG_DEBUG`` AFTER import arms the emitter but the
|
|
tier stays frozen at the import-time ``setup()`` value, silently no-op'ing
|
|
every verbose/debug-gated boundary. We reuse the bootstrap's
|
|
``_resolve_tier`` so the §6 fail-closed DEBUG guard still applies (a
|
|
production / unresolved DEBUG request raises → degrade to the frozen tier).
|
|
"""
|
|
try:
|
|
return _resolve_tier(dict(os.environ))
|
|
except RuntimeError:
|
|
# Fail-closed DEBUG refusal: fall back to the import-time resolved tier
|
|
# (never silently escalate to debug in production).
|
|
return current_tier()
|
|
|
|
|
|
def _tier_permits(boundary: str) -> bool:
|
|
"""True iff the active tier is at-or-above the boundary's minimum tier."""
|
|
need = _TIER_RANK.get(_BOUNDARY_TIER.get(boundary, "default"), 0)
|
|
have = _TIER_RANK.get(_active_tier(), 0)
|
|
return have >= need
|
|
|
|
|
|
# ── Edge headers (spec §5 — 9-key allow-list + 12-name deny-list) ───────────
|
|
|
|
# The closed 9-key edge-header allow-list. Always-present in the envelope;
|
|
# absent header → ``None``.
|
|
_EDGE_ALLOW = (
|
|
"cf-ray",
|
|
"cf-mitigated",
|
|
"cf-cache-status",
|
|
"x-railway-edge",
|
|
"x-railway-request-id",
|
|
"x-hikari-trace",
|
|
"retry-after",
|
|
"via",
|
|
"server",
|
|
)
|
|
|
|
# Exact-match deny-list (spec §5). REJECTED even if accidentally present in the
|
|
# allow-list — these carry client IP / geo PII and must never round-trip.
|
|
_EDGE_DENY = frozenset(
|
|
{
|
|
"cf-ipcountry",
|
|
"cf-connecting-ip",
|
|
"cf-ipcity",
|
|
"cf-iplatitude",
|
|
"cf-iplongitude",
|
|
"cf-iptimezone",
|
|
"cf-visitor",
|
|
"cf-worker",
|
|
"true-client-ip",
|
|
"x-forwarded-for",
|
|
"x-real-ip",
|
|
"forwarded",
|
|
}
|
|
)
|
|
|
|
|
|
def extract_edge_headers(headers: Any) -> Dict[str, Optional[str]]:
|
|
"""Build the closed 9-key ``edge_headers`` bag from a headers mapping.
|
|
|
|
All nine keys are ALWAYS present; an absent (or deny-listed) header maps to
|
|
``None``. ``headers`` is any case-insensitive mapping exposing ``.get`` /
|
|
iteration of ``(name, value)`` pairs (Starlette ``Headers``, httpx, dict).
|
|
"""
|
|
bag: Dict[str, Optional[str]] = {k: None for k in _EDGE_ALLOW}
|
|
if headers is None:
|
|
return bag
|
|
try:
|
|
getter = headers.get
|
|
except AttributeError:
|
|
return bag
|
|
for key in _EDGE_ALLOW:
|
|
if key in _EDGE_DENY: # belt-and-braces: never emit a deny-listed key
|
|
continue
|
|
val = getter(key)
|
|
if val is not None:
|
|
bag[key] = str(val)
|
|
return bag
|
|
|
|
|
|
# ── PII scrub (spec §6) ──────────────────────────────────────────────────────
|
|
|
|
# Bearer tokens, OpenAI/Stripe-style secret keys, publishable keys, and URL
|
|
# userinfo. Applied to any captured free-text metadata value
|
|
# (``message_scrubbed``, stack frames) before it is emitted. The ``sk-``/``pk-``
|
|
# key bodies allow hyphens/underscores so test-style keys such as the spec
|
|
# regression fixture ``sk-test-12345`` are redacted alongside real production
|
|
# keys (``sk-<48+ base62>``).
|
|
#
|
|
# Parity with the canonical TS scrubber (``harness/src/cvdiag/scrub.ts``):
|
|
# * Bearer — grabs the WHOLE token (``\S+``) to match TS ``Bearer\s+\S+``;
|
|
# the legacy ``[A-Za-z0-9._\-]+`` stopped at ``/``/``+``/``=`` and left an
|
|
# un-redacted JWT tail (e.g. ``Bearer a.b.c/sig+more=`` → ``…/sig+more=``).
|
|
# * URL userinfo — redacts BOTH ``scheme://user:pw@host`` AND colon-less
|
|
# ``scheme://token@host`` (TS ``([scheme]://)[^/\s?#]*@``); the legacy
|
|
# ``[^/\s:@]+:[^/\s@]+@`` required a mandatory ``:`` so a bare-token
|
|
# authority such as ``https://ghp_xxx@host`` LEAKED. The userinfo class
|
|
# excludes ``?``/``#`` so the match never crosses into the query/fragment.
|
|
_SCRUB_PATTERNS = (
|
|
re.compile(r"Bearer\s+\S+", re.IGNORECASE),
|
|
re.compile(r"\bsk-[A-Za-z0-9][A-Za-z0-9_-]{3,}"),
|
|
re.compile(r"\bpk-[A-Za-z0-9][A-Za-z0-9_-]{3,}"),
|
|
re.compile(r"(?P<scheme>[a-z][a-z0-9+.\-]*://)[^/\s?#]*@", re.IGNORECASE),
|
|
)
|
|
|
|
# Per-event field byte caps (spec §5). message_scrubbed ≤512B.
|
|
_MESSAGE_CAP = 512
|
|
|
|
# Hard input-size guard (mirrors TS ``SCRUB_MAX_SCAN_LEN``): no regex ever runs
|
|
# on a string longer than this. A longer value has only its bounded prefix
|
|
# scanned and a self-describing ``…[unscanned:<N>]`` marker records the dropped
|
|
# tail length, so an adversarial multi-KB string can never make the regex
|
|
# engine scan unbounded input. 2 KB covers any legitimate metadata value with
|
|
# headroom. Set below the byte cap so the marker survives the §5 byte clamp.
|
|
_SCRUB_MAX_SCAN_LEN = 400
|
|
|
|
|
|
def _run_scrub_regexes(s: str) -> str:
|
|
"""Apply the secret regexes in sequence (TS ``runScrubRegexes`` parity)."""
|
|
for pat in _SCRUB_PATTERNS:
|
|
if pat.groupindex.get("scheme"):
|
|
s = pat.sub(r"\g<scheme>[REDACTED]@", s)
|
|
else:
|
|
s = pat.sub("[REDACTED]", s)
|
|
return s
|
|
|
|
|
|
def scrub(text: Any) -> str:
|
|
"""Redact secrets from a free-text value and cap it at 512 bytes.
|
|
|
|
Returns ``"[REDACTED]"`` substitutions for any matched secret pattern so a
|
|
synthetic ``sk-test-12345`` in an exception message can never reach the
|
|
emitted envelope. A value longer than ``_SCRUB_MAX_SCAN_LEN`` has only its
|
|
bounded prefix scanned, with an ``…[unscanned:<N>]`` marker (TS parity).
|
|
"""
|
|
if text is None:
|
|
return ""
|
|
s = str(text)
|
|
if len(s) > _SCRUB_MAX_SCAN_LEN:
|
|
dropped_tail = len(s) - _SCRUB_MAX_SCAN_LEN
|
|
s = f"{_run_scrub_regexes(s[:_SCRUB_MAX_SCAN_LEN])}…[unscanned:{dropped_tail}]"
|
|
else:
|
|
s = _run_scrub_regexes(s)
|
|
encoded = s.encode("utf-8")
|
|
if len(encoded) > _MESSAGE_CAP:
|
|
s = encoded[:_MESSAGE_CAP].decode("utf-8", errors="ignore")
|
|
return s
|
|
|
|
|
|
# ── Envelope construction ──────────────────────────────────────────────────
|
|
|
|
_TEST_ID_HEADER = "x-test-id"
|
|
_AIMOCK_CONTEXT_HEADER = "x-aimock-context"
|
|
# UUIDv7 variant/version nibbles (RFC 9562) the schema regex requires.
|
|
_SLUG_FALLBACK = "unknown"
|
|
_DEMO_FALLBACK = "default"
|
|
|
|
|
|
def _uuid7() -> str:
|
|
"""Generate a lowercase-hyphenated UUIDv7 (RFC 9562) string.
|
|
|
|
48-bit Unix-ms timestamp, version nibble 7, variant 10 — matches the
|
|
schema ``TEST_ID_PATTERN``. Used as the fallback ``test_id`` when no
|
|
inbound ``x-test-id`` correlation header is present.
|
|
"""
|
|
unix_ms = int(time.time() * 1000) & ((1 << 48) - 1)
|
|
rand_a = secrets.randbits(12)
|
|
rand_b = secrets.randbits(62)
|
|
msb = (unix_ms << 16) | (0x7 << 12) | rand_a
|
|
lsb = (0b10 << 62) | rand_b
|
|
return str(uuid.UUID(int=(msb << 64) | lsb))
|
|
|
|
|
|
_UUID7_RE = re.compile(
|
|
r"^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$"
|
|
)
|
|
|
|
|
|
def normalize_test_id(raw: Optional[str]) -> str:
|
|
"""Return a schema-valid lowercased UUIDv7, minting one if ``raw`` is
|
|
absent or not a well-formed UUIDv7."""
|
|
if raw:
|
|
candidate = raw.strip().lower()
|
|
if _UUID7_RE.match(candidate):
|
|
return candidate
|
|
return _uuid7()
|
|
|
|
|
|
def _span_id() -> str:
|
|
"""16-hex span id, unique per emit (schema ``SPAN_ID_PATTERN``)."""
|
|
return secrets.token_hex(8)
|
|
|
|
|
|
_SLUG_RE = re.compile(r"^[a-z][a-z0-9-]{0,63}$")
|
|
|
|
|
|
def _normalize_slug(raw: Optional[str]) -> str:
|
|
"""Coerce the inbound ``x-aimock-context`` slug into the closed slug shape
|
|
(``^[a-z][a-z0-9-]{0,63}$``), falling back to ``unknown`` when unusable."""
|
|
if raw:
|
|
candidate = raw.strip().lower()
|
|
if _SLUG_RE.match(candidate):
|
|
return candidate
|
|
return _SLUG_FALLBACK
|
|
|
|
|
|
def build_envelope(
|
|
*,
|
|
boundary: str,
|
|
outcome: str,
|
|
test_id: str,
|
|
slug: str,
|
|
demo: str,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
edge_headers: Optional[Dict[str, Optional[str]]] = None,
|
|
duration_ms: Optional[int] = None,
|
|
parent_span_id: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Assemble a schema-v1 backend envelope (``layer=backend``).
|
|
|
|
All envelope-required fields are populated; ``edge_headers`` defaults to the
|
|
closed 9-key all-null bag when not supplied. ``metadata`` is passed through
|
|
verbatim — unknown keys are stamped ``_metadata_dropped`` by the schema
|
|
validator inside ``emit_cvdiag``.
|
|
"""
|
|
return {
|
|
"schema_version": 1,
|
|
"test_id": test_id,
|
|
"trace_id": test_id,
|
|
"span_id": _span_id(),
|
|
"parent_span_id": parent_span_id,
|
|
"layer": "backend",
|
|
"boundary": boundary,
|
|
"slug": slug,
|
|
"demo": demo,
|
|
"ts": _now_iso(),
|
|
"mono_ns": time.monotonic_ns(),
|
|
"duration_ms": duration_ms,
|
|
"outcome": outcome,
|
|
"edge_headers": edge_headers or {k: None for k in _EDGE_ALLOW},
|
|
"metadata": metadata or {},
|
|
}
|
|
|
|
|
|
def _now_iso() -> str:
|
|
"""ISO-8601 millisecond-precision timestamp with a ``Z`` suffix."""
|
|
# ``time.gmtime`` + manual ms keeps this dependency-free and 3.9-safe.
|
|
now = time.time()
|
|
secs = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(now))
|
|
ms = int((now - int(now)) * 1000)
|
|
return f"{secs}.{ms:03d}Z"
|
|
|
|
|
|
def emit_backend_boundary(
|
|
boundary: str,
|
|
*,
|
|
outcome: str = "info",
|
|
test_id: str,
|
|
slug: str,
|
|
demo: str,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
edge_headers: Optional[Dict[str, Optional[str]]] = None,
|
|
duration_ms: Optional[int] = None,
|
|
parent_span_id: Optional[str] = None,
|
|
) -> None:
|
|
"""Emit one backend boundary envelope, honoring the env gate + tier matrix.
|
|
|
|
No-op when the emitter is disabled or the active tier does not permit this
|
|
boundary. Never raises into the caller.
|
|
"""
|
|
if not cvdiag_backend_enabled():
|
|
return
|
|
if not _tier_permits(boundary):
|
|
return
|
|
try:
|
|
envelope = build_envelope(
|
|
boundary=boundary,
|
|
outcome=outcome,
|
|
test_id=test_id,
|
|
slug=slug,
|
|
demo=demo,
|
|
metadata=metadata,
|
|
edge_headers=edge_headers,
|
|
duration_ms=duration_ms,
|
|
parent_span_id=parent_span_id,
|
|
)
|
|
emit_cvdiag(envelope)
|
|
except Exception as err: # noqa: BLE001 - instrumentation must not throw
|
|
logger.warning("CVDIAG backend emit-failed boundary=%s error=%s", boundary, err)
|
|
|
|
|
|
# ── Per-request correlation context ─────────────────────────────────────────
|
|
|
|
|
|
class _RequestCtx:
|
|
"""Holds the per-request correlation identity + timing the boundaries share.
|
|
|
|
Carried on ``request.state`` so the middleware, the LLM hook, and the agent
|
|
hooks all stamp the same ``test_id`` / ``slug`` / ``demo`` onto their
|
|
envelopes.
|
|
"""
|
|
|
|
__slots__ = (
|
|
"test_id",
|
|
"slug",
|
|
"demo",
|
|
"ingress_mono_ns",
|
|
"sse_seq",
|
|
"first_byte_emitted",
|
|
"bytes_streamed",
|
|
)
|
|
|
|
def __init__(self, *, test_id: str, slug: str, demo: str) -> None:
|
|
self.test_id = test_id
|
|
self.slug = slug
|
|
self.demo = demo
|
|
self.ingress_mono_ns = time.monotonic_ns()
|
|
self.sse_seq = 0
|
|
self.first_byte_emitted = False
|
|
self.bytes_streamed = 0
|
|
|
|
|
|
def _demo_from_path(path: str) -> str:
|
|
"""Derive the ``demo`` label from the mounted sub-app path.
|
|
|
|
Each demo is mounted at ``/<demo>`` (e.g. ``/voice``, ``/byoc-hashbrown``);
|
|
the root agent serves the default demo. Strip the leading slash and any
|
|
trailing AG-UI segment so ``/byoc-hashbrown/`` → ``byoc-hashbrown`` and
|
|
``/`` → ``default``.
|
|
"""
|
|
trimmed = path.strip("/")
|
|
if not trimmed:
|
|
return _DEMO_FALLBACK
|
|
return trimmed.split("/", 1)[0] or _DEMO_FALLBACK
|
|
|
|
|
|
# ── HTTP middleware: ingress / first_byte / sse.event / sse.aborted /
|
|
# response.complete / error.caught ─────────────────────────────────────────
|
|
|
|
|
|
class CvdiagBackendMiddleware(BaseHTTPMiddleware):
|
|
"""Starlette middleware emitting the HTTP-observable backend boundaries.
|
|
|
|
Wires six of the eleven boundaries around the request lifecycle:
|
|
|
|
* ``backend.request.ingress`` on entry
|
|
* ``backend.sse.first_byte`` on the first streamed chunk
|
|
* ``backend.sse.event`` per streamed chunk (debug tier)
|
|
* ``backend.sse.aborted`` on premature stream termination
|
|
* ``backend.response.complete`` on clean stream close
|
|
* ``backend.error.caught`` on any exception escaping the inner app
|
|
|
|
The agent/LLM boundaries (``agent.enter``, ``llm.call.*``, ``agent.exit``)
|
|
are emitted by the agent hooks / LLM httpx hook installed separately, all
|
|
keyed on the same ``test_id`` this middleware stamps onto ``request.state``.
|
|
|
|
Inert when ``CVDIAG_BACKEND_EMITTER`` is off: the dispatch fast-paths to a
|
|
bare ``call_next`` with no envelope construction and no response wrapping.
|
|
"""
|
|
|
|
async def dispatch(self, request: Request, call_next) -> Response:
|
|
if not cvdiag_backend_enabled():
|
|
return await call_next(request)
|
|
|
|
headers = request.headers
|
|
ctx = _RequestCtx(
|
|
test_id=normalize_test_id(headers.get(_TEST_ID_HEADER)),
|
|
slug=_normalize_slug(headers.get(_AIMOCK_CONTEXT_HEADER)),
|
|
demo=_demo_from_path(request.url.path),
|
|
)
|
|
request.state.cvdiag = ctx
|
|
|
|
emit_backend_boundary(
|
|
"backend.request.ingress",
|
|
outcome="info",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
edge_headers=extract_edge_headers(headers),
|
|
metadata={
|
|
"method": request.method,
|
|
"path": request.url.path,
|
|
"content_length": _int_or_none(headers.get("content-length")),
|
|
},
|
|
)
|
|
|
|
try:
|
|
response = await call_next(request)
|
|
except Exception as exc: # noqa: BLE001 - observe then re-raise
|
|
emit_backend_boundary(
|
|
"backend.error.caught",
|
|
outcome="err",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
metadata={
|
|
"exception_type": type(exc).__name__,
|
|
"message_scrubbed": scrub(str(exc)),
|
|
"stack_brief": [],
|
|
"truncated": False,
|
|
},
|
|
)
|
|
raise
|
|
|
|
return self._wrap_response(request, response, ctx)
|
|
|
|
def _wrap_response(
|
|
self, request: Request, response: Response, ctx: "_RequestCtx"
|
|
) -> Response:
|
|
"""Wrap a streaming response so SSE boundaries fire as chunks flow.
|
|
|
|
Non-streaming responses are returned unwrapped after emitting
|
|
``backend.response.complete`` directly.
|
|
|
|
NOTE: ``BaseHTTPMiddleware`` re-wraps the inner ``StreamingResponse`` as
|
|
a private ``_StreamingResponse`` before it reaches us, so an
|
|
``isinstance(response, StreamingResponse)`` check is always False here.
|
|
Detect streaming by the presence of a ``body_iterator`` (which both the
|
|
public and the private response carry) instead.
|
|
"""
|
|
if not hasattr(response, "body_iterator"):
|
|
emit_backend_boundary(
|
|
"backend.response.complete",
|
|
outcome="ok",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
duration_ms=_elapsed_ms(ctx.ingress_mono_ns),
|
|
edge_headers=extract_edge_headers(response.headers),
|
|
metadata={
|
|
"http_status": response.status_code,
|
|
"content_length": _int_or_none(
|
|
response.headers.get("content-length")
|
|
),
|
|
"total_duration_ms": _elapsed_ms(ctx.ingress_mono_ns),
|
|
"sse_event_count": ctx.sse_seq,
|
|
},
|
|
)
|
|
return response
|
|
|
|
inner = response.body_iterator
|
|
edge = extract_edge_headers(response.headers)
|
|
status = response.status_code
|
|
|
|
async def _instrumented():
|
|
# ``completed`` distinguishes a clean stream exhaustion (→
|
|
# response.complete) from an early termination (→ sse.aborted).
|
|
#
|
|
# IMPORTANT (Starlette ``BaseHTTPMiddleware`` quirk): when the INNER
|
|
# endpoint generator raises mid-stream, Starlette swallows the error
|
|
# internally and our ``async for`` simply ends — we never see an
|
|
# exception there. The abort surface we CAN observe is the consumer
|
|
# tearing the stream down early (client disconnect), which closes
|
|
# this generator and raises ``GeneratorExit`` / ``CancelledError``
|
|
# into it. We therefore catch ``BaseException`` (not just
|
|
# ``Exception``) so a disconnect-driven abort is captured, and emit
|
|
# ``backend.response.complete`` only on a clean exhaustion.
|
|
completed = False
|
|
terminated_kind = "rst"
|
|
try:
|
|
async for chunk in inner:
|
|
ctx.bytes_streamed += len(chunk) if chunk else 0
|
|
if not ctx.first_byte_emitted:
|
|
ctx.first_byte_emitted = True
|
|
emit_backend_boundary(
|
|
"backend.sse.first_byte",
|
|
outcome="info",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
edge_headers=edge,
|
|
metadata={
|
|
"delta_ms_from_ingress": _elapsed_ms(
|
|
ctx.ingress_mono_ns
|
|
)
|
|
},
|
|
)
|
|
emit_backend_boundary(
|
|
"backend.sse.event",
|
|
outcome="info",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
metadata={
|
|
"event_type": "chunk",
|
|
"payload_size_bytes": len(chunk) if chunk else 0,
|
|
"sequence_num": ctx.sse_seq,
|
|
},
|
|
)
|
|
ctx.sse_seq += 1
|
|
yield chunk
|
|
completed = True
|
|
except BaseException as exc: # noqa: BLE001 - observe abort then re-raise
|
|
# GeneratorExit (disconnect) and CancelledError carry no
|
|
# message; an in-iterator error would. Pick a termination_kind.
|
|
terminated_kind = (
|
|
"rst"
|
|
if isinstance(exc, (GeneratorExit,))
|
|
else (
|
|
"timeout"
|
|
if isinstance(exc, asyncio.CancelledError)
|
|
else "chunk_error"
|
|
)
|
|
)
|
|
raise
|
|
finally:
|
|
if completed:
|
|
emit_backend_boundary(
|
|
"backend.response.complete",
|
|
outcome="ok",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
duration_ms=_elapsed_ms(ctx.ingress_mono_ns),
|
|
edge_headers=edge,
|
|
metadata={
|
|
"http_status": status,
|
|
"content_length": ctx.bytes_streamed,
|
|
"total_duration_ms": _elapsed_ms(ctx.ingress_mono_ns),
|
|
"sse_event_count": ctx.sse_seq,
|
|
},
|
|
)
|
|
else:
|
|
emit_backend_boundary(
|
|
"backend.sse.aborted",
|
|
outcome="err",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
edge_headers=edge,
|
|
metadata={
|
|
"termination_kind": terminated_kind,
|
|
"bytes_before_abort": ctx.bytes_streamed,
|
|
},
|
|
)
|
|
|
|
response.body_iterator = _instrumented()
|
|
return response
|
|
|
|
|
|
def _int_or_none(raw: Any) -> Optional[int]:
|
|
"""Parse an int header value, returning ``None`` on absence / malformed."""
|
|
if raw is None:
|
|
return None
|
|
try:
|
|
return int(raw)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _elapsed_ms(start_mono_ns: int) -> int:
|
|
"""Whole milliseconds elapsed since a ``time.monotonic_ns`` start mark."""
|
|
return max(0, (time.monotonic_ns() - start_mono_ns) // 1_000_000)
|
|
|
|
|
|
# ── Agent + LLM boundaries ──────────────────────────────────────────────────
|
|
|
|
# The LLM-call boundaries (start / heartbeat / response) and the agent
|
|
# enter/exit boundaries are emitted via the explicit helpers below. They are
|
|
# called from the agent factory's hook points (strands ``HookProvider``) and
|
|
# from the outbound httpx event hook, all keyed on the request ``ctx``.
|
|
|
|
|
|
def emit_agent_enter(ctx: "_RequestCtx", *, agent_name: str, model_id: str) -> None:
|
|
"""Emit ``backend.agent.enter`` (default tier)."""
|
|
emit_backend_boundary(
|
|
"backend.agent.enter",
|
|
outcome="info",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
metadata={"agent_name": agent_name, "model_id": model_id},
|
|
)
|
|
|
|
|
|
def emit_agent_exit(
|
|
ctx: "_RequestCtx", *, terminal_outcome: str, total_duration_ms: int
|
|
) -> None:
|
|
"""Emit ``backend.agent.exit`` (default tier)."""
|
|
emit_backend_boundary(
|
|
"backend.agent.exit",
|
|
outcome="ok" if terminal_outcome == "ok" else "err",
|
|
test_id=ctx.test_id,
|
|
slug=ctx.slug,
|
|
demo=ctx.demo,
|
|
duration_ms=total_duration_ms,
|
|
metadata={
|
|
"terminal_outcome": terminal_outcome,
|
|
"total_duration_ms": total_duration_ms,
|
|
},
|
|
)
|
|
|
|
|
|
class LlmCallScope:
|
|
"""Async context manager spanning one outbound LLM call.
|
|
|
|
On ``__aenter__`` emits ``backend.llm.call.start`` and launches a heartbeat
|
|
task that emits ``backend.llm.call.heartbeat`` every ``interval_s`` (≈10s)
|
|
while the call is outstanding (verbose tier). On ``__aexit__`` emits
|
|
``backend.llm.call.response`` with the measured latency.
|
|
|
|
All emission is gated/tiered through ``emit_backend_boundary``, so with the
|
|
emitter off or at default tier this scope is effectively free (the
|
|
heartbeat task still ticks but every emit is suppressed; callers that want
|
|
zero task overhead can skip the scope when ``cvdiag_backend_enabled()`` is
|
|
false).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
ctx: "_RequestCtx",
|
|
*,
|
|
provider: str,
|
|
model: str,
|
|
prompt_token_count_estimate: int = 0,
|
|
interval_s: float = 10.0,
|
|
) -> None:
|
|
self._ctx = ctx
|
|
self._provider = provider
|
|
self._model = model
|
|
self._prompt_tokens = prompt_token_count_estimate
|
|
self._interval_s = interval_s
|
|
self._start_mono_ns = 0
|
|
self._hb_task: Optional[asyncio.Task] = None
|
|
|
|
async def __aenter__(self) -> "LlmCallScope":
|
|
self._start_mono_ns = time.monotonic_ns()
|
|
emit_backend_boundary(
|
|
"backend.llm.call.start",
|
|
outcome="info",
|
|
test_id=self._ctx.test_id,
|
|
slug=self._ctx.slug,
|
|
demo=self._ctx.demo,
|
|
metadata={
|
|
"provider": self._provider,
|
|
"model": self._model,
|
|
"prompt_token_count_estimate": self._prompt_tokens,
|
|
},
|
|
)
|
|
self._hb_task = asyncio.ensure_future(self._heartbeat())
|
|
return self
|
|
|
|
async def _heartbeat(self) -> None:
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(self._interval_s)
|
|
emit_backend_boundary(
|
|
"backend.llm.call.heartbeat",
|
|
outcome="info",
|
|
test_id=self._ctx.test_id,
|
|
slug=self._ctx.slug,
|
|
demo=self._ctx.demo,
|
|
metadata={
|
|
"elapsed_ms_since_start": _elapsed_ms(self._start_mono_ns)
|
|
},
|
|
)
|
|
except asyncio.CancelledError: # normal shutdown on call completion
|
|
return
|
|
|
|
async def __aexit__(self, exc_type, exc, tb) -> bool:
|
|
if self._hb_task is not None:
|
|
hb_task = self._hb_task
|
|
self._hb_task = None
|
|
hb_task.cancel()
|
|
try:
|
|
await hb_task
|
|
except asyncio.CancelledError:
|
|
# Cooperative cancellation (was ``except (CancelledError,
|
|
# Exception)``, which swallowed the CALLER's cancel and broke
|
|
# cooperative cancellation). Suppress ONLY the heartbeat task's
|
|
# OWN cancellation — the one we just requested. If THIS task is
|
|
# being cancelled by the caller (a pending cancellation request,
|
|
# ``current_task().cancelling() > 0``), the CancelledError is the
|
|
# caller's and MUST propagate. ``Task.cancelling()`` is 3.11+
|
|
# (production runs 3.12); on older runtimes the attribute is
|
|
# absent and we degrade to suppressing (the legacy behavior).
|
|
current = asyncio.current_task()
|
|
cancelling = getattr(current, "cancelling", None)
|
|
if current is not None and cancelling is not None and cancelling() > 0:
|
|
raise
|
|
except Exception: # noqa: BLE001 - heartbeat body must never throw out
|
|
pass
|
|
emit_backend_boundary(
|
|
"backend.llm.call.response",
|
|
outcome="err" if exc_type is not None else "ok",
|
|
test_id=self._ctx.test_id,
|
|
slug=self._ctx.slug,
|
|
demo=self._ctx.demo,
|
|
duration_ms=_elapsed_ms(self._start_mono_ns),
|
|
metadata={
|
|
"provider": self._provider,
|
|
"model": self._model,
|
|
"response_token_count": None,
|
|
"latency_ms": _elapsed_ms(self._start_mono_ns),
|
|
"error_class": type(exc).__name__ if exc is not None else None,
|
|
},
|
|
)
|
|
return False # never suppress the underlying exception
|