diff --git a/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/CHANGELOG.md b/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/CHANGELOG.md index eaf43047e..8dd6c0387 100644 --- a/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/CHANGELOG.md +++ b/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +### Fixed + +- Propagate the active OpenTelemetry context into wrapped tool functions so + automatic function calls executed in a ``ThreadPoolExecutor`` / + ``run_in_executor`` worker attach to the agent trace instead of starting a + new root trace (issue #38). + ## Version 0.9.0 (2026-09-07) ### Added diff --git a/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/tool_call_wrapper.py b/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/tool_call_wrapper.py index e1bd6ea85..12dc3933c 100644 --- a/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/tool_call_wrapper.py +++ b/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/tool_call_wrapper.py @@ -27,6 +27,8 @@ ToolOrDict, ) +from opentelemetry import context as otel_context +from opentelemetry.trace import INVALID_SPAN, get_current_span from opentelemetry.util.genai import hook_advice from ._compat import TelemetryHandler, ToolInvocation @@ -133,20 +135,53 @@ def _fail_tool_advice( state.invocation.fail(error) +def _capture_parent_context() -> Optional[otel_context.Context]: + """Snapshot the OTel context at tool-wrapping time when it carries a span. + + ``wrapped_tool`` runs while the agent/LLM span is active (see + ``generate_content._wrapped_config_with_tools``). The wrapped tool itself, + however, is executed by the Google GenAI SDK's automatic function calling + -- and by agent frameworks -- inside a ``ThreadPoolExecutor`` / + ``run_in_executor`` worker. Worker threads do not inherit ``contextvars``, + so ``start_execute_tool`` in the worker sees an empty context and parents + every tool span to nothing, fragmenting one logical trace into several + (issue #38). + + Capturing the context here lets each tool call re-attach it before the + invocation span is created, so the tool span becomes a child of the span + that was active where the tool was wrapped. Returns ``None`` when no span + is active, so normal single-threaded execution is left untouched. + """ + if get_current_span(otel_context.get_current()) is INVALID_SPAN: + return None + return otel_context.get_current() + + def _wrap_tool_function( tool_function: ToolFunction, telemetry_handler: TelemetryHandler, ): + parent_context = _capture_parent_context() + if inspect.iscoroutinefunction(tool_function): @functools.wraps(tool_function) async def wrapped_function(*args, **kwargs): - state = _prepare_tool_advice( - tool_function, - telemetry_handler, - args, - kwargs, + token = ( + otel_context.attach(parent_context) + if parent_context is not None + else None ) + try: + state = _prepare_tool_advice( + tool_function, + telemetry_handler, + args, + kwargs, + ) + finally: + if token is not None: + otel_context.detach(token) try: result = await tool_function(*args, **kwargs) except BaseException as error: @@ -160,12 +195,21 @@ async def wrapped_function(*args, **kwargs): @functools.wraps(tool_function) def wrapped_function(*args, **kwargs): - state = _prepare_tool_advice( - tool_function, - telemetry_handler, - args, - kwargs, + token = ( + otel_context.attach(parent_context) + if parent_context is not None + else None ) + try: + state = _prepare_tool_advice( + tool_function, + telemetry_handler, + args, + kwargs, + ) + finally: + if token is not None: + otel_context.detach(token) try: result = tool_function(*args, **kwargs) except BaseException as error: diff --git a/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/tests/utils/test_tool_call_wrapper.py b/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/tests/utils/test_tool_call_wrapper.py index db0e5b372..9eb0cf0ca 100644 --- a/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/tests/utils/test_tool_call_wrapper.py +++ b/instrumentation-loongsuite/loongsuite-instrumentation-google-genai/tests/utils/test_tool_call_wrapper.py @@ -16,6 +16,7 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +import concurrent.futures import json import unittest from unittest.mock import patch @@ -244,3 +245,71 @@ def somefunction(arg=None): except Exception: span = self.otel.get_span_named("execute_tool somefunction") self.assertEqual(span.attributes["error.type"], "Exception") + + def test_parallel_tool_calls_share_parent_trace(self): + # Regression for #38: an agent runs wrapped tools concurrently in a + # ThreadPoolExecutor. Worker threads do not inherit contextvars, so + # without context propagation each tool span starts its own root trace + # instead of joining the active agent span's trace. + tracer = get_tracer_provider().get_tracer("test-#38") + + def get_weather(): + pass + + def get_stock(): + pass + + with tracer.start_as_current_span("invoke_agent") as parent: + parent_trace_id = parent.get_span_context().trace_id + wrapped_weather = self.wrap(get_weather) + wrapped_stock = self.wrap(get_stock) + with concurrent.futures.ThreadPoolExecutor( + max_workers=2 + ) as executor: + futures = [ + executor.submit(wrapped_weather), + executor.submit(wrapped_stock), + ] + for future in futures: + future.result() + + weather_span = self.otel.get_span_named("execute_tool get_weather") + stock_span = self.otel.get_span_named("execute_tool get_stock") + # Both tool spans must belong to the agent's trace, not new roots. + self.assertEqual( + weather_span.context.trace_id, + parent_trace_id, + "get_weather tool span started a new trace (context lost across " + "the executor worker)", + ) + self.assertEqual( + stock_span.context.trace_id, + parent_trace_id, + "get_stock tool span started a new trace (context lost across " + "the executor worker)", + ) + + def test_run_in_executor_tool_call_shares_parent_trace(self): + # Regression for #38 via the asyncio.run_in_executor path named in the + # issue: the coroutine offloads a sync tool to the default executor. + tracer = get_tracer_provider().get_tracer("test-#38-async") + + def get_weather(): + pass + + async def drive(): + loop = asyncio.get_event_loop() + wrapped_weather = self.wrap(get_weather) + await loop.run_in_executor(None, wrapped_weather) + + with tracer.start_as_current_span("invoke_agent") as parent: + parent_trace_id = parent.get_span_context().trace_id + asyncio.run(drive()) + + weather_span = self.otel.get_span_named("execute_tool get_weather") + self.assertEqual( + weather_span.context.trace_id, + parent_trace_id, + "run_in_executor tool span started a new trace (context lost " + "across the executor worker)", + )