From 6718ba8bdfd54764c00149bf60518a25bb0e4770 Mon Sep 17 00:00:00 2001
From: u7-u7 <116776533+u7-u7@users.noreply.github.com>
Date: Tue, 27 Oct 2026 10:41:13 +0800
Subject: [PATCH] fix(async): preserve context in guarded thread calls

---
 hy_memory/utils/async_executor.py    |  5 ++-
 tests/test_async_executor_context.py | 51 ++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_async_executor_context.py

diff --git a/hy_memory/utils/async_executor.py b/hy_memory/utils/async_executor.py
index 0ab0c7e..28afa74 100644
--- a/hy_memory/utils/async_executor.py
+++ b/hy_memory/utils/async_executor.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import contextvars
 import functools
 from concurrent.futures import Executor
 from typing import Any, Callable, Optional
@@ -39,7 +40,9 @@ async def run_in_executor_guarded(
 
 async def to_thread_guarded(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
     """Guarded equivalent of ``asyncio.to_thread`` using the default pool."""
-    return await run_in_executor_guarded(None, func, *args, **kwargs)
+    context = contextvars.copy_context()
+    call = functools.partial(func, *args, **kwargs)
+    return await run_in_executor_guarded(None, context.run, call)
 
 
 def concurrent_result_guarded(
diff --git a/tests/test_async_executor_context.py b/tests/test_async_executor_context.py
new file mode 100644
index 0000000..5fb6a47
--- /dev/null
+++ b/tests/test_async_executor_context.py
@@ -0,0 +1,51 @@
+"""The guarded to_thread adapter preserves asyncio.to_thread context semantics."""
+import asyncio
+import contextvars
+import importlib.util
+from pathlib import Path
+import unittest
+
+
+# This helper is dependency-free; exercise it without initializing the SDK.
+_spec = importlib.util.spec_from_file_location(
+    "async_executor_under_test",
+    Path(__file__).resolve().parents[1] / "hy_memory/utils/async_executor.py",
+)
+_executor = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_executor)
+
+
+class ThreadContextTests(unittest.IsolatedAsyncioTestCase):
+    async def test_copies_context_and_does_not_propagate_worker_changes(self):
+        request_id = contextvars.ContextVar("request_id", default="missing")
+        request_id.set("caller")
+
+        def work(prefix, *, suffix):
+            value = request_id.get()
+            request_id.set("worker")
+            return prefix + value + suffix
+
+        result = await _executor.to_thread_guarded(work, "[", suffix="]")
+        self.assertEqual(result, "[caller]")
+        self.assertEqual(request_id.get(), "caller")
+
+    async def test_concurrent_callers_keep_separate_contexts(self):
+        request_id = contextvars.ContextVar("request_id", default="missing")
+
+        async def call(value):
+            request_id.set(value)
+            await asyncio.sleep(0)
+            return await _executor.to_thread_guarded(request_id.get)
+
+        self.assertEqual(await asyncio.gather(call("a"), call("b")), ["a", "b"])
+
+    async def test_context_is_copied_when_coroutine_runs(self):
+        request_id = contextvars.ContextVar("request_id", default="missing")
+        request_id.set("created")
+        pending = _executor.to_thread_guarded(request_id.get)
+        request_id.set("awaited")
+        self.assertEqual(await pending, "awaited")
+
+
+if __name__ == "__main__":
+    unittest.main()