fix(google-genai): propagate context to threaded tool calls (#38)

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

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:

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)",
+        )