From 5662d9703622b9cd4ac5cffc554f0438b4b6da46 Mon Sep 17 00:00:00 2001 From: Tong Chen Date: Fri, 9 Oct 2026 17:45:34 +0800 Subject: [PATCH] fix: release Single Agent toolkits when startup fails or is stopped (#2046) --- backend/app/agent/factory/single_agent.py | 48 ++- .../app/agent/factory/toolkit_assembler.py | 54 ++- backend/app/service/single_agent_service.py | 8 +- .../factory/test_toolkit_assembly_rollback.py | 314 ++++++++++++++++++ .../app/service/test_single_agent_service.py | 39 +++ src/components/Folder/FilePreview.tsx | 2 +- src/components/Folder/index.tsx | 2 +- src/service/artifactAssetApi.test.ts | 11 +- src/service/artifactAssetApi.ts | 8 +- .../components/Folder/FilePreview.test.tsx | 21 ++ 10 files changed, 470 insertions(+), 37 deletions(-) create mode 100644 backend/tests/app/agent/factory/test_toolkit_assembly_rollback.py diff --git a/backend/app/agent/factory/single_agent.py b/backend/app/agent/factory/single_agent.py index 80cd559ef..e74b8b44f 100644 --- a/backend/app/agent/factory/single_agent.py +++ b/backend/app/agent/factory/single_agent.py @@ -103,22 +103,38 @@ async def single_agent( if bundle_context: system_message = f"{system_message}\n\n{bundle_context}" - agent = agent_model( - Agents.single_agent, - BaseMessage.make_assistant_message( - role_name="Single Agent", - content=system_message, - ), - options, - assembly.tools, - tool_names=assembly.tool_names, - toolkits_to_register_agent=assembly.toolkits_to_register_agent, - **( - {"managed_resources": managed_execution.model_resources} - if managed_execution is not None - else {} - ), - ) + try: + agent = agent_model( + Agents.single_agent, + BaseMessage.make_assistant_message( + role_name="Single Agent", + content=system_message, + ), + options, + assembly.tools, + tool_names=assembly.tool_names, + toolkits_to_register_agent=assembly.toolkits_to_register_agent, + **( + {"managed_resources": managed_execution.model_resources} + if managed_execution is not None + else {} + ), + ) + except BaseException: + # Without an Agent, nothing would release the assembled Browser + # reservation, Terminal sessions or MCP servers. + if managed_execution is None: + from app.agent.factory.toolkit_assembler import ( + _rollback_runtime_assembly, + ) + + await _rollback_runtime_assembly( + assembly, + project_id=options.project_id, + options=options, + hands=hands, + ) + raise if pause_event is not None: agent.pause_event = pause_event if assembly.observable_todo_toolkit is not None: diff --git a/backend/app/agent/factory/toolkit_assembler.py b/backend/app/agent/factory/toolkit_assembler.py index db61d3c25..7382bd042 100644 --- a/backend/app/agent/factory/toolkit_assembler.py +++ b/backend/app/agent/factory/toolkit_assembler.py @@ -12,6 +12,7 @@ # limitations under the License. # ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. ========= +import asyncio import inspect import logging import os @@ -171,7 +172,7 @@ async def _rollback_runtime_assembly( options: Chat, hands: IHands | None, ) -> None: - """Best-effort rollback when fail-closed Bundle assembly aborts.""" + """Best-effort rollback when assembly or Agent creation aborts.""" candidates = list(assembly.cleanup_toolkits) if ( @@ -188,7 +189,13 @@ async def _rollback_runtime_assembly( if cleanup is None: cleanup = getattr(toolkit, "cleanup", None) if cleanup is not None: - outcome = cleanup() + # Terminal cleanup waits for process groups to stop. Keep that + # wait off the event loop shared by every other Session. + outcome = ( + cleanup() + if inspect.iscoroutinefunction(cleanup) + else await asyncio.to_thread(cleanup) + ) if inspect.isawaitable(outcome): await outcome disposed.append(toolkit) @@ -386,8 +393,30 @@ def _mcp_config( async def assemble_single_agent_toolkits( + options: Chat, + **kwargs, +) -> ToolkitAssembly: + assembly = ToolkitAssembly() + try: + return await _assemble_single_agent_toolkits( + options, assembly=assembly, **kwargs + ) + except BaseException: + # A failure or cancellation can arrive after the Browser reservation, + # Terminal or MCP servers were created. Nothing else owns them yet. + await _rollback_runtime_assembly( + assembly, + project_id=options.project_id, + options=options, + hands=kwargs.get("hands"), + ) + raise + + +async def _assemble_single_agent_toolkits( options: Chat, *, + assembly: ToolkitAssembly, task_id: str, working_directory: str, hands: IHands | None, @@ -397,7 +426,6 @@ async def assemble_single_agent_toolkits( runtime_environment: ResolvedRuntimeEnvironment | None = None, ) -> ToolkitAssembly: config = _merged_config(options) - assembly = ToolkitAssembly() pinned_skill_sources = ( runtime_environment.pinned_skill_sources(Agents.single_agent) if runtime_environment is not None @@ -622,15 +650,17 @@ async def assemble_single_agent_toolkits( "owned_target_url": owned_target_url, "stealth": not bool(owned_target_url), } + # Record the reservation first, so rollback can return it even + # if the toolkit cannot be constructed. + assembly.browser_port = selected_port + assembly.browser_cdp_url = cdp_url + assembly.browser_session_id = toolkit_session_id + assembly.browser_owned_by_hands = cdp_owned_by_hands toolkit = HybridBrowserToolkit( options.project_id, **browser_options ) toolkit.agent_name = Agents.single_agent assembly.browser_toolkit = toolkit - assembly.browser_port = selected_port - assembly.browser_cdp_url = cdp_url - assembly.browser_session_id = toolkit_session_id - assembly.browser_owned_by_hands = cdp_owned_by_hands assembly.toolkits_to_register_agent.append(toolkit) registered = message_integration.register_toolkits(toolkit) assembly.add_tools( @@ -723,18 +753,15 @@ async def assemble_single_agent_toolkits( EnvironmentSetupRequiredError, ) - await _rollback_runtime_assembly( - assembly, - project_id=options.project_id, - options=options, - hands=hands, - ) raise EnvironmentSetupRequiredError( ["bundle_mcp_start_failed"] ) from exc assembly.add_tools(toolkit.get_tools(), "MCPToolkit") else: toolkit = MCPToolkit(**mcp_options) + # Servers that started before connect() failed or was + # cancelled must still be stopped by rollback or disposal. + assembly.cleanup_toolkits.append(toolkit) try: await toolkit.connect() except Exception: @@ -743,7 +770,6 @@ async def assemble_single_agent_toolkits( exc_info=True, ) else: - assembly.cleanup_toolkits.append(toolkit) assembly.add_tools(toolkit.get_tools(), "MCPToolkit") if _enabled(config, "agent") and can_delegate: diff --git a/backend/app/service/single_agent_service.py b/backend/app/service/single_agent_service.py index f4483aa75..21b3671a4 100644 --- a/backend/app/service/single_agent_service.py +++ b/backend/app/service/single_agent_service.py @@ -118,7 +118,13 @@ async def _dispose_stale_agent_runtime( if cleanup is None: cleanup = getattr(toolkit, "cleanup", None) if cleanup is not None: - outcome = cleanup() + # Terminal cleanup waits for process groups to stop. Keep that + # wait off the event loop shared by every other Session. + outcome = ( + cleanup() + if inspect.iscoroutinefunction(cleanup) + else await asyncio.to_thread(cleanup) + ) if inspect.isawaitable(outcome): await outcome disposed = True diff --git a/backend/tests/app/agent/factory/test_toolkit_assembly_rollback.py b/backend/tests/app/agent/factory/test_toolkit_assembly_rollback.py new file mode 100644 index 000000000..276bb0240 --- /dev/null +++ b/backend/tests/app/agent/factory/test_toolkit_assembly_rollback.py @@ -0,0 +1,314 @@ +# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. ========= +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. ========= + +import asyncio +from typing import Any + +import pytest + +import app.agent.factory.toolkit_assembler as assembler +from app.agent.factory import browser as browser_factory +from app.agent.factory.toolkit_assembler import assemble_single_agent_toolkits +from app.model.chat import Chat +from app.service.task import TaskLock, task_locks + +pytestmark = pytest.mark.unit + +_BROWSER = {"port": 9555} + + +class DummyTaskLock: + def add_human_input_listen(self, agent_name: str) -> None: + pass + + async def put_queue(self, data: Any) -> None: + pass + + async def get_human_input(self, agent_name: str) -> str: + return "" + + +@pytest.fixture(autouse=True) +def cdp_pool(monkeypatch): + pool = browser_factory.CdpBrowserPoolManager() + monkeypatch.setattr(browser_factory, "_cdp_pool_manager", pool) + return pool + + +def _options(sample_chat_data, *enabled: str) -> Chat: + return Chat( + **{ + **sample_chat_data, + "cdp_browsers": [dict(_BROWSER)], + "installed_mcp": { + "mcpServers": {"local": {"command": "python", "args": ["-V"]}} + }, + "toolkit_config": { + name: {"enabled": name in enabled} + for name in assembler.DEFAULT_SINGLE_AGENT_TOOLKIT_CONFIG + }, + } + ) + + +def _install_toolkits(monkeypatch, tmp_path, project_id, events, cdp_pool): + import app.agent.toolkit.human_toolkit as human_toolkit + + monkeypatch.setattr( + human_toolkit, "get_task_lock", lambda api_task_id: DummyTaskLock() + ) + monkeypatch.setitem( + task_locks, project_id, TaskLock(project_id, asyncio.Queue(), {}) + ) + monkeypatch.setenv("file_save_path", str(tmp_path)) + monkeypatch.delenv("EIGENT_RUNTIME", raising=False) + + class FakeTerminalToolkit: + def __init__(self, *args, **kwargs): + pass + + @classmethod + def toolkit_name(cls): + return "TerminalToolkit" + + def get_tools(self): + return [] + + def cleanup(self): + events.append("terminal.cleanup") + + class FakeMCPToolkit: + def __init__(self, **kwargs): + pass + + async def connect(self): + events.append( + "mcp.started:browser-reserved" + if cdp_pool._occupied_browsers + else "mcp.started" + ) + + async def disconnect(self): + events.append("mcp.disconnect") + + def get_tools(self): + return [] + + monkeypatch.setattr(assembler, "TerminalToolkit", FakeTerminalToolkit) + monkeypatch.setattr(assembler, "MCPToolkit", FakeMCPToolkit) + return FakeMCPToolkit + + +def _browser_is_free(cdp_pool) -> bool: + selected = cdp_pool.acquire_browser([dict(_BROWSER)], "next-session") + return selected is not None + + +@pytest.mark.asyncio +async def test_late_assembly_failure_releases_started_toolkits( + sample_chat_data, monkeypatch, tmp_path, cdp_pool +): + options = _options(sample_chat_data, "browser", "terminal", "mcp", "agent") + events: list[str] = [] + _install_toolkits( + monkeypatch, tmp_path, options.project_id, events, cdp_pool + ) + + class FailingAgentToolkit: + def __init__(self, **_kwargs): + raise RuntimeError("delegation toolkit failed") + + monkeypatch.setattr( + assembler, "DepthLimitedAgentToolkit", FailingAgentToolkit + ) + + with pytest.raises(RuntimeError, match="delegation toolkit failed"): + await assemble_single_agent_toolkits( + options, + task_id=options.task_id, + working_directory=str(tmp_path), + hands=None, + can_delegate=True, + ) + + assert events == [ + "mcp.started:browser-reserved", + "mcp.disconnect", + "terminal.cleanup", + ] + assert _browser_is_free(cdp_pool) + + +@pytest.mark.asyncio +async def test_browser_toolkit_failure_releases_the_browser_slot( + sample_chat_data, monkeypatch, tmp_path, cdp_pool +): + options = _options(sample_chat_data, "browser") + _install_toolkits(monkeypatch, tmp_path, options.project_id, [], cdp_pool) + + class FailingBrowserToolkit: + def __init__(self, *_args, **_kwargs): + raise TypeError("unsupported browser option") + + monkeypatch.setattr( + assembler, "HybridBrowserToolkit", FailingBrowserToolkit + ) + + with pytest.raises(TypeError, match="unsupported browser option"): + await assemble_single_agent_toolkits( + options, + task_id=options.task_id, + working_directory=str(tmp_path), + hands=None, + can_delegate=False, + ) + + assert _browser_is_free(cdp_pool) + + +@pytest.mark.asyncio +async def test_agent_creation_failure_releases_assembled_toolkits( + sample_chat_data, monkeypatch, tmp_path, cdp_pool +): + import app.agent.factory.single_agent as single_agent_factory + + options = _options(sample_chat_data, "browser", "terminal", "mcp") + events: list[str] = [] + _install_toolkits( + monkeypatch, tmp_path, options.project_id, events, cdp_pool + ) + + def failing_agent_model(*_args, **_kwargs): + raise RuntimeError("model unavailable") + + monkeypatch.setattr( + single_agent_factory, "agent_model", failing_agent_model + ) + + with pytest.raises(RuntimeError, match="model unavailable"): + await single_agent_factory.single_agent(options) + + assert events == [ + "mcp.started:browser-reserved", + "mcp.disconnect", + "terminal.cleanup", + ] + assert _browser_is_free(cdp_pool) + + +@pytest.mark.asyncio +async def test_cancelled_mcp_startup_stops_servers_that_started( + sample_chat_data, monkeypatch, tmp_path, cdp_pool +): + options = _options(sample_chat_data, "browser", "terminal", "mcp") + events: list[str] = [] + fake_mcp = _install_toolkits( + monkeypatch, tmp_path, options.project_id, events, cdp_pool + ) + started = asyncio.Event() + + class SlowMCPToolkit(fake_mcp): + async def connect(self): + await super().connect() + started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(assembler, "MCPToolkit", SlowMCPToolkit) + assembly = asyncio.create_task( + assemble_single_agent_toolkits( + options, + task_id=options.task_id, + working_directory=str(tmp_path), + hands=None, + can_delegate=False, + ) + ) + await asyncio.wait_for(started.wait(), timeout=5) + assembly.cancel() + + with pytest.raises(asyncio.CancelledError): + await assembly + + assert events == [ + "mcp.started:browser-reserved", + "mcp.disconnect", + "terminal.cleanup", + ] + assert _browser_is_free(cdp_pool) + + +@pytest.mark.asyncio +async def test_failed_mcp_startup_remains_owned_by_the_agent( + sample_chat_data, monkeypatch, tmp_path, cdp_pool +): + options = _options(sample_chat_data, "mcp") + events: list[str] = [] + fake_mcp = _install_toolkits( + monkeypatch, tmp_path, options.project_id, events, cdp_pool + ) + + class PartlyFailingMCPToolkit(fake_mcp): + async def connect(self): + await super().connect() + raise RuntimeError("one MCP server failed") + + monkeypatch.setattr(assembler, "MCPToolkit", PartlyFailingMCPToolkit) + + assembly = await assemble_single_agent_toolkits( + options, + task_id=options.task_id, + working_directory=str(tmp_path), + hands=None, + can_delegate=False, + ) + + assert "MCPToolkit" not in assembly.tool_names + assert [type(toolkit) for toolkit in assembly.cleanup_toolkits] == [ + PartlyFailingMCPToolkit + ] + + +@pytest.mark.asyncio +async def test_rollback_runs_sync_cleanup_off_the_event_loop( + sample_chat_data, +): + import threading + + loop_thread = threading.get_ident() + cleanup_threads: list[int] = [] + disconnect_threads: list[int] = [] + + class BlockingCleanupToolkit: + def cleanup(self): + cleanup_threads.append(threading.get_ident()) + + class AsyncDisconnectToolkit: + async def disconnect(self): + disconnect_threads.append(threading.get_ident()) + + options = Chat(**sample_chat_data) + assembly = assembler.ToolkitAssembly( + cleanup_toolkits=[BlockingCleanupToolkit(), AsyncDisconnectToolkit()] + ) + + await assembler._rollback_runtime_assembly( + assembly, + project_id=options.project_id, + options=options, + hands=None, + ) + + assert len(cleanup_threads) == 1 + assert cleanup_threads[0] != loop_thread + assert disconnect_threads == [loop_thread] diff --git a/backend/tests/app/service/test_single_agent_service.py b/backend/tests/app/service/test_single_agent_service.py index 9a8a13950..d504d456e 100644 --- a/backend/tests/app/service/test_single_agent_service.py +++ b/backend/tests/app/service/test_single_agent_service.py @@ -680,3 +680,42 @@ async def test_stale_runtime_cleanup_failure_blocks_replacement(): assert stale_toolkit in task_lock.registered_toolkits assert agent._runtime_cleanup_toolkits == (stale_toolkit,) + + +@pytest.mark.asyncio +async def test_stale_runtime_sync_cleanup_runs_off_the_event_loop(): + import threading + + from app.service.single_agent_service import _dispose_stale_agent_runtime + + loop_thread = threading.get_ident() + cleanup_threads: list[int] = [] + disconnect_threads: list[int] = [] + + class BlockingCleanupToolkit: + def cleanup(self): + cleanup_threads.append(threading.get_ident()) + + class AsyncDisconnectToolkit: + async def disconnect(self): + disconnect_threads.append(threading.get_ident()) + + blocking = BlockingCleanupToolkit() + async_toolkit = AsyncDisconnectToolkit() + agent = MagicMock() + agent._cdp_release_callback = None + agent._runtime_cleanup_toolkits = (blocking, async_toolkit) + task_lock = MagicMock() + task_lock.registered_toolkits = [blocking, async_toolkit] + + await _dispose_stale_agent_runtime( + agent, + task_lock, + task_id="run-runtime-switch", + ) + + assert len(cleanup_threads) == 1 + assert cleanup_threads[0] != loop_thread + assert disconnect_threads == [loop_thread] + assert task_lock.registered_toolkits == [] + assert agent._runtime_cleanup_toolkits == () diff --git a/src/components/Folder/FilePreview.tsx b/src/components/Folder/FilePreview.tsx index 8545a7108..928de256f 100644 --- a/src/components/Folder/FilePreview.tsx +++ b/src/components/Folder/FilePreview.tsx @@ -84,7 +84,7 @@ export function FilePreview({ previewRequestRef.current = controller; setSelectedFile(target); setLoading(true); - void resolveArtifactAssetFile(target) + void resolveArtifactAssetFile(target, controller.signal) .then((resolved) => loadFilePreview(resolved, { ipcRenderer, diff --git a/src/components/Folder/index.tsx b/src/components/Folder/index.tsx index 7844bfb21..f57bf38d5 100644 --- a/src/components/Folder/index.tsx +++ b/src/components/Folder/index.tsx @@ -1142,7 +1142,7 @@ export default function Folder({ data: _data, spaceId }: FolderProps) { setIsShowSourceCode(Boolean(isShowSourceCode)); } setLoading(true); - void resolveArtifactAssetFile(file) + void resolveArtifactAssetFile(file, controller.signal) .then((resolvedFile) => loadFilePreview(resolvedFile, { ipcRenderer, diff --git a/src/service/artifactAssetApi.test.ts b/src/service/artifactAssetApi.test.ts index 0f46ad4f4..13cbacc10 100644 --- a/src/service/artifactAssetApi.test.ts +++ b/src/service/artifactAssetApi.test.ts @@ -58,15 +58,22 @@ describe('Artifact asset resolution', () => { contentType: 'text/csv', }, } as FileInfo; + const controller = new AbortController(); - await expect(resolveArtifactAssetFile(file)).resolves.toMatchObject({ + await expect( + resolveArtifactAssetFile(file, controller.signal) + ).resolves.toMatchObject({ path: 'https://assets.example/report.csv?signature=one', isRemote: true, size: 123, mimeType: 'text/csv', }); + // The preview owner can cancel the lookup when it moves to another file. expect(proxyFetchGetMock).toHaveBeenCalledWith( - '/api/v1/chat/files/73/download' + '/api/v1/chat/files/73/download', + undefined, + undefined, + { signal: controller.signal } ); }); diff --git a/src/service/artifactAssetApi.ts b/src/service/artifactAssetApi.ts index 722d9d20e..8d153f7ec 100644 --- a/src/service/artifactAssetApi.ts +++ b/src/service/artifactAssetApi.ts @@ -21,7 +21,8 @@ type ArtifactDownloadResponse = { /** Resolve a Cloud-restored Artifact into the same bounded preview pipeline. */ export async function resolveArtifactAssetFile( - file: FileInfo + file: FileInfo, + signal?: AbortSignal ): Promise { if (file.localPathAvailable !== false) return file; if (/^https?:\/\//i.test(file.path)) return file; @@ -30,7 +31,10 @@ export async function resolveArtifactAssetFile( throw new Error('This Artifact has not finished uploading yet.'); } const result = (await proxyFetchGet( - `/api/v1/chat/files/${encodeURIComponent(chatFileId)}/download` + `/api/v1/chat/files/${encodeURIComponent(chatFileId)}/download`, + undefined, + undefined, + { signal } )) as ArtifactDownloadResponse; if (!result?.download_url) { throw new Error('Artifact download URL is unavailable.'); diff --git a/test/unit/components/Folder/FilePreview.test.tsx b/test/unit/components/Folder/FilePreview.test.tsx index 6b896132d..82c0e91e5 100644 --- a/test/unit/components/Folder/FilePreview.test.tsx +++ b/test/unit/components/Folder/FilePreview.test.tsx @@ -199,6 +199,27 @@ describe('file selection and loading recovery', () => { await waitFor(() => expect(screen.queryByRole('alert')).toBeNull()); expect(loadFilePreviewMock).toHaveBeenCalledTimes(2); }); + it('cancels the Artifact lookup of a preview that was replaced', async () => { + resolveArtifactAssetFileMock.mockImplementationOnce( + () => new Promise(() => {}) + ); + const next = { ...file, name: 'next.zip', path: '/workspace/next.zip' }; + const { rerender } = render(); + await waitFor(() => + expect(resolveArtifactAssetFileMock).toHaveBeenCalledTimes(1) + ); + const [, signal] = resolveArtifactAssetFileMock.mock.calls[0]; + expect(signal).toBeInstanceOf(AbortSignal); + expect(signal.aborted).toBe(false); + + rerender(); + + await waitFor(() => + expect(resolveArtifactAssetFileMock).toHaveBeenCalledTimes(2) + ); + expect(signal.aborted).toBe(true); + expect(resolveArtifactAssetFileMock.mock.calls[1][1].aborted).toBe(false); + }); it('does not restore a cleared selection when an old request finishes', async () => { let resolve!: (value: FileInfo) => void; loadFilePreviewMock.mockImplementationOnce(