From 14079ed7ca0bd583001fcf97cb8e0541fc9e4e6b Mon Sep 17 00:00:00 2001 From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:15:43 -0700 Subject: [PATCH] feat(hooks): correlate llm:request/llm:response with a request_id `llm:request` and `llm:response` shared no field identifying the call they belonged to, so every consumer had to pair them positionally (FIFO over the event stream). That silently mis-attributes whenever a second LLM call is in flight -- a background summarizer, a session-naming hook, a forked sub-agent. Both events still parse and the counts still look plausible, so the error is invisible. Measured on real captures, FIFO crossed 12-31 pairs per run and put a summarizer's cost on the agent; the resulting figure survived six probes before anyone noticed it was wrong. `HookRegistry.emit()` now stamps a client-generated correlation id: llm:request -> generates request_id (uuid4), opens the call llm:response, provider:error -> echo it, then close the call provider:retry, provider:throttle -> echo it, call stays open everything else -> untouched The error path matters as much as the happy one: a call that times out never emits a response at all, which is precisely the case positional pairing gets most wrong. `provider:error` now carries the id of the request it belongs to. Scoping is by contextvars, so the id follows the async task that issued the call -- two concurrent calls hold two independent slots. The stamp happens in the PyO3 bridge, on the caller's Python stack, because the spawned future runs off-thread and no longer has the emitting task's context. Providers need no change. The policy lives in `amplifier_core.correlation` and is the single authority on which events carry an id; the Rust bridge only applies it, behind a cheap `llm:`/`provider:` prefix filter. Backward compatible in both directions: - Consumers that ignore `request_id` see an otherwise identical payload; nothing was renamed, moved, or removed. - Captures already on disk have no `request_id` and remain readable exactly as before -- the kernel does not rewrite history. Consumers treat the field as optional and keep their prior heuristic as the fallback. - A provider supplying its own `request_id` keeps it; the kernel adopts that value for the rest of the call and never overwrites it. - When a response fires with no matching request in its context, no id is stamped. An absent id is always preferable to a wrong one. Out-of-process (gRPC/WASM) providers do not share the kernel's Python context and must supply `request_id` themselves; documented in PROVIDER_CONTRACT.md. --- bindings/python/src/correlation.rs | 98 ++++++++ bindings/python/src/hooks.rs | 10 + bindings/python/src/lib.rs | 1 + docs/HOOKS_API.md | 65 ++++++ docs/contracts/PROVIDER_CONTRACT.md | 27 +++ docs/specs/PROVIDER_SPECIFICATION.md | 9 + python/amplifier_core/correlation.py | 143 ++++++++++++ tests/test_hooks_request_id.py | 338 +++++++++++++++++++++++++++ 8 files changed, 691 insertions(+) create mode 100644 bindings/python/src/correlation.rs create mode 100644 python/amplifier_core/correlation.py create mode 100644 tests/test_hooks_request_id.py diff --git a/bindings/python/src/correlation.rs b/bindings/python/src/correlation.rs new file mode 100644 index 00000000..6d45fa78 --- /dev/null +++ b/bindings/python/src/correlation.rs @@ -0,0 +1,98 @@ +// --------------------------------------------------------------------------- +// Correlation id stamping for LLM call events +// --------------------------------------------------------------------------- +// +// `llm:request` and `llm:response` carried no field in common that identified +// the call they belonged to, so every consumer paired them *positionally*. +// Any concurrently-issued call (a background summarizer, a session-naming +// hook, a forked sub-agent) silently mis-files the pairing: both events parse, +// the counts look right, and the cost attribution is wrong. +// +// The kernel closes that gap on the emit path, so providers need no change: +// `llm:request` gets a generated `request_id` and the matching terminal event +// echoes it. The *policy* -- which events carry an id, and how the in-flight +// call is scoped -- lives in `amplifier_core.correlation` (pure Python), which +// is authoritative. This module is the thin bridge that applies it to the +// serialized event payload. +// +// Scoping is by `contextvars`, so the id follows the async task that issued +// the call. Concurrent calls run in separate tasks and hold separate slots -- +// which is precisely the case positional pairing gets wrong. + +use pyo3::prelude::*; +use serde_json::Value; + +/// Event-data field carrying the correlation id. +/// +/// Mirrors `amplifier_core.correlation.REQUEST_ID_FIELD`; the two are pinned +/// together by `tests/test_hooks_request_id.py`. +pub(crate) const REQUEST_ID_FIELD: &str = "request_id"; + +/// Cheap pre-filter: only the `llm:` and `provider:` families can carry a +/// correlation id, so every other event skips the Python round-trip entirely. +/// +/// This is deliberately *broader* than the real set -- `resolve_request_id` +/// in `amplifier_core.correlation` remains the single authority on which +/// events actually get stamped, so the two cannot drift into disagreement. +fn maybe_correlated(event: &str) -> bool { + event.starts_with("llm:") || event.starts_with("provider:") +} + +/// Stamp the correlation id onto `data` for correlated events. +/// +/// Returns `data` unchanged for every other event, for non-object payloads, +/// and whenever the policy declines to supply an id (e.g. a response with no +/// matching request in this context -- an absent id is always preferable to a +/// wrong one). +/// +/// An explicit `request_id` already present in `data` always wins: the policy +/// adopts it as the id of the call in flight and the value is left untouched. +pub(crate) fn stamp_request_id(py: Python<'_>, event: &str, mut data: Value) -> PyResult { + if !maybe_correlated(event) { + return Ok(data); + } + + let Value::Object(ref mut map) = data else { + return Ok(data); + }; + + let explicit: Option = map + .get(REQUEST_ID_FIELD) + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) + .map(str::to_owned); + + let resolved: Option = py + .import("amplifier_core.correlation")? + .getattr("resolve_request_id")? + .call1((event, explicit.clone()))? + .extract()?; + + if let Some(request_id) = resolved { + if explicit.as_deref() != Some(request_id.as_str()) { + map.insert(REQUEST_ID_FIELD.to_string(), Value::String(request_id)); + } + } + + Ok(data) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prefilter_admits_llm_and_provider_families() { + assert!(maybe_correlated("llm:request")); + assert!(maybe_correlated("llm:response")); + assert!(maybe_correlated("provider:error")); + assert!(maybe_correlated("provider:retry")); + } + + #[test] + fn prefilter_rejects_unrelated_events() { + assert!(!maybe_correlated("tool:pre")); + assert!(!maybe_correlated("session:start")); + assert!(!maybe_correlated("content_block:delta")); + } +} diff --git a/bindings/python/src/hooks.rs b/bindings/python/src/hooks.rs index ba0a7d3a..b9dec4cc 100644 --- a/bindings/python/src/hooks.rs +++ b/bindings/python/src/hooks.rs @@ -121,6 +121,11 @@ impl PyHookRegistry { /// Emit an event and return the aggregated result as a JSON string. /// /// Calls all registered handlers for the event in priority order. + /// + /// For the LLM call event family this also stamps the correlation id + /// (`request_id`) so `llm:request` and its matching terminal event can be + /// paired by identity instead of by position. See + /// [`crate::correlation`] for the policy and its scoping rules. fn emit<'py>( &self, py: Python<'py>, @@ -133,6 +138,11 @@ impl PyHookRegistry { let json_str: String = json_dumps_safe(py, &serializable)?; let value: Value = serde_json::from_str(&json_str) .map_err(|e| PyErr::new::(format!("Invalid JSON: {e}")))?; + // Stamp the correlation id before handlers see the event. Must happen + // here, on the caller's Python stack, because the in-flight call is + // scoped by contextvars -- the spawned future below runs off-thread + // and no longer has the emitting task's context. + let value = crate::correlation::stamp_request_id(py, &event, value)?; wrap_future_as_coroutine( py, diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index c33e0fe2..d2419039 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -23,6 +23,7 @@ use pyo3::prelude::*; mod bridges; mod cancellation; mod coordinator; +mod correlation; mod errors; mod helpers; mod hooks; diff --git a/docs/HOOKS_API.md b/docs/HOOKS_API.md index c68dae99..5f9000c1 100644 --- a/docs/HOOKS_API.md +++ b/docs/HOOKS_API.md @@ -217,6 +217,71 @@ provider on 2026-08-28. --- +## Infrastructure-Owned Event Fields + +`HookRegistry.emit()` stamps a small number of fields onto event data before +any handler sees it. Handlers and event consumers can rely on them being +present without any provider or module doing anything. + +| Field | Events | Owner | Notes | +|-------|--------|-------|-------| +| `timestamp` | all | infrastructure | UTC ISO-8601. Callers cannot omit or override it. | +| `request_id` | LLM call family (below) | infrastructure, caller may override | Correlates one LLM call's events. | + +### `request_id` — LLM call correlation + +`llm:request` and `llm:response` previously shared no field identifying the +call they belonged to, so consumers had to pair them **positionally** (FIFO +over the event stream). That silently mis-attributes whenever a second call is +in flight — a background summarizer, a session-naming hook, a forked +sub-agent. Both events still parse and the counts still look plausible, so the +error is invisible: measured on real captures, positional pairing crossed +12–31 pairs per run and put a summarizer's cost on the agent. + +The kernel now stamps a correlation id on the emit path: + +| Event | Behaviour | +|-------|-----------| +| `llm:request` | Generates a `request_id` (uuid4) and opens the call. | +| `llm:response` | Echoes it exactly, then closes the call. | +| `provider:error` | Echoes it, then closes the call — this is how a call that **times out** stays attributable even though no response ever arrives. | +| `provider:retry`, `provider:throttle` | Echo it without closing the call. | +| everything else | Untouched. | + +**Scoping.** The in-flight call is held in a `contextvars.ContextVar`, so the +id follows the *async task* that issued the call. Two concurrent calls run in +two tasks (`asyncio.gather`, `create_task`, a forked session, a worker thread) +and therefore hold two independent slots — which is exactly the case +positional pairing gets wrong. A provider whose request and response are +emitted from *different* tasks (e.g. a streaming callback on its own task) +will not correlate automatically; such a provider should pass `request_id` +explicitly. + +**Explicit ids win.** A provider that puts its own `request_id` in the event +data keeps it; the kernel adopts that value for the rest of the call and never +overwrites it. + +**Absence is meaningful.** If a response-family event fires with no matching +request in its context, no `request_id` is stamped at all. An absent id is +always preferable to a wrong one. + +**Backward compatibility.** The field is purely additive: + +- Consumers that ignore `request_id` see an otherwise identical payload — + nothing was renamed, moved, or removed. +- Event streams captured before this change carry no `request_id`. They remain + readable exactly as before; the kernel does not rewrite history. Consumers + must treat the field as **optional** (`data.get("request_id")`) and keep + their prior pairing heuristic as the fallback path. +- Providers that already emit their own correlation id are unaffected. + +The policy lives in `amplifier_core.correlation` and is the single authority +on which events carry an id. Module authors can call +`amplifier_core.correlation.current_request_id()` to tag their own logs or +custom events with the enclosing LLM call. + +--- + ## Hook Registration Register hooks to handle specific events. diff --git a/docs/contracts/PROVIDER_CONTRACT.md b/docs/contracts/PROVIDER_CONTRACT.md index 2317b1e1..3817d3f2 100644 --- a/docs/contracts/PROVIDER_CONTRACT.md +++ b/docs/contracts/PROVIDER_CONTRACT.md @@ -171,6 +171,33 @@ coordinator.register_contributor( See [CONTRIBUTION_CHANNELS.md](../specs/CONTRIBUTION_CHANNELS.md) for the pattern. +### `request_id` — Call Correlation (Kernel-Supplied) + +Providers **SHOULD NOT** do anything. The kernel stamps `request_id` onto +`llm:request` on the emit path and echoes the same value onto the matching +`llm:response` (and onto `provider:error` when the call fails or times out), +so a consumer can pair a call's events by identity instead of by position. + +Two requirements fall on providers: + +1. **Emit `llm:request` before the call and `llm:response`/`provider:error` + after it, from the same async task.** Correlation is scoped by + `contextvars`, which is what keeps two concurrent calls apart. A provider + that emits the response from a *different* task (e.g. a streaming callback + scheduled separately) must pass `request_id` explicitly instead. + +2. **Do not overwrite it.** If a provider has a meaningful upstream id of its + own it may put `request_id` in the event data — an explicit value always + wins and the kernel adopts it for the rest of the call. Otherwise leave the + field alone. + +Out-of-process (gRPC/WASM) providers do not share the kernel's Python context +and **MUST** supply `request_id` themselves on both events if they want their +calls correlated. + +See [HOOKS_API.md](../HOOKS_API.md#request_id--llm-call-correlation) for the +full policy and its backward-compatibility guarantees. + ### `llm:response` Event — `usage` Payload Schema Providers **MUST** emit `llm:response` with the following `usage` payload. Key names are normative — derived from the kernel `Usage` struct (`crates/amplifier-core/src/messages.rs`): diff --git a/docs/specs/PROVIDER_SPECIFICATION.md b/docs/specs/PROVIDER_SPECIFICATION.md index cabca8d7..486c517a 100644 --- a/docs/specs/PROVIDER_SPECIFICATION.md +++ b/docs/specs/PROVIDER_SPECIFICATION.md @@ -159,6 +159,15 @@ coordinator.register_contributor( ) ``` +### Call Correlation + +The kernel stamps `request_id` onto `llm:request` and echoes it onto the +matching `llm:response` (and `provider:error`), so consumers pair a call's +events by identity rather than by position. In-process providers get this for +free as long as both events are emitted from the same async task; gRPC/WASM +providers must supply `request_id` themselves. See +[PROVIDER_CONTRACT.md](../contracts/PROVIDER_CONTRACT.md#request_id--call-correlation-kernel-supplied). + ### Debug Levels Support via config flags: diff --git a/python/amplifier_core/correlation.py b/python/amplifier_core/correlation.py new file mode 100644 index 00000000..f1f24ab5 --- /dev/null +++ b/python/amplifier_core/correlation.py @@ -0,0 +1,143 @@ +"""Correlation identity for LLM call events (infrastructure-owned). + +`llm:request` and `llm:response` historically shared **no** field that +identified the call they belonged to. Consumers were forced to pair them +*positionally* (FIFO over the event stream), which silently mis-attributes +every time a second LLM call is in flight concurrently -- a background +summarizer, a session-naming hook, a forked sub-agent. The mis-pairing is +invisible: both events parse, the counts look plausible, and the resulting +per-caller cost attribution is simply wrong. + +This module defines the correlation policy the kernel applies on the emit +path so that **every** provider gets correct pairing without changing a line +of provider code: + +* `llm:request` carries a generated ``request_id``. +* The terminal event of the same call (`llm:response`, or `provider:error` + when the call fails or times out) echoes that exact value. +* `provider:retry` / `provider:throttle` echo it too, without ending the call. + +Scoping is by :mod:`contextvars`, so the id follows the *async task* that +issued the call. Two concurrent calls live in two tasks (``asyncio.gather``, +``create_task``, a forked session, a thread) and therefore hold two +independent slots -- which is exactly the case positional pairing gets wrong. + +Backward compatibility +---------------------- +The field is **additive**. Consumers that ignore ``request_id`` see an +otherwise identical payload. Event streams captured before this change carry +no ``request_id`` at all, so consumers must treat it as *optional* +(``data.get("request_id")``) and keep whatever pairing heuristic they used +before as the fallback. A provider that supplies its own ``request_id`` +always wins -- the kernel never overwrites an explicit value. +""" + +from __future__ import annotations + +import uuid +from contextvars import ContextVar + +__all__ = [ + "REQUEST_ID_FIELD", + "REQUEST_EVENTS", + "TERMINAL_EVENTS", + "INTERIM_EVENTS", + "new_request_id", + "current_request_id", + "resolve_request_id", + "reset_request_id", +] + +#: Name of the correlation field stamped onto event data. +REQUEST_ID_FIELD = "request_id" + +#: Events that *open* a call and generate the correlation id. +REQUEST_EVENTS = frozenset({"llm:request"}) + +#: Events that *close* a call. They echo the id, then end the call so a later +#: unrelated event in the same task cannot inherit a stale id. +TERMINAL_EVENTS = frozenset({"llm:response", "provider:error"}) + +#: Events that echo the id of an in-flight call without ending it. +INTERIM_EVENTS = frozenset({"provider:retry", "provider:throttle"}) + +# (request_id, in_flight). `in_flight` is True between the request event and +# the terminal event of the same call. A closed slot is never read again -- +# absence of a correlation id is always preferable to a wrong one. +_CALL: ContextVar[tuple[str, bool] | None] = ContextVar( + "amplifier_core_llm_call", default=None +) + + +def new_request_id() -> str: + """Generate a fresh correlation id. + + Client-generated (uuid4) on purpose: a provider-assigned id only exists + *after* the response comes back, which is far too late to stamp onto the + request -- and is absent entirely when the call times out. + """ + return str(uuid.uuid4()) + + +def current_request_id() -> str | None: + """Correlation id of the in-flight LLM call in this context, if any. + + Returns ``None`` when no call is in flight (including after the call's + terminal event). Useful for module authors who want to tag their own + logs or custom events with the enclosing call. + """ + call = _CALL.get() + if call is None or not call[1]: + return None + return call[0] + + +def reset_request_id() -> None: + """Clear the correlation slot for this context (test/teardown helper).""" + _CALL.set(None) + + +def resolve_request_id(event: str, explicit: str | None = None) -> str | None: + """Return the correlation id to stamp on ``event``, or ``None``. + + This is the whole policy, and it is deliberately the *only* place that + decides. The kernel emit path calls it for every event in the ``llm:`` + and ``provider:`` families; this function is authoritative about which + of those actually carry a correlation id. + + Args: + event: Event name being emitted. + explicit: A ``request_id`` the caller already put in the event data. + An explicit value always wins and is adopted as the id of the + call in flight. + + Returns: + The id to stamp, or ``None`` when this event carries no correlation + id (unknown event, or a response with no matching request in this + context). + """ + if event in REQUEST_EVENTS: + if explicit: + _CALL.set((explicit, True)) + return explicit + request_id = new_request_id() + _CALL.set((request_id, True)) + return request_id + + if event in TERMINAL_EVENTS: + if explicit: + _CALL.set((explicit, False)) + return explicit + request_id = current_request_id() + if request_id is None: + return None + _CALL.set((request_id, False)) + return request_id + + if event in INTERIM_EVENTS: + if explicit: + _CALL.set((explicit, True)) + return explicit + return current_request_id() + + return None diff --git a/tests/test_hooks_request_id.py b/tests/test_hooks_request_id.py new file mode 100644 index 00000000..89657a23 --- /dev/null +++ b/tests/test_hooks_request_id.py @@ -0,0 +1,338 @@ +"""Tests for LLM call correlation (`request_id`) stamping in HookRegistry.emit(). + +`llm:request` and `llm:response` shared no field identifying the call they +belonged to, so consumers paired them positionally (FIFO). Measured on real +captures, any concurrently-issued call -- a background summarizer, a +session-naming hook -- crosses the pairing silently: both events parse and the +resulting cost attribution is simply wrong. + +These tests pin the fix at both levels: + +* the policy in `amplifier_core.correlation` (pure Python, exhaustive), and +* the emit path in `HookRegistry.emit()` (end-to-end, including the + concurrency case FIFO gets wrong). +""" + +import asyncio +import uuid + +import pytest +from amplifier_core import correlation +from amplifier_core.correlation import REQUEST_ID_FIELD +from amplifier_core.correlation import current_request_id +from amplifier_core.correlation import new_request_id +from amplifier_core.correlation import resolve_request_id +from amplifier_core.hooks import HookRegistry +from amplifier_core.models import HookResult + + +@pytest.fixture(autouse=True) +def _clean_correlation_slot(): + """Each test starts with no call in flight.""" + correlation.reset_request_id() + yield + correlation.reset_request_id() + + +def _recorder(sink): + async def handler(event, data): + sink.append((event, dict(data))) + return HookResult(action="continue") + + return handler + + +# --------------------------------------------------------------------------- +# Policy: amplifier_core.correlation +# --------------------------------------------------------------------------- + + +def test_field_name_is_request_id(): + """The field name is contract surface -- pinned here and in the Rust bridge.""" + assert REQUEST_ID_FIELD == "request_id" + + +def test_new_request_id_is_a_uuid4_string(): + value = new_request_id() + assert isinstance(value, str) + parsed = uuid.UUID(value) + assert parsed.version == 4 + assert str(parsed) == value + + +def test_request_generates_and_response_echoes(): + request_id = resolve_request_id("llm:request") + assert request_id + assert resolve_request_id("llm:response") == request_id + + +def test_each_request_gets_a_distinct_id(): + first = resolve_request_id("llm:request") + resolve_request_id("llm:response") + second = resolve_request_id("llm:request") + assert first != second + + +def test_explicit_request_id_wins_and_is_adopted(): + assert resolve_request_id("llm:request", "provider-supplied") == "provider-supplied" + assert resolve_request_id("llm:response") == "provider-supplied" + + +def test_response_without_a_request_has_no_id(): + """Absent is always better than wrong -- never invent a pairing.""" + assert resolve_request_id("llm:response") is None + + +def test_closed_call_is_not_reused_by_a_later_event(): + """A stale id would silently mis-attribute; the slot closes on the terminal event.""" + request_id = resolve_request_id("llm:request") + assert resolve_request_id("llm:response") == request_id + assert resolve_request_id("llm:response") is None + assert resolve_request_id("provider:throttle") is None + + +def test_error_path_echoes_the_request_id(): + """A call that times out has no response -- provider:error carries the id instead.""" + request_id = resolve_request_id("llm:request") + assert resolve_request_id("provider:error") == request_id + + +def test_retry_and_throttle_echo_without_ending_the_call(): + request_id = resolve_request_id("llm:request") + assert resolve_request_id("provider:retry") == request_id + assert resolve_request_id("provider:throttle") == request_id + assert resolve_request_id("llm:response") == request_id + + +def test_unrelated_events_carry_no_correlation_id(): + resolve_request_id("llm:request") + for event in ("tool:pre", "session:start", "content_block:delta", "provider:resolve"): + assert resolve_request_id(event) is None + + +def test_current_request_id_tracks_the_in_flight_call(): + assert current_request_id() is None + request_id = resolve_request_id("llm:request") + assert current_request_id() == request_id + resolve_request_id("llm:response") + assert current_request_id() is None + + +# --------------------------------------------------------------------------- +# Emit path: HookRegistry.emit() +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_emit_stamps_request_id_on_llm_request(): + registry = HookRegistry() + seen = [] + registry.register("llm:request", _recorder(seen), name="capture") + + await registry.emit("llm:request", {"provider": "anthropic", "model": "m"}) + + assert len(seen) == 1 + request_id = seen[0][1].get(REQUEST_ID_FIELD) + assert request_id, "llm:request must carry a non-null request_id" + assert uuid.UUID(request_id).version == 4 + + +@pytest.mark.asyncio +async def test_emit_pairs_request_and_response_by_id(): + registry = HookRegistry() + seen = [] + registry.register("llm:request", _recorder(seen), name="req") + registry.register("llm:response", _recorder(seen), name="resp") + + await registry.emit("llm:request", {"provider": "anthropic"}) + await registry.emit("llm:response", {"provider": "anthropic", "usage": {}}) + + request_id = seen[0][1][REQUEST_ID_FIELD] + assert request_id + assert seen[1][1][REQUEST_ID_FIELD] == request_id + + +@pytest.mark.asyncio +async def test_emit_preserves_an_explicit_request_id(): + registry = HookRegistry() + seen = [] + registry.register("llm:request", _recorder(seen), name="req") + registry.register("llm:response", _recorder(seen), name="resp") + + await registry.emit("llm:request", {REQUEST_ID_FIELD: "upstream-42"}) + await registry.emit("llm:response", {}) + + assert seen[0][1][REQUEST_ID_FIELD] == "upstream-42" + assert seen[1][1][REQUEST_ID_FIELD] == "upstream-42" + + +@pytest.mark.asyncio +async def test_emit_leaves_unrelated_events_untouched(): + """Consumers of every other event see a byte-identical payload.""" + registry = HookRegistry() + seen = [] + registry.register("tool:pre", _recorder(seen), name="tool") + + await registry.emit("tool:pre", {"tool": "read_file"}) + + assert REQUEST_ID_FIELD not in seen[0][1] + + +@pytest.mark.asyncio +async def test_emit_omits_request_id_when_no_call_is_in_flight(): + """An event stream with no request still emits cleanly -- no invented id.""" + registry = HookRegistry() + seen = [] + registry.register("llm:response", _recorder(seen), name="resp") + + await registry.emit("llm:response", {"provider": "anthropic"}) + + assert REQUEST_ID_FIELD not in seen[0][1] + assert seen[0][1]["provider"] == "anthropic" + + +@pytest.mark.asyncio +async def test_emit_carries_request_id_onto_the_error_path(): + """The timeout case: a request with no response must still be attributable.""" + registry = HookRegistry() + seen = [] + registry.register("llm:request", _recorder(seen), name="req") + registry.register("provider:error", _recorder(seen), name="err") + + await registry.emit("llm:request", {"provider": "anthropic"}) + await registry.emit("provider:error", {"error": "timeout after 10s"}) + + assert seen[1][1][REQUEST_ID_FIELD] == seen[0][1][REQUEST_ID_FIELD] + + +@pytest.mark.asyncio +async def test_concurrent_calls_get_distinct_ids_and_pair_correctly(): + """The case FIFO gets wrong. + + Interleaving is forced to reproduce the measured trace (agent request, + summarizer request, summarizer response, agent response) where positional + pairing charges each response to the other caller. + """ + registry = HookRegistry() + seen = [] + for event in ("llm:request", "llm:response"): + registry.register(event, _recorder(seen), name=f"cap-{event}") + + summarizer_requested = asyncio.Event() + summarizer_responded = asyncio.Event() + + async def agent_call(): + await registry.emit("llm:request", {"caller": "agent"}) + await summarizer_responded.wait() + await registry.emit("llm:response", {"caller": "agent"}) + + async def summarizer_call(): + await summarizer_requested.wait() + await registry.emit("llm:request", {"caller": "summarizer"}) + await registry.emit("llm:response", {"caller": "summarizer"}) + summarizer_responded.set() + + async def run(): + task_agent = asyncio.create_task(agent_call()) + task_summarizer = asyncio.create_task(summarizer_call()) + await asyncio.sleep(0) + summarizer_requested.set() + await asyncio.gather(task_agent, task_summarizer) + + await run() + + order = [(event, payload["caller"]) for event, payload in seen] + assert order == [ + ("llm:request", "agent"), + ("llm:request", "summarizer"), + ("llm:response", "summarizer"), + ("llm:response", "agent"), + ], "expected the interleaving that defeats positional pairing" + + by_caller = {} + for event, payload in seen: + by_caller.setdefault(payload["caller"], {})[event] = payload[REQUEST_ID_FIELD] + + agent = by_caller["agent"] + summarizer = by_caller["summarizer"] + + assert agent["llm:request"] == agent["llm:response"] + assert summarizer["llm:request"] == summarizer["llm:response"] + assert agent["llm:request"] != summarizer["llm:request"] + + # And the positional pairing this replaces would have crossed them: + # FIFO joins the first request to the first response, which here belong + # to different callers. + requests = [p[REQUEST_ID_FIELD] for e, p in seen if e == "llm:request"] + responses = [p[REQUEST_ID_FIELD] for e, p in seen if e == "llm:response"] + assert requests[0] != responses[0], "expected FIFO to mis-pair this trace" + + +@pytest.mark.asyncio +async def test_many_concurrent_calls_all_pair_correctly(): + registry = HookRegistry() + seen = [] + for event in ("llm:request", "llm:response"): + registry.register(event, _recorder(seen), name=f"cap-{event}") + + async def one_call(index: int): + await registry.emit("llm:request", {"caller": index}) + # Yield control so every call is genuinely in flight at once. + await asyncio.sleep(0.01) + await registry.emit("llm:response", {"caller": index}) + + await asyncio.gather(*(one_call(i) for i in range(20))) + + ids = {} + for event, payload in seen: + ids.setdefault(payload["caller"], {})[event] = payload[REQUEST_ID_FIELD] + + assert len(ids) == 20 + for caller, pair in ids.items(): + assert pair["llm:request"] == pair["llm:response"], f"caller {caller} mis-paired" + assert len({pair["llm:request"] for pair in ids.values()}) == 20 + + +# --------------------------------------------------------------------------- +# Backward compatibility +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_id_ignoring_consumer_sees_an_otherwise_identical_payload(): + """The change is additive: nothing existing is renamed, moved, or dropped.""" + registry = HookRegistry() + seen = [] + registry.register("llm:response", _recorder(seen), name="resp") + + payload = {"provider": "anthropic", "model": "m", "usage": {"input_tokens": 7}} + await registry.emit("llm:request", {"provider": "anthropic"}) + await registry.emit("llm:response", dict(payload)) + + observed = seen[0][1] + for key, value in payload.items(): + assert observed[key] == value + assert set(observed) - set(payload) <= {REQUEST_ID_FIELD, "timestamp"} + + +def test_historical_capture_without_request_id_still_parses(): + """Every capture already on disk has no request_id -- it must stay readable. + + Consumers treat the field as optional and fall back to their prior + heuristic; the kernel never rewrites history. + """ + historical = [ + {"event": "llm:request", "data": {"provider": "openai", "model": "gpt-5"}}, + {"event": "llm:response", "data": {"provider": "openai", "usage": {}}}, + ] + + for record in historical: + assert record["data"].get(REQUEST_ID_FIELD) is None + + paired = list( + zip( + [r for r in historical if r["event"] == "llm:request"], + [r for r in historical if r["event"] == "llm:response"], + ) + ) + assert len(paired) == 1