Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion src/google/adk/flows/llm_flows/base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
from typing import TYPE_CHECKING

from google.genai import types
from opentelemetry import context as otel_context
from opentelemetry import trace
from websockets.exceptions import ConnectionClosed
from websockets.exceptions import ConnectionClosedOK

Expand All @@ -44,6 +46,7 @@
from ...models.base_llm_connection import BaseLlmConnection
from ...models.llm_request import LlmRequest
from ...models.llm_response import LlmResponse

from ...telemetry import tracing
from ...telemetry.tracing import trace_call_llm
from ...telemetry.tracing import trace_send_data
Expand Down Expand Up @@ -1127,7 +1130,17 @@ async def _call_llm_async(
llm = self.__get_llm(invocation_context)

async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
with tracer.start_as_current_span('call_llm') as span:
# Use explicit span management instead of start_as_current_span context
# manager to ensure span.end() is always called. In multi-agent scenarios
# with transfer_to_agent, the async generator may receive GeneratorExit
# after an async context switch (sub-agent execution). This causes
# context.detach() to raise ValueError (stale contextvars token), which
# prevents span.end() from being reached when using the context manager.
# See: https://github.com/google/adk-python/issues/4715
span = tracer.start_span('call_llm')
ctx = trace.set_span_in_context(span)
token = otel_context.attach(ctx)
try:
if invocation_context.run_config.support_cfc:
invocation_context.live_request_queue = LiveRequestQueue()
responses_generator = self.run_live(invocation_context)
Expand Down Expand Up @@ -1187,6 +1200,12 @@ async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
llm_response = altered_llm_response

yield llm_response
finally:
try:
otel_context.detach(token)
except ValueError:
pass
span.end()

async with Aclosing(_call_llm_with_tracing()) as agen:
async for event in agen:
Expand Down
3 changes: 3 additions & 0 deletions tests/unittests/telemetry/test_functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,9 @@ def do_replace(tracer):
monkeypatch.setattr(
tracer, 'start_as_current_span', real_tracer.start_as_current_span
)
monkeypatch.setattr(
tracer, 'start_span', real_tracer.start_span
)

do_replace(tracing.tracer)
do_replace(base_agent.tracer)
Expand Down