From 32aa81b1f13caf795a1bfb3c72d1d1d34b4ddb9a Mon Sep 17 00:00:00 2001 From: Max Ind Date: Mon, 2 Nov 2026 06:21:13 -0800 Subject: [PATCH] fix(telemetry): Record the chunks of a stream cut short before its aggregate Since the inference span records a stream's `partial=False` aggregate rather than its `partial=True` chunks, a stream that failed before the aggregate arrived recorded no output message at all. Previously its chunks were recorded. When the model call ends with the span still open, the skipped chunks are now recorded, as one output message. Reported in https://github.com/google/adk-python/pull/7815#issuecomment-6504671789. Co-authored-by: Max Ind PiperOrigin-RevId: 1012119875 --- src/google/adk/telemetry/_instrumentation.py | 15 +++++++- .../telemetry/test_instrumentation.py | 36 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 270cd732cc9..36cb5b70cfa 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -611,6 +611,16 @@ def record_llm_response( exit_stack.close() self._inference_span_ended = True + def _record_unfinished_stream( + self, invocation_context: InvocationContext + ) -> None: + """Records the chunks of a stream cut short before its aggregate.""" + if self._inference_span_ended: + return + for response in self._llm_responses: + if response.partial: + tracing.trace_inference_result(invocation_context, self.span, response) + def _record_chunk_timing(self, response: LlmResponse) -> None: """Records the wait for ``response``, if it is a chunk being timed. @@ -1139,7 +1149,10 @@ async def record_inference_telemetry( else: # A span that is not recorded has no start time. tel_ctx.span_started_ns = time.time_ns() - yield tel_ctx + try: + yield tel_ctx + finally: + tel_ctx._record_unfinished_stream(invocation_context) # pylint: disable=protected-access finally: context_api.detach(token) inference_error = sys.exc_info()[1] diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index 381fae3ec98..8d69aacf73f 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/tests/unittests/telemetry/test_instrumentation.py @@ -1793,6 +1793,42 @@ async def test_chunk_timing_is_left_to_the_genai_instrumentor( assert not telemetry.points(_TIME_PER_OUTPUT_CHUNK) +@pytest.mark.asyncio +async def test_stream_cut_short_before_its_aggregate_records_its_chunks( + telemetry: _Telemetry, monkeypatch: pytest.MonkeyPatch +): + """Chunks wait for the aggregate; without one they are all there is.""" + monkeypatch.setenv( + "OTEL_SEMCONV_STABILITY_OPT_IN", "gen_ai_latest_experimental" + ) + ctx = await _invocation_context(_agent()) + model_response_event = mock.MagicMock() + model_response_event.id = "event-1" + + with pytest.raises(RuntimeError): + async with _instrumentation.record_inference_telemetry( + LlmRequest(model="some-model"), ctx, model_response_event + ) as tel_ctx: + for text in ("text ", "response"): + tel_ctx.record_llm_response( + ctx, + _llm_response( + content=types.Content( + role="model", parts=[types.Part(text=text)] + ), + partial=True, + finish_reason=None, + ), + ) + raise RuntimeError("stream interrupted") + + messages = tel_ctx.span.operation_details_attributes["gen_ai.output.messages"] + assert [part["content"] for part in messages[0]["parts"]] == [ + "text ", + "response", + ] + + # --- record_invocation -----------------------------------------------------