feat: Propagate context to thread pools

PiperOrigin-RevId: 896063584

Author: google-genai-bot

src/google/adk/flows/llm_flows/functions.py

@@ -20,6 +20,7 @@
 import base64
 import binascii
 from concurrent.futures import ThreadPoolExecutor
+import contextvars
 import copy
 import functools
 import inspect
@@ -140,6 +141,7 @@ async def _call_tool_in_thread_pool(
   """
   from ...tools.function_tool import FunctionTool
 
+  ctx = contextvars.copy_context()
   loop = asyncio.get_running_loop()
   executor = _get_tool_thread_pool(max_workers)
 
@@ -160,7 +162,9 @@ def run_sync_tool():
         # For other sync tool types, we can't easily run them in thread pool
         return None
 
-    result = await loop.run_in_executor(executor, run_sync_tool)
+    result = await loop.run_in_executor(
+        executor, lambda: ctx.run(run_sync_tool)
+    )
     if result is not None:
       return result
   else:
@@ -171,7 +175,9 @@ def run_async_tool_in_new_loop():
       # Create a new event loop for this thread
       return asyncio.run(tool.run_async(args=args, tool_context=tool_context))
 
-    return await loop.run_in_executor(executor, run_async_tool_in_new_loop)
+    return await loop.run_in_executor(
+        executor, lambda: ctx.run(run_async_tool_in_new_loop)
+    )
 
   # Fall back to normal async execution for non-FunctionTool sync tools
   return await tool.run_async(args=args, tool_context=tool_context)

src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
 import logging
 from typing import Any
 from typing import Optional
@@ -298,7 +299,8 @@ def run_gepa():
 
     _logger.info("Running the GEPA optimizer...")
 
-    gepa_results = await loop.run_in_executor(None, run_gepa)
+    ctx = contextvars.copy_context()
+    gepa_results = await loop.run_in_executor(None, lambda: ctx.run(run_gepa))
 
     _logger.info("GEPA optimization finished. Preparing final results...")
 

tests/unittests/flows/llm_flows/test_functions_thread_pool.py

@@ -15,6 +15,7 @@
 """Tests for thread pool execution of tools in Live API mode."""
 
 import asyncio
+import contextvars
 import threading
 import time
 
@@ -349,6 +350,54 @@ def sync_func() -> dict:
     pool = _get_tool_thread_pool(max_workers=12)
     assert pool is not None
 
+  @pytest.mark.asyncio
+  async def test_contextvars_propagation_sync_tool(self):
+    """Test that contextvars propagate to sync tools in thread pool."""
+    test_var = contextvars.ContextVar('test_var', default='default')
+    test_var.set('main_thread_value')
+
+    def sync_func() -> dict[str, str]:
+      return {'value': test_var.get()}
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'value': 'main_thread_value'}
+
+  @pytest.mark.asyncio
+  async def test_contextvars_propagation_async_tool(self):
+    """Test that contextvars propagate to async tools in thread pool."""
+    test_var = contextvars.ContextVar('test_var', default='default')
+    test_var.set('main_thread_value')
+
+    async def async_func() -> dict[str, str]:
+      return {'value': test_var.get()}
+
+    tool = FunctionTool(async_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'value': 'main_thread_value'}
+
 
 class TestToolThreadPoolConfig:
   """Tests for the tool_thread_pool_config in RunConfig."""