fix(telemetry): Record the chunks of a stream cut short before its aggregate

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]

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 -----------------------------------------------------