fix: propagate contextvars to tool executor threads

[Hundao]

Description

Fix MCP tool context params (data_dir, etc.) being silently lost when tool executors run in thread pool workers, causing save_data, serve_file_to_user, and other data tools to fail with "Missing required argument: data_dir".

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring (no functional changes)

Related Issues

Fixes #6853

Changes Made

• core/framework/graph/event_loop/tool_result_handler.py: Wrap tool_executor call in contextvars.copy_context().run() so execution context (data_dir, etc.) propagates into worker threads spawned by run_in_executor
• core/tests/test_tool_context_propagation.py: Regression test verifying contextvars are visible inside tool executors running in thread pool

Testing

• Unit tests pass (cd core && pytest tests/)
• Lint passes (cd core && ruff check .)
• Manual testing performed
• E2E: ran vulnerability_assessment agent, confirmed save_data writes to session data dir instead of falling back to /tmp
• Regression test: verifies _execution_context is readable inside run_in_executor worker thread

Checklist

• My code follows the project's style guidelines
• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Summary by CodeRabbit

• Bug Fixes

  • Context variables are now correctly preserved and available when tools run in background threads, ensuring execution-context data (e.g., runtime settings) is accessible during tool execution.

• Tests

  • Added regression tests that validate context propagation for tool execution in thread-pool scenarios, covering both populated and empty context cases to prevent regressions.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: db40d7d1-37da-4fe2-8f56-89940bb9e7e8

📥 Commits

Reviewing files that changed from the base of the PR and between 6dcf9ae and 51344fc.

📒 Files selected for processing (1)

• core/tests/test_tool_context_propagation.py

🚧 Files skipped from review as they are similar to previous changes (1)

• core/tests/test_tool_context_propagation.py

---

📝 Walkthrough

Walkthrough

The PR ensures contextvars set in the async execution context (e.g., data_dir) are preserved when offloading tool executors to a thread pool by running the executor under contextvars.copy_context().run(). It adds tests validating propagation and absence behavior.

Changes

Cohort / File(s)	Summary
Context Propagation Fix core/framework/graph/event_loop/tool_result_handler.py	Replace loop.run_in_executor(None, tool_executor, tool_use) with a context-preserving call using ctx = contextvars.copy_context() and ctx.run(...) so contextvars (like data_dir) are available inside worker threads.
Regression Tests core/tests/test_tool_context_propagation.py	Add two async tests: one asserting a contextvars value (with data_dir) is visible inside the thread-executed tool executor, and one asserting no context yields None inside the executor.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 I hopped into threads where context slipped away,
I wrapped it in a copy so it now can stay,
data_dir travels safe to the tool's little den,
Tests nod and whisper, "No loss now, my friend!" 🥕

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'fix: propagate contextvars to tool executor threads' accurately and concisely summarizes the main change: enabling context variable propagation to thread-pool worker threads handling tool execution.
Linked Issues check	✅ Passed	The PR implementation fully addresses issue #6853: it wraps tool_executor invocation with contextvars.copy_context().run() to propagate execution context into worker threads, includes regression tests verifying _execution_context visibility, and validates the fix through E2E testing.
Out of Scope Changes check	✅ Passed	All changes are directly related to fixing the contextvars propagation issue: the core fix in tool_result_handler.py and regression tests in test_tool_context_propagation.py are both scoped to the stated objective.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

core/tests/test_tool_context_propagation.py (1)

16-17: Prefer public ToolRegistry context helpers in test setup/teardown.

Lines 45-55 directly mutate private _execution_context; using ToolRegistry.set_execution_context/reset_execution_context better matches production flow and reduces coupling.

💡 Proposed refactor

-from framework.runner.tool_registry import _execution_context
+from framework.runner.tool_registry import ToolRegistry, _execution_context
@@
-    token = _execution_context.set({"data_dir": "/tmp/test_data"})
+    token = ToolRegistry.set_execution_context(data_dir="/tmp/test_data")
@@
-        _execution_context.reset(token)
+        ToolRegistry.reset_execution_context(token)

Also applies to: 45-55

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/tests/test_tool_context_propagation.py` around lines 16 - 17, Tests are
directly importing and mutating the private _execution_context; replace those
direct mutations with the public API by importing ToolRegistry and calling
ToolRegistry.set_execution_context(...) in the setup and
ToolRegistry.reset_execution_context() in teardown (or the equivalent methods on
ToolRegistry) so the test uses the same context management as production and
avoids touching _execution_context directly.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/tests/test_tool_context_propagation.py`:
- Line 12: Add "from __future__ import annotations" at the top of
core/tests/test_tool_context_propagation.py and add explicit return type
annotations (-> None) to the test function signatures mentioned in the review
(the test functions defined at the commented lines 22, 29, and 65, e.g., the
test_* functions that exercise tool context propagation). Ensure each function
signature includes the return type annotation (-> None) and keep existing
parameter types/annotations intact so the file complies with the repository
typing rules.

---

Nitpick comments:
In `@core/tests/test_tool_context_propagation.py`:
- Around line 16-17: Tests are directly importing and mutating the private
_execution_context; replace those direct mutations with the public API by
importing ToolRegistry and calling ToolRegistry.set_execution_context(...) in
the setup and ToolRegistry.reset_execution_context() in teardown (or the
equivalent methods on ToolRegistry) so the test uses the same context management
as production and avoids touching _execution_context directly.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 01997f23-f525-4ced-abbb-171c32a53ab3

📥 Commits

Reviewing files that changed from the base of the PR and between eba7524 and 6dcf9ae.

📒 Files selected for processing (2)

• core/framework/graph/event_loop/tool_result_handler.py
• core/tests/test_tool_context_propagation.py