fix(runtime): 隔离线程任务上下文 (#6420)

* fix(runtime): isolate thread task contexts

* test: use canonical message schema import
This commit is contained in:
InfinityPacer
2026-08-23 20:10:23 +08:00
committed by GitHub
parent 45b41a5caa
commit 146f8649f6
8 changed files with 154 additions and 16 deletions
+11 -2
View File
@@ -3,6 +3,7 @@ import threading
import time import time
import uuid import uuid
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from contextvars import Context, copy_context
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any, Callable, Optional, Protocol from typing import Any, Callable, Optional, Protocol
from urllib.parse import urlparse from urllib.parse import urlparse
@@ -288,7 +289,11 @@ class BrowserSessionHelper:
executor = cls._get_existing_session_executor(session_key) executor = cls._get_existing_session_executor(session_key)
if executor: if executor:
future = executor.submit( context = copy_context()
# 会话线程保持空底层上下文,每次操作只使用当前调用快照。
future = Context().run(
executor.submit,
context.run,
cls._run_session_task, cls._run_session_task,
session_key, session_key,
cls._close_session_in_thread, cls._close_session_in_thread,
@@ -387,7 +392,11 @@ class BrowserSessionHelper:
for _ in range(2): for _ in range(2):
executor = cls._get_session_executor(session_key) executor = cls._get_session_executor(session_key)
try: try:
future = executor.submit( context = copy_context()
# 会话线程保持空底层上下文,每次操作只使用当前调用快照。
future = Context().run(
executor.submit,
context.run,
cls._run_session_task, cls._run_session_task,
session_key, session_key,
callback, callback,
+8 -1
View File
@@ -4,6 +4,7 @@ import json
import threading import threading
from abc import ABCMeta, abstractmethod from abc import ABCMeta, abstractmethod
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from contextvars import Context, copy_context
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Protocol from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Protocol
@@ -225,7 +226,13 @@ async def run_agent_blocking(
await semaphore.acquire() await semaphore.acquire()
try: 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: except Exception:
semaphore.release() semaphore.release()
raise raise
+11 -1
View File
@@ -7,7 +7,9 @@ import queue
import re import re
import threading import threading
import time import time
from contextvars import Context, copy_context
from datetime import datetime from datetime import datetime
from functools import partial
from typing import Any, Literal, Optional, List, Dict, Protocol, Union from typing import Any, Literal, Optional, List, Dict, Protocol, Union
from typing import Callable from typing import Callable
@@ -959,8 +961,16 @@ class MessageQueueManager(metaclass=SingletonClass):
if immediately or self._is_in_scheduled_time(datetime.now()): if immediately or self._is_in_scheduled_time(datetime.now()):
# _send 会执行具体渠道回调,可能包含网络 IO;放到 executor # _send 会执行具体渠道回调,可能包含网络 IO;放到 executor
# 避免 async 调用方所在事件循环被同步发送阻塞。 # 避免 async 调用方所在事件循环被同步发送阻塞。
context = copy_context()
call = partial(self._send, *args, **kwargs)
loop = asyncio.get_running_loop() 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 return
self.queue.put({ self.queue.put({
"args": args, "args": args,
+9 -4
View File
@@ -6,7 +6,9 @@ import pickle
import threading import threading
from collections import defaultdict, deque from collections import defaultdict, deque
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from contextvars import Context, copy_context
from datetime import date, datetime from datetime import date, datetime
from functools import partial
from time import sleep from time import sleep
from typing import Any, Callable, List, Optional, Tuple from typing import Any, Callable, List, Optional, Tuple
@@ -371,14 +373,17 @@ class WorkflowExecutor:
if not node_id: if not node_id:
continue continue
# 提交任务到线程池,每个节点使用上下文快照,避免并行节点互相修改同一个对象 # 节点分别复制业务上下文和调用上下文,避免共享可变状态或丢失触发链路
future = self.executor.submit( context = copy_context()
future = Context().run(
self.executor.submit,
context.run,
self.execute_node, self.execute_node,
self.workflow.id, self.workflow.id,
node_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: finally:
self.executor.shutdown(wait=True, cancel_futures=True) self.executor.shutdown(wait=True, cancel_futures=True)
+15
View File
@@ -11,6 +11,7 @@ from app.agent.tools.base import (
shutdown_blocking_executors, shutdown_blocking_executors,
) )
from app.agent.tools.manager import MoviePilotToolsManager from app.agent.tools.manager import MoviePilotToolsManager
from app.runtime.correlation import correlation_scope, get_correlation_id
class SlowAgentTool(MoviePilotTool): class SlowAgentTool(MoviePilotTool):
@@ -99,6 +100,20 @@ def test_run_blocking_keeps_bucket_slot_until_worker_finishes():
asyncio.run(_run_scenario()) 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(): def test_shutdown_blocking_executors_clears_agent_tool_workers():
"""测试结束清理应关闭 Agent 工具阻塞线程池,避免全量测试退出时等待 worker。""" """测试结束清理应关闭 Agent 工具阻塞线程池,避免全量测试退出时等待 worker。"""
+18
View File
@@ -18,6 +18,7 @@ from app.adapters.network.browser import (
launch_browser_context, launch_browser_context,
launch_browser_context_async, launch_browser_context_async,
) )
from app.runtime.correlation import correlation_scope, get_correlation_id
class _FakeResponse: 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 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(): def test_browser_session_helper_closes_session_on_worker_thread():
"""关闭会话时应在创建浏览器对象的工作线程内释放资源。""" """关闭会话时应在创建浏览器对象的工作线程内释放资源。"""
page = _FakePage() page = _FakePage()
+28 -7
View File
@@ -12,7 +12,8 @@ sys.modules.setdefault("psutil", ModuleType("psutil"))
from app.chain.message import MessageChain from app.chain.message import MessageChain
from app.application.messaging.message import MessageQueueManager 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 ( from app.foundation.identity import (
SYSTEM_INTERNAL_USER_ID, SYSTEM_INTERNAL_USER_ID,
is_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): def test_async_send_message_uses_executor_for_immediate_send(self):
"""异步立即发送不能在事件循环里直接执行同步渠道回调。""" """异步立即发送不能在事件循环里直接执行同步渠道回调。"""
class _FakeLoop: class _FakeLoop:
def __init__(self): def __init__(self):
self.called = False self.called = False
async def run_in_executor(self, executor, func): async def run_in_executor(self, _executor, func, *args):
self.called = True self.called = True
func() func(*args)
async def _run(): async def _run():
manager = MessageQueueManager() manager = MessageQueueManager()
fake_loop = _FakeLoop() fake_loop = _FakeLoop()
with patch("asyncio.get_running_loop", return_value=fake_loop), patch.object( with patch(
manager, "_send" "asyncio.get_running_loop",
) as send: return_value=fake_loop,
), patch.object(manager, "_send") as send:
await manager.async_send_message("payload", immediately=True) await manager.async_send_message("payload", immediately=True)
self.assertTrue(fake_loop.called) self.assertTrue(fake_loop.called)
send.assert_called_once_with("payload") send.assert_called_once_with("payload")
asyncio.run(_run()) 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"])
+54 -1
View File
@@ -5,8 +5,9 @@ import time
from types import SimpleNamespace from types import SimpleNamespace
from app.chain import workflow as workflow_module 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.types import EventType
from app.schemas.workflow import Action, ActionContext, ActionResult
from app import workflow as workflow_package from app import workflow as workflow_package
@@ -127,6 +128,58 @@ class _OpaqueValue:
return "opaque-value" 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): def test_workflow_executor_resumes_downstream_nodes(monkeypatch):
"""恢复执行时应释放已完成节点的后继节点。""" """恢复执行时应释放已完成节点的后继节点。"""
calls = [] calls = []