magic_coder_types/

tools.rs

use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[cfg(feature = "schemars")]
use schemars::{JsonSchema, schema_for};
#[cfg(feature = "schemars")]
use serde_json::Value;

/// Convert a Rust struct schema into an OpenAI tool `parameters` object.
///
/// - remove `$schema` and `title`
/// - convert `definitions` to `$defs`
/// - convert `oneOf` to `anyOf`
#[cfg(feature = "schemars")]
pub fn tool_parameters<T: JsonSchema>() -> Value {
    let mut v = serde_json::to_value(schema_for!(T)).expect("can't parse value from schema");

    // remove the $schema and title fields
    if let Some(value) = v.as_object_mut() {
        value.remove("$schema");
        value.remove("title");
    }

    let mut v_str = serde_json::to_string(&v).unwrap();
    v_str = v_str
        .replace("/definitions/", "/$defs/")
        .replace("\"definitions\":", "\"$defs\":");

    // Replace oneOf with anyOf, because it's better supported by the LLMs
    v_str = v_str.replace("\"oneOf\":", "\"anyOf\":");

    let mut v: Value = serde_json::from_str(&v_str).expect("can't parse value from updated schema");
    enforce_openai_strict_schema(&mut v);
    v
}

#[cfg(feature = "schemars")]
fn enforce_openai_strict_schema(v: &mut Value) {
    match v {
        Value::Object(map) => {
            // Recurse first so we fix nested schemas too.
            for (_k, child) in map.iter_mut() {
                enforce_openai_strict_schema(child);
            }

            // If this looks like an object schema, enforce strict rules.
            let is_object = map
                .get("type")
                .and_then(|t| t.as_str())
                .is_some_and(|t| t == "object");
            let has_props = map.get("properties").is_some();
            if is_object || has_props {
                map.entry("additionalProperties".to_string())
                    .or_insert(Value::Bool(false));

                if let Some(Value::Object(props)) = map.get("properties") {
                    let mut keys: Vec<String> = props.keys().cloned().collect();
                    keys.sort();
                    map.insert(
                        "required".to_string(),
                        Value::Array(keys.into_iter().map(Value::String).collect()),
                    );
                }
            }
        }
        Value::Array(arr) => {
            for child in arr.iter_mut() {
                enforce_openai_strict_schema(child);
            }
        }
        _ => {}
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct ReadFileArgs {
    /// Path to file.
    pub path: String,
    /// Optional starting line (0-based).
    pub offset: Option<usize>,
    /// Optional maximum number of lines to read.
    pub limit: Option<usize>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct ListDirArgs {
    /// Directory path to list.
    pub path: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum GlobKind {
    Files,
    Dirs,
    All,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct GlobArgs {
    /// Glob pattern to match. Supports `*`, `**`, `?`, and character classes.
    pub pattern: String,
    /// Optional directory root to search under. Defaults to `"."`.
    pub path: Option<String>,
    /// Optional maximum number of returned paths. Defaults to `50`.
    pub limit: Option<usize>,
    /// Optional match kind. Defaults to `files`.
    pub kind: Option<GlobKind>,
    /// Optional exclude patterns. Defaults to an empty list.
    #[serde(default)]
    pub exclude: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct GrepArgs {
    /// Regex pattern to search for.
    pub pattern: String,
    /// Optional path (file or directory) to search in.
    pub path: Option<String>,
    /// Optional glob filter, e.g. `"*.rs"`.
    pub glob: Option<String>,
    /// Optional limit for returned matches.
    pub head_limit: Option<usize>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct RunShellArgs {
    /// Shell command line to run (executed via `bash -lc`), supports pipes/redirection.
    pub command: String,
    /// Optional working directory.
    pub cwd: Option<String>,
    /// Optional timeout in seconds for foreground execution.
    /// Omit to use the default 30 second timeout.
    /// Must be omitted when `bg=true`.
    /// For longer-running work like model training, set a larger value up front on the safe side to avoid retries.
    pub timeout_seconds: Option<u64>,
    /// Optional maximum captured bytes per stream (stdout/stderr) for foreground execution.
    /// Must be omitted when `bg=true`.
    ///
    /// Truncated output keeps roughly the first 30% and last 70%, so very large
    /// values are usually unnecessary; prefer a few KB or low tens of KB and only
    /// increase if needed.
    pub max_output_bytes: Option<u64>,
    /// When true, spawn the shell in the background and return immediately with a shell id.
    #[serde(default)]
    pub bg: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct ReadShellOutputArgs {
    /// Background shell id returned by `run_shell` with `bg=true`.
    pub shell_id: String,
    /// When true, read from the start of the log. Defaults to `false` meaning read from the end.
    #[serde(default)]
    pub from_start: bool,
    /// Optional 0-based line offset from the selected side. Defaults to `0`.
    pub offset: Option<usize>,
    /// Optional maximum number of lines to read. Defaults to `200`, max `1000`.
    pub limit: Option<usize>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct StopShellArgs {
    /// Background shell id returned by `run_shell` with `bg=true`.
    pub shell_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct SleepArgs {
    /// Sleep duration in seconds. Clients may clamp this to a supported range.
    pub seconds: u64,
    /// Background shell ids to watch. Use an empty array for a plain timer.
    /// If any watched shell exits early, the sleep may end early.
    #[serde(default)]
    pub shell_ids: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct ApplyDiffArgs {
    /// Wrap file blocks between `*** Begin Patch` and `*** End Patch`.
    ///
    /// Start each block with `*** Add File: <path>`,
    /// `*** Update File: <path>`, or `*** Delete File: <path>`;
    /// `*** Move to: <path>` may follow Update. Add lines start `+`.
    /// `@@ <unchanged anchor>` starts an update hunk and searches forward after
    /// that line; following space/`-` lines must match the current file, while
    /// `+` lines are inserted. Include enough context to identify one location.
    pub diff: String,

    /// Optional working directory used as the root for relative patch paths.
    pub cwd: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct DeleteFilesArgs {
    /// Paths to delete (relative to project root; no absolute paths; no `..`).
    pub paths: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(JsonSchema))]
pub struct CallMcpToolArgs {
    /// Connected local MCP server UUID.
    pub server_id: Uuid,
    /// Name of the tool to call on the selected MCP server.
    pub tool_name: String,
    /// Arguments object matching the selected MCP tool's advertised input schema.
    #[serde(default)]
    pub arguments: serde_json::Map<String, serde_json::Value>,
}

/// Tools (function definitions) to send to the OpenAI Responses API.
#[cfg(feature = "schemars")]
pub fn openai_tools() -> Vec<Value> {
    vec![
        serde_json::json!({
            "type": "function",
            "name": "read_file",
            "description": "Read a local file (by path), optionally with offset/limit. Returns plain text with a frontmatter block containing `path`, `offset`, `limit`, `total_lines`, `truncated`, and `returned_lines`, followed by the requested line-numbered file contents.",
            "strict": true,
            "parameters": tool_parameters::<ReadFileArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "list_dir",
            "description": "List a local directory (by path). Returns plain text with `Path`, `Entries`, and one entry per line similar to `ls`; directories end with `/`.",
            "strict": true,
            "parameters": tool_parameters::<ListDirArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "glob",
            "description": "Find local file or directory paths using a glob pattern under a search root. Use this for path discovery when you need matching paths, not file contents. Returns plain text with Returned, Total, and one relative path per line.",
            "strict": true,
            "parameters": tool_parameters::<GlobArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "grep",
            "description": "Search for a regex pattern in files. Returns a JSON string including matches (file, line_number, line). May be truncated to head_limit.",
            "strict": true,
            "parameters": tool_parameters::<GrepArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "run_shell",
            "description": "Run a shell command via `bash -lc` (supports pipes/redirection). Requires user confirmation unless the client auto-approves it. Use `max_output_bytes` intentionally for foreground runs: prefer the smallest limit that answers the question, and increase only when needed. Oversize output is cut from the middle, preserving roughly the first 30% and last 70%, so large requests are rarely necessary just to inspect the tail. Set `bg=true` to start a background shell that returns immediately with a shell id. When `bg=true`, omit `timeout_seconds` and omit `max_output_bytes`.",
            "parameters": tool_parameters::<RunShellArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "read_shell_output",
            "description": "Read captured output from a background shell started with `run_shell(bg=true)`. Output is line-oriented. By default it reads from the end; set `from_start=true` to read from the beginning.",
            "strict": true,
            "parameters": tool_parameters::<ReadShellOutputArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "stop_shell",
            "description": "Stop a background shell started with `run_shell(bg=true)` and discard its retained state and logs.",
            "strict": true,
            "parameters": tool_parameters::<StopShellArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "sleep",
            "description": "Wait for 15 to 275 seconds. Provide `shell_ids` to return early when any watched background shell exits. Use `shell_ids: []` for a plain timer.",
            "strict": true,
            "parameters": tool_parameters::<SleepArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "apply_diff",
            "description": "Apply one ApplyPatch document to the local working tree. Files commit independently; returns applied changes and per-file failures.",
            "strict": true,
            "parameters": tool_parameters::<ApplyDiffArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "delete_files",
            "description": "Delete one or more files by path (relative to project root). Returns a JSON string listing deleted and missing paths.",
            "strict": true,
            "parameters": tool_parameters::<DeleteFilesArgs>(),
        }),
        serde_json::json!({
            "type": "function",
            "name": "call_mcp_tool",
            "description": "Call one tool from a connected local MCP server by `server_id` and `tool_name`. The `arguments` field must be a JSON object matching that tool's advertised input schema.",
            "parameters": tool_parameters::<CallMcpToolArgs>(),
        }),
    ]
}

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

    #[test]
    fn glob_args_default_exclude_to_empty_list() {
        let args: GlobArgs = serde_json::from_value(json!({
            "pattern": "**/*.rs",
            "path": "src",
            "limit": 50,
            "kind": "files",
        }))
        .expect("glob args");

        assert_eq!(args.pattern, "**/*.rs");
        assert_eq!(args.path.as_deref(), Some("src"));
        assert_eq!(args.limit, Some(50));
        assert_eq!(args.kind, Some(GlobKind::Files));
        assert!(args.exclude.is_empty());
    }

    #[cfg(feature = "schemars")]
    #[test]
    fn run_shell_tool_schema_encourages_small_output_limits() {
        let run_shell = openai_tools()
            .into_iter()
            .find(|tool| tool.get("name").and_then(Value::as_str) == Some("run_shell"))
            .expect("run_shell tool");

        let description = run_shell
            .get("description")
            .and_then(Value::as_str)
            .expect("run_shell description");
        assert!(description.contains("max_output_bytes"));
        assert!(description.contains("30%"));
        assert!(description.contains("70%"));
        assert!(description.contains("smallest limit"));
        assert!(description.contains("bg=true"));
        assert!(description.contains("omit `timeout_seconds`"));
        assert!(description.contains("omit `max_output_bytes`"));

        let timeout_description = run_shell
            .get("parameters")
            .and_then(|value| value.get("properties"))
            .and_then(|value| value.get("timeout_seconds"))
            .and_then(|value| value.get("description"))
            .and_then(Value::as_str)
            .expect("timeout_seconds description");
        assert!(timeout_description.contains("30 second timeout"));
        assert!(timeout_description.contains("model training"));
        assert!(timeout_description.contains("safe side"));
        assert!(timeout_description.contains("Must be omitted when `bg=true`"));

        let max_output_description = run_shell
            .get("parameters")
            .and_then(|value| value.get("properties"))
            .and_then(|value| value.get("max_output_bytes"))
            .and_then(|value| value.get("description"))
            .and_then(Value::as_str)
            .expect("max_output_bytes description");
        assert!(max_output_description.contains("30%"));
        assert!(max_output_description.contains("70%"));
        assert!(max_output_description.contains("few KB"));
        assert!(max_output_description.contains("Must be omitted when `bg=true`"));

        let properties = run_shell
            .get("parameters")
            .and_then(|value| value.get("properties"))
            .and_then(Value::as_object)
            .expect("run_shell parameters");
        assert!(properties.contains_key("bg"));
    }

    #[cfg(feature = "schemars")]
    #[test]
    fn openai_tools_include_glob_tool() {
        let glob_tool = openai_tools()
            .into_iter()
            .find(|tool| tool.get("name").and_then(Value::as_str) == Some("glob"))
            .expect("glob tool");

        let description = glob_tool
            .get("description")
            .and_then(Value::as_str)
            .expect("glob description");
        assert!(description.contains("path discovery"));
        assert!(description.contains("Returned"));
        assert!(description.contains("Total"));

        let properties = glob_tool
            .get("parameters")
            .and_then(|value| value.get("properties"))
            .and_then(Value::as_object)
            .expect("glob parameters");
        assert!(properties.contains_key("pattern"));
        assert!(properties.contains_key("path"));
        assert!(properties.contains_key("limit"));
        assert!(properties.contains_key("kind"));
        assert!(properties.contains_key("exclude"));
    }

    #[cfg(feature = "schemars")]
    #[test]
    fn openai_tools_describe_plaintext_file_and_directory_reads() {
        let tools = openai_tools();

        let read_file = tools
            .iter()
            .find(|tool| tool.get("name").and_then(Value::as_str) == Some("read_file"))
            .expect("read_file tool");
        let read_file_description = read_file
            .get("description")
            .and_then(Value::as_str)
            .expect("read_file description");
        assert!(read_file_description.contains("Returns plain text"));
        assert!(read_file_description.contains("frontmatter"));
        assert!(read_file_description.contains("returned_lines"));
        assert!(read_file_description.contains("line-numbered"));

        let list_dir = tools
            .iter()
            .find(|tool| tool.get("name").and_then(Value::as_str) == Some("list_dir"))
            .expect("list_dir tool");
        let list_dir_description = list_dir
            .get("description")
            .and_then(Value::as_str)
            .expect("list_dir description");
        assert!(list_dir_description.contains("Returns plain text"));
        assert!(list_dir_description.contains("Path"));
        assert!(list_dir_description.contains("Entries"));
        assert!(list_dir_description.contains("similar to `ls`"));
        assert!(list_dir_description.contains("directories end with `/`"));
    }

    #[cfg(feature = "schemars")]
    #[test]
    fn openai_tools_describe_apply_diff_contract() {
        let apply_diff = openai_tools()
            .into_iter()
            .find(|tool| tool.get("name").and_then(Value::as_str) == Some("apply_diff"))
            .expect("apply_diff tool");

        let description = apply_diff
            .get("description")
            .and_then(Value::as_str)
            .expect("apply_diff description");
        let diff_description = apply_diff
            .pointer("/parameters/properties/diff/description")
            .and_then(Value::as_str)
            .expect("apply_diff diff description");

        assert!(description.contains("per-file failures"));
        assert!(diff_description.contains("*** Begin Patch"));
        assert!(diff_description.contains("*** Add File: <path>"));
        assert!(diff_description.contains("*** Update File: <path>"));
        assert!(diff_description.contains("*** Delete File: <path>"));
        assert!(diff_description.contains("*** Move to: <path>"));
        assert!(diff_description.contains("@@ <unchanged anchor>"));
        assert!(diff_description.contains("must match the current file"));
        assert!(diff_description.contains("identify one location"));
        assert!(description.len() + diff_description.len() <= 750);
    }

    #[test]
    fn background_shell_tool_args_default_to_tail_reads() {
        let read_shell_output: ReadShellOutputArgs = serde_json::from_value(json!({
            "shell_id": "bg_123"
        }))
        .expect("read_shell_output args");
        assert_eq!(read_shell_output.shell_id, "bg_123");
        assert!(!read_shell_output.from_start);
        assert_eq!(read_shell_output.offset, None);
        assert_eq!(read_shell_output.limit, None);

        let run_shell: RunShellArgs = serde_json::from_value(json!({
            "command": "echo hi"
        }))
        .expect("run_shell args");
        assert_eq!(run_shell.command, "echo hi");
        assert!(!run_shell.bg);

        let sleep: SleepArgs = serde_json::from_value(json!({
            "seconds": 30,
            "shell_ids": ["bg_123", "bg_456"]
        }))
        .expect("sleep args");
        assert_eq!(sleep.seconds, 30);
        assert_eq!(sleep.shell_ids, vec!["bg_123", "bg_456"]);

        let timer_only_sleep: SleepArgs = serde_json::from_value(json!({
            "seconds": 15
        }))
        .expect("timer-only sleep args");
        assert_eq!(timer_only_sleep.seconds, 15);
        assert!(timer_only_sleep.shell_ids.is_empty());
    }

    #[cfg(feature = "schemars")]
    #[test]
    fn openai_tools_include_background_shell_tools() {
        let tools = openai_tools();
        let names = tools
            .iter()
            .filter_map(|tool| tool.get("name").and_then(Value::as_str))
            .collect::<Vec<_>>();

        assert!(names.contains(&"read_shell_output"));
        assert!(names.contains(&"stop_shell"));
        assert!(names.contains(&"sleep"));
    }

    #[test]
    fn call_mcp_tool_args_default_arguments_to_empty_object() {
        let args: CallMcpToolArgs = serde_json::from_value(json!({
            "server_id": "00000000-0000-0000-0000-000000000000",
            "tool_name": "create_page"
        }))
        .expect("call_mcp_tool args");

        assert_eq!(args.server_id, Uuid::nil());
        assert_eq!(args.tool_name, "create_page");
        assert!(args.arguments.is_empty());
    }

    #[cfg(feature = "schemars")]
    #[test]
    fn openai_tools_include_call_mcp_tool() {
        let tool = openai_tools()
            .into_iter()
            .find(|tool| tool.get("name").and_then(Value::as_str) == Some("call_mcp_tool"))
            .expect("call_mcp_tool");

        let description = tool
            .get("description")
            .and_then(Value::as_str)
            .expect("call_mcp_tool description");
        assert!(description.contains("connected local MCP server"));
        assert!(description.contains("server_id"));
        assert!(description.contains("tool_name"));

        let properties = tool
            .get("parameters")
            .and_then(|value| value.get("properties"))
            .and_then(Value::as_object)
            .expect("call_mcp_tool parameters");
        assert!(properties.contains_key("server_id"));
        assert!(properties.contains_key("tool_name"));
        assert!(properties.contains_key("arguments"));
        assert_eq!(
            properties["arguments"].get("additionalProperties"),
            Some(&Value::Bool(true))
        );
        assert!(tool.get("strict").is_none());
    }
}