feat(tools): optimize coding workflows
This commit is contained in:
Xubin Ren
@@ -0,0 +1,409 @@
|
||||
"""Session support for long-running exec workflows."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shutil
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from nanobot.agent.tools.base import Tool, tool_parameters
|
||||
from nanobot.agent.tools.schema import BooleanSchema, IntegerSchema, StringSchema, tool_parameters_schema
|
||||
|
||||
|
||||
DEFAULT_YIELD_MS = 1000
|
||||
MAX_YIELD_MS = 30_000
|
||||
DEFAULT_MAX_OUTPUT_CHARS = 10_000
|
||||
MAX_OUTPUT_CHARS = 50_000
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _SessionPoll:
|
||||
output: str
|
||||
done: bool
|
||||
exit_code: int | None
|
||||
timed_out: bool = False
|
||||
terminated: bool = False
|
||||
stdin_closed: bool = False
|
||||
truncated_chars: int = 0
|
||||
|
||||
|
||||
class _ExecSession:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
process: asyncio.subprocess.Process,
|
||||
timeout: int,
|
||||
) -> None:
|
||||
self.session_id = session_id
|
||||
self.process = process
|
||||
self.deadline = time.monotonic() + timeout
|
||||
self.last_access = time.monotonic()
|
||||
self._chunks: list[str] = []
|
||||
self._lock = asyncio.Lock()
|
||||
self._timed_out = False
|
||||
self._stdout_task = asyncio.create_task(self._read_stream(process.stdout, ""))
|
||||
self._stderr_task = asyncio.create_task(self._read_stream(process.stderr, "STDERR:\n"))
|
||||
|
||||
async def _read_stream(
|
||||
self,
|
||||
stream: asyncio.StreamReader | None,
|
||||
prefix: str,
|
||||
) -> None:
|
||||
if stream is None:
|
||||
return
|
||||
first = True
|
||||
while True:
|
||||
chunk = await stream.read(4096)
|
||||
if not chunk:
|
||||
break
|
||||
text = chunk.decode("utf-8", errors="replace")
|
||||
if prefix and first:
|
||||
text = prefix + text
|
||||
first = False
|
||||
async with self._lock:
|
||||
self._chunks.append(text)
|
||||
|
||||
async def write(self, chars: str) -> str | None:
|
||||
if self.process.returncode is not None:
|
||||
return "session has already exited"
|
||||
if self.process.stdin is None:
|
||||
return "session stdin is not available"
|
||||
try:
|
||||
self.process.stdin.write(chars.encode("utf-8"))
|
||||
await self.process.stdin.drain()
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
return "session stdin is closed"
|
||||
return None
|
||||
|
||||
async def close_stdin(self) -> str | None:
|
||||
if self.process.returncode is not None:
|
||||
return "session has already exited"
|
||||
if self.process.stdin is None:
|
||||
return "session stdin is not available"
|
||||
self.process.stdin.close()
|
||||
with suppress(BrokenPipeError, ConnectionResetError):
|
||||
await self.process.stdin.wait_closed()
|
||||
return None
|
||||
|
||||
async def poll(
|
||||
self,
|
||||
yield_time_ms: int,
|
||||
max_output_chars: int,
|
||||
*,
|
||||
terminated: bool = False,
|
||||
stdin_closed: bool = False,
|
||||
) -> _SessionPoll:
|
||||
self.last_access = time.monotonic()
|
||||
if yield_time_ms > 0 and self.process.returncode is None:
|
||||
await asyncio.sleep(min(yield_time_ms, MAX_YIELD_MS) / 1000)
|
||||
|
||||
if self.process.returncode is None and time.monotonic() >= self.deadline:
|
||||
self._timed_out = True
|
||||
await self.kill()
|
||||
|
||||
if self.process.returncode is not None:
|
||||
with suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(self._stdout_task, self._stderr_task),
|
||||
timeout=2.0,
|
||||
)
|
||||
|
||||
async with self._lock:
|
||||
output = "".join(self._chunks)
|
||||
self._chunks.clear()
|
||||
|
||||
output, truncated = _truncate_output(output, max_output_chars)
|
||||
return _SessionPoll(
|
||||
output=output,
|
||||
done=self.process.returncode is not None,
|
||||
exit_code=self.process.returncode,
|
||||
timed_out=self._timed_out,
|
||||
terminated=terminated,
|
||||
stdin_closed=stdin_closed,
|
||||
truncated_chars=truncated,
|
||||
)
|
||||
|
||||
async def kill(self) -> None:
|
||||
if self.process.returncode is not None:
|
||||
return
|
||||
self.process.kill()
|
||||
with suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(self.process.wait(), timeout=5.0)
|
||||
|
||||
|
||||
class ExecSessionManager:
|
||||
def __init__(self, *, max_sessions: int = 8, idle_timeout: int = 1800) -> None:
|
||||
self.max_sessions = max_sessions
|
||||
self.idle_timeout = idle_timeout
|
||||
self._sessions: dict[str, _ExecSession] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def start(
|
||||
self,
|
||||
*,
|
||||
command: str,
|
||||
cwd: str,
|
||||
env: dict[str, str],
|
||||
timeout: int,
|
||||
shell_program: str | None,
|
||||
login: bool,
|
||||
yield_time_ms: int,
|
||||
max_output_chars: int,
|
||||
) -> tuple[str, _SessionPoll]:
|
||||
async with self._lock:
|
||||
await self._cleanup_locked()
|
||||
if len(self._sessions) >= self.max_sessions:
|
||||
raise RuntimeError(f"maximum exec sessions reached ({self.max_sessions})")
|
||||
process = await self._spawn(command, cwd, env, shell_program, login)
|
||||
session_id = uuid.uuid4().hex[:12]
|
||||
session = _ExecSession(session_id=session_id, process=process, timeout=timeout)
|
||||
self._sessions[session_id] = session
|
||||
|
||||
poll = await session.poll(yield_time_ms, max_output_chars)
|
||||
if poll.done:
|
||||
async with self._lock:
|
||||
self._sessions.pop(session_id, None)
|
||||
return session_id, poll
|
||||
|
||||
async def write(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
chars: str | None,
|
||||
close_stdin: bool,
|
||||
terminate: bool,
|
||||
yield_time_ms: int,
|
||||
max_output_chars: int,
|
||||
) -> _SessionPoll:
|
||||
async with self._lock:
|
||||
await self._cleanup_locked()
|
||||
session = self._sessions.get(session_id)
|
||||
if session is None:
|
||||
raise KeyError(session_id)
|
||||
|
||||
if chars is not None:
|
||||
error = await session.write(chars)
|
||||
if error:
|
||||
raise RuntimeError(error)
|
||||
stdin_closed = False
|
||||
if close_stdin:
|
||||
error = await session.close_stdin()
|
||||
if error:
|
||||
raise RuntimeError(error)
|
||||
stdin_closed = True
|
||||
if terminate:
|
||||
await session.kill()
|
||||
poll = await session.poll(
|
||||
yield_time_ms,
|
||||
max_output_chars,
|
||||
terminated=terminate,
|
||||
stdin_closed=stdin_closed,
|
||||
)
|
||||
if poll.done:
|
||||
async with self._lock:
|
||||
self._sessions.pop(session_id, None)
|
||||
return poll
|
||||
|
||||
async def _cleanup_locked(self) -> None:
|
||||
now = time.monotonic()
|
||||
stale = [
|
||||
session_id
|
||||
for session_id, session in self._sessions.items()
|
||||
if session.process.returncode is not None
|
||||
or now - session.last_access > self.idle_timeout
|
||||
]
|
||||
for session_id in stale:
|
||||
session = self._sessions.pop(session_id)
|
||||
await session.kill()
|
||||
|
||||
async def _spawn(
|
||||
self,
|
||||
command: str,
|
||||
cwd: str,
|
||||
env: dict[str, str],
|
||||
shell_program: str | None,
|
||||
login: bool,
|
||||
) -> asyncio.subprocess.Process:
|
||||
from nanobot.agent.tools import shell
|
||||
|
||||
if shell._IS_WINDOWS:
|
||||
return await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
)
|
||||
shell_program = shell_program or shutil.which("bash") or "/bin/bash"
|
||||
args = [shell_program]
|
||||
if login and shell_program.rsplit("/", 1)[-1] in {"bash", "zsh"}:
|
||||
args.append("-l")
|
||||
args.extend(["-c", command])
|
||||
return await asyncio.create_subprocess_exec(
|
||||
*args,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_EXEC_SESSION_MANAGER = ExecSessionManager()
|
||||
|
||||
|
||||
def clamp_session_int(value: int | None, default: int, minimum: int, maximum: int) -> int:
|
||||
if value is None:
|
||||
return default
|
||||
return min(max(value, minimum), maximum)
|
||||
|
||||
|
||||
def _truncate_output(output: str, max_output_chars: int) -> tuple[str, int]:
|
||||
if len(output) <= max_output_chars:
|
||||
return output, 0
|
||||
half = max_output_chars // 2
|
||||
omitted = len(output) - max_output_chars
|
||||
return (
|
||||
output[:half]
|
||||
+ f"\n\n... ({omitted:,} chars truncated) ...\n\n"
|
||||
+ output[-half:],
|
||||
omitted,
|
||||
)
|
||||
|
||||
|
||||
def format_session_poll(session_id: str, poll: _SessionPoll) -> str:
|
||||
parts = [poll.output] if poll.output else []
|
||||
if poll.truncated_chars:
|
||||
parts.append(f"(output truncated by {poll.truncated_chars:,} chars)")
|
||||
if poll.timed_out:
|
||||
parts.append("Error: Command timed out; session was terminated.")
|
||||
if poll.terminated and not poll.timed_out:
|
||||
parts.append("Session terminated.")
|
||||
if poll.stdin_closed:
|
||||
parts.append("Stdin closed.")
|
||||
if poll.done:
|
||||
parts.append(f"Exit code: {poll.exit_code}")
|
||||
else:
|
||||
parts.append(f"Process running. session_id: {session_id}")
|
||||
return "\n".join(parts) if parts else "(no output yet)"
|
||||
|
||||
|
||||
@tool_parameters(
|
||||
tool_parameters_schema(
|
||||
session_id=StringSchema("Session id returned by exec when yield_time_ms is used."),
|
||||
chars=StringSchema(
|
||||
"Bytes/text to write to stdin. Omit or pass an empty string to only poll recent output.",
|
||||
nullable=True,
|
||||
),
|
||||
close_stdin=BooleanSchema(
|
||||
description="Close stdin after writing chars. Useful for commands waiting for EOF.",
|
||||
default=False,
|
||||
),
|
||||
terminate=BooleanSchema(
|
||||
description="Terminate the running exec session.",
|
||||
default=False,
|
||||
),
|
||||
yield_time_ms=IntegerSchema(
|
||||
DEFAULT_YIELD_MS,
|
||||
description="Milliseconds to wait before returning recent output (default 1000, max 30000).",
|
||||
minimum=0,
|
||||
maximum=MAX_YIELD_MS,
|
||||
),
|
||||
max_output_chars=IntegerSchema(
|
||||
DEFAULT_MAX_OUTPUT_CHARS,
|
||||
description="Maximum output characters to return from this poll (default 10000, max 50000).",
|
||||
minimum=1000,
|
||||
maximum=MAX_OUTPUT_CHARS,
|
||||
),
|
||||
max_output_tokens=IntegerSchema(
|
||||
DEFAULT_MAX_OUTPUT_CHARS,
|
||||
description="Compatibility alias for max_output_chars. The current runtime uses a character budget.",
|
||||
minimum=1000,
|
||||
maximum=MAX_OUTPUT_CHARS,
|
||||
nullable=True,
|
||||
),
|
||||
required=["session_id"],
|
||||
)
|
||||
)
|
||||
class WriteStdinTool(Tool):
|
||||
"""Write to or poll a running exec session."""
|
||||
|
||||
_scopes = {"core", "subagent"}
|
||||
config_key = "exec"
|
||||
|
||||
@classmethod
|
||||
def config_cls(cls):
|
||||
from nanobot.agent.tools.shell import ExecToolConfig
|
||||
|
||||
return ExecToolConfig
|
||||
|
||||
@classmethod
|
||||
def enabled(cls, ctx: Any) -> bool:
|
||||
return ctx.config.exec.enable
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
manager: ExecSessionManager | None = None,
|
||||
) -> None:
|
||||
self._manager = manager or DEFAULT_EXEC_SESSION_MANAGER
|
||||
|
||||
@classmethod
|
||||
def create(cls, ctx: Any) -> Tool:
|
||||
return cls()
|
||||
|
||||
@property
|
||||
def exclusive(self) -> bool:
|
||||
return True
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "write_stdin"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return (
|
||||
"Write text to a running exec session and return recent output. "
|
||||
"Use chars='' to poll without writing. Set close_stdin=true to send EOF, "
|
||||
"or terminate=true to stop the session. Sessions finish automatically "
|
||||
"when their process exits."
|
||||
)
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
session_id: str,
|
||||
chars: str | None = None,
|
||||
close_stdin: bool = False,
|
||||
terminate: bool = False,
|
||||
yield_time_ms: int | None = None,
|
||||
max_output_chars: int | None = None,
|
||||
max_output_tokens: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
try:
|
||||
if max_output_chars is None:
|
||||
max_output_chars = max_output_tokens
|
||||
poll = await self._manager.write(
|
||||
session_id=session_id,
|
||||
chars=chars,
|
||||
close_stdin=close_stdin,
|
||||
terminate=terminate,
|
||||
yield_time_ms=clamp_session_int(yield_time_ms, DEFAULT_YIELD_MS, 0, MAX_YIELD_MS),
|
||||
max_output_chars=clamp_session_int(
|
||||
max_output_chars,
|
||||
DEFAULT_MAX_OUTPUT_CHARS,
|
||||
1000,
|
||||
MAX_OUTPUT_CHARS,
|
||||
),
|
||||
)
|
||||
return format_session_poll(session_id, poll)
|
||||
except KeyError:
|
||||
return f"Error: exec session not found: {session_id}"
|
||||
except Exception as exc:
|
||||
return f"Error writing to exec session: {exc}"
|
||||
Reference in New Issue
Block a user