diff --git a/astrbot/core/computer/booters/local.py b/astrbot/core/computer/booters/local.py index 30de507ba..4edeb040e 100644 --- a/astrbot/core/computer/booters/local.py +++ b/astrbot/core/computer/booters/local.py @@ -118,18 +118,43 @@ class LocalShellComponent(ShellComponent): # `command` is intentionally executed through the current shell so # local computer-use behavior matches existing tool semantics. # Safety relies on `_is_safe_command()` and the allowed-root checks. - result = subprocess.run( # noqa: S602 # nosemgrep: python.lang.security.audit.dangerous-subprocess-use-audit + proc = subprocess.Popen( # noqa: S602 # nosemgrep: python.lang.security.audit.dangerous-subprocess-use-audit command, shell=shell, cwd=working_dir, env=run_env, - timeout=timeout or 300, - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, ) + try: + stdout, stderr = proc.communicate(timeout=timeout or 300) + except subprocess.TimeoutExpired: + should_kill_parent = sys.platform != "win32" + if sys.platform == "win32": + try: + taskkill_result = subprocess.run( + ["taskkill", "/F", "/T", "/PID", str(proc.pid)], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=5, + ) + should_kill_parent = taskkill_result.returncode != 0 + except Exception: + should_kill_parent = True + if should_kill_parent: + try: + proc.kill() + except Exception: + pass + try: + proc.wait(timeout=5) + except Exception: + pass + raise return { - "stdout": _decode_shell_output(result.stdout), - "stderr": _decode_shell_output(result.stderr), - "exit_code": result.returncode, + "stdout": _decode_shell_output(stdout), + "stderr": _decode_shell_output(stderr), + "exit_code": proc.returncode, } return await asyncio.to_thread(_run) diff --git a/tests/test_local_shell_component.py b/tests/test_local_shell_component.py index d1369091e..48eeda7be 100644 --- a/tests/test_local_shell_component.py +++ b/tests/test_local_shell_component.py @@ -3,23 +3,37 @@ from __future__ import annotations import asyncio import subprocess +import pytest + from astrbot.core.computer.booters import local as local_booter from astrbot.core.computer.booters.local import LocalShellComponent -class _FakeCompletedProcess: +class _FakePopen: def __init__(self, stdout: bytes, stderr: bytes = b"", returncode: int = 0): - self.stdout = stdout - self.stderr = stderr + self._stdout = stdout + self._stderr = stderr + self.returncode = returncode + self.pid = 12345 + + def communicate(self, timeout=None): + return self._stdout, self._stderr + + def wait(self, timeout=None): + pass + + +class _FakeTaskkillResult: + def __init__(self, returncode: int): self.returncode = returncode def test_local_shell_component_decodes_utf8_output(monkeypatch): def fake_run(*args, **kwargs): _ = args, kwargs - return _FakeCompletedProcess(stdout="技能内容".encode()) + return _FakePopen(stdout="技能内容".encode()) - monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr(subprocess, "Popen", fake_run) result = asyncio.run(LocalShellComponent().exec("dummy")) @@ -33,9 +47,9 @@ def test_local_shell_component_prefers_utf8_before_windows_locale( ): def fake_run(*args, **kwargs): _ = args, kwargs - return _FakeCompletedProcess(stdout="技能内容".encode()) + return _FakePopen(stdout="技能内容".encode()) - monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr(subprocess, "Popen", fake_run) monkeypatch.setattr(local_booter.os, "name", "nt", raising=False) monkeypatch.setattr( local_booter.locale, @@ -53,9 +67,9 @@ def test_local_shell_component_prefers_utf8_before_windows_locale( def test_local_shell_component_falls_back_to_gbk_on_windows(monkeypatch): def fake_run(*args, **kwargs): _ = args, kwargs - return _FakeCompletedProcess(stdout="微博热搜".encode("gbk")) + return _FakePopen(stdout="微博热搜".encode("gbk")) - monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr(subprocess, "Popen", fake_run) monkeypatch.setattr(local_booter.os, "name", "nt", raising=False) monkeypatch.setattr( local_booter.locale, @@ -73,9 +87,9 @@ def test_local_shell_component_falls_back_to_gbk_on_windows(monkeypatch): def test_local_shell_component_falls_back_to_utf8_replace(monkeypatch): def fake_run(*args, **kwargs): _ = args, kwargs - return _FakeCompletedProcess(stdout=b"\xffabc") + return _FakePopen(stdout=b"\xffabc") - monkeypatch.setattr(subprocess, "run", fake_run) + monkeypatch.setattr(subprocess, "Popen", fake_run) monkeypatch.setattr(local_booter.os, "name", "posix", raising=False) monkeypatch.setattr( local_booter.locale, @@ -86,3 +100,37 @@ def test_local_shell_component_falls_back_to_utf8_replace(monkeypatch): result = asyncio.run(LocalShellComponent().exec("dummy")) assert result["stdout"] == "\ufffdabc" + + +def test_local_shell_component_falls_back_when_windows_taskkill_fails(monkeypatch): + class TimeoutPopen: + pid = 12345 + + def __init__(self): + self.killed = False + self.wait_timeout = None + + def communicate(self, timeout=None): + raise subprocess.TimeoutExpired(cmd="dummy", timeout=timeout) + + def kill(self): + self.killed = True + + def wait(self, timeout=None): + self.wait_timeout = timeout + + proc = TimeoutPopen() + + monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: proc) + monkeypatch.setattr( + subprocess, + "run", + lambda *_args, **_kwargs: _FakeTaskkillResult(returncode=1), + ) + monkeypatch.setattr(local_booter.sys, "platform", "win32") + + with pytest.raises(subprocess.TimeoutExpired): + asyncio.run(LocalShellComponent().exec("dummy", timeout=1)) + + assert proc.killed + assert proc.wait_timeout == 5