diff --git a/app/adapters/network/browser.py b/app/adapters/network/browser.py index ae6a4c1a2..8f9c5488c 100644 --- a/app/adapters/network/browser.py +++ b/app/adapters/network/browser.py @@ -3,6 +3,7 @@ import threading import time import uuid from concurrent.futures import ThreadPoolExecutor +from contextvars import Context, copy_context from dataclasses import dataclass, field from typing import Any, Callable, Optional, Protocol from urllib.parse import urlparse @@ -288,7 +289,11 @@ class BrowserSessionHelper: executor = cls._get_existing_session_executor(session_key) if executor: - future = executor.submit( + context = copy_context() + # 会话线程保持空底层上下文，每次操作只使用当前调用快照。 + future = Context().run( + executor.submit, + context.run, cls._run_session_task, session_key, cls._close_session_in_thread, @@ -387,7 +392,11 @@ class BrowserSessionHelper: for _ in range(2): executor = cls._get_session_executor(session_key) try: - future = executor.submit( + context = copy_context() + # 会话线程保持空底层上下文，每次操作只使用当前调用快照。 + future = Context().run( + executor.submit, + context.run, cls._run_session_task, session_key, callback, diff --git a/app/agent/tools/base.py b/app/agent/tools/base.py index 0c0a9c988..845c4198a 100644 --- a/app/agent/tools/base.py +++ b/app/agent/tools/base.py @@ -4,6 +4,7 @@ import json import threading from abc import ABCMeta, abstractmethod from concurrent.futures import ThreadPoolExecutor +from contextvars import Context, copy_context from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Protocol @@ -225,7 +226,13 @@ async def run_agent_blocking( await semaphore.acquire() try: - future = _get_blocking_executor(bucket_name).submit(bound_call) + context = copy_context() + # 长期 worker 保持空底层上下文，每个任务只在自己的调用快照内运行。 + future = Context().run( + _get_blocking_executor(bucket_name).submit, + context.run, + bound_call, + ) except Exception: semaphore.release() raise diff --git a/app/application/messaging/message.py b/app/application/messaging/message.py index ae3c02f44..69ded86ff 100644 --- a/app/application/messaging/message.py +++ b/app/application/messaging/message.py @@ -7,7 +7,9 @@ import queue import re import threading import time +from contextvars import Context, copy_context from datetime import datetime +from functools import partial from typing import Any, Literal, Optional, List, Dict, Protocol, Union from typing import Callable @@ -959,8 +961,16 @@ class MessageQueueManager(metaclass=SingletonClass): if immediately or self._is_in_scheduled_time(datetime.now()): # _send 会执行具体渠道回调，可能包含网络 IO；放到 executor # 避免 async 调用方所在事件循环被同步发送阻塞。 + context = copy_context() + call = partial(self._send, *args, **kwargs) loop = asyncio.get_running_loop() - await loop.run_in_executor(None, lambda: self._send(*args, **kwargs)) + # 默认执行器保持空底层上下文，渠道调用只使用当前消息快照。 + await Context().run( + loop.run_in_executor, + None, + context.run, + call, + ) return self.queue.put({ "args": args, diff --git a/app/chain/workflow.py b/app/chain/workflow.py index a52a5b0d1..78e344e42 100644 --- a/app/chain/workflow.py +++ b/app/chain/workflow.py @@ -6,7 +6,9 @@ import pickle import threading from collections import defaultdict, deque from concurrent.futures import ThreadPoolExecutor +from contextvars import Context, copy_context from datetime import date, datetime +from functools import partial from time import sleep from typing import Any, Callable, List, Optional, Tuple @@ -371,14 +373,17 @@ class WorkflowExecutor: if not node_id: continue - # 提交任务到线程池，每个节点使用上下文快照，避免并行节点互相修改同一个对象。 - future = self.executor.submit( + # 节点分别复制业务上下文和调用上下文，避免共享可变状态或丢失触发链路。 + context = copy_context() + future = Context().run( + self.executor.submit, + context.run, self.execute_node, self.workflow.id, node_id, - copy.deepcopy(self.context) + copy.deepcopy(self.context), ) - future.add_done_callback(self.on_node_complete) + future.add_done_callback(partial(context.run, self.on_node_complete)) finally: self.executor.shutdown(wait=True, cancel_futures=True) diff --git a/tests/test_agent_tool_timeouts.py b/tests/test_agent_tool_timeouts.py index 18dae839d..fb48b7f92 100644 --- a/tests/test_agent_tool_timeouts.py +++ b/tests/test_agent_tool_timeouts.py @@ -11,6 +11,7 @@ from app.agent.tools.base import ( shutdown_blocking_executors, ) from app.agent.tools.manager import MoviePilotToolsManager +from app.runtime.correlation import correlation_scope, get_correlation_id class SlowAgentTool(MoviePilotTool): @@ -99,6 +100,20 @@ def test_run_blocking_keeps_bucket_slot_until_worker_finishes(): asyncio.run(_run_scenario()) +def test_run_blocking_preserves_each_call_context(): + """长期复用的工具线程必须读取当前调用，而不是首个调用的上下文。""" + async def _run_scenario(): + observed = [] + for correlation_id in ("request-one", "request-two"): + with correlation_scope(correlation_id): + observed.append( + await MoviePilotTool.run_blocking("web", get_correlation_id) + ) + return observed + + assert asyncio.run(_run_scenario()) == ["request-one", "request-two"] + + def test_shutdown_blocking_executors_clears_agent_tool_workers(): """测试结束清理应关闭 Agent 工具阻塞线程池，避免全量测试退出时等待 worker。""" diff --git a/tests/test_browser_helper.py b/tests/test_browser_helper.py index 7e66aa6b9..e9179fc4e 100644 --- a/tests/test_browser_helper.py +++ b/tests/test_browser_helper.py @@ -18,6 +18,7 @@ from app.adapters.network.browser import ( launch_browser_context, launch_browser_context_async, ) +from app.runtime.correlation import correlation_scope, get_correlation_id class _FakeResponse: @@ -348,6 +349,23 @@ def test_browser_session_helper_runs_same_session_on_one_worker_thread(): assert session_thread_ids[0] not in caller_thread_ids +def test_browser_session_helper_preserves_each_call_context(): + """会话固定线程必须使用每次操作的上下文，不能保留首次请求状态。""" + page = _FakePage() + context = _FakeContext([page]) + helper = BrowserSessionHelper() + observed = [] + + with patch.object(BrowserSessionHelper, "_launch_context", return_value=context): + for correlation_id in ("request-one", "request-two"): + with correlation_scope(correlation_id): + observed.append( + helper.with_session("session-1", lambda _session: get_correlation_id()) + ) + + assert observed == ["request-one", "request-two"] + + def test_browser_session_helper_closes_session_on_worker_thread(): """关闭会话时应在创建浏览器对象的工作线程内释放资源。""" page = _FakePage() diff --git a/tests/test_system_notification_dispatch.py b/tests/test_system_notification_dispatch.py index f6852bd5e..24e14fab2 100644 --- a/tests/test_system_notification_dispatch.py +++ b/tests/test_system_notification_dispatch.py @@ -12,7 +12,8 @@ sys.modules.setdefault("psutil", ModuleType("psutil")) from app.chain.message import MessageChain from app.application.messaging.message import MessageQueueManager -from app.schemas import Message +from app.schemas.message import Message +from app.runtime.correlation import correlation_scope, get_correlation_id from app.foundation.identity import ( SYSTEM_INTERNAL_USER_ID, is_internal_user_id, @@ -69,23 +70,43 @@ class TestSystemNotificationDispatch(unittest.TestCase): def test_async_send_message_uses_executor_for_immediate_send(self): """异步立即发送不能在事件循环里直接执行同步渠道回调。""" - class _FakeLoop: def __init__(self): self.called = False - async def run_in_executor(self, executor, func): + async def run_in_executor(self, _executor, func, *args): self.called = True - func() + func(*args) async def _run(): manager = MessageQueueManager() fake_loop = _FakeLoop() - with patch("asyncio.get_running_loop", return_value=fake_loop), patch.object( - manager, "_send" - ) as send: + with patch( + "asyncio.get_running_loop", + return_value=fake_loop, + ), patch.object(manager, "_send") as send: await manager.async_send_message("payload", immediately=True) self.assertTrue(fake_loop.called) send.assert_called_once_with("payload") asyncio.run(_run()) + + def test_async_send_message_preserves_call_context(self): + """异步立即发送的同步渠道回调应保留当前请求关联 ID。""" + observed = [] + + async def _run(): + manager = MessageQueueManager() + with patch.object( + manager, + "_send", + side_effect=lambda *_args, **_kwargs: observed.append( + get_correlation_id() + ), + ): + with correlation_scope("message-request"): + await manager.async_send_message("payload", immediately=True) + + asyncio.run(_run()) + + self.assertEqual(observed, ["message-request"]) diff --git a/tests/test_workflow_execution.py b/tests/test_workflow_execution.py index df0af96bf..e0e68c19b 100644 --- a/tests/test_workflow_execution.py +++ b/tests/test_workflow_execution.py @@ -5,8 +5,9 @@ import time from types import SimpleNamespace from app.chain import workflow as workflow_module -from app.schemas import Action, ActionContext, ActionResult +from app.runtime.correlation import correlation_scope, get_correlation_id from app.schemas.types import EventType +from app.schemas.workflow import Action, ActionContext, ActionResult from app import workflow as workflow_package @@ -127,6 +128,58 @@ class _OpaqueValue: return "opaque-value" +def test_workflow_executor_preserves_trigger_context(monkeypatch): + """工作流节点及其完成回调应保留触发链路的关联 ID。""" + observed = [] + release = threading.Event() + + def run_action(_action, context): + observed.append(("node", get_correlation_id())) + assert release.wait(timeout=1) + return ActionResult(success=True, message="ok", context=context) + + fake_manager = _FakeWorkflowManager( + [], + results={"A": run_action}, + ) + monkeypatch.setattr(workflow_module, "WorkFlowManager", lambda: fake_manager) + monkeypatch.setattr( + workflow_module.global_vars, + "workflow_resume", + lambda _workflow_id: None, + ) + monkeypatch.setattr( + workflow_module.global_vars, + "is_workflow_stopped", + lambda _workflow_id: False, + ) + + executor = workflow_module.WorkflowExecutor( + _build_workflow( + actions=[ + {"id": "A", "type": "FakeAction", "name": "动作A", "data": {}} + ], + flows=[], + ), + step_callback=lambda _action, _context: observed.append( + ("completion", get_correlation_id()) + ), + ) + timer = threading.Timer(0.05, release.set) + try: + with correlation_scope("workflow-request"): + timer.start() + executor.execute() + finally: + release.set() + timer.join(timeout=1) + + assert observed == [ + ("node", "workflow-request"), + ("completion", "workflow-request"), + ] + + def test_workflow_executor_resumes_downstream_nodes(monkeypatch): """恢复执行时应释放已完成节点的后继节点。""" calls = []