johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! Code mode: every MCP tool call is driven through Monty.
//!
//! The model sees one tool, `run_code`. Its Python calls the MCP tools as
//! `async def` functions; Monty executes the program and only its result
//! reaches the model. See `docs/adrs/0004-code-mode.md`.

pub mod builtins;
pub mod catalog;
pub mod convert;
pub mod dispatch;
mod engine;
mod pool;
mod protocol;
pub mod scope;
pub mod watch;

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use futures::stream::{FuturesUnordered, StreamExt};
use serde_json::{json, Value};
use tokio::sync::Mutex;

use self::builtins::Builtins;
use self::catalog::Catalog;
use self::engine::Engine;
use self::pool::PoolSession;
pub use self::protocol::Completion;
use self::protocol::{CallOutcome, HostEvent, Job, SandboxSession};
pub use self::watch::{WatchLimits, WatchManager, WatchNotice};
use crate::tools::{content_for_log, ToolCallOutput, ToolDefinition, ToolError, ToolExecutor};

/// Name of the single model-facing tool.
pub const RUN_CODE_TOOL: &str = "run_code";

/// Python helpers fed into every new REPL scope.
pub const HELPERS: &str = include_str!("helpers.py");

const PREVIEW_CHARS: usize = 2_000;

/// Sandbox limits for each `run_code` snippet.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CodeModeLimits {
    /// Python execution time per snippet, excluding time awaiting tools.
    pub max_duration: Duration,
    /// Wall-clock time per `run_code`, including tool calls.
    pub max_wall: Duration,
    /// Time per nested tool call.
    pub tool_timeout: Duration,
    /// Nested tool calls per snippet.
    pub max_tool_calls: usize,
    /// Host suspensions (calls, lookups, future resolutions) per snippet.
    pub max_suspensions: usize,
    /// Collected `print()` output per snippet.
    pub max_print_bytes: usize,
    /// Characters of `run_code` output returned to the model.
    pub max_result_chars: usize,
    /// Worker heap cap (pool sandbox only).
    pub max_memory: usize,
}

impl Default for CodeModeLimits {
    fn default() -> Self {
        Self {
            max_duration: Duration::from_secs(30),
            max_wall: Duration::from_secs(300),
            tool_timeout: Duration::from_secs(120),
            max_tool_calls: 100,
            max_suspensions: 10_000,
            max_print_bytes: 1024 * 1024,
            max_result_chars: 20_000,
            max_memory: 4 * 1024 * 1024 * 1024,
        }
    }
}

/// Where snippets run.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum SandboxKind {
    /// Pool when a worker binary is found, otherwise in-process.
    Auto,
    /// Engine thread inside the agent; no heap cap.
    #[default]
    InProcess,
    /// `monty` worker subprocesses; crash isolation and a heap cap.
    Pool,
}

impl std::str::FromStr for SandboxKind {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "auto" => Ok(Self::Auto),
            "in-process" => Ok(Self::InProcess),
            "pool" => Ok(Self::Pool),
            other => Err(format!(
                "unknown sandbox {other:?}; use auto, in-process, or pool"
            )),
        }
    }
}

/// Code mode configuration.
#[derive(Clone, Debug, Default)]
pub struct CodeModeOptions {
    pub limits: CodeModeLimits,
    /// Log every nested call (mutating calls are always logged).
    pub verbose: bool,
    pub sandbox: SandboxKind,
    /// Worker binary; falls back to `MONTY_WORKER`, then `monty` on `PATH`.
    pub monty_worker: Option<PathBuf>,
    /// Operator Python run in every new scope after the helpers.
    pub preamble: Option<String>,
    /// MCP resources; enables `list_resources`/`read_resource` and watches.
    pub resources: Option<Arc<dyn crate::resources::ResourceSource>>,
    pub watch_limits: WatchLimits,
    /// Saved watch specs; restored by `restore_watches`.
    pub watch_state: Option<PathBuf>,
}

impl std::fmt::Debug for dyn crate::resources::ResourceSource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("ResourceSource")
    }
}

#[derive(Clone)]
enum Sandbox {
    InProcess,
    Pool(Arc<monty_pool::Pool>),
}

pub(crate) struct Shared {
    inner: Arc<dyn ToolExecutor>,
    catalog: Arc<Catalog>,
    definition: ToolDefinition,
    limits: CodeModeLimits,
    verbose: bool,
    sandbox: Sandbox,
    preamble: Arc<String>,
    builtins: Arc<Builtins>,
}

/// Thread REPLs kept alive across runs; least recently used are dropped.
const MAX_THREAD_SCOPES: usize = 8;

/// Wraps an executor (normally `McpToolExecutor`) behind a single `run_code`
/// tool. `thread_scope` keeps one REPL per conversation thread, so variables
/// survive across that thread's runs; `run_scope` is a fresh REPL.
pub struct CodeModeToolExecutor {
    shared: Arc<Shared>,
    threads: std::sync::Mutex<Vec<(String, Arc<CodeModeRun>)>>,
    notices: std::sync::Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<WatchNotice>>>,
}

impl CodeModeToolExecutor {
    /// In-process sandbox with default limits.
    pub async fn new(inner: Arc<dyn ToolExecutor>) -> Result<Self, ToolError> {
        Self::with_options(inner, CodeModeOptions::default()).await
    }

    pub async fn with_limits(
        inner: Arc<dyn ToolExecutor>,
        limits: CodeModeLimits,
        verbose: bool,
    ) -> Result<Self, ToolError> {
        Self::with_options(
            inner,
            CodeModeOptions {
                limits,
                verbose,
                ..CodeModeOptions::default()
            },
        )
        .await
    }

    pub async fn with_options(
        inner: Arc<dyn ToolExecutor>,
        options: CodeModeOptions,
    ) -> Result<Self, ToolError> {
        let tools = inner.list_tools().await?;
        let sandbox = select_sandbox(&options).await?;
        let mut preamble = HELPERS.to_string();
        if let Some(extra) = &options.preamble {
            preamble.push('\n');
            preamble.push_str(extra);
        }
        let preamble = Arc::new(preamble);
        let (notices_tx, notices_rx) = tokio::sync::mpsc::unbounded_channel();

        // Watch filters: read-only MCP tools plus list/read built-ins, tighter
        // limits, own scope per watch.
        let watches = options.resources.as_ref().map(|resources| {
            let read_builtins = Builtins {
                resources: Some(resources.clone()),
                watches: None,
            };
            let read_only_tools: Vec<ToolDefinition> = tools
                .iter()
                .filter(|tool| !tool.is_mutating())
                .cloned()
                .collect();
            let filter_catalog = Arc::new(Catalog::with_builtins(
                &read_builtins.definitions(),
                &read_only_tools,
            ));
            let filter_limits = CodeModeLimits {
                max_duration: Duration::from_secs(2),
                max_wall: Duration::from_secs(30),
                max_tool_calls: 10,
                ..options.limits.clone()
            };
            let filter_shared = Arc::new(Shared {
                inner: inner.clone(),
                definition: run_code_definition(&filter_catalog, &filter_limits),
                catalog: filter_catalog,
                limits: filter_limits,
                verbose: options.verbose,
                sandbox: sandbox.clone(),
                preamble: preamble.clone(),
                builtins: Arc::new(read_builtins),
            });
            Arc::new(WatchManager::new(
                resources.clone(),
                filter_shared,
                notices_tx,
                options.watch_limits.clone(),
                options.verbose,
                options.watch_state.clone(),
            ))
        });

        let builtins = Builtins {
            resources: options.resources.clone(),
            watches,
        };
        let catalog = Arc::new(Catalog::with_builtins(&builtins.definitions(), &tools));
        let definition = run_code_definition(&catalog, &options.limits);
        Ok(Self {
            shared: Arc::new(Shared {
                inner,
                catalog,
                definition,
                limits: options.limits,
                verbose: options.verbose,
                sandbox,
                preamble,
                builtins: Arc::new(builtins),
            }),
            threads: std::sync::Mutex::new(Vec::new()),
            notices: std::sync::Mutex::new(Some(notices_rx)),
        })
    }

    /// Watch notices for the agent loop (once; later calls get `None`).
    pub fn take_watch_notices(&self) -> Option<tokio::sync::mpsc::UnboundedReceiver<WatchNotice>> {
        self.notices.lock().expect("notice lock").take()
    }

    /// Restarts watches saved under `CodeModeOptions::watch_state`.
    pub async fn restore_watches(&self) -> usize {
        match &self.shared.builtins.watches {
            Some(watches) => watches.restore().await,
            None => 0,
        }
    }

    pub fn catalog(&self) -> &Catalog {
        &self.shared.catalog
    }

    /// `"pool"` or `"in-process"`.
    pub fn sandbox_name(&self) -> &'static str {
        match self.shared.sandbox {
            Sandbox::InProcess => "in-process",
            Sandbox::Pool(_) => "pool",
        }
    }

    /// A fresh REPL scope.
    pub fn new_run(&self) -> Result<CodeModeRun, ToolError> {
        CodeModeRun::new(self.shared.clone())
    }
}

async fn select_sandbox(options: &CodeModeOptions) -> Result<Sandbox, ToolError> {
    let worker = || find_worker(options.monty_worker.clone());
    match options.sandbox {
        SandboxKind::InProcess => Ok(Sandbox::InProcess),
        SandboxKind::Pool => {
            let path = worker().ok_or_else(|| {
                ToolError::new(
                    "--code-sandbox pool: no monty worker found; install with `cargo install monty-runtime --locked` or pass --monty-worker",
                )
            })?;
            pool::connect(path, &options.limits)
                .await
                .map(Sandbox::Pool)
                .map_err(ToolError::new)
        }
        SandboxKind::Auto => match worker() {
            Some(path) => match pool::connect(path.clone(), &options.limits).await {
                Ok(pool) => Ok(Sandbox::Pool(pool)),
                Err(error) => {
                    warn_no_heap_cap(&format!("{} ({error})", path.display()));
                    Ok(Sandbox::InProcess)
                }
            },
            None => {
                warn_no_heap_cap("no monty worker found");
                Ok(Sandbox::InProcess)
            }
        },
    }
}

fn warn_no_heap_cap(reason: &str) {
    eprintln!(
        "{}",
        json!({
            "level": "warn",
            "event": "code_mode.sandbox",
            "sandbox": "in-process",
            "heap_cap": false,
            "reason": reason,
            "hint": "cargo install monty-runtime --locked",
        })
    );
}

/// `explicit`, then `MONTY_WORKER`, then `monty` on `PATH`.
pub fn find_worker(explicit: Option<PathBuf>) -> Option<PathBuf> {
    if let Some(path) = explicit {
        return path.is_file().then_some(path);
    }
    if let Some(path) = std::env::var_os("MONTY_WORKER").map(PathBuf::from) {
        return path.is_file().then_some(path);
    }
    std::env::var_os("PATH").and_then(|paths| {
        std::env::split_paths(&paths)
            .map(|dir| dir.join("monty"))
            .find(|candidate| candidate.is_file())
    })
}

#[async_trait]
impl ToolExecutor for CodeModeToolExecutor {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        Ok(vec![self.shared.definition.clone()])
    }

    /// Without a run scope every call gets a fresh REPL.
    async fn call_tool(&self, name: &str, arguments: Value) -> Result<ToolCallOutput, ToolError> {
        self.new_run()?.call_tool(name, arguments).await
    }

    fn run_scope(&self) -> Option<Arc<dyn ToolExecutor>> {
        match self.new_run() {
            Ok(run) => Some(Arc::new(run)),
            Err(error) => {
                crate::log::emit(
                    "error",
                    "code_mode.engine_failed",
                    json!({"error": error.to_string()}),
                );
                None
            }
        }
    }

    /// One REPL per thread. The runtime never overlaps runs on a thread, so
    /// the scope needs no extra locking beyond its own session mutex.
    fn thread_scope(&self, thread_id: &str) -> Option<Arc<dyn ToolExecutor>> {
        let mut threads = self.threads.lock().expect("thread scope lock");
        if let Some(index) = threads.iter().position(|(id, _)| id == thread_id) {
            let entry = threads.remove(index);
            let run = entry.1.clone();
            threads.push(entry);
            return Some(run);
        }
        let run = match self.new_run() {
            Ok(run) => Arc::new(run),
            Err(error) => {
                crate::log::emit(
                    "error",
                    "code_mode.engine_failed",
                    json!({"error": error.to_string(), "thread_id": thread_id}),
                );
                return None;
            }
        };
        threads.push((thread_id.to_string(), run.clone()));
        if threads.len() > MAX_THREAD_SCOPES {
            threads.remove(0);
        }
        Some(run)
    }
}

/// One REPL scope: variables persist across `run_code` calls until the scope
/// is dropped or `restart` is passed.
pub struct CodeModeRun {
    shared: Arc<Shared>,
    session: Mutex<Box<dyn SandboxSession>>,
    /// Names bound by earlier snippets and their last description.
    variables: std::sync::Mutex<Variables>,
}

#[derive(Default)]
struct Variables {
    candidates: Vec<String>,
    described: Vec<(String, String)>,
}

impl CodeModeRun {
    pub(crate) fn new(shared: Arc<Shared>) -> Result<Self, ToolError> {
        let session: Box<dyn SandboxSession> = match &shared.sandbox {
            Sandbox::InProcess => Box::new(
                Engine::spawn(
                    shared.catalog.clone(),
                    shared.limits.clone(),
                    shared.preamble.clone(),
                )
                .map_err(|error| {
                    ToolError::new(format!("failed to start code mode engine: {error}"))
                })?,
            ),
            Sandbox::Pool(pool) => Box::new(PoolSession::spawn(
                pool.clone(),
                shared.catalog.clone(),
                shared.limits.clone(),
                shared.preamble.clone(),
            )),
        };
        Ok(Self {
            shared,
            session: Mutex::new(session),
            variables: std::sync::Mutex::default(),
        })
    }

    /// Runs one snippet, servicing its nested tool calls concurrently.
    pub async fn run_code(&self, code: String, restart: bool) -> ToolCallOutput {
        if let Some(helper) = scope::shadowed_helper(&code) {
            return ToolCallOutput::error(format!(
                "{RUN_CODE_TOOL} rejected: `{helper}` is a built-in helper; pick another variable name. Nothing ran."
            ));
        }
        if restart {
            *self.variables.lock().expect("variables lock") = Variables::default();
        }
        let bound = scope::assigned_names(&code);
        let (outcome, audit) = self.execute(code, restart).await;
        if !matches!(outcome, Err(Failure::Host(_))) {
            self.refresh_variables(bound).await;
        }
        let mut output = match outcome {
            Ok(completion) => render_outcome(Ok(completion), &self.shared.limits),
            Err(Failure::Python(traceback)) => render_outcome(Err(traceback), &self.shared.limits),
            Err(Failure::Host(message)) => ToolCallOutput::error(message),
        };
        output.audit = audit;
        output
    }

    /// Re-describes the tracked variables by running a small probe in the
    /// same scope; names no longer bound drop out.
    async fn refresh_variables(&self, bound: Vec<String>) {
        let candidates = {
            let mut variables = self.variables.lock().expect("variables lock");
            for name in bound.into_iter().filter(|name| scope::trackable(name)) {
                if !variables.candidates.contains(&name) {
                    variables.candidates.push(name);
                }
            }
            variables.candidates.clone()
        };
        if candidates.is_empty() {
            return;
        }
        let described: Vec<(String, String)> = match self
            .execute(scope::describe_program(&candidates), false)
            .await
        {
            (
                Ok(Completion {
                    value: Some(Value::Object(map)),
                    ..
                }),
                _,
            ) => candidates
                .iter()
                .filter_map(|name| {
                    map.get(name)
                        .and_then(Value::as_str)
                        .map(|kind| (name.clone(), kind.to_string()))
                })
                .collect(),
            _ => return,
        };
        let mut variables = self.variables.lock().expect("variables lock");
        variables
            .candidates
            .retain(|name| described.iter().any(|(known, _)| known == name));
        variables.described = described;
    }

    /// Variables still bound in this scope, as `(name, "list[80]")`.
    pub fn variables(&self) -> Vec<(String, String)> {
        self.variables
            .lock()
            .expect("variables lock")
            .described
            .clone()
    }

    /// Runs one snippet and returns its raw outcome (for watch filters and
    /// `run_code`), servicing nested calls concurrently.
    pub(crate) async fn execute(
        &self,
        code: String,
        restart: bool,
    ) -> (Result<Completion, Failure>, Option<Value>) {
        let session = self.session.lock().await;
        let limits = &self.shared.limits;
        let (events_tx, mut events) = tokio::sync::mpsc::unbounded_channel();
        let (results_tx, results_rx) = tokio::sync::mpsc::unbounded_channel();
        if let Err(message) = session.submit(Job {
            code,
            restart,
            events: events_tx,
            results: results_rx,
        }) {
            return (Err(Failure::Host(message)), None);
        }

        let mut audit = Audit::default();
        let deadline = tokio::time::sleep(limits.max_wall);
        tokio::pin!(deadline);
        let mut in_flight = FuturesUnordered::new();
        loop {
            tokio::select! {
                _ = &mut deadline => {
                    let message = format!(
                        "{RUN_CODE_TOOL} exceeded its {:?} wall-clock limit; in-flight tool calls were cancelled",
                        limits.max_wall
                    );
                    return (Err(Failure::Host(message)), audit.into_value());
                }
                event = events.recv() => match event {
                    Some(HostEvent::Call { call_id, py_name, mcp_name, arguments, mutating }) => {
                        self.log_call(&py_name, &mcp_name, &arguments, mutating);
                        audit.start(call_id, &mcp_name, &arguments, mutating);
                        let inner = self.shared.inner.clone();
                        let builtins = self.shared.builtins.clone();
                        let timeout = limits.tool_timeout;
                        in_flight.push(async move {
                            let call = async {
                                match mcp_name.strip_prefix(builtins::PREFIX) {
                                    Some(builtin) => builtins.call(builtin, arguments).await,
                                    None => inner.call_tool(&mcp_name, arguments).await,
                                }
                            };
                            let outcome = match tokio::time::timeout(timeout, call).await {
                                Ok(result) => result.map_err(|error| error.to_string()),
                                Err(_) => Err(format!(
                                    "tool {mcp_name} timed out after {timeout:?}"
                                )),
                            };
                            (call_id, mcp_name, mutating, outcome)
                        });
                    }
                    Some(HostEvent::Done(outcome)) => {
                        return (outcome.map_err(Failure::Python), audit.into_value());
                    }
                    None => return (Err(Failure::Host("code mode engine stopped".to_string())), audit.into_value()),
                },
                Some((call_id, mcp_name, mutating, outcome)) = in_flight.next(), if !in_flight.is_empty() => {
                    self.log_result(&mcp_name, mutating, &outcome);
                    audit.finish(call_id, &outcome);
                    let _ = results_tx.send((call_id, outcome));
                }
            }
        }
    }

    fn log_call(&self, py_name: &str, mcp_name: &str, arguments: &Value, mutating: bool) {
        if !(self.shared.verbose || mutating) {
            return;
        }
        let mut record = json!({
            "level": if mutating { "warn" } else { "info" },
            "event": "mcp.tool_call",
            "via": "code_mode",
            "tool": mcp_name,
            "function": py_name,
            "arguments": arguments,
        });
        if mutating {
            record["mutating"] = json!(true);
        }
        eprintln!("{record}");
    }

    fn log_result(&self, mcp_name: &str, mutating: bool, outcome: &CallOutcome) {
        if !(self.shared.verbose || mutating) {
            return;
        }
        let (is_error, bytes, content) = match outcome {
            Ok(output) => (
                output.is_error,
                output.content.len(),
                content_for_log(&output.content, 500),
            ),
            Err(message) => (true, message.len(), Value::String(message.clone())),
        };
        let level = match (is_error, mutating) {
            (true, _) => "error",
            (false, true) => "warn",
            (false, false) => "info",
        };
        let mut record = json!({
            "level": level,
            "event": "mcp.tool_result",
            "via": "code_mode",
            "tool": mcp_name,
            "is_error": is_error,
            "bytes": bytes,
            "content": content,
        });
        if mutating {
            record["mutating"] = json!(true);
        }
        eprintln!("{record}");
    }
}

/// Why a snippet produced no completion.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum Failure {
    /// Python raised; the traceback.
    Python(String),
    /// The host stopped it (wall deadline, engine gone).
    Host(String),
}

/// Nested calls made by one `run_code`, for session records and metrics.
#[derive(Default)]
struct Audit {
    calls: Vec<(u32, Value)>,
}

impl Audit {
    fn start(&mut self, call_id: u32, tool: &str, arguments: &Value, mutating: bool) {
        self.calls.push((
            call_id,
            json!({
                "tool": tool,
                "arguments": arguments,
                "mutating": mutating,
                "finished": false,
            }),
        ));
    }

    fn finish(&mut self, call_id: u32, outcome: &CallOutcome) {
        if let Some((_, record)) = self.calls.iter_mut().find(|(id, _)| *id == call_id) {
            let (is_error, bytes) = match outcome {
                Ok(output) => (output.is_error, output.content.len()),
                Err(message) => (true, message.len()),
            };
            record["finished"] = json!(true);
            record["is_error"] = json!(is_error);
            record["bytes"] = json!(bytes);
        }
    }

    fn into_value(self) -> Option<Value> {
        if self.calls.is_empty() {
            return None;
        }
        Some(json!({
            "nested_calls": self.calls.into_iter().map(|(_, record)| record).collect::<Vec<_>>(),
        }))
    }
}

#[async_trait]
impl ToolExecutor for CodeModeRun {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        Ok(vec![self.shared.definition.clone()])
    }

    async fn call_tool(&self, name: &str, arguments: Value) -> Result<ToolCallOutput, ToolError> {
        if name != RUN_CODE_TOOL {
            return Ok(ToolCallOutput::error(format!(
                "unknown tool {name:?}: call MCP tools from Python inside {RUN_CODE_TOOL}"
            )));
        }
        let code = arguments
            .get("code")
            .and_then(Value::as_str)
            .ok_or_else(|| ToolError::new(format!("{RUN_CODE_TOOL} requires a string `code`")))?
            .to_string();
        let restart = arguments
            .get("restart")
            .and_then(Value::as_bool)
            .unwrap_or(false);
        Ok(self.run_code(code, restart).await)
    }

    fn context_note(&self) -> Option<String> {
        scope::render_note(&self.variables())
    }
}

/// Applies the return contract and the model-facing size guard.
pub fn render_outcome(
    outcome: Result<Completion, String>,
    limits: &CodeModeLimits,
) -> ToolCallOutput {
    let completion = match outcome {
        Ok(completion) => completion,
        Err(traceback) => {
            let mut message = format!("{RUN_CODE_TOOL} raised:\n{traceback}");
            if let Some(hint) = hint_for(&traceback) {
                message.push_str("\nHint: ");
                message.push_str(hint);
            }
            return ToolCallOutput::error(message);
        }
    };

    if let Some(Value::String(text)) = &completion.value {
        if is_awaitable_repr(text) {
            return ToolCallOutput::error(format!(
                "{RUN_CODE_TOOL} result is an un-awaited awaitable ({text}); add `await`, e.g. `await asyncio.gather(...)`"
            ));
        }
    }

    let result = match (completion.prints.is_empty(), completion.value) {
        (true, Some(value)) => value,
        (true, None) => json!(EMPTY_RESULT),
        (false, None) => json!({"output": completion.prints}),
        (false, Some(value)) => json!({"output": completion.prints, "result": value}),
    };
    let content = match result {
        Value::String(text) => text,
        other => other.to_string(),
    };

    let length = content.chars().count();
    if length > limits.max_result_chars {
        let preview: String = content.chars().take(PREVIEW_CHARS).collect();
        return ToolCallOutput::error(format!(
            "{RUN_CODE_TOOL} result is {length} characters, over the {} character limit. \
             Assign the data to a variable and return a filtered subset; \
             variables persist between {RUN_CODE_TOOL} calls.\npreview:\n{preview}",
            limits.max_result_chars
        ));
    }

    ToolCallOutput::text(content)
}

/// Model-facing text for a snippet with no prints and a `None` result.
pub const EMPTY_RESULT: &str = "(no output; the last expression was None)";

/// Known Monty gaps and common slips, matched against a traceback.
const HINTS: &[(&[&str], &str)] = &[
    (
        &["unexpected keyword argument 'default'", "JSONEncoder"],
        "Monty's json.dumps has no default=/cls=; use to_json(x) or plain json.dumps(x).",
    ),
    (
        &["unhashable type"],
        "a dict, list, or set cannot be a dict key, set element, or Counter key; convert it first (tuple(...), str(...), or to_json(...)).",
    ),
];

fn hint_for(traceback: &str) -> Option<&'static str> {
    HINTS
        .iter()
        .find(|(needles, _)| needles.iter().all(|needle| traceback.contains(needle)))
        .map(|(_, hint)| *hint)
}

/// `<coroutine ...>`, `<gather(4)>`, and Monty's `Repr('<...>')` wrapping.
fn is_awaitable_repr(text: &str) -> bool {
    let inner = text
        .strip_prefix("Repr('")
        .and_then(|rest| rest.strip_suffix("')"))
        .unwrap_or(text);
    [
        "<coroutine",
        "<gather",
        "<future",
        "<Future",
        "<Task",
        "<async_generator",
    ]
    .iter()
    .any(|prefix| inner.starts_with(prefix))
}

fn run_code_definition(catalog: &Catalog, limits: &CodeModeLimits) -> ToolDefinition {
    let modules = monty::builtin_module_names().join(", ");
    let functions = if catalog.is_empty() {
        "# no tools available".to_string()
    } else {
        catalog.render()
    };
    let description = format!(
        "Run Python in a sandbox (Monty, a Python subset). Every tool is a Python \
function, callable only from here; there are no other tools.\n\
\n\
Plan first, then write ONE complete program that fetches, filters, and returns the answer.\n\
\n\
Rules:\n\
- Call functions with `await` and keyword arguments only.\n\
- Run independent calls concurrently: `a, b = await asyncio.gather(f(x=1), g(y=2))`.\n\
- Results arrive parsed: JSON becomes dict/list, other text becomes str. A failing tool raises RuntimeError.\n\
- `Returns:` in each docstring shows the result shape. Only for `-> Any` tools, or after a failure, inspect a small sample (type(x), list(x)[:5]).\n\
- Results can be very large. Filter, count, or summarize in Python and return only what is needed.\n\
- The value of the last expression is returned. print() output is returned too.\n\
- Variables persist across {RUN_CODE_TOOL} calls for this whole conversation; pass restart=true to reset. Earlier tool results are NOT in your message history, so keep data you may need later in variables. The variables still in scope are listed with each new message.\n\
- Output over {max} characters is rejected: keep data in a variable and return a smaller subset.\n\
- Importable modules: {modules}. No files, network, environment, subprocesses, or third-party packages.\n\
- Monty gaps: json.dumps has no default=/cls= (use to_json). An un-awaited coroutine or gather is an error.\n\
\n\
Pre-imported: asyncio, json, re. Built-in helpers (reserved names; do not assign to them):\n\
- items(x): x[\"items\"] for a Kubernetes list, x for a list\n\
- names(x): metadata.name (or name) of every item\n\
- dig(obj, \"a.b.c\", default=None): nested lookup; dig(obj, \"a\", \"b.c\") also works\n\
- count_by(x, \"status.phase\"): counts by a path\n\
- to_json(x, indent=None): json.dumps that stringifies unknown types\n\
\n\
Example:\n\
```python\n\
data = await some_tool(name=\"value\")\n\
{{\"total\": len(items(data)), \"by_phase\": count_by(data, \"status.phase\")}}\n\
```\n\
\n\
Functions:\n\
```python\n\
{functions}\n\
```",
        max = limits.max_result_chars,
    );

    ToolDefinition::new(
        RUN_CODE_TOOL,
        Some(description),
        json!({
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Python to run. The last expression is the result.",
                },
                "restart": {
                    "type": "boolean",
                    "description": "Reset variables before running.",
                },
            },
            "required": ["code"],
        }),
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    fn done(value: Option<Value>, prints: &str) -> Result<Completion, String> {
        Ok(Completion {
            value,
            prints: prints.to_string(),
        })
    }

    #[test]
    fn return_contract_matches_code_mode() {
        let limits = CodeModeLimits::default();
        assert_eq!(
            render_outcome(done(Some(json!([1, 2])), ""), &limits).content,
            "[1,2]"
        );
        assert_eq!(
            render_outcome(done(Some(json!("text")), ""), &limits).content,
            "text"
        );
        assert_eq!(
            render_outcome(done(None, ""), &limits).content,
            EMPTY_RESULT
        );
        assert_eq!(
            render_outcome(done(None, "hi\n"), &limits).content,
            r#"{"output":"hi\n"}"#
        );
        let both = render_outcome(done(Some(json!(3)), "hi\n"), &limits);
        let parsed: Value = serde_json::from_str(&both.content).unwrap();
        assert_eq!(parsed, json!({"output": "hi\n", "result": 3}));
        assert!(!both.is_error);

        let failed = render_outcome(Err("NameError: x".to_string()), &limits);
        assert!(failed.is_error);
        assert!(failed.content.contains("NameError: x"));
    }

    #[test]
    fn unawaited_results_fail_and_known_gaps_get_hints() {
        let limits = CodeModeLimits::default();
        for repr in ["Repr('<gather(4)>')", "<coroutine object f at 0x1>"] {
            let output = render_outcome(done(Some(json!(repr)), ""), &limits);
            assert!(output.is_error, "{repr}");
            assert!(output.content.contains("add `await`"));
        }
        assert!(!render_outcome(done(Some(json!("<b>html</b>")), ""), &limits).is_error);

        let failed = render_outcome(
            Err(
                "TypeError: JSONEncoder.__init__() got an unexpected keyword argument 'default'"
                    .into(),
            ),
            &limits,
        );
        assert!(failed.content.contains("Hint: Monty's json.dumps"));
        let plain = render_outcome(Err("NameError: x".into()), &limits);
        assert!(!plain.content.contains("Hint"));
    }

    #[test]
    fn oversized_results_are_rejected_with_a_preview() {
        let limits = CodeModeLimits {
            max_result_chars: 10,
            ..CodeModeLimits::default()
        };
        let output = render_outcome(done(Some(json!("abcdefghijklmnop")), ""), &limits);
        assert!(output.is_error);
        assert!(output.content.contains("16 characters"));
        assert!(output.content.contains("abcdefghijklmnop"));
    }

    #[test]
    fn definition_lists_functions_helpers_and_modules() {
        let catalog = Catalog::new(&[ToolDefinition::new(
            "http_get",
            Some("GET a URL."),
            json!({"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}),
        )]);
        let definition = run_code_definition(&catalog, &CodeModeLimits::default());
        let description = definition.description.unwrap();
        assert_eq!(definition.name, "run_code");
        assert!(description.contains("async def http_get(*, url: str) -> Any:"));
        assert!(description.contains("json"));
        assert!(description.contains("count_by(x, \"status.phase\")"));
        assert!(description.contains("Plan first"));
        assert_eq!(definition.input_schema["required"], json!(["code"]));
    }

    #[test]
    fn sandbox_kind_parses() {
        assert_eq!("auto".parse::<SandboxKind>().unwrap(), SandboxKind::Auto);
        assert_eq!("pool".parse::<SandboxKind>().unwrap(), SandboxKind::Pool);
        assert_eq!(
            "in-process".parse::<SandboxKind>().unwrap(),
            SandboxKind::InProcess
        );
        assert!("docker".parse::<SandboxKind>().is_err());
    }

    #[test]
    fn missing_explicit_worker_is_not_found() {
        assert_eq!(find_worker(Some(PathBuf::from("/nonexistent/monty"))), None);
    }
}