feat(tools): optimize coding workflows

nanobot/agent/tools/shell.py

@@ -8,6 +8,7 @@ import re
 import shutil
 import sys
 from contextlib import suppress
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
@@ -15,8 +16,17 @@ from loguru import logger
 from pydantic import Field
 
 from nanobot.agent.tools.base import Tool, tool_parameters
+from nanobot.agent.tools.exec_session import (
+    DEFAULT_MAX_OUTPUT_CHARS,
+    DEFAULT_YIELD_MS,
+    DEFAULT_EXEC_SESSION_MANAGER,
+    MAX_OUTPUT_CHARS,
+    MAX_YIELD_MS,
+    clamp_session_int,
+    format_session_poll,
+)
 from nanobot.agent.tools.sandbox import wrap_command
-from nanobot.agent.tools.schema import IntegerSchema, StringSchema, tool_parameters_schema
+from nanobot.agent.tools.schema import BooleanSchema, IntegerSchema, StringSchema, tool_parameters_schema
 from nanobot.config.paths import get_media_dir
 from nanobot.config.schema import Base
 
@@ -44,10 +54,22 @@ class ExecToolConfig(Base):
     deny_patterns: list[str] = Field(default_factory=list)
 
 
+@dataclass(slots=True)
+class _PreparedCommand:
+    command: str
+    cwd: str
+    env: dict[str, str]
+    timeout: int
+    shell_program: str | None
+    login: bool
+
+
 @tool_parameters(
     tool_parameters_schema(
         command=StringSchema("The shell command to execute"),
+        cmd=StringSchema("Compatibility alias for command"),
         working_dir=StringSchema("Optional working directory for the command"),
+        workdir=StringSchema("Compatibility alias for working_dir"),
         timeout=IntegerSchema(
             60,
             description=(
@@ -57,7 +79,44 @@ class ExecToolConfig(Base):
             minimum=1,
             maximum=600,
         ),
-        required=["command"],
+        shell=StringSchema(
+            "Optional shell binary to launch. On Unix, supports sh, bash, or zsh.",
+            nullable=True,
+        ),
+        login=BooleanSchema(
+            description="Whether to run bash/zsh with login shell semantics (default true).",
+            default=True,
+            nullable=True,
+        ),
+        yield_time_ms=IntegerSchema(
+            description=(
+                "Optional milliseconds to wait before returning output. "
+                "When set, a still-running command returns a session_id that "
+                "can be polled or written to with write_stdin. Omit this field "
+                "to keep one-shot exec behavior."
+            ),
+            minimum=0,
+            maximum=MAX_YIELD_MS,
+            nullable=True,
+        ),
+        max_output_chars=IntegerSchema(
+            description=(
+                "Maximum output characters to return when yield_time_ms is used "
+                "(default 10000, max 50000)."
+            ),
+            minimum=1000,
+            maximum=MAX_OUTPUT_CHARS,
+            nullable=True,
+        ),
+        max_output_tokens=IntegerSchema(
+            description=(
+                "Compatibility alias for max_output_chars. The current runtime "
+                "uses a character budget."
+            ),
+            minimum=1000,
+            maximum=MAX_OUTPUT_CHARS,
+            nullable=True,
+        ),
     )
 )
 class ExecTool(Tool):
@@ -98,6 +157,7 @@ class ExecTool(Tool):
         sandbox: str = "",
         path_append: str = "",
         allowed_env_keys: list[str] | None = None,
+        session_manager: Any | None = None,
     ):
         self.timeout = timeout
         self.working_dir = working_dir
@@ -125,6 +185,7 @@ class ExecTool(Tool):
         self.restrict_to_workspace = restrict_to_workspace
         self.path_append = path_append
         self.allowed_env_keys = allowed_env_keys or []
+        self._session_manager = session_manager or DEFAULT_EXEC_SESSION_MANAGER
 
     @property
     def name(self) -> str:
@@ -153,7 +214,10 @@ class ExecTool(Tool):
             "Prefer read_file/write_file/edit_file over cat/echo/sed, "
             "and grep/glob over shell find/grep. "
             "Use -y or --yes flags to avoid interactive prompts. "
-            "Output is truncated at 10 000 chars; timeout defaults to 60s."
+            "For long-running or interactive commands, pass yield_time_ms; "
+            "if the command keeps running, exec returns a session_id that can "
+            "be polled or written to with write_stdin. Output is truncated at "
+            "10 000 chars; timeout defaults to 60s."
         )
 
     @property
@@ -161,9 +225,111 @@ class ExecTool(Tool):
         return True
 
     async def execute(
-        self, command: str, working_dir: str | None = None,
-        timeout: int | None = None, **kwargs: Any,
+        self, command: str | None = None, cmd: str | None = None,
+        working_dir: str | None = None, workdir: str | None = None,
+        timeout: int | None = None, shell: str | None = None,
+        login: bool | None = None, yield_time_ms: int | None = None,
+        max_output_chars: int | None = None,
+        max_output_tokens: int | None = None,
+        **kwargs: Any,
     ) -> str:
+        command = command or cmd
+        working_dir = working_dir or workdir
+        if not command:
+            return "Error: Missing command. Provide command or cmd."
+        if max_output_chars is None:
+            max_output_chars = max_output_tokens
+
+        prepared = self._prepare_command(command, working_dir, timeout, shell, login)
+        if isinstance(prepared, str):
+            return prepared
+
+        if yield_time_ms is not None:
+            return await self._execute_session(prepared, yield_time_ms, max_output_chars)
+
+        try:
+            process = await self._spawn(
+                prepared.command,
+                prepared.cwd,
+                prepared.env,
+                prepared.shell_program,
+                prepared.login,
+            )
+
+            try:
+                stdout, stderr = await asyncio.wait_for(
+                    process.communicate(),
+                    timeout=prepared.timeout,
+                )
+            except asyncio.TimeoutError:
+                await self._kill_process(process)
+                return f"Error: Command timed out after {prepared.timeout} seconds"
+            except asyncio.CancelledError:
+                await self._kill_process(process)
+                raise
+
+            output_parts = []
+
+            if stdout:
+                output_parts.append(stdout.decode("utf-8", errors="replace"))
+
+            if stderr:
+                stderr_text = stderr.decode("utf-8", errors="replace")
+                if stderr_text.strip():
+                    output_parts.append(f"STDERR:\n{stderr_text}")
+
+            output_parts.append(f"\nExit code: {process.returncode}")
+
+            result = "\n".join(output_parts) if output_parts else "(no output)"
+
+            max_len = clamp_session_int(max_output_chars, self._MAX_OUTPUT, 1000, MAX_OUTPUT_CHARS)
+            if len(result) > max_len:
+                half = max_len // 2
+                result = (
+                    result[:half]
+                    + f"\n\n... ({len(result) - max_len:,} chars truncated) ...\n\n"
+                    + result[-half:]
+                )
+
+            return result
+
+        except Exception as e:
+            return f"Error executing command: {str(e)}"
+
+    async def _execute_session(
+        self,
+        prepared: _PreparedCommand,
+        yield_time_ms: int | None,
+        max_output_chars: int | None,
+    ) -> str:
+        try:
+            session_id, poll = await self._session_manager.start(
+                command=prepared.command,
+                cwd=prepared.cwd,
+                env=prepared.env,
+                timeout=prepared.timeout,
+                shell_program=prepared.shell_program,
+                login=prepared.login,
+                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 Exception as exc:
+            return f"Error executing command: {exc}"
+
+    def _prepare_command(
+        self,
+        command: str,
+        working_dir: str | None = None,
+        timeout: int | None = None,
+        shell: str | None = None,
+        login: bool | None = None,
+    ) -> _PreparedCommand | str:
         cwd = working_dir or self.working_dir or os.getcwd()
 
         # Prevent an LLM-supplied working_dir from escaping the configured
@@ -211,52 +377,24 @@ class ExecTool(Tool):
                 env["NANOBOT_PATH_APPEND"] = self.path_append
                 command = f'export PATH="$PATH{os.pathsep}$NANOBOT_PATH_APPEND"; {command}'
 
-        try:
-            process = await self._spawn(command, cwd, env)
+        shell_program, shell_error = self._resolve_shell(shell)
+        if shell_error:
+            return shell_error
 
-            try:
-                stdout, stderr = await asyncio.wait_for(
-                    process.communicate(),
-                    timeout=effective_timeout,
-                )
-            except asyncio.TimeoutError:
-                await self._kill_process(process)
-                return f"Error: Command timed out after {effective_timeout} seconds"
-            except asyncio.CancelledError:
-                await self._kill_process(process)
-                raise
-
-            output_parts = []
-
-            if stdout:
-                output_parts.append(stdout.decode("utf-8", errors="replace"))
-
-            if stderr:
-                stderr_text = stderr.decode("utf-8", errors="replace")
-                if stderr_text.strip():
-                    output_parts.append(f"STDERR:\n{stderr_text}")
-
-            output_parts.append(f"\nExit code: {process.returncode}")
-
-            result = "\n".join(output_parts) if output_parts else "(no output)"
-
-            max_len = self._MAX_OUTPUT
-            if len(result) > max_len:
-                half = max_len // 2
-                result = (
-                    result[:half]
-                    + f"\n\n... ({len(result) - max_len:,} chars truncated) ...\n\n"
-                    + result[-half:]
-                )
-
-            return result
-
-        except Exception as e:
-            return f"Error executing command: {str(e)}"
+        return _PreparedCommand(
+            command=command,
+            cwd=cwd,
+            env=env,
+            timeout=effective_timeout,
+            shell_program=shell_program,
+            login=True if login is None else login,
+        )
 
     @staticmethod
     async def _spawn(
         command: str, cwd: str, env: dict[str, str],
+        shell_program: str | None = None,
+        login: bool = True,
     ) -> asyncio.subprocess.Process:
         """Launch *command* in a platform-appropriate shell."""
         if _IS_WINDOWS:
@@ -272,9 +410,13 @@ class ExecTool(Tool):
                 cwd=cwd,
                 env=env,
             )
-        bash = shutil.which("bash") or "/bin/bash"
+        shell_program = shell_program or shutil.which("bash") or "/bin/bash"
+        args = [shell_program]
+        if login and Path(shell_program).name in {"bash", "zsh"}:
+            args.append("-l")
+        args.extend(["-c", command])
         return await asyncio.create_subprocess_exec(
-            bash, "-l", "-c", command,
+            *args,
             stdin=asyncio.subprocess.DEVNULL,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
@@ -282,6 +424,31 @@ class ExecTool(Tool):
             env=env,
         )
 
+    @staticmethod
+    def _resolve_shell(shell: str | None) -> tuple[str | None, str | None]:
+        if not shell:
+            return None, None
+        if _IS_WINDOWS:
+            return None, "Error: shell parameter is not supported on Windows"
+        if "\0" in shell or "\n" in shell or "\r" in shell:
+            return None, "Error: shell contains invalid characters"
+        allowed = {"sh", "bash", "zsh"}
+        path = Path(shell).expanduser()
+        if path.is_absolute():
+            if path.name not in allowed:
+                return None, f"Error: unsupported shell {shell!r}. Allowed: bash, sh, zsh"
+            if not path.is_file() or not os.access(path, os.X_OK):
+                return None, f"Error: shell is not executable: {shell}"
+            return str(path), None
+        if "/" in shell or "\\" in shell:
+            return None, "Error: shell must be a shell name or absolute path"
+        if shell not in allowed:
+            return None, f"Error: unsupported shell {shell!r}. Allowed: bash, sh, zsh"
+        resolved = shutil.which(shell)
+        if not resolved:
+            return None, f"Error: shell not found: {shell}"
+        return resolved, None
+
     @staticmethod
     async def _kill_process(process: asyncio.subprocess.Process) -> None:
         """Kill a subprocess and reap it to prevent zombies."""