feat: stream Bash progress and guard message actions during replies

This commit is contained in:
yangziyu-hhh
2026-06-10 13:46:50 +08:00
parent f5caba81d6
commit 402e2bfee0
6 changed files with 207 additions and 31 deletions

View File

@@ -1174,10 +1174,21 @@ class AgentStreamExecutor:
# Set tool context
tool.model = self.model
tool.context = self.agent
tool.progress_callback = lambda message: self._emit_event(
"tool_execution_progress",
{
"tool_call_id": tool_id,
"tool_name": tool_name,
"message": message,
}
)
# Execute tool
start_time = time.time()
result: ToolResult = tool.execute_tool(arguments)
try:
result: ToolResult = tool.execute_tool(arguments)
finally:
tool.progress_callback = None
execution_time = time.time() - start_time
result_dict = {
@@ -1719,4 +1730,4 @@ class AgentStreamExecutor:
not as a message. The AgentLLMModel will handle this.
"""
# Don't add system message here - it will be handled separately by the LLM adapter
return self.messages
return self.messages

View File

@@ -38,6 +38,16 @@ class BaseTool:
description: str = "Base tool"
params: dict = {} # Store JSON Schema
model: Optional[Any] = None # LLM model instance, type depends on bot implementation
progress_callback = None
def report_progress(self, message: str):
callback = getattr(self, "progress_callback", None)
if not callback:
return
try:
callback(str(message))
except Exception as e:
logger.debug(f"[{self.name}] progress callback failed: {e}")
@classmethod
def get_json_schema(cls) -> dict:

View File

@@ -4,9 +4,12 @@ Bash tool - Execute bash commands
import os
import re
import signal
import sys
import subprocess
import tempfile
import threading
import time
from typing import Dict, Any
from agent.tools.base_tool import BaseTool, ToolResult
@@ -19,6 +22,8 @@ class Bash(BaseTool):
"""Tool for executing bash commands"""
_IS_WIN = sys.platform == "win32"
_PROGRESS_MAX_BYTES = 32 * 1024
_PROGRESS_INTERVAL = 0.5
name: str = "bash"
description: str = f"""Execute a bash command in the current working directory. Returns stdout and stderr. Output is truncated to last {DEFAULT_MAX_LINES} lines or {DEFAULT_MAX_BYTES // 1024}KB (whichever is hit first). If truncated, full output is saved to a temp file.
@@ -113,17 +118,11 @@ SAFETY:
if command and not command.strip().lower().startswith("chcp"):
command = f"chcp 65001 >nul 2>&1 && {command}"
result = subprocess.run(
result = self._run_streaming(
command,
shell=True,
cwd=self.cwd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
encoding="utf-8",
errors="replace",
timeout=timeout,
env=env,
timeout,
env,
dotenv_vars,
)
logger.debug(f"[Bash] Exit code: {result.returncode}")
@@ -236,6 +235,93 @@ SAFETY:
except Exception as e:
return ToolResult.fail(f"Error executing command: {str(e)}")
def _run_streaming(self, command: str, timeout: int, env: dict, dotenv_vars: dict):
process = subprocess.Popen(
command,
shell=True,
cwd=self.cwd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=env,
start_new_session=not self._IS_WIN,
)
stdout_chunks, stderr_chunks = [], []
recent = bytearray()
recent_lock = threading.Lock()
def drain(stream, chunks):
while True:
chunk = os.read(stream.fileno(), 4096)
if not chunk:
break
chunks.append(chunk)
with recent_lock:
recent.extend(chunk)
if len(recent) > self._PROGRESS_MAX_BYTES:
del recent[:-self._PROGRESS_MAX_BYTES]
readers = [
threading.Thread(target=drain, args=(process.stdout, stdout_chunks), daemon=True),
threading.Thread(target=drain, args=(process.stderr, stderr_chunks), daemon=True),
]
for reader in readers:
reader.start()
started = time.monotonic()
last_reported_at = started
last_snapshot = None
try:
while process.poll() is None:
now = time.monotonic()
elapsed = now - started
if elapsed >= timeout:
self._kill_process(process)
raise subprocess.TimeoutExpired(command, timeout)
if elapsed >= self._PROGRESS_INTERVAL and now - last_reported_at >= self._PROGRESS_INTERVAL:
with recent_lock:
snapshot = bytes(recent).decode("utf-8", errors="replace")
snapshot = self._redact_progress(snapshot, dotenv_vars)
if snapshot and snapshot != last_snapshot:
self.report_progress(snapshot)
last_snapshot = snapshot
last_reported_at = now
time.sleep(0.1)
finally:
if process.poll() is None:
self._kill_process(process)
process.wait()
for reader in readers:
reader.join()
from types import SimpleNamespace
return SimpleNamespace(
returncode=process.returncode,
stdout=b"".join(stdout_chunks).decode("utf-8", errors="replace"),
stderr=b"".join(stderr_chunks).decode("utf-8", errors="replace"),
)
def _kill_process(self, process):
if self._IS_WIN:
process.kill()
else:
try:
os.killpg(process.pid, signal.SIGKILL)
except PermissionError:
process.kill()
@staticmethod
def _redact_progress(text: str, dotenv_vars: dict) -> str:
text = re.sub(
r'(?i)\b(API_KEY|TOKEN|PASSWORD|AUTHORIZATION)\s*=\s*[^\s]+',
lambda match: f"{match.group(1)}=[REDACTED]",
text,
)
for value in dotenv_vars.values():
value = str(value or "")
if len(value) >= 6:
text = text.replace(value, "[REDACTED]")
return text
def _get_safety_warning(self, command: str) -> str:
"""
Get safety warning for absolutely catastrophic commands only.