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

[RichardoMrMu]

What

Fixes #38. Parallel/concurrent tool calls in an agent produce multiple
disconnected traces instead of one trace with multiple spans.

Root cause

tool_call_wrapper.wrapped_tool() wraps each tool while the agent/LLM span is
active (generate_content._wrapped_config_with_tools). The wrapped tool builds
its ToolInvocation -> start_execute_tool, whose span is parented to whatever
is in the OpenTelemetry context at call time. The Google GenAI SDK's
automatic function calling (and agent frameworks) execute those wrapped tools in
a ThreadPoolExecutor / run_in_executor worker. Worker threads do not inherit
contextvars, so the worker sees an empty context and each tool span starts a
new root trace.

This matches the maintainer's confirmed analysis on #38, and is the same class
of issue already worked around per-instrumentation in bfclv4
(threading_propagation.py) and openhands (session_context.py).

Fix

tool_call_wrapper now captures the active context when a tool is wrapped
(only when a span is active) and re-attaches it around invocation creation in
both the sync and async wrappers, then detaches immediately. The tool function
itself runs in its normal context; only the invocation/span creation is
re-parented. When no span is active the capture returns None and behavior is
unchanged, so single-threaded execution is unaffected.

Only tool_call_wrapper.py changes on the source side; _compat.py and
generate_content.py are untouched.

Tests

Added two regression tests in tests/utils/test_tool_call_wrapper.py:

• test_parallel_tool_calls_share_parent_trace runs two wrapped tools in a
  ThreadPoolExecutor under an agent span and asserts both execute_tool spans
  share the agent's trace_id.
• test_run_in_executor_tool_call_shares_parent_trace covers the
  asyncio.run_in_executor path named in the issue.

Both fail on the current wrapper (tool spans land on new traces) and pass with
the fix. The existing single-threaded tool-wrapper tests are unchanged and still
pass.

[Copilot]

Copilot review overview

🟡 Changes recommended

Out-of-order context-token detachment leaks the captured trace context into reused workers.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1

Open (1)

• Detach captured parent token after invocation completes · New

What changed in this PR

Propagates OpenTelemetry context to threaded Google GenAI tool calls.

Changes:

• Captures and reattaches parent context for tool spans.
• Adds thread-pool regression tests.
• Documents the fix.

File	Description
tool_call_wrapper.py	Adds context propagation.
test_tool_call_wrapper.py	Tests threaded tool calls.
CHANGELOG.md	Records the fix.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[sipercai]

Thanks for working on cross-thread tool-call tracing. We tested the two regression tests added in this PR against current main (2872e4f) using loongsuite-instrument:

• With opentelemetry-instrumentation-threading installed and automatically enabled, both tests pass without this PR.
• With threading instrumentation explicitly disabled, both tests fail, reproducing the context-propagation issue.

LoongSuite already includes threading instrumentation in its default bootstrap registry. For manual instrumentation, it can be enabled with:

from opentelemetry.instrumentation.threading import ThreadingInstrumentor

ThreadingInstrumentor().instrument()

For the tested ThreadPoolExecutor and asyncio.run_in_executor paths, we prefer letting the two instrumentations work together: threading instrumentation propagates the active context at task submission, while the Google GenAI instrumentation and GenAI util manage tool spans and their lifecycle. This also supports other instrumentations used inside the tool and avoids duplicating parent-context capture in individual framework wrappers.

We also checked child spans created inside a tool. On main with threading enabled, the hierarchy is request → tool → child. With this PR, the child becomes a sibling of the tool because the captured context is detached before the tool body runs. The two added tests only check trace IDs, so they do not catch this regression.

Could you share a minimal reproduction that still loses context with threading instrumentation enabled, including the package versions and startup command? Based on the current reproduction, enabling the existing threading instrumentation appears sufficient, and we would prefer documenting/verifying that integration over adding another context bridge in the Google tool wrapper.

Test environment: Python 3.12.13, OpenTelemetry SDK 1.39.1, threading instrumentation 0.60b1, Google GenAI SDK 2.25.0, and LoongSuite distro/Google instrumentation/GenAI util from source. The Google SDK automatic-tool-call check used a mocked HTTP transport; the thread propagation checks did not require a live model service.

[RichardoMrMu]

Thanks @sipercai for the thorough investigation — you're right on both counts, and I'm closing this PR.

I re-checked against the opentelemetry-instrumentation-threading source: it already wraps futures.ThreadPoolExecutor.submit (and Thread.start/run), capturing context.get_current() at submit time and re-attaching it around the entire task body in the worker. That covers exactly the ThreadPoolExecutor / asyncio.run_in_executor paths this PR targeted, so with LoongSuite's default threading instrumentation enabled the tool span already joins the agent trace — no separate context bridge is needed in the Google tool wrapper.

You also correctly spotted the regression: my change detaches the captured context right after the tool span is created but before the tool body runs, so a child span created inside the tool becomes a sibling of the tool span instead of its child (request → tool → child degrades to request → {tool, child}). My two tests only asserted trace IDs, which are fixed at attach time, so they passed while missing the broken parent-child hierarchy — thanks for catching that.

Since the existing threading instrumentation is the right layer for this and my wrapper-level bridge both duplicates it and breaks span nesting, closing is the correct call. If it'd be useful, I'm happy to instead contribute a short docs note + a hierarchy-asserting test that verifies the threading-instrumentation integration for the automatic-function-calling path. Thanks again for the careful review.