refactor(telemetry): migrate stream_with_chunking telemetry to plugin hooks

[ajbozarth]

Pull Request

Issue

Fixes #1290

Description

Migrates stream_with_chunking telemetry to plugin hooks as part of the Phase 2 migration: moves the orchestration span and its streaming metrics off direct mellea.telemetry calls and onto plugin hooks. With this change, no async path in mellea/stdlib/ calls the old trace_application / set_span_* API, so the deprecated helpers can finally be removed.

Span and metrics on hooks. New streaming hook types in mellea/plugins/types.py (with payload dataclasses in mellea/plugins/hooks/streaming.py) carry the data previously stamped inline — chunk index, requirement counts, pass/fail, exceptions — keyed by a streaming_id so pre/post hooks pair cleanly across the _run_async_in_thread tasks. A new StreamingTracingPlugin subscribes to these hooks and emits the orchestration span via internal helpers in mellea/telemetry/tracing.py; the existing metrics plugins gain streaming subscriptions, replacing the inline record_* calls. Metrics record fire-and-forget so they never block the stream. Spans, events, and metric names/attributes match main — no observability regression — and the span still nests correctly under the action span when called via a session.

Cleanup. With streaming migrated, the four deprecated public tracing helpers (trace_application, set_span_attribute, set_span_error, set_span_status_error) are removed; the internal setter logic moves to _tracing_setters.py.

Bug fix surfaced during migration. The backend chat span attaches in the caller task but finishes in the orchestration task that drains the MOT, so its OTel context detach crossed asyncio tasks and failed — leaving the generation span ambient and mis-nesting validation chat spans under generation instead of as its sibling. The fix adds orchestration-task hooks that re-attach the stream_with_chunking span as ambient context for the drain/validate loop, and _safe_detach now skips cross-task detaches within a reattach scope. See 6c3de9e5 for details.

Follow-up bug fix. functional.py now catches BaseException (not just Exception) in the component error path, so CancelledError/KeyboardInterrupt fire COMPONENT_POST_ERROR (closing the action span) before propagating. Caught while writing this PR; the handler re-raises, so nothing is swallowed.

Docs. docs/docs/observability/tracing.md documents the stream_with_chunking span (attributes + lifecycle events) and shows the two chat spans as siblings under it.

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[ajbozarth]

I investigated the flaky tests, and I'm pretty sure they stem from the same issue as #1367 which I fixed in #1369

If they are still flaky in the future I will address them then.

[ajbozarth]

FYI I'm working on solving the CI failure. tl;dr python 3.11 wait_for() creates a child task for the asynchronous function it calls, 3.12 and later do not. This breaks nesting trace spans using hooks, so I'm writing up a "fix" for 3.11

[ajbozarth]

Python 3.11 tracing fix

The 3.11 CI job was failing on the streaming tracing tests. Root cause: cpex runs every plugin hook through asyncio.wait_for, which on Python ≤3.11 runs the hook in a child task with a copied context. So any OTel context.attach() done inside a hook is lost when the hook returns, and the later detach fails cross-task (the "Failed to detach context" errors). Python 3.12 changed wait_for to run in the caller's task, so it only breaks on 3.11.

Since we can't attach context from inside a hook on 3.11, I gate that attach/detach off there — spans still emit, just flattened instead of nested. Directly-attached spans (like session) are unaffected, since normal parent→child context propagation still works. The gate lives in the plugin layer that actually knows it's running in a hook; the span helpers stay version-agnostic.

I also added a session → action → chat nesting test through the real hook path. The existing nesting tests call the hook methods directly, so they never hit cpex's wait_for boundary — which is why this slipped through.

[psschwei]

seems test is still failing.

Codex summary (in case it helps):

• [High] Python 3.11 CI is failing because span context attach is disabled too broadly.

  • Location: mellea/telemetry/tracing_plugins.py:22, lines 22-27 and 77-87.

  • Snippet:

    _CONTEXT_ATTACH_SUPPORTED: bool = sys.version_info >= (3, 12)
    ...
    attach_context=_CONTEXT_ATTACH_SUPPORTED,

  • Why it matters: on Python 3.11, BackendTracingPlugin.on_pre_call() starts and stashes a backend span but never makes it current context. The failing test then starts nested spans and correctly observes parent is None.

  • Evidence: CI and local 3.11 both fail test/telemetry/test_tracing_plugins.py::test_nested_span_during_call_parents_under_backend_span at line 581. Local 3.12 passes; local 3.11 module run is 1 failed, 31 passed.

  • Suggested fix: either restore context attachment where the hook is not crossing cpex’s wait_for() boundary, or make the 3.11 flattening behavior explicit and update this test. If flattening is intentional, the PR should treat it as a supported-runtime observability regression, not only a test expectation issue.

Why Tests are failing: The active CI failure is Python 3.11-specific. The PR added _CONTEXT_ATTACH_SUPPORTED = sys.version_info >= (3, 12), so on 3.11 the backend span is not attached as current context. The test still expects nested spans to parent under that backend span, so by_name["nested-caller-task"].parent is None.

[psschwei]

And from claude:

CI failure root cause

Only code-checks / quality (3.11) actually failed. The 3.12 and 3.13 matrix jobs show CANCELLED — that's fail-fast killing them after 3.11 failed, not independent failures. Everything else (docs build, actionlint, DCO, label, hold) is green.

One test failed:

FAILED test/telemetry/test_tracing_plugins.py::test_nested_span_during_call_parents_under_backend_span
  - assert None is not None
  +  where None = <ReadableSpan ...>.parent
= 1 failed, 3324 passed, 169 skipped ...

Root cause — regression from the last commit

1005eb14 fix(telemetry): gate hook-scoped span context attach on Python 3.12+ added this gate in mellea/telemetry/tracing_plugins.py:

_CONTEXT_ATTACH_SUPPORTED: bool = sys.version_info >= (3, 12)

and made BackendTracingPlugin.on_pre_call pass attach_context=_CONTEXT_ATTACH_SUPPORTED into start_backend_span. On Python 3.11 that value is False, so start_backend_span skips attaching the chat span as ambient OTel context.

The failing test (test/telemetry/test_tracing_plugins.py:539) starts two nested spans between on_pre_call and on_post_call and asserts they parent under the backend span:

assert by_name["nested-caller-task"].parent is not None          # ← fails on 3.11
assert by_name["nested-caller-task"].parent.span_id == backend_span_id

With no ambient attach on 3.11, those nested spans have no parent → parent is None → assert None is not None.

Why the gate is too broad for this test

The commit's rationale is correct for the real hook path — cpex runs hooks through asyncio.wait_for, which on ≤3.11 copies the contextvars Context into a child task, so an in-hook attach is lost. But this test calls backend_plugin.on_pre_call(...) directly (line 555), not through cpex's wrapper. In a direct call the attach would propagate fine on 3.11 — yet the version gate unconditionally disables it, so the test that specifically asserts nesting breaks.

Confirmed locally: _CONTEXT_ATTACH_SUPPORTED is False for (3, 11), and the test passes on 3.12 (where the gate is True). It was added in bbddd0ea (already on main) and passed until this branch's gate commit.

Fix options

The behavior change (flat spans on ≤3.11) is intentional, so the test should reflect it:

1. Guard the test's nesting assertions by version — on ≤3.11 assert the spans emit but are flat (parent is None); on ≥3.12 assert they nest. Matches the documented intent ("hook spans are emitted flat (no nesting)" on ≤3.11) and keeps coverage on both. Preferred, since the commit explicitly promises "spans still emit, just flat."
2. Skip the test on <3.12 with @pytest.mark.skipif(sys.version_info < (3, 12), reason=...) — simpler, but loses the flat-emission assertion.

One thing worth flagging: it's slightly surprising that a directly-invoked on_pre_call also loses the attach. If any production path invokes these hooks directly (rather than via cpex's wait_for), it would get flat spans on 3.11 too. If that's a real path, gating inside start_backend_span on "am I running inside a copied child task" rather than on Python version would be more precise — but that's a design consideration, not required to unblock CI.

[ajbozarth]

Pushed the CI fix, I missed a test in the plugins file that tested nesting when updating all the nesting test in my last commit

[ajbozarth]

@psschwei I believe all of your comments thus far have been addressed now

[psschwei]

This came up on the last review. I don't know the hooks/telemetry code well enough to be able to really evaluate the problem, but it did survive a few rounds of pushback, so I'm going to post it.

---

Telemetry regression: stream_with_chunking span + metrics leak when acomplete() isn't called

On main, _orchestrate_streaming wrapped its body in with trace_application("stream_with_chunking") as span:, so the span always closed inside the orchestration task's own finally — the telemetry lifecycle was bounded by the task, independent of how the caller consumed the result.

This PR moves STREAMING_END (the only thing that calls finish_streaming_span — which ends the span and pops it from the process-global _in_flight_spans — and fires the outcome/error metrics) to fire only from acomplete(). So a caller that drains astream() or watches events() to the terminal CompletedEvent but never calls acomplete():

• leaks the span in _in_flight_spans forever (unbounded, one entry per run — bad in a long-lived process),
• never exports the span (silently missing from traces),
• under-counts outcome/error metrics,
• and on 3.12+ leaves the caller-task OTel token attached, so later spans mis-parent.

This is a plausible pattern, not misuse — a token-streaming UI genuinely doesn't need acomplete()'s return values, and the astream() docstring's "effectively a no-op" wording implies it's skippable. All the shipped examples happen to call acomplete(), but nothing enforces it, and no test covers the drain-without-acomplete() path.

Fix

Fire STREAMING_END from the orchestrator's finally instead (or in addition, deduped by the existing _streaming_end_fired flag), placed after STREAMING_ORCHESTRATION_END so the reattached span is released before the span is finished. That runs on the orchestration task, on a live loop, at deterministic run-end — restoring the pre-PR guarantee and fixing the metrics under-count for free. acomplete() then just calls the same one-shot helper as a fallback.

All the payload data (completed, full_text, final_validations, _streaming_failure_reason) is already populated before the finally runs. The setup-path STREAMING_END in stream_with_chunking()'s except stays as-is (no orchestration task exists there, so no double-fire).

1. Extract a one-shot helper on StreamChunkingResult so the orchestrator and acomplete() share one firing site:

async def _fire_streaming_end(self) -> None:
    """Fire STREAMING_END exactly once. Safe to call from the orchestrator
    finally and from acomplete(); the first caller wins."""
    if self._streaming_end_fired:
        return
    self._streaming_end_fired = True
    if has_plugins(HookType.STREAMING_END):
        from ..plugins.hooks.streaming import StreamingEndPayload

        await invoke_hook(
            HookType.STREAMING_END,
            StreamingEndPayload(
                streaming_id=self._streaming_id,
                success=self.completed,
                failure_reason=self._streaming_failure_reason,
                exception=self._streaming_end_exception,
                model=self._mot.generation.model,
                provider=self._mot.generation.provider,
                full_text_length=len(self.full_text),
            ),
        )

2. Call it from the orchestrator finally, shielded, after STREAMING_ORCHESTRATION_END and before the terminal queue bookkeeping:

        # ... existing STREAMING_ORCHESTRATION_END invoke_hook (shielded) ...

        # Fire STREAMING_END here (not only in acomplete()) so the span always
        # closes when the orchestration task finishes, regardless of whether the
        # caller drains astream()/events() without calling acomplete(). Shielded
        # like ORCHESTRATION_END above. Runs AFTER ORCHESTRATION_END so the
        # reattached span is released before the span is finished
        # (StreamingTracingPlugin relies on that ordering).
        try:
            await result._fire_streaming_end()
        except BaseException:
            pass

        completed_ev = CompletedEvent(
            success=result.completed, full_text=result.full_text, attempts_used=1
        )
        # ... existing terminal put_nowait / _done.set() ...

3. Simplify acomplete() to call the same helper (now an idempotent fallback):

    async def acomplete(self) -> None:
        ...
        await self._done.wait()
        await self._fire_streaming_end()   # no-op if the orchestrator already fired
        # ... existing raise-once logic ...

Plus a regression test for the drain-astream()-without-acomplete() path asserting the span is exported and _in_flight_spans is empty.

[psschwei]

The other thing that's gnawing at me a little bit (though I know the reasons for it and don't have an alternative to propose) is that we have different behavior in different versions of Python. That just doesn't feel right on some level, but it may be something that just is what it is too.

[ajbozarth]

stream_with_chunking span + metrics leak when acomplete() isn't called

This is actually just Claude not reading deep enough into the code, I have hit this misunderstanding multiple times. acomplete() should always be called, not doing so is bad practice and not the intended use of the API, without calling it other issues than just the telemetry will not be finalized. It is even documented, but due to the obfuscated-ness of the streaming design Claude can't follow it and treats the function as optional.

he other thing that's gnawing at me a little bit (though I know the reasons for it and don't have an alternative to propose) is that we have different behavior in different versions of Python. That just doesn't feel right on some level, but it may be something that just is what it is too.

I agree, but in this case the only "solution" would be to change how the cpex library works to avoid the python version implemented difference, but given apex is in the middle of replacing it's python implementation with rust (which would indirectly solve this issue) I decided to introduce this workaround and remove it when we update cpex to 0.2 in the future (contextforge-org/cpex#20)

[ajbozarth]

I did a sanity check on my above statement with Claude and it went down the same rabbit hole again, so I just added another comment and validated that it prevented Claude from going down that hole again

[ajbozarth]

I also added a comment on the streaming phase 2 epic issue detailing the edge case that Claude keeps getting caught on so I make sure to address it in one of the eventual sub-issues on that epic when I plan it. #1013 (comment)

[psschwei]

Honestly, I think some of my unease stems from hooks not being a great fit for how streaming is usually instrumented: typically you use one span for the whole response and close it when the stream ends (it's what the OTel GenAI semconv does and roughly where #444 is already heading). Because plugin hooks are decoupled from the streaming code, we end up managing the span's lifecycle from the outside instead. Which works... but it isn't ideal.

I don't want to block this PR on that (it's a non-trivial redesign) but I think we should strongly consider "one span per stream" as the target design for streaming, and maybe capture that on #444 so it doesn't get lost. WDYT?

[ajbozarth]

Talked with @psschwei offline and we've decided that any improvements to the design should land with the streaming validation feature itself and not the telemetry. I've outline our discussion in #1013 (comment) and further design improvements would be bundled with that work

[psschwei]

As discussed, deferring design questions to the parent epic.

One small, non-blocking test issue noticed when running the tests locally: There are new "coroutine never awaited" warnings in the streaming test suite. main emits 0; this branch emits 2 (PluginExecutor._with_semaphore / execute_plugin), attributed to test_no_requirements_events_omits_full_validation_event but originating from the new FIRE_AND_FORGET streaming metrics hooks in an earlier test. This is test hygiene, not a production defect — in production cpex schedules these via asyncio.create_task() and they run; the warning only appears because the test loop tears down first without draining. The streaming stdlib tests that fire metrics hooks could drain_background_tasks() / discard_background_tasks() (as test_tracing_backend.py already does) to keep the suite clean.

[ajbozarth]

@psschwei I looked into that nit and it was not what you thought it was. I addressed it in bbdd31f see that commit msg for clarification

[nrfulton]

FYI @jakelorocco maybe also take a look when you get back.