diff --git a/.github/instructions/instrumentation.instructions.md b/.github/instructions/instrumentation.instructions.md index 6e56f01c0..07baefae1 100644 --- a/.github/instructions/instrumentation.instructions.md +++ b/.github/instructions/instrumentation.instructions.md @@ -59,8 +59,8 @@ prefer opt-in or additive. Breaking changes need explicit justification in the P See [request property capture](../../AGENTS.md#capturing-properties-that-may-contain-iterators). - Streaming responses must be instrumented by subclassing the util's `SyncStreamWrapper` / `AsyncStreamWrapper` (`opentelemetry.util.genai.stream`). Flag hand-rolled stream wrappers, and - invocation-backed wrappers that do not pass the invocation to `super().__init__(stream, invocation)` - or do not `invocation.suspend()` before the stream is returned. Flag an `_execution_context()` + invocation-backed wrappers that do not pass the invocation to `super().__init__(stream, invocation)`, + which suspends it before the stream is returned. Flag an `_execution_context()` override on an invocation-backed wrapper: the base wrapper already activates a suspended invocation during reads, and an override bypasses that guard. - Instrumentation should not change what a call returns or when its work happens. Flag: work the SDK diff --git a/AGENTS.md b/AGENTS.md index 9373ef629..12125724b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -264,16 +264,16 @@ finally: A streamed response only finishes once the caller has drained the stream, so the invocation must stay open until then. Do **not** call `invocation.stop()` when the SDK returns the stream — the span would close before any chunks arrive. The invocation's span must also not stay current in -the caller's context while the stream is unconsumed: call `invocation.suspend()` before returning -the stream, and re-activate the invocation only while a chunk is being read. +the caller's context while the stream is unconsumed. The shared stream wrappers suspend the +invocation before returning the stream and re-activate it only during reads and cleanup. Instrument streams by subclassing `SyncStreamWrapper` / `AsyncStreamWrapper` from `opentelemetry.util.genai.stream` (the public, supported helpers). The base class proxies the underlying SDK stream, drives iteration, and finalizes telemetry exactly once on success, error, or `close()`. Subclasses pass the SDK stream and the invocation to -`super().__init__(stream, invocation)`, which also marks the invocation as a streamed request -(`gen_ai.request.stream`), call `invocation.suspend()` before returning the stream, and implement -three hooks. The base `_execution_context()` makes a suspended invocation current again for each +`super().__init__(stream, invocation)`, which suspends it and marks it as a streamed request +(`gen_ai.request.stream`), and implement three hooks. +The base `_execution_context()` makes a suspended invocation current again for each read and cleanup operation and restores the context before the chunk is returned, so invocation-backed wrappers don't override it. @@ -287,7 +287,6 @@ invocation-backed wrappers don't override it. class MyStreamWrapper(SyncStreamWrapper[Chunk]): def __init__(self, stream, invocation, capture_content): super().__init__(stream, invocation) - invocation.suspend() ... def _process_chunk(self, chunk): ... # accumulate state diff --git a/instrumentation/opentelemetry-instrumentation-google-genai/.changelog/899.fixed b/instrumentation/opentelemetry-instrumentation-google-genai/.changelog/899.fixed new file mode 100644 index 000000000..7217cd767 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-google-genai/.changelog/899.fixed @@ -0,0 +1 @@ +Restore caller context between Google GenAI stream reads while preserving invocation context during iteration and cleanup. diff --git a/instrumentation/opentelemetry-instrumentation-google-genai/tests/test_stream_context.py b/instrumentation/opentelemetry-instrumentation-google-genai/tests/test_stream_context.py new file mode 100644 index 000000000..1e6ae4144 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-google-genai/tests/test_stream_context.py @@ -0,0 +1,314 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Iterator +from contextlib import nullcontext +from functools import wraps + +import pytest +from google.genai import Client +from google.genai.models import AsyncModels, Models +from google.genai.types import GenerateContentResponse + +from opentelemetry import context, trace +from opentelemetry.instrumentation.google_genai import ( + GoogleGenAiSdkInstrumentor, +) +from opentelemetry.instrumentation.google_genai.interactions import ( + _HAS_INTERACTIONS, + AsyncInteractionsResource, + InteractionsResource, +) +from opentelemetry.sdk._logs import LoggerProvider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.semconv._incubating.attributes import gen_ai_attributes +from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE +from opentelemetry.test_util_genai.instrumentor import instrument +from opentelemetry.trace import Span, StatusCode + +REQUIRES_INTERACTIONS = pytest.mark.skipif( + not _HAS_INTERACTIONS, + reason="Interactions are not supported in this version of google-genai", +) + + +pytestmark = [ + pytest.mark.parametrize( + "expected_operation", + [ + pytest.param( + "generate_content", id="models.generate_content_stream" + ), + pytest.param( + "interactions.create", + id="interactions.create-model", + marks=REQUIRES_INTERACTIONS, + ), + pytest.param( + "invoke_agent", + id="interactions.create-agent", + marks=REQUIRES_INTERACTIONS, + ), + pytest.param( + "fetch_response", + id="interactions.get", + marks=REQUIRES_INTERACTIONS, + ), + ], + ), + pytest.mark.parametrize( + "ending", ["drain", "caller_error", "stream_error"] + ), +] + + +def _assert_finished_spans( + span_exporter: InMemorySpanExporter, + sdk_span: Span, + caller: Span, + expected_operation: str, + ending: str, +) -> None: + invocation_span, caller_span = span_exporter.get_finished_spans() + assert invocation_span.context == sdk_span.get_span_context() + assert invocation_span.parent == caller.get_span_context() + assert caller_span.context == caller.get_span_context() + assert invocation_span.attributes is not None + assert ( + invocation_span.attributes[gen_ai_attributes.GEN_AI_OPERATION_NAME] + == expected_operation + ) + failed = ending != "drain" + assert invocation_span.status.status_code == ( + StatusCode.ERROR if failed else StatusCode.UNSET + ) + if failed: + assert invocation_span.attributes[ERROR_TYPE] == "ConnectionError" + else: + assert ERROR_TYPE not in invocation_span.attributes + + +def test_sync_stream_context( + monkeypatch: pytest.MonkeyPatch, + tracer_provider: TracerProvider, + meter_provider: MeterProvider, + logger_provider: LoggerProvider, + span_exporter: InMemorySpanExporter, + expected_operation: str, + ending: str, +) -> None: + tracer = tracer_provider.get_tracer(__name__) + sdk_span: Span | None = None + error = ConnectionError("stream failed") + chunk = ( + GenerateContentResponse(response_id="response-id", candidates=[]) + if expected_operation == "generate_content" + else {"event_type": "content"} + ) + + def produce() -> Iterator[object]: + for _ in range(2): + assert trace.get_current_span() is sdk_span + yield chunk + if ending == "stream_error": + raise error + + resource_type = ( + Models + if expected_operation == "generate_content" + else InteractionsResource + ) + method_name = ( + "generate_content_stream" + if expected_operation == "generate_content" + else "get" + if expected_operation == "fetch_response" + else "create" + ) + original = getattr(resource_type, method_name) + source = produce() + + @wraps(original) + def call(*args: object, **kwargs: object) -> Iterator[object]: + nonlocal sdk_span + sdk_span = trace.get_current_span() + return source + + monkeypatch.setattr(resource_type, method_name, call) + # The instrumentor sets this independently of the shared content setting. + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_EMIT_EVENT", "true") + with ( + instrument( + GoogleGenAiSdkInstrumentor(), + tracer_provider=tracer_provider, + meter_provider=meter_provider, + logger_provider=logger_provider, + content_capture="NO_CONTENT", + ), + tracer.start_as_current_span("caller") as caller, + ): + client = Client(api_key="test-key") + before = context.get_current() + if expected_operation == "generate_content": + stream = client.models.generate_content_stream( + model="gemini-2.5-flash", contents="Hello" + ) + elif expected_operation == "fetch_response": + stream = client.interactions.get("response-id", stream=True) + elif expected_operation == "invoke_agent": + stream = client.interactions.create( + agent="test-agent", input="Hello", stream=True + ) + else: + stream = client.interactions.create( + model="gemini-2.5-flash", input="Hello", stream=True + ) + assert sdk_span is not None + assert sdk_span.is_recording() + assert context.get_current() is before + received: list[object] = [] + + with ( + pytest.raises(ConnectionError) + if ending != "drain" + else nullcontext() + ) as raised: + with stream: + for item in stream: + assert item is chunk + assert context.get_current() is before + assert sdk_span.is_recording() + received.append(item) + if ending == "caller_error": + raise error + + if raised is not None: + assert raised.value is error + assert len(received) == (2 if ending == "drain" else 1) + assert context.get_current() is before + assert not sdk_span.is_recording() + + _assert_finished_spans( + span_exporter, sdk_span, caller, expected_operation, ending + ) + + +@pytest.mark.asyncio +async def test_async_stream_context( + monkeypatch: pytest.MonkeyPatch, + tracer_provider: TracerProvider, + meter_provider: MeterProvider, + logger_provider: LoggerProvider, + span_exporter: InMemorySpanExporter, + expected_operation: str, + ending: str, +) -> None: + tracer = tracer_provider.get_tracer(__name__) + sdk_span: Span | None = None + error = ConnectionError("stream failed") + chunk = ( + GenerateContentResponse(response_id="response-id", candidates=[]) + if expected_operation == "generate_content" + else {"event_type": "content"} + ) + + async def produce() -> AsyncIterator[object]: + for _ in range(2): + assert trace.get_current_span() is sdk_span + yield chunk + if ending == "stream_error": + raise error + + resource_type = ( + AsyncModels + if expected_operation == "generate_content" + else AsyncInteractionsResource + ) + method_name = ( + "generate_content_stream" + if expected_operation == "generate_content" + else "get" + if expected_operation == "fetch_response" + else "create" + ) + original = getattr(resource_type, method_name) + source = produce() + + @wraps(original) + async def call(*args: object, **kwargs: object) -> AsyncIterator[object]: + nonlocal sdk_span + sdk_span = trace.get_current_span() + return source + + monkeypatch.setattr(resource_type, method_name, call) + # The instrumentor sets this independently of the shared content setting. + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_EMIT_EVENT", "true") + with ( + instrument( + GoogleGenAiSdkInstrumentor(), + tracer_provider=tracer_provider, + meter_provider=meter_provider, + logger_provider=logger_provider, + content_capture="NO_CONTENT", + ), + tracer.start_as_current_span("caller") as caller, + ): + client = Client(api_key="test-key") + before = context.get_current() + if expected_operation == "generate_content": + stream = await client.aio.models.generate_content_stream( + model="gemini-2.5-flash", contents="Hello" + ) + elif expected_operation == "fetch_response": + stream = await client.aio.interactions.get( + "response-id", stream=True + ) + elif expected_operation == "invoke_agent": + stream = await client.aio.interactions.create( + agent="test-agent", input="Hello", stream=True + ) + else: + stream = await client.aio.interactions.create( + model="gemini-2.5-flash", input="Hello", stream=True + ) + assert sdk_span is not None + assert sdk_span.is_recording() + assert context.get_current() is before + received: list[object] = [] + + async def consume() -> None: + try: + async with stream: + async for item in stream: + assert item is chunk + assert context.get_current() is before + assert sdk_span.is_recording() + received.append(item) + if ending == "caller_error": + raise error + finally: + assert context.get_current() is before + + with ( + pytest.raises(ConnectionError) + if ending != "drain" + else nullcontext() + ) as raised: + await asyncio.create_task(consume()) + + if raised is not None: + assert raised.value is error + assert len(received) == (2 if ending == "drain" else 1) + assert context.get_current() is before + assert not sdk_span.is_recording() + + _assert_finished_spans( + span_exporter, sdk_span, caller, expected_operation, ending + ) diff --git a/tox.ini b/tox.ini index f2de43091..b9c58eecc 100644 --- a/tox.ini +++ b/tox.ini @@ -311,7 +311,7 @@ commands = test-instrumentation-genai-openai_agents-conformance: pytest {toxinidir}/instrumentation/opentelemetry-instrumentation-genai-openai-agents/tests/test_conformance.py --vcr-record=none {posargs} test-instrumentation-google-genai-{oldest,latest}: pytest --ignore={toxinidir}/instrumentation/opentelemetry-instrumentation-google-genai/tests/test_conformance.py {toxinidir}/instrumentation/opentelemetry-instrumentation-google-genai/tests --vcr-record=none {posargs} - test-instrumentation-google-genai-v2: pytest {toxinidir}/instrumentation/opentelemetry-instrumentation-google-genai/tests/interactions --vcr-record=none {posargs} + test-instrumentation-google-genai-v2: pytest {toxinidir}/instrumentation/opentelemetry-instrumentation-google-genai/tests/interactions {toxinidir}/instrumentation/opentelemetry-instrumentation-google-genai/tests/test_stream_context.py --vcr-record=none {posargs} test-instrumentation-google-genai-conformance: pytest {toxinidir}/instrumentation/opentelemetry-instrumentation-google-genai/tests/test_conformance.py --vcr-record=none {posargs} test-instrumentation-genai-agno-{oldest,latest}: pytest --ignore={toxinidir}/instrumentation/opentelemetry-instrumentation-genai-agno/tests/test_conformance.py {toxinidir}/instrumentation/opentelemetry-instrumentation-genai-agno/tests --vcr-record=none {posargs} diff --git a/util/opentelemetry-util-genai/.changelog/899.fixed b/util/opentelemetry-util-genai/.changelog/899.fixed new file mode 100644 index 000000000..f11a5d5aa --- /dev/null +++ b/util/opentelemetry-util-genai/.changelog/899.fixed @@ -0,0 +1 @@ +Suspend invocations in the shared stream wrapper constructors so caller context is restored between stream reads. diff --git a/util/opentelemetry-util-genai/AGENTS.md b/util/opentelemetry-util-genai/AGENTS.md index 091248c4b..1139da852 100644 --- a/util/opentelemetry-util-genai/AGENTS.md +++ b/util/opentelemetry-util-genai/AGENTS.md @@ -71,16 +71,16 @@ is usually hardcoded in specific invocation and does not need to be passed. A streamed response only finishes once the caller has drained the stream, so the invocation must stay open until then. Do **not** call `invocation.stop()` when the SDK returns the stream — the span would close before any chunks arrive. The invocation's span must also not stay current in -the caller's context while the stream is unconsumed: call `invocation.suspend()` before returning -the stream, and re-activate the invocation only while a chunk is being read. +the caller's context while the stream is unconsumed. The shared stream wrappers suspend the +invocation before returning the stream and re-activate it only during reads and cleanup. Instrument streams by subclassing `SyncStreamWrapper` / `AsyncStreamWrapper` from `opentelemetry.util.genai.stream` (the public, supported helpers). The base class proxies the underlying SDK stream, drives iteration, and finalizes telemetry exactly once on success, error, or `close()`. Subclasses pass the SDK stream and the invocation to -`super().__init__(stream, invocation)`, which also marks the invocation as a streamed request -(`gen_ai.request.stream`), call `invocation.suspend()` before returning the stream, and implement -three hooks. The base `_execution_context()` makes a suspended invocation current again for each +`super().__init__(stream, invocation)`, which suspends it and marks it as a streamed request +(`gen_ai.request.stream`), and implement three hooks. +The base `_execution_context()` makes a suspended invocation current again for each read and cleanup operation and restores the context before the chunk is returned, so invocation-backed wrappers don't override it. @@ -94,7 +94,6 @@ invocation-backed wrappers don't override it. class MyStreamWrapper(SyncStreamWrapper[Chunk]): def __init__(self, stream, invocation, capture_content): super().__init__(stream, invocation) - invocation.suspend() ... def _process_chunk(self, chunk): ... # accumulate state diff --git a/util/opentelemetry-util-genai/README.rst b/util/opentelemetry-util-genai/README.rst index 7d8be4e0c..eb77b72ab 100644 --- a/util/opentelemetry-util-genai/README.rst +++ b/util/opentelemetry-util-genai/README.rst @@ -44,8 +44,9 @@ returning a chunk to the consumer. Generator ``send``/``throw`` and ``asend``/``athrow`` also use this scope when the underlying stream supports them. For invocation-backed streams, pass the invocation to the base constructor -(``super().__init__(stream, invocation)``) and call ``invocation.suspend()`` -before returning the stream. The base wrapper then makes the suspended invocation +(``super().__init__(stream, invocation)``). It calls ``invocation.suspend()`` +automatically when available, restoring the caller's context before the stream +is returned. The base wrapper then makes the suspended invocation current during reads, cleanup and finalization of an abandoned stream, so no ``_execution_context()`` override is needed. Keep finalization in the existing ``_on_stream_end`` and ``_on_stream_error`` hooks. diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py index 11f6f1734..adbea3e78 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py @@ -99,11 +99,10 @@ class _StreamTelemetry(Generic[ChunkT], metaclass=ABCMeta): def _execution_context(self) -> AbstractContextManager[None]: """Scope stream reads and cleanup, restoring context before returning a chunk. - By default a suspended invocation is made current again for each read: - an instrumentation that hands the stream back to the caller calls - ``invocation.suspend()`` and gets the span restored around tool or - provider code without overriding this hook. Invocations that are still - attached, or that never attached, leave the context alone. + The constructor suspends invocations that support context management. + Their spans are made current again around tool or provider code without + overriding this hook. Invocations that never attached leave the context + alone. """ invocation = getattr(self, "_self_invocation", None) if invocation is None or not getattr(invocation, "_suspended", False): @@ -169,8 +168,8 @@ class SyncStreamWrapper( ``super().__init__(stream, invocation)`` and implement the three telemetry hooks: ``_process_chunk`` for per-chunk state, ``_on_stream_end`` for successful finalization, and ``_on_stream_error`` for failure finalization. - ``_execution_context`` scopes each read and cleanup; by default it makes a - suspended invocation current again and can be overridden. + The constructor suspends the invocation when it supports ``suspend()``. + ``_execution_context`` makes it current again for each read and cleanup. Users should consume subclasses as normal streams, for example with ``for chunk in wrapper`` or ``with wrapper``. The hook methods are called @@ -188,6 +187,9 @@ def __init__( if invocation is not None and hasattr(invocation, "_request_stream"): setattr(invocation, "_request_stream", True) self._bind_stream(stream) + suspend = getattr(invocation, "suspend", None) + if suspend is not None: + suspend() # The SDK stream, held loosely typed: subclasses re-expose it through a # ``stream`` property carrying the SDK's own type. @@ -296,8 +298,8 @@ class AsyncStreamWrapper( ``super().__init__(stream, invocation)`` and implement the three telemetry hooks: ``_process_chunk`` for per-chunk state, ``_on_stream_end`` for successful finalization, and ``_on_stream_error`` for failure finalization. - ``_execution_context`` scopes each read and cleanup; by default it makes a - suspended invocation current again and can be overridden. + The constructor suspends the invocation when it supports ``suspend()``. + ``_execution_context`` makes it current again for each read and cleanup. Users should consume subclasses as normal async streams, for example with ``async for chunk in wrapper`` or ``async with wrapper``. The hook methods @@ -316,6 +318,9 @@ def __init__( if invocation is not None and hasattr(invocation, "_request_stream"): setattr(invocation, "_request_stream", True) self._bind_stream(stream) + suspend = getattr(invocation, "suspend", None) + if suspend is not None: + suspend() # See ``SyncStreamWrapper._self_stream``. _self_stream: Any @@ -517,7 +522,6 @@ def __init__( ) -> None: super().__init__(stream, invocation) self._self_tool_invocation = invocation - invocation.suspend() self._self_chunks: list[Any] = [] def __del__(self) -> None: @@ -563,7 +567,6 @@ def __init__( ) -> None: super().__init__(stream, invocation) self._self_tool_invocation = invocation - invocation.suspend() self._self_chunks: list[Any] = [] def __del__(self) -> None: diff --git a/util/opentelemetry-util-genai/tests/test_stream.py b/util/opentelemetry-util-genai/tests/test_stream.py index 8ac0aaeea..3fb0848e4 100644 --- a/util/opentelemetry-util-genai/tests/test_stream.py +++ b/util/opentelemetry-util-genai/tests/test_stream.py @@ -2331,60 +2331,100 @@ def _parent_and_attached_invocation(): return parent, token, invocation -@pytest.mark.parametrize("ending", ["exhaust", "error", "close"]) -def test_default_execution_context_leaves_unsuspended_invocation_alone(ending): - """An instrumentation that never calls suspend() must still end with the parent current.""" +@pytest.mark.parametrize( + "ending", ["exhaust", "error", "caller_error", "close"] +) +def test_sync_wrapper_suspends_invocation(ending: str) -> None: parent, token, invocation = _parent_and_attached_invocation() + error = RuntimeError("boom") + cleanup_spans = [] try: - def chunks(): - yield "a" - if ending == "error": - raise RuntimeError("boom") - yield "b" + def chunks() -> Iterator[str]: + try: + assert get_current_span() is invocation.span + yield "a" + assert get_current_span() is invocation.span + if ending == "error": + raise error + yield "b" + finally: + cleanup_spans.append(get_current_span()) wrapper = _TestSyncStreamWrapper(chunks(), invocation) + assert get_current_span() is parent + assert invocation.span.is_recording() + assert next(wrapper) == "a" + assert get_current_span() is parent + assert invocation.span.is_recording() if ending == "exhaust": - assert list(wrapper) == ["a", "b"] + assert list(wrapper) == ["b"] elif ending == "error": - with pytest.raises(RuntimeError): + with pytest.raises(RuntimeError) as raised: list(wrapper) + assert raised.value is error + elif ending == "caller_error": + with pytest.raises(RuntimeError) as raised, wrapper: + raise error + assert raised.value is error else: - assert next(iter(wrapper)) == "a" wrapper.close() + assert cleanup_spans == [invocation.span] assert wrapper._self_finalized + assert not invocation.span.is_recording() assert get_current_span() is parent finally: + invocation.stop() detach(token) parent.end() @pytest.mark.asyncio -@pytest.mark.parametrize("ending", ["exhaust", "error", "close"]) -async def test_default_execution_context_leaves_unsuspended_invocation_alone_async( - ending, -): +@pytest.mark.parametrize( + "ending", ["exhaust", "error", "caller_error", "close"] +) +async def test_async_wrapper_suspends_invocation(ending: str) -> None: parent, token, invocation = _parent_and_attached_invocation() + error = RuntimeError("boom") + cleanup_spans = [] try: - async def chunks(): - yield "a" - if ending == "error": - raise RuntimeError("boom") - yield "b" + async def chunks() -> AsyncGenerator[str]: + try: + assert get_current_span() is invocation.span + yield "a" + assert get_current_span() is invocation.span + if ending == "error": + raise error + yield "b" + finally: + cleanup_spans.append(get_current_span()) wrapper = _TestAsyncStreamWrapper(chunks(), invocation) + assert get_current_span() is parent + assert invocation.span.is_recording() + assert await wrapper.__anext__() == "a" + assert get_current_span() is parent + assert invocation.span.is_recording() if ending == "exhaust": - assert [c async for c in wrapper] == ["a", "b"] + assert [c async for c in wrapper] == ["b"] elif ending == "error": - with pytest.raises(RuntimeError): + with pytest.raises(RuntimeError) as raised: _ = [c async for c in wrapper] + assert raised.value is error + elif ending == "caller_error": + with pytest.raises(RuntimeError) as raised: + async with wrapper: + raise error + assert raised.value is error else: - assert await wrapper.__anext__() == "a" await wrapper.aclose() + assert cleanup_spans == [invocation.span] assert wrapper._self_finalized + assert not invocation.span.is_recording() assert get_current_span() is parent finally: + invocation.stop() detach(token) parent.end()