mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-09-05 23:47:41 +08:00
fix(agent): preserve full command output in temp logs
Return only a 10KB preview to the agent so large command results do not overwhelm conversations while keeping the full output available for follow-up reads. Add pytest to the project dependencies to make the regression tests runnable in the project venv.
This commit is contained in:
@@ -5,7 +5,8 @@ import os
|
|||||||
import signal
|
import signal
|
||||||
import subprocess
|
import subprocess
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Optional, Type
|
from tempfile import NamedTemporaryFile
|
||||||
|
from typing import Optional, TextIO, Type
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
@@ -15,7 +16,7 @@ from app.log import logger
|
|||||||
|
|
||||||
DEFAULT_TIMEOUT_SECONDS = 60
|
DEFAULT_TIMEOUT_SECONDS = 60
|
||||||
MAX_TIMEOUT_SECONDS = 300
|
MAX_TIMEOUT_SECONDS = 300
|
||||||
MAX_OUTPUT_CHARS = 6000
|
MAX_OUTPUT_PREVIEW_BYTES = 10 * 1024
|
||||||
READ_CHUNK_SIZE = 4096
|
READ_CHUNK_SIZE = 4096
|
||||||
KILL_GRACE_SECONDS = 3
|
KILL_GRACE_SECONDS = 3
|
||||||
COMMAND_CONCURRENCY_LIMIT = 2
|
COMMAND_CONCURRENCY_LIMIT = 2
|
||||||
@@ -25,40 +26,93 @@ _command_semaphore = asyncio.Semaphore(COMMAND_CONCURRENCY_LIMIT)
|
|||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class _CommandOutput:
|
class _CommandOutput:
|
||||||
"""保存受限命令输出,避免大输出一次性进入内存。"""
|
"""保存前 10KB 预览,并在超限时将完整输出写入临时文件。"""
|
||||||
|
|
||||||
limit: int
|
preview_limit_bytes: int
|
||||||
stdout_chunks: list[str] = field(default_factory=list)
|
preview_entries: list[tuple[str, str]] = field(default_factory=list)
|
||||||
stderr_chunks: list[str] = field(default_factory=list)
|
captured_bytes: int = 0
|
||||||
captured_chars: int = 0
|
preview_truncated: bool = False
|
||||||
truncated: bool = False
|
temp_file_path: Optional[str] = None
|
||||||
|
temp_file_handle: Optional[TextIO] = None
|
||||||
|
last_written_stream: Optional[str] = None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _clip_text_to_bytes(text: str, byte_limit: int) -> str:
|
||||||
|
if byte_limit <= 0:
|
||||||
|
return ""
|
||||||
|
return text.encode("utf-8")[:byte_limit].decode("utf-8", errors="ignore")
|
||||||
|
|
||||||
|
def _write_chunk(self, stream_name: str, text: str) -> None:
|
||||||
|
if not self.temp_file_handle or not text:
|
||||||
|
return
|
||||||
|
|
||||||
|
if self.last_written_stream != stream_name:
|
||||||
|
if self.temp_file_handle.tell() > 0:
|
||||||
|
self.temp_file_handle.write("\n")
|
||||||
|
title = "标准输出" if stream_name == "stdout" else "错误输出"
|
||||||
|
self.temp_file_handle.write(f"[{title}]\n")
|
||||||
|
self.last_written_stream = stream_name
|
||||||
|
|
||||||
|
self.temp_file_handle.write(text)
|
||||||
|
|
||||||
|
def _ensure_temp_file(self) -> None:
|
||||||
|
if self.temp_file_handle:
|
||||||
|
return
|
||||||
|
|
||||||
|
temp_file = NamedTemporaryFile(
|
||||||
|
mode="w",
|
||||||
|
encoding="utf-8",
|
||||||
|
suffix=".log",
|
||||||
|
prefix="moviepilot-command-",
|
||||||
|
delete=False,
|
||||||
|
)
|
||||||
|
self.temp_file_path = temp_file.name
|
||||||
|
self.temp_file_handle = temp_file
|
||||||
|
for stream_name, chunk in self.preview_entries:
|
||||||
|
self._write_chunk(stream_name, chunk)
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
if not self.temp_file_handle:
|
||||||
|
return
|
||||||
|
self.temp_file_handle.flush()
|
||||||
|
self.temp_file_handle.close()
|
||||||
|
self.temp_file_handle = None
|
||||||
|
|
||||||
def append(self, stream_name: str, text: str) -> None:
|
def append(self, stream_name: str, text: str) -> None:
|
||||||
if not text:
|
if not text:
|
||||||
return
|
return
|
||||||
|
|
||||||
remaining = self.limit - self.captured_chars
|
if self.temp_file_handle:
|
||||||
if remaining <= 0:
|
self._write_chunk(stream_name, text)
|
||||||
self.truncated = True
|
|
||||||
return
|
return
|
||||||
|
|
||||||
captured = text[:remaining]
|
chunk_bytes = len(text.encode("utf-8"))
|
||||||
if stream_name == "stdout":
|
remaining = self.preview_limit_bytes - self.captured_bytes
|
||||||
self.stdout_chunks.append(captured)
|
if chunk_bytes <= remaining:
|
||||||
else:
|
self.preview_entries.append((stream_name, text))
|
||||||
self.stderr_chunks.append(captured)
|
self.captured_bytes += chunk_bytes
|
||||||
|
return
|
||||||
|
|
||||||
self.captured_chars += len(captured)
|
self.preview_truncated = True
|
||||||
if len(text) > remaining:
|
self._ensure_temp_file()
|
||||||
self.truncated = True
|
self._write_chunk(stream_name, text)
|
||||||
|
|
||||||
|
preview = self._clip_text_to_bytes(text, remaining)
|
||||||
|
if preview:
|
||||||
|
self.preview_entries.append((stream_name, preview))
|
||||||
|
self.captured_bytes += len(preview.encode("utf-8"))
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def stdout(self) -> str:
|
def stdout(self) -> str:
|
||||||
return "".join(self.stdout_chunks).strip()
|
return "".join(
|
||||||
|
text for stream_name, text in self.preview_entries if stream_name == "stdout"
|
||||||
|
).strip()
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def stderr(self) -> str:
|
def stderr(self) -> str:
|
||||||
return "".join(self.stderr_chunks).strip()
|
return "".join(
|
||||||
|
text for stream_name, text in self.preview_entries if stream_name == "stderr"
|
||||||
|
).strip()
|
||||||
|
|
||||||
|
|
||||||
class ExecuteCommandInput(BaseModel):
|
class ExecuteCommandInput(BaseModel):
|
||||||
@@ -78,7 +132,7 @@ class ExecuteCommandTool(MoviePilotTool):
|
|||||||
description: str = (
|
description: str = (
|
||||||
"Safely execute shell commands on the server. Useful for system "
|
"Safely execute shell commands on the server. Useful for system "
|
||||||
"maintenance, checking status, or running custom scripts. Includes "
|
"maintenance, checking status, or running custom scripts. Includes "
|
||||||
"timeout, concurrency, and hard output limits."
|
"timeout, concurrency, and output preview limits."
|
||||||
)
|
)
|
||||||
args_schema: Type[BaseModel] = ExecuteCommandInput
|
args_schema: Type[BaseModel] = ExecuteCommandInput
|
||||||
require_admin: bool = True
|
require_admin: bool = True
|
||||||
@@ -107,7 +161,7 @@ class ExecuteCommandTool(MoviePilotTool):
|
|||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _subprocess_kwargs() -> dict:
|
def _subprocess_kwargs() -> dict:
|
||||||
"""为子进程创建独立进程组,便于超时或输出过大时清理整棵子进程。"""
|
"""为子进程创建独立进程组,便于超时场景清理整棵子进程。"""
|
||||||
kwargs = {
|
kwargs = {
|
||||||
"stdin": subprocess.DEVNULL,
|
"stdin": subprocess.DEVNULL,
|
||||||
"stdout": asyncio.subprocess.PIPE,
|
"stdout": asyncio.subprocess.PIPE,
|
||||||
@@ -124,23 +178,14 @@ class ExecuteCommandTool(MoviePilotTool):
|
|||||||
stream: asyncio.StreamReader,
|
stream: asyncio.StreamReader,
|
||||||
stream_name: str,
|
stream_name: str,
|
||||||
output: _CommandOutput,
|
output: _CommandOutput,
|
||||||
limit_reached: asyncio.Event,
|
|
||||||
) -> None:
|
) -> None:
|
||||||
"""按块读取输出,达到上限后通知主流程终止命令。"""
|
"""按块读取输出,始终只把前 10KB 保留在返回结果中。"""
|
||||||
while True:
|
while True:
|
||||||
chunk = await stream.read(READ_CHUNK_SIZE)
|
chunk = await stream.read(READ_CHUNK_SIZE)
|
||||||
if not chunk:
|
if not chunk:
|
||||||
break
|
break
|
||||||
|
|
||||||
if output.truncated:
|
|
||||||
limit_reached.set()
|
|
||||||
continue
|
|
||||||
|
|
||||||
output.append(stream_name, chunk.decode("utf-8", errors="replace"))
|
output.append(stream_name, chunk.decode("utf-8", errors="replace"))
|
||||||
if output.truncated:
|
|
||||||
limit_reached.set()
|
|
||||||
# 达到上限后继续排空管道但不再保存内容,避免子进程因 pipe 反压卡住。
|
|
||||||
continue
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _terminate_process(process: asyncio.subprocess.Process, sig: int):
|
def _terminate_process(process: asyncio.subprocess.Process, sig: int):
|
||||||
@@ -205,27 +250,33 @@ class ExecuteCommandTool(MoviePilotTool):
|
|||||||
output: _CommandOutput,
|
output: _CommandOutput,
|
||||||
timeout: int,
|
timeout: int,
|
||||||
timed_out: bool,
|
timed_out: bool,
|
||||||
output_limited: bool,
|
|
||||||
timeout_note: Optional[str],
|
timeout_note: Optional[str],
|
||||||
) -> str:
|
) -> str:
|
||||||
if timed_out:
|
if timed_out:
|
||||||
result = f"命令执行超时 (限制: {timeout}秒,已终止进程)"
|
result = f"命令执行超时 (限制: {timeout}秒,已终止进程)"
|
||||||
elif output_limited:
|
|
||||||
result = (
|
|
||||||
f"命令输出超过限制 (限制: {MAX_OUTPUT_CHARS}字符,"
|
|
||||||
f"已截断并终止进程,退出码: {exit_code})"
|
|
||||||
)
|
|
||||||
else:
|
else:
|
||||||
result = f"命令执行完成 (退出码: {exit_code})"
|
result = f"命令执行完成 (退出码: {exit_code})"
|
||||||
|
|
||||||
if timeout_note:
|
if timeout_note:
|
||||||
result += f"\n\n提示:\n{timeout_note}"
|
result += f"\n\n提示:\n{timeout_note}"
|
||||||
|
if output.temp_file_path:
|
||||||
|
file_note = (
|
||||||
|
"截至命令终止前的完整输出"
|
||||||
|
if timed_out
|
||||||
|
else "完整输出"
|
||||||
|
)
|
||||||
|
result += (
|
||||||
|
"\n\n提示:\n"
|
||||||
|
f"命令输出超过 10KB,仅返回前 {MAX_OUTPUT_PREVIEW_BYTES} 字节内容。\n"
|
||||||
|
f"{file_note}已写入临时文件: {output.temp_file_path}\n"
|
||||||
|
"如需完整内容,请继续读取该文件。"
|
||||||
|
)
|
||||||
if output.stdout:
|
if output.stdout:
|
||||||
result += f"\n\n标准输出:\n{output.stdout}"
|
result += f"\n\n标准输出:\n{output.stdout}"
|
||||||
if output.stderr:
|
if output.stderr:
|
||||||
result += f"\n\n错误输出:\n{output.stderr}"
|
result += f"\n\n错误输出:\n{output.stderr}"
|
||||||
if output.truncated:
|
if output.preview_truncated:
|
||||||
result += "\n\n...(输出内容过长,已截断)"
|
result += "\n\n...(仅展示前 10KB 内容)"
|
||||||
if not output.stdout and not output.stderr:
|
if not output.stdout and not output.stderr:
|
||||||
result += "\n\n(无输出内容)"
|
result += "\n\n(无输出内容)"
|
||||||
return result
|
return result
|
||||||
@@ -252,51 +303,40 @@ class ExecuteCommandTool(MoviePilotTool):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
async with _command_semaphore:
|
async with _command_semaphore:
|
||||||
# 命令输出可能非常大,必须边读边截断,不能使用 communicate() 一次性收集。
|
# 命令输出可能非常大,必须边读边落盘,不能使用 communicate() 一次性收集。
|
||||||
process = await asyncio.create_subprocess_shell(
|
process = await asyncio.create_subprocess_shell(
|
||||||
command, **self._subprocess_kwargs()
|
command, **self._subprocess_kwargs()
|
||||||
)
|
)
|
||||||
output = _CommandOutput(limit=MAX_OUTPUT_CHARS)
|
output = _CommandOutput(preview_limit_bytes=MAX_OUTPUT_PREVIEW_BYTES)
|
||||||
limit_reached = asyncio.Event()
|
|
||||||
wait_task = asyncio.create_task(process.wait())
|
wait_task = asyncio.create_task(process.wait())
|
||||||
limit_task = asyncio.create_task(limit_reached.wait())
|
|
||||||
reader_tasks = [
|
reader_tasks = [
|
||||||
asyncio.create_task(
|
asyncio.create_task(
|
||||||
self._read_stream(
|
self._read_stream(process.stdout, "stdout", output)
|
||||||
process.stdout, "stdout", output, limit_reached
|
|
||||||
)
|
|
||||||
),
|
),
|
||||||
asyncio.create_task(
|
asyncio.create_task(
|
||||||
self._read_stream(
|
self._read_stream(process.stderr, "stderr", output)
|
||||||
process.stderr, "stderr", output, limit_reached
|
|
||||||
)
|
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
timed_out = False
|
timed_out = False
|
||||||
output_limited = False
|
try:
|
||||||
done, _ = await asyncio.wait(
|
await asyncio.wait_for(
|
||||||
{wait_task, limit_task},
|
asyncio.shield(wait_task), timeout=normalized_timeout
|
||||||
timeout=normalized_timeout,
|
)
|
||||||
return_when=asyncio.FIRST_COMPLETED,
|
except asyncio.TimeoutError:
|
||||||
)
|
timed_out = True
|
||||||
|
|
||||||
if wait_task not in done:
|
|
||||||
if limit_task in done:
|
|
||||||
output_limited = True
|
|
||||||
else:
|
|
||||||
timed_out = True
|
|
||||||
await self._cleanup_process(process, wait_task)
|
await self._cleanup_process(process, wait_task)
|
||||||
|
|
||||||
limit_task.cancel()
|
try:
|
||||||
await self._finish_reader_tasks(reader_tasks)
|
await self._finish_reader_tasks(reader_tasks)
|
||||||
|
finally:
|
||||||
|
output.close()
|
||||||
|
|
||||||
return self._format_result(
|
return self._format_result(
|
||||||
exit_code=process.returncode,
|
exit_code=process.returncode,
|
||||||
output=output,
|
output=output,
|
||||||
timeout=normalized_timeout,
|
timeout=normalized_timeout,
|
||||||
timed_out=timed_out,
|
timed_out=timed_out,
|
||||||
output_limited=output_limited,
|
|
||||||
timeout_note=timeout_note,
|
timeout_note=timeout_note,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -89,3 +89,4 @@ openai~=2.33.0
|
|||||||
google-genai~=1.74.0
|
google-genai~=1.74.0
|
||||||
ddgs~=9.10.0
|
ddgs~=9.10.0
|
||||||
websocket-client~=1.8.0
|
websocket-client~=1.8.0
|
||||||
|
pytest~=8.4.0
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
|
import re
|
||||||
import shlex
|
import shlex
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
@@ -8,7 +9,7 @@ import unittest
|
|||||||
|
|
||||||
from app.agent.tools.impl.execute_command import (
|
from app.agent.tools.impl.execute_command import (
|
||||||
ExecuteCommandTool,
|
ExecuteCommandTool,
|
||||||
MAX_OUTPUT_CHARS,
|
MAX_OUTPUT_PREVIEW_BYTES,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -21,6 +22,11 @@ def _python_command(code: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
class TestExecuteCommandTool(unittest.TestCase):
|
class TestExecuteCommandTool(unittest.TestCase):
|
||||||
|
def _temp_file_path_from_result(self, result: str) -> str:
|
||||||
|
match = re.search(r"临时文件: (.+)", result)
|
||||||
|
self.assertIsNotNone(match)
|
||||||
|
return match.group(1).strip()
|
||||||
|
|
||||||
def _run_command(self, command: str, timeout: int = 60) -> str:
|
def _run_command(self, command: str, timeout: int = 60) -> str:
|
||||||
tool = ExecuteCommandTool(session_id="session-1", user_id="10001")
|
tool = ExecuteCommandTool(session_id="session-1", user_id="10001")
|
||||||
return asyncio.run(tool.run(command=command, timeout=timeout))
|
return asyncio.run(tool.run(command=command, timeout=timeout))
|
||||||
@@ -31,9 +37,19 @@ class TestExecuteCommandTool(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
result = self._run_command(command)
|
result = self._run_command(command)
|
||||||
|
temp_file_path = self._temp_file_path_from_result(result)
|
||||||
|
|
||||||
self.assertIn("输出内容过长,已截断", result)
|
self.addCleanup(lambda: os.path.exists(temp_file_path) and os.unlink(temp_file_path))
|
||||||
self.assertLess(len(result), MAX_OUTPUT_CHARS + 500)
|
self.assertIn("命令输出超过 10KB", result)
|
||||||
|
self.assertIn("仅展示前 10KB 内容", result)
|
||||||
|
self.assertIn("如需完整内容,请继续读取该文件", result)
|
||||||
|
self.assertLess(len(result), MAX_OUTPUT_PREVIEW_BYTES + 600)
|
||||||
|
|
||||||
|
with open(temp_file_path, encoding="utf-8") as file_handle:
|
||||||
|
file_content = file_handle.read()
|
||||||
|
|
||||||
|
self.assertIn("[标准输出]", file_content)
|
||||||
|
self.assertGreater(len(file_content), 100000)
|
||||||
|
|
||||||
def test_timeout_returns_partial_output_promptly(self):
|
def test_timeout_returns_partial_output_promptly(self):
|
||||||
command = _python_command(
|
command = _python_command(
|
||||||
@@ -48,6 +64,24 @@ class TestExecuteCommandTool(unittest.TestCase):
|
|||||||
self.assertIn("命令执行超时", result)
|
self.assertIn("命令执行超时", result)
|
||||||
self.assertIn("started", result)
|
self.assertIn("started", result)
|
||||||
|
|
||||||
|
def test_timeout_with_large_output_writes_partial_full_log_to_temp_file(self):
|
||||||
|
command = _python_command(
|
||||||
|
"import sys, time; sys.stdout.write('x' * 20000); sys.stdout.flush(); time.sleep(5)"
|
||||||
|
)
|
||||||
|
|
||||||
|
result = self._run_command(command, timeout=1)
|
||||||
|
temp_file_path = self._temp_file_path_from_result(result)
|
||||||
|
|
||||||
|
self.addCleanup(lambda: os.path.exists(temp_file_path) and os.unlink(temp_file_path))
|
||||||
|
self.assertIn("命令执行超时", result)
|
||||||
|
self.assertIn("截至命令终止前的完整输出已写入临时文件", result)
|
||||||
|
|
||||||
|
with open(temp_file_path, encoding="utf-8") as file_handle:
|
||||||
|
file_content = file_handle.read()
|
||||||
|
|
||||||
|
self.assertIn("[标准输出]", file_content)
|
||||||
|
self.assertGreaterEqual(file_content.count("x"), 20000)
|
||||||
|
|
||||||
def test_timeout_is_capped(self):
|
def test_timeout_is_capped(self):
|
||||||
command = _python_command("print('ok')")
|
command = _python_command("print('ok')")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user