mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-12 23:29:16 +03:00
fix(exec): clean up failed job assignment
This commit is contained in:
@@ -181,18 +181,17 @@ class WindowsJob:
|
|||||||
self.close()
|
self.close()
|
||||||
raise error
|
raise error
|
||||||
|
|
||||||
assigned = False
|
if not _kernel32.AssignProcessToJobObject(self._handle, process):
|
||||||
|
error = _win_error("AssignProcessToJobObject")
|
||||||
|
_kernel32.TerminateProcess(process, 1)
|
||||||
|
_close_handle(process)
|
||||||
|
self.close()
|
||||||
|
raise error
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if not _kernel32.AssignProcessToJobObject(self._handle, process):
|
|
||||||
raise _win_error("AssignProcessToJobObject")
|
|
||||||
assigned = True
|
|
||||||
_resume_primary_thread(pid)
|
_resume_primary_thread(pid)
|
||||||
except Exception:
|
except Exception:
|
||||||
if assigned:
|
self.terminate()
|
||||||
self.terminate()
|
|
||||||
else:
|
|
||||||
_kernel32.TerminateProcess(process, 1)
|
|
||||||
self.close()
|
|
||||||
raise
|
raise
|
||||||
finally:
|
finally:
|
||||||
_close_handle(process)
|
_close_handle(process)
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ import sys
|
|||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path, PureWindowsPath
|
from pathlib import Path, PureWindowsPath
|
||||||
from typing import Any, Protocol, cast, runtime_checkable
|
from typing import Any, Protocol, cast
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
from pydantic import Field
|
from pydantic import Field
|
||||||
@@ -45,14 +45,11 @@ _IS_WINDOWS = sys.platform == "win32"
|
|||||||
_PROCESS_TREE_OWNER_ATTR = "_nanobot_process_tree_owner"
|
_PROCESS_TREE_OWNER_ATTR = "_nanobot_process_tree_owner"
|
||||||
|
|
||||||
|
|
||||||
@runtime_checkable
|
|
||||||
class _ProcessTreeOwner(Protocol):
|
class _ProcessTreeOwner(Protocol):
|
||||||
creation_flags: int
|
creation_flags: int
|
||||||
|
|
||||||
def assign_and_resume(self, pid: int) -> None: ...
|
def assign_and_resume(self, pid: int) -> None: ...
|
||||||
|
|
||||||
def close(self) -> None: ...
|
|
||||||
|
|
||||||
def release(self) -> None: ...
|
def release(self) -> None: ...
|
||||||
|
|
||||||
def terminate(self) -> None: ...
|
def terminate(self) -> None: ...
|
||||||
@@ -554,6 +551,7 @@ class ExecTool(Tool):
|
|||||||
"""Launch *command* in a platform-appropriate shell."""
|
"""Launch *command* in a platform-appropriate shell."""
|
||||||
if _IS_WINDOWS:
|
if _IS_WINDOWS:
|
||||||
windows_job = None
|
windows_job = None
|
||||||
|
process = None
|
||||||
creation_flags = 0
|
creation_flags = 0
|
||||||
if process_tree and sys.platform == "win32":
|
if process_tree and sys.platform == "win32":
|
||||||
windows_job = ExecTool._create_windows_job()
|
windows_job = ExecTool._create_windows_job()
|
||||||
@@ -601,6 +599,8 @@ class ExecTool(Tool):
|
|||||||
except BaseException:
|
except BaseException:
|
||||||
if windows_job is not None:
|
if windows_job is not None:
|
||||||
windows_job.terminate()
|
windows_job.terminate()
|
||||||
|
if process is not None:
|
||||||
|
await ExecTool._kill_process(process)
|
||||||
raise
|
raise
|
||||||
shell_program = shell_program or shutil.which("bash") or "/bin/bash"
|
shell_program = shell_program or shutil.which("bash") or "/bin/bash"
|
||||||
args: list[str] = [shell_program]
|
args: list[str] = [shell_program]
|
||||||
@@ -757,11 +757,8 @@ class ExecTool(Tool):
|
|||||||
def _process_tree_owner(
|
def _process_tree_owner(
|
||||||
process: asyncio.subprocess.Process,
|
process: asyncio.subprocess.Process,
|
||||||
) -> _ProcessTreeOwner | None:
|
) -> _ProcessTreeOwner | None:
|
||||||
attributes = getattr(process, "__dict__", None)
|
# _spawn is the only writer for this private ownership marker.
|
||||||
if not isinstance(attributes, dict):
|
return cast(_ProcessTreeOwner | None, vars(process).get(_PROCESS_TREE_OWNER_ATTR))
|
||||||
return None
|
|
||||||
owner = cast(dict[str, object], attributes).get(_PROCESS_TREE_OWNER_ATTR)
|
|
||||||
return owner if isinstance(owner, _ProcessTreeOwner) else None
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _create_windows_job() -> _ProcessTreeOwner:
|
def _create_windows_job() -> _ProcessTreeOwner:
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ platform-specific binaries (all subprocess calls are mocked).
|
|||||||
import asyncio
|
import asyncio
|
||||||
import shutil
|
import shutil
|
||||||
import sys
|
import sys
|
||||||
from unittest.mock import AsyncMock, patch
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
@@ -27,9 +27,6 @@ class _FakeWindowsJob:
|
|||||||
def assign_and_resume(self, pid: int) -> None:
|
def assign_and_resume(self, pid: int) -> None:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def close(self) -> None:
|
|
||||||
pass
|
|
||||||
|
|
||||||
def release(self) -> None:
|
def release(self) -> None:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@@ -149,6 +146,32 @@ class TestSpawnUnix:
|
|||||||
|
|
||||||
class TestSpawnWindows:
|
class TestSpawnWindows:
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_job_assignment_failure_kills_suspended_process(self):
|
||||||
|
env = {"PATH": ""}
|
||||||
|
process = AsyncMock()
|
||||||
|
process.pid = 123
|
||||||
|
process.returncode = None
|
||||||
|
process.kill = MagicMock()
|
||||||
|
process.wait.return_value = -9
|
||||||
|
job = MagicMock(spec=_FakeWindowsJob)
|
||||||
|
job.creation_flags = 0x4
|
||||||
|
job.assign_and_resume.side_effect = OSError("OpenProcess failed")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
|
||||||
|
patch("nanobot.agent.tools.shell.sys.platform", "win32"),
|
||||||
|
patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec,
|
||||||
|
patch.object(ExecTool, "_create_windows_job", return_value=job),
|
||||||
|
pytest.raises(OSError, match="OpenProcess failed"),
|
||||||
|
):
|
||||||
|
mock_exec.return_value = process
|
||||||
|
await ExecTool._spawn("echo hi", r"C:\work", env, process_tree=True)
|
||||||
|
|
||||||
|
job.terminate.assert_called_once_with()
|
||||||
|
process.kill.assert_called_once_with()
|
||||||
|
process.wait.assert_awaited_once_with()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_single_line_uses_powershell(self):
|
async def test_single_line_uses_powershell(self):
|
||||||
"""Single-line commands on Windows now route through PowerShell."""
|
"""Single-line commands on Windows now route through PowerShell."""
|
||||||
|
|||||||
Reference in New Issue
Block a user