fix(smolagents): suspend stream invocations between reads

instrumentation/opentelemetry-instrumentation-genai-smolagents/.changelog/1006.fixed

@@ -0,0 +1 @@
+Restore the caller's context while a model or agent run stream is unconsumed, so caller work is not parented under the `chat` or `invoke_agent` span.

instrumentation/opentelemetry-instrumentation-genai-smolagents/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
   "opentelemetry-api ~= 1.43",
   "opentelemetry-instrumentation >= 0.64b0, <1",
   "opentelemetry-semantic-conventions >= 0.64b0, <1",
-  "opentelemetry-util-genai >= 1.2b0, <2",
+  "opentelemetry-util-genai >= 1.3b0, <2",
 ]
 
 [project.optional-dependencies]

instrumentation/opentelemetry-instrumentation-genai-smolagents/src/opentelemetry/instrumentation/genai/smolagents/patch.py

@@ -20,7 +20,7 @@
 
 import logging
 from collections.abc import Callable, Generator, Mapping
-from contextlib import ExitStack
+from contextlib import AbstractContextManager, ExitStack
 from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
 
 from smolagents import AgentMaxStepsError
@@ -297,6 +297,7 @@ def __init__(
     ) -> None:
         super().__init__(stream, invocation=invocation)
         self._self_inference = invocation
+        invocation.suspend()
         self._self_capture_content = handler.should_capture_content()
         self._self_content: list[str] = []
         self._self_input_tokens = 0
@@ -305,7 +306,8 @@ def __init__(
 
     def __del__(self) -> None:
         try:
-            self._finalize_success()
+            with self._execution_context():
+                self._finalize_success()
         except BaseException:  # pylint: disable=broad-except
             pass
 
@@ -424,15 +426,23 @@ def __init__(
         # invocation to the base class.
         super().__init__(stream)
         self._self_agent_invocation = invocation
+        invocation.suspend()
         self._self_agent = agent
         self._self_capture_content = capture_content
         self._self_final_output: object = None
         self._self_saw_final = False
         self._self_finished = False
 
+    def _execution_context(self) -> AbstractContextManager[None]:
+        return self._self_agent_invocation.activate()
+
     def __del__(self) -> None:
         try:
-            self._finalize_success()
+            with self._execution_context():
+                try:
+                    self._self_stream.close()
+                finally:
+                    self._finalize_success()
         except BaseException:  # pylint: disable=broad-except
             pass
 

instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_agents.py

@@ -7,6 +7,7 @@
 import inspect
 import json
 import pathlib
+import sys
 import tempfile
 import weakref
 from collections.abc import Generator
@@ -36,7 +37,7 @@
 )
 from opentelemetry.semconv._incubating.metrics import gen_ai_metrics
 from opentelemetry.semconv.attributes import error_attributes
-from opentelemetry.trace import StatusCode
+from opentelemetry.trace import INVALID_SPAN, StatusCode, get_current_span
 
 from .test_utils import (
     BrokenTool,
@@ -309,6 +310,18 @@ def test_abandoned_streaming_run_does_not_leak_context(
     assert unrelated.parent is None
 
 
+def test_abandoned_streaming_run_closes_the_run(
+    instrument_with_content, span_exporter, monkeypatch
+) -> None:
+    unraisable: list[Any] = []
+    monkeypatch.setattr(sys, "unraisablehook", unraisable.append)
+    agent = CodeAgent(tools=[], model=FakeCodeModel(), max_steps=3)
+    _first_step_only(agent)
+    gc.collect()
+
+    assert unraisable == []
+
+
 def test_code_agent_streaming_is_lazy_until_drained(
     instrument_with_content, span_exporter
 ) -> None:
@@ -330,6 +343,42 @@ def test_code_agent_streaming_is_lazy_until_drained(
     assert attr(agent_span, error_attributes.ERROR_TYPE) is None
 
 
+class _SpanningCodeModel(FakeCodeModel):
+    def __init__(self, tracer: Any) -> None:
+        super().__init__()
+        self._tracer = tracer
+
+    def generate(self, *args: Any, **kwargs: Any) -> Any:
+        with self._tracer.start_as_current_span("model work"):
+            return super().generate(*args, **kwargs)
+
+
+def test_streaming_run_is_current_only_while_read(
+    instrument_with_content, span_exporter, tracer_provider
+) -> None:
+    tracer = tracer_provider.get_tracer("test")
+    agent = CodeAgent(tools=[], model=_SpanningCodeModel(tracer), max_steps=3)
+    stream = agent.run("Test question", stream=True)
+    assert get_current_span() is INVALID_SPAN
+
+    next(stream)
+    assert get_current_span() is INVALID_SPAN
+    with tracer.start_as_current_span("caller"):
+        pass
+    list(stream)
+
+    spans = span_exporter.get_finished_spans()
+    (agent_span,) = spans_by_operation(spans, "invoke_agent")
+    (caller,) = [span for span in spans if span.name == "caller"]
+    model_work = [span for span in spans if span.name == "model work"]
+    assert caller.parent is None
+    assert model_work
+    assert all(
+        span.parent.span_id == agent_span.context.span_id
+        for span in model_work
+    )
+
+
 def test_streaming_run_stays_a_generator(
     instrument_with_content, span_exporter
 ) -> None:

instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_models.py

@@ -44,7 +44,7 @@
     error_attributes,
     server_attributes,
 )
-from opentelemetry.trace import StatusCode
+from opentelemetry.trace import INVALID_SPAN, StatusCode, get_current_span
 from opentelemetry.util.genai.types import BlobPart, TextPart, UriPart
 
 from .test_utils import (
@@ -652,6 +652,45 @@ def test_generate_stream_is_lazy_and_records_the_drained_response(
     assert attr(span, GenAI.GEN_AI_RESPONSE_FINISH_REASONS) is None
 
 
+def test_generate_stream_is_current_only_while_read(
+    instrument_with_content, span_exporter, tracer_provider
+) -> None:
+    model = transformers_model(stream_chunks=["In ", "Paris"])
+    stream = model.generate_stream(messages=MESSAGES)
+    assert get_current_span() is INVALID_SPAN
+
+    next(stream)
+    assert get_current_span() is INVALID_SPAN
+    with tracer_provider.get_tracer("test").start_as_current_span("caller"):
+        pass
+    list(stream)
+
+    (caller,) = [
+        span
+        for span in span_exporter.get_finished_spans()
+        if span.name == "caller"
+    ]
+    assert caller.parent is None
+
+
+def test_generate_stream_reads_under_the_chat_span(
+    instrument_with_content, span_exporter
+) -> None:
+    model = transformers_model()
+    seen = []
+
+    def chunks():
+        for text in ["In ", "Paris"]:
+            seen.append(get_current_span().get_span_context().span_id)
+            yield text
+
+    model.streamer = chunks()
+    list(model.generate_stream(messages=MESSAGES))
+
+    (chat,) = spans_by_operation(span_exporter.get_finished_spans(), "chat")
+    assert seen == [chat.context.span_id] * 2
+
+
 def test_generate_stream_stays_a_generator(
     instrument_with_content, span_exporter
 ) -> None: