Code MVP: a tool registry and a shell tool

"""
chapter 06: tools and tool execution.
A tool registry with schema validation, plus a shell tool that returns a
STRUCTURED, TRUNCATED result (exit code, timing, head+tail of output).
This is the agent's hands; the capstone plugs Chapter 7's approvals in front.
"""
import subprocess, time
from dataclasses import dataclass
from typing import Callable

# --- Truncation: keep the head and tail, elide the noisy middle ---
def truncate_output(text: str, head: int = 50, tail: int = 50) -> str:
    lines = text.splitlines()
    if len(lines) <= head + tail:
        return text
    omitted = len(lines) - head - tail
    return "\n".join(lines[:head] + [f"... ({omitted} lines omitted) ..."] + lines[-tail:])

# --- A structured result the model can actually reason about ---
def format_result(exit_code: int, output: str, seconds: float) -> str:
    total_lines = len(output.splitlines())
    return (f"Exit code: {exit_code}\n"
            f"Wall time: {seconds:.2f} seconds\n"
            f"Total output lines: {total_lines}\n"
            f"Output:\n{truncate_output(output)}")

@dataclass
class Tool:
    name: str
    description: str
    parameters: dict          # JSON schema (Chapter 3)
    handler: Callable         # the function that does the work

class ToolRegistry:
    def __init__(self):
        self._tools: dict[str, Tool] = {}

    def register(self, tool: Tool):
        self._tools[tool.name] = tool

    def schemas(self) -> list:
        """The 'tools' field for the prompt (Chapter 3)."""
        return [{"type": "function", "name": t.name,
                 "description": t.description, "parameters": t.parameters}
                for t in self._tools.values()]

    def validate(self, name: str, args: dict) -> str | None:
        """Cheap schema check: required keys present? Returns an error or None."""
        tool = self._tools.get(name)
        if tool is None:
            return f"Unknown tool: {name}"
        required = tool.parameters.get("required", [])
        missing = [k for k in required if k not in args]
        return f"Missing required args: {missing}" if missing else None

    def call(self, name: str, args: dict) -> str:
        error = self.validate(name, args)           # step 1 & 2: parse + validate
        if error:
            return f"Exit code: 1\nError: {error}"
        # step 3 (permissions) is inserted here by Chapter 7.
        return self._tools[name].handler(**args)    # step 4: execute

# --- The shell tool itself: execute, capture, time, format ---
def shell_handler(command: str, timeout_ms: int = 10000) -> str:
    start = time.time()
    try:
        proc = subprocess.run(command, shell=True, capture_output=True,
                              text=True, timeout=timeout_ms / 1000)
        combined = proc.stdout + proc.stderr            # step 5: capture
        return format_result(proc.returncode, combined, time.time() - start)
    except subprocess.TimeoutExpired:
        return format_result(124, "Command timed out", time.time() - start)

if __name__ == "__main__":
    registry = ToolRegistry()
    registry.register(Tool(
        name="shell", description="Run a shell command",
        parameters={"type": "object",
                    "properties": {"command": {"type": "string"},
                                   "timeout_ms": {"type": "number"}},
                    "required": ["command"]},
        handler=shell_handler))

    print(registry.call("shell", {"command": "echo hi && seq 1 200"}))
    print("---")
    print(registry.call("shell", {}))   # missing required arg -> clean error

Run it and you will see a structured result with an exit code, a wall time, and the 200-line output neatly truncated to its head and tail. Call it with bad arguments and you get a clean validation error instead of a crash. That robustness is the whole point.