From ebeb9c9b7dced0d1845402e4ec13d659449c0e3d Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 10:20:49 -0600 Subject: [PATCH 01/12] refactor: python tools now use tree-sitter queries instead of AST --- Cargo.lock | 355 ++--------------- Cargo.toml | 4 +- src/parsers/python.rs | 897 +++++++++++++++++++++++++++++++++--------- 3 files changed, 749 insertions(+), 507 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e2741b3..162b10d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -859,7 +859,7 @@ dependencies = [ "proc-macro2", "quote", "regex", - "rustc-hash 2.1.2", + "rustc-hash", "shlex", "syn", ] @@ -1388,12 +1388,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "crunchy" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" - [[package]] name = "crypto-common" version = "0.1.7" @@ -1518,34 +1512,13 @@ dependencies = [ "syn", ] -[[package]] -name = "derive_more" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a9b99b9cbbe49445b21764dc0625032a89b145a2642e67603e1c936f5458d05" -dependencies = [ - "derive_more-impl 1.0.0", -] - [[package]] name = "derive_more" version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134" dependencies = [ - "derive_more-impl 2.1.1", -] - -[[package]] -name = "derive_more-impl" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb7330aeadfbe296029522e6c40f315320aba36fc43a5b3632f3795348f3bd22" -dependencies = [ - "proc-macro2", - "quote", - "syn", - "unicode-xid", + "derive_more-impl", ] [[package]] @@ -2105,15 +2078,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "getopts" -version = "0.2.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" -dependencies = [ - "unicode-width", -] - [[package]] name = "getrandom" version = "0.2.17" @@ -2258,15 +2222,6 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" -[[package]] -name = "hashbrown" -version = "0.14.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" -dependencies = [ - "ahash", -] - [[package]] name = "hashbrown" version = "0.15.5" @@ -2822,18 +2777,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "is-macro" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d57a3e447e24c22647738e4607f1df1e0ec6f72e16182c4cd199f647cdfb0e4" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "is-terminal" version = "0.4.17" @@ -2870,15 +2813,6 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" -[[package]] -name = "itertools" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.13.0" @@ -2987,12 +2921,6 @@ dependencies = [ "simple_asn1", ] -[[package]] -name = "lalrpop-util" -version = "0.20.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553" - [[package]] name = "lazy_static" version = "1.5.0" @@ -3021,12 +2949,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "libm" -version = "0.2.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" - [[package]] name = "libredox" version = "0.1.15" @@ -3171,8 +3093,6 @@ dependencies = [ "reqwest-eventsource", "rmcp", "rust-embed", - "rustpython-ast", - "rustpython-parser", "scraper", "serde", "serde_json", @@ -3188,6 +3108,8 @@ dependencies = [ "tokio", "tokio-graceful", "tokio-stream", + "tree-sitter", + "tree-sitter-python", "unicode-segmentation", "unicode-width", "url", @@ -3233,64 +3155,6 @@ dependencies = [ "libc", ] -[[package]] -name = "malachite" -version = "0.4.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fbdf9cb251732db30a7200ebb6ae5d22fe8e11397364416617d2c2cf0c51cb5" -dependencies = [ - "malachite-base", - "malachite-nz", - "malachite-q", -] - -[[package]] -name = "malachite-base" -version = "0.4.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ea0ed76adf7defc1a92240b5c36d5368cfe9251640dcce5bd2d0b7c1fd87aeb" -dependencies = [ - "hashbrown 0.14.5", - "itertools 0.11.0", - "libm", - "ryu", -] - -[[package]] -name = "malachite-bigint" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d149aaa2965d70381709d9df4c7ee1fc0de1c614a4efc2ee356f5e43d68749f8" -dependencies = [ - "derive_more 1.0.0", - "malachite", - "num-integer", - "num-traits", - "paste", -] - -[[package]] -name = "malachite-nz" -version = "0.4.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34a79feebb2bc9aa7762047c8e5495269a367da6b5a90a99882a0aeeac1841f7" -dependencies = [ - "itertools 0.11.0", - "libm", - "malachite-base", -] - -[[package]] -name = "malachite-q" -version = "0.4.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50f235d5747b1256b47620f5640c2a17a88c7569eebdf27cd9cb130e1a619191" -dependencies = [ - "itertools 0.11.0", - "malachite-base", - "malachite-nz", -] - [[package]] name = "markup5ever" version = "0.12.1" @@ -3945,12 +3809,6 @@ dependencies = [ "subtle", ] -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "pastey" version = "0.2.1" @@ -4293,7 +4151,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.2", + "rustc-hash", "rustls 0.23.37", "socket2 0.6.3", "thiserror 2.0.18", @@ -4313,7 +4171,7 @@ dependencies = [ "lru-slab", "rand 0.9.2", "ring", - "rustc-hash 2.1.2", + "rustc-hash", "rustls 0.23.37", "rustls-pki-types", "slab", @@ -4364,8 +4222,6 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ - "libc", - "rand_chacha 0.3.1", "rand_core 0.6.4", ] @@ -4375,20 +4231,10 @@ version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ - "rand_chacha 0.9.0", + "rand_chacha", "rand_core 0.9.5", ] -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core 0.6.4", -] - [[package]] name = "rand_chacha" version = "0.9.0" @@ -4732,12 +4578,6 @@ version = "0.1.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - [[package]] name = "rustc-hash" version = "2.1.2" @@ -4851,63 +4691,6 @@ dependencies = [ "untrusted", ] -[[package]] -name = "rustpython-ast" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5" -dependencies = [ - "is-macro", - "malachite-bigint", - "rustpython-parser-core", - "static_assertions", -] - -[[package]] -name = "rustpython-parser" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb" -dependencies = [ - "anyhow", - "is-macro", - "itertools 0.11.0", - "lalrpop-util", - "log", - "malachite-bigint", - "num-traits", - "phf", - "phf_codegen", - "rustc-hash 1.1.0", - "rustpython-ast", - "rustpython-parser-core", - "tiny-keccak", - "unic-emoji-char", - "unic-ucd-ident", - "unicode_names2", -] - -[[package]] -name = "rustpython-parser-core" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad" -dependencies = [ - "is-macro", - "memchr", - "rustpython-parser-vendored", -] - -[[package]] -name = "rustpython-parser-vendored" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6" -dependencies = [ - "memchr", - "once_cell", -] - [[package]] name = "rustversion" version = "1.0.22" @@ -5388,12 +5171,6 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - [[package]] name = "stop-words" version = "0.9.0" @@ -5403,6 +5180,12 @@ dependencies = [ "serde_json", ] +[[package]] +name = "streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b2231b7c3057d5e4ad0156fb3dc807d900806020c5ffa3ee6ff2c8c76fb8520" + [[package]] name = "string_cache" version = "0.8.9" @@ -5739,15 +5522,6 @@ dependencies = [ "time-core", ] -[[package]] -name = "tiny-keccak" -version = "2.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" -dependencies = [ - "crunchy", -] - [[package]] name = "tinystr" version = "0.8.3" @@ -6053,6 +5827,35 @@ dependencies = [ "tracing-log", ] +[[package]] +name = "tree-sitter" +version = "0.24.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5387dffa7ffc7d2dae12b50c6f7aab8ff79d6210147c6613561fc3d474c6f75" +dependencies = [ + "cc", + "regex", + "regex-syntax", + "streaming-iterator", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-language" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" + +[[package]] +name = "tree-sitter-python" +version = "0.23.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d065aaa27f3aaceaf60c1f0e0ac09e1cb9eb8ed28e7bcdaa52129cffc7f4b04" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree_magic_mini" version = "3.2.2" @@ -6136,58 +5939,6 @@ dependencies = [ "syn", ] -[[package]] -name = "unic-char-property" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221" -dependencies = [ - "unic-char-range", -] - -[[package]] -name = "unic-char-range" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc" - -[[package]] -name = "unic-common" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc" - -[[package]] -name = "unic-emoji-char" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d" -dependencies = [ - "unic-char-property", - "unic-char-range", - "unic-ucd-version", -] - -[[package]] -name = "unic-ucd-ident" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987" -dependencies = [ - "unic-char-property", - "unic-char-range", - "unic-ucd-version", -] - -[[package]] -name = "unic-ucd-version" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4" -dependencies = [ - "unic-common", -] - [[package]] name = "unicase" version = "2.9.0" @@ -6224,28 +5975,6 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" -[[package]] -name = "unicode_names2" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd" -dependencies = [ - "phf", - "unicode_names2_generator", -] - -[[package]] -name = "unicode_names2_generator" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e" -dependencies = [ - "getopts", - "log", - "phf_codegen", - "rand 0.8.5", -] - [[package]] name = "universal-hash" version = "0.5.1" diff --git a/Cargo.toml b/Cargo.toml index 837f646..8b2177b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -91,8 +91,8 @@ strum_macros = "0.27.2" indoc = "2.0.6" rmcp = { version = "0.16.0", features = ["client", "transport-child-process"] } num_cpus = "1.17.0" -rustpython-parser = "0.4.0" -rustpython-ast = "0.4.0" +tree-sitter = "0.24" +tree-sitter-python = "0.23" colored = "3.0.0" clap_complete = { version = "4.5.58", features = ["unstable-dynamic"] } gman = "0.3.0" diff --git a/src/parsers/python.rs b/src/parsers/python.rs index 147ae9b..ff411b9 100644 --- a/src/parsers/python.rs +++ b/src/parsers/python.rs @@ -1,13 +1,11 @@ use crate::function::{FunctionDeclaration, JsonSchema}; -use anyhow::{Context, Result, bail}; -use ast::{Stmt, StmtFunctionDef}; +use anyhow::{Context, Result, anyhow, bail}; use indexmap::IndexMap; -use rustpython_ast::{Constant, Expr, UnaryOp}; -use rustpython_parser::{Mode, ast}; use serde_json::Value; use std::fs::File; use std::io::Read; use std::path::Path; +use tree_sitter::{Node, Parser, Tree}; #[derive(Debug)] struct Param { @@ -28,178 +26,197 @@ pub fn generate_python_declarations( tool_file .read_to_string(&mut src) .with_context(|| format!("Failed to load script at '{tool_file:?}'"))?; - let suite = parse_suite(&src, file_name)?; + let tree = parse_tree(&src, file_name)?; let is_tool = parent .and_then(|p| p.file_name()) .is_some_and(|n| n == "tools"); - let mut declarations = python_to_function_declarations(file_name, &suite, is_tool)?; - if is_tool { - for d in &mut declarations { - d.agent = true; - } - } - - Ok(declarations) + python_to_function_declarations(file_name, &src, &tree, is_tool) } -fn parse_suite(src: &str, filename: &str) -> Result { - let mod_ast = - rustpython_parser::parse(src, Mode::Module, filename).context("failed to parse python")?; +fn parse_tree(src: &str, filename: &str) -> Result { + let mut parser = Parser::new(); + let language = tree_sitter_python::LANGUAGE.into(); + parser + .set_language(&language) + .context("failed to initialize python tree-sitter parser")?; - let suite = match mod_ast { - ast::Mod::Module(m) => m.body, - ast::Mod::Interactive(m) => m.body, - ast::Mod::Expression(_) => bail!("expected a module; got a single expression"), - _ => bail!("unexpected parse mode/AST variant"), - }; + let tree = parser + .parse(src.as_bytes(), None) + .ok_or_else(|| anyhow!("failed to parse python: {filename}"))?; - Ok(suite) + if tree.root_node().has_error() { + bail!("failed to parse python: syntax error in {filename}"); + } + + Ok(tree) } fn python_to_function_declarations( file_name: &str, - module: &ast::Suite, + src: &str, + tree: &Tree, is_tool: bool, ) -> Result> { let mut out = Vec::new(); + let root = tree.root_node(); + let mut cursor = root.walk(); - for stmt in module { - if let Stmt::FunctionDef(fd) = stmt { - let func_name = fd.name.to_string(); + for stmt in root.named_children(&mut cursor) { + let Some(fd) = unwrap_function_definition(stmt) else { + continue; + }; - if func_name.starts_with('_') && func_name != "_instructions" { - continue; - } + let func_name = function_name(fd, src)?.to_string(); - if is_tool && func_name != "run" { - continue; - } - - let description = get_docstring_from_body(&fd.body).unwrap_or_default(); - let params = collect_params(fd); - let schema = build_parameters_schema(¶ms, &description); - let name = if is_tool && func_name == "run" { - underscore(file_name) - } else { - underscore(&func_name) - }; - let desc_trim = description.trim().to_string(); - if desc_trim.is_empty() { - bail!("Missing or empty description on function: {func_name}"); - } - - out.push(FunctionDeclaration { - name, - description: desc_trim, - parameters: schema, - agent: !is_tool, - }); + if func_name.starts_with('_') && func_name != "_instructions" { + continue; } + + if is_tool && func_name != "run" { + continue; + } + + let description = get_docstring_from_function(fd, src).unwrap_or_default(); + let params = collect_params(fd, src)?; + let schema = build_parameters_schema(¶ms, &description); + let name = if is_tool && func_name == "run" { + underscore(file_name) + } else { + underscore(&func_name) + }; + let desc_trim = description.trim().to_string(); + if desc_trim.is_empty() { + bail!("Missing or empty description on function: {func_name}"); + } + + out.push(FunctionDeclaration { + name, + description: desc_trim, + parameters: schema, + agent: !is_tool, + }); } Ok(out) } -fn get_docstring_from_body(body: &[Stmt]) -> Option { - let first = body.first()?; - if let Stmt::Expr(expr_stmt) = first - && let Expr::Constant(constant) = &*expr_stmt.value - && let Constant::Str(s) = &constant.value - { - return Some(s.clone()); +fn unwrap_function_definition(node: Node<'_>) -> Option> { + match node.kind() { + "function_definition" => Some(node), + "decorated_definition" => { + let mut cursor = node.walk(); + node.named_children(&mut cursor) + .find(|child| child.kind() == "function_definition") + } + _ => None, } - None } -fn collect_params(fd: &StmtFunctionDef) -> Vec { +fn function_name<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { + let name_node = node + .child_by_field_name("name") + .ok_or_else(|| anyhow!("function_definition missing name"))?; + node_text(name_node, src) +} + +fn get_docstring_from_function(node: Node<'_>, src: &str) -> Option { + let body = node.child_by_field_name("body")?; + let mut cursor = body.walk(); + let first = body.named_children(&mut cursor).next()?; + if first.kind() != "expression_statement" { + return None; + } + + let mut expr_cursor = first.walk(); + let expr = first.named_children(&mut expr_cursor).next()?; + if expr.kind() != "string" { + return None; + } + + let text = node_text(expr, src).ok()?; + strip_string_quotes(text) +} + +fn strip_string_quotes(text: &str) -> Option { + let quote_offset = text + .char_indices() + .find_map(|(idx, ch)| (ch == '\'' || ch == '"').then_some(idx))?; + let prefix = &text[..quote_offset]; + if !prefix.chars().all(|ch| ch.is_ascii_alphabetic()) { + return None; + } + if prefix.chars().any(|ch| ch == 'f' || ch == 'F') { + return None; + } + + let literal = &text[quote_offset..]; + let quote = if literal.starts_with("\"\"\"") { + "\"\"\"" + } else if literal.starts_with("'''") { + "'''" + } else if literal.starts_with('"') { + "\"" + } else if literal.starts_with('\'') { + "'" + } else { + return None; + }; + + if literal.len() < quote.len() * 2 || !literal.ends_with(quote) { + return None; + } + + Some(literal[quote.len()..literal.len() - quote.len()].to_string()) +} + +fn collect_params(node: Node<'_>, src: &str) -> Result> { + let parameters = node + .child_by_field_name("parameters") + .ok_or_else(|| anyhow!("function_definition missing parameters"))?; let mut out = Vec::new(); + let mut cursor = parameters.walk(); - for a in fd.args.posonlyargs.iter().chain(fd.args.args.iter()) { - let name = a.def.arg.to_string(); - let mut ty = get_arg_type(a.def.annotation.as_deref()); - let mut required = a.default.is_none(); - - if ty.ends_with('?') { - ty.pop(); - required = false; + for param in parameters.named_children(&mut cursor) { + match param.kind() { + "identifier" => out.push(Param { + name: node_text(param, src)?.to_string(), + ty_hint: String::new(), + required: true, + default: None, + doc_type: None, + doc_desc: None, + }), + "typed_parameter" => out.push(build_param( + parameter_name(param, src)?, + get_arg_type(param.child_by_field_name("type"), src)?, + true, + None, + )), + "default_parameter" => out.push(build_param( + parameter_name(param, src)?, + String::new(), + false, + Some(Value::Null), + )), + "typed_default_parameter" => out.push(build_param( + parameter_name(param, src)?, + get_arg_type(param.child_by_field_name("type"), src)?, + false, + Some(Value::Null), + )), + "list_splat_pattern" | "dictionary_splat_pattern" | "positional_separator" => { + bail!( + "Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions" + ) + } + "keyword_separator" => continue, + other => bail!("Unsupported parameter type: {other}"), } - - let default = if a.default.is_some() { - Some(Value::Null) - } else { - None - }; - - out.push(Param { - name, - ty_hint: ty, - required, - default, - doc_type: None, - doc_desc: None, - }); } - for a in &fd.args.kwonlyargs { - let name = a.def.arg.to_string(); - let mut ty = get_arg_type(a.def.annotation.as_deref()); - let mut required = a.default.is_none(); - - if ty.ends_with('?') { - ty.pop(); - required = false; - } - - let default = if a.default.is_some() { - Some(Value::Null) - } else { - None - }; - - out.push(Param { - name, - ty_hint: ty, - required, - default, - doc_type: None, - doc_desc: None, - }); - } - - if let Some(vararg) = &fd.args.vararg { - let name = vararg.arg.to_string(); - let inner = get_arg_type(vararg.annotation.as_deref()); - let ty = if inner.is_empty() { - "list[str]".into() - } else { - format!("list[{inner}]") - }; - - out.push(Param { - name, - ty_hint: ty, - required: false, - default: None, - doc_type: None, - doc_desc: None, - }); - } - - if let Some(kwarg) = &fd.args.kwarg { - let name = kwarg.arg.to_string(); - out.push(Param { - name, - ty_hint: "object".into(), - required: false, - default: None, - doc_type: None, - doc_desc: None, - }); - } - - if let Some(doc) = get_docstring_from_body(&fd.body) { + if let Some(doc) = get_docstring_from_function(node, src) { let meta = parse_docstring_args(&doc); for p in &mut out { if let Some((t, d)) = meta.get(&p.name) { @@ -218,69 +235,155 @@ fn collect_params(fd: &StmtFunctionDef) -> Vec { } } - out + Ok(out) } -fn get_arg_type(annotation: Option<&Expr>) -> String { - match annotation { - None => "".to_string(), - Some(Expr::Name(n)) => n.id.to_string(), - Some(Expr::Subscript(sub)) => match &*sub.value { - Expr::Name(name) if &name.id == "Optional" => { - let inner = get_arg_type(Some(&sub.slice)); - format!("{inner}?") - } - Expr::Name(name) if &name.id == "List" => { - let inner = get_arg_type(Some(&sub.slice)); - format!("list[{inner}]") - } - Expr::Name(name) if &name.id == "Literal" => { - let vals = literal_members(&sub.slice); - format!("literal:{}", vals.join("|")) - } - _ => "any".to_string(), - }, - _ => "any".to_string(), +fn build_param(name: &str, mut ty: String, mut required: bool, default: Option) -> Param { + if ty.ends_with('?') { + ty.pop(); + required = false; + } + + Param { + name: name.to_string(), + ty_hint: ty, + required, + default, + doc_type: None, + doc_desc: None, } } -fn expr_to_str(e: &Expr) -> String { - match e { - Expr::Constant(c) => match &c.value { - Constant::Str(s) => s.clone(), - Constant::Int(i) => i.to_string(), - Constant::Float(f) => f.to_string(), - Constant::Bool(b) => b.to_string(), - Constant::None => "None".to_string(), - Constant::Ellipsis => "...".to_string(), - Constant::Bytes(b) => String::from_utf8_lossy(b).into_owned(), - Constant::Complex { real, imag } => format!("{real}+{imag}j"), - _ => "any".to_string(), - }, +fn parameter_name<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { + if let Some(name) = node.child_by_field_name("name") { + return node_text(name, src); + } - Expr::Name(n) => n.id.to_string(), + let mut cursor = node.walk(); + node.named_children(&mut cursor) + .find(|child| child.kind() == "identifier") + .ok_or_else(|| anyhow!("parameter missing name")) + .and_then(|name| node_text(name, src)) +} - Expr::UnaryOp(u) => { - if matches!(u.op, UnaryOp::USub) { - let inner = expr_to_str(&u.operand); - if inner.parse::().is_ok() || inner.chars().all(|c| c.is_ascii_digit()) { - return format!("-{inner}"); - } +fn get_arg_type(annotation: Option>, src: &str) -> Result { + let Some(annotation) = annotation else { + return Ok(String::new()); + }; + + match annotation.kind() { + "type" => get_arg_type(named_child(annotation, 0), src), + "generic_type" => { + let value = annotation + .child_by_field_name("type") + .or_else(|| named_child(annotation, 0)) + .ok_or_else(|| anyhow!("generic_type missing value"))?; + let value_name = if value.kind() == "identifier" { + node_text(value, src)? + } else { + return Ok("any".to_string()); + }; + + let inner = annotation + .child_by_field_name("type_parameter") + .or_else(|| annotation.child_by_field_name("parameters")) + .or_else(|| named_child(annotation, 1)) + .ok_or_else(|| anyhow!("generic_type missing inner type"))?; + + match value_name { + "Optional" => Ok(format!("{}?", generic_inner_type(inner, src)?)), + "List" => Ok(format!("list[{}]", generic_inner_type(inner, src)?)), + "Literal" => Ok(format!( + "literal:{}", + literal_members(inner, src)?.join("|") + )), + _ => Ok("any".to_string()), } - "any".to_string() } + "identifier" => Ok(node_text(annotation, src)?.to_string()), + "subscript" => { + let value = annotation + .child_by_field_name("value") + .or_else(|| named_child(annotation, 0)) + .ok_or_else(|| anyhow!("subscript missing value"))?; + let value_name = if value.kind() == "identifier" { + node_text(value, src)? + } else { + return Ok("any".to_string()); + }; - Expr::Tuple(t) => t.elts.iter().map(expr_to_str).collect::>().join(","), - - _ => "any".to_string(), + let inner = annotation + .child_by_field_name("subscript") + .or_else(|| annotation.child_by_field_name("slice")) + .or_else(|| named_child(annotation, 1)) + .ok_or_else(|| anyhow!("subscript missing inner type"))?; + match value_name { + "Optional" => Ok(format!("{}?", get_arg_type(Some(inner), src)?)), + "List" => Ok(format!("list[{}]", get_arg_type(Some(inner), src)?)), + "Literal" => Ok(format!( + "literal:{}", + literal_members(inner, src)?.join("|") + )), + _ => Ok("any".to_string()), + } + } + _ => Ok("any".to_string()), } } -fn literal_members(e: &Expr) -> Vec { - match e { - Expr::Tuple(t) => t.elts.iter().map(expr_to_str).collect(), - _ => vec![expr_to_str(e)], +fn generic_inner_type(node: Node<'_>, src: &str) -> Result { + if node.kind() == "type_parameter" { + return get_arg_type(named_child(node, 0), src); } + + get_arg_type(Some(node), src) +} + +fn literal_members(node: Node<'_>, src: &str) -> Result> { + if node.kind() == "type" { + return literal_members( + named_child(node, 0).ok_or_else(|| anyhow!("type missing inner literal"))?, + src, + ); + } + + if node.kind() == "tuple" || node.kind() == "type_parameter" { + let mut cursor = node.walk(); + let members = node + .named_children(&mut cursor) + .map(|child| expr_to_str(child, src)) + .collect::>>()?; + + return Ok(if members.is_empty() { + vec!["any".to_string()] + } else { + members + }); + } + + Ok(vec![expr_to_str(node, src)?]) +} + +fn expr_to_str(node: Node<'_>, src: &str) -> Result { + match node.kind() { + "type" => expr_to_str( + named_child(node, 0).ok_or_else(|| anyhow!("type missing expression"))?, + src, + ), + "string" | "integer" | "float" | "true" | "false" | "none" | "identifier" + | "unary_operator" => Ok(node_text(node, src)?.trim().to_string()), + _ => Ok("any".to_string()), + } +} + +fn named_child(node: Node<'_>, index: usize) -> Option> { + let mut cursor = node.walk(); + node.named_children(&mut cursor).nth(index) +} + +fn node_text<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { + node.utf8_text(src.as_bytes()) + .map_err(|err| anyhow!("invalid utf-8 in python source: {err}")) } fn parse_docstring_args(doc: &str) -> IndexMap { @@ -417,3 +520,413 @@ fn apply_type_to_schema(ty: &str, s: &mut JsonSchema) { .into(), ); } + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::time::{SystemTime, UNIX_EPOCH}; + + fn parse_source( + source: &str, + file_name: &str, + parent: &Path, + ) -> Result> { + let unique = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time went backwards") + .as_nanos(); + let path = std::env::temp_dir().join(format!("loki_python_parser_{file_name}_{unique}.py")); + fs::write(&path, source).expect("failed to write temp python source"); + let file = File::open(&path).expect("failed to open temp python source"); + let result = generate_python_declarations(file, file_name, Some(parent)); + let _ = fs::remove_file(&path); + result + } + + fn properties(schema: &JsonSchema) -> &IndexMap { + schema + .properties + .as_ref() + .expect("missing schema properties") + } + + fn property<'a>(schema: &'a JsonSchema, name: &str) -> &'a JsonSchema { + properties(schema) + .get(name) + .unwrap_or_else(|| panic!("missing property: {name}")) + } + + #[test] + fn test_tool_demo_py() { + let source = r#" +import os +from typing import List, Literal, Optional + +def run( + string: str, + string_enum: Literal["foo", "bar"], + boolean: bool, + integer: int, + number: float, + array: List[str], + string_optional: Optional[str] = None, + array_optional: Optional[List[str]] = None, +): + """Demonstrates how to create a tool using Python and how to use comments. + Args: + string: Define a required string property + string_enum: Define a required string property with enum + boolean: Define a required boolean property + integer: Define a required integer property + number: Define a required number property + array: Define a required string array property + string_optional: Define an optional string property + array_optional: Define an optional string array property + """ + output = f"""string: {string} +string_enum: {string_enum} +string_optional: {string_optional} +boolean: {boolean} +integer: {integer} +number: {number} +array: {array} +array_optional: {array_optional}""" + + for key, value in os.environ.items(): + if key.startswith("LLM_"): + output = f"{output}\n{key}: {value}" + + return output +"#; + + let declarations = parse_source(source, "demo_py", Path::new("tools")).unwrap(); + assert_eq!(declarations.len(), 1); + + let decl = &declarations[0]; + assert_eq!(decl.name, "demo_py"); + assert!(!decl.agent); + assert!(decl.description.starts_with("Demonstrates how to create")); + + let params = &decl.parameters; + assert_eq!(params.type_value.as_deref(), Some("object")); + assert_eq!( + params.required.as_ref().unwrap(), + &vec![ + "string".to_string(), + "string_enum".to_string(), + "boolean".to_string(), + "integer".to_string(), + "number".to_string(), + "array".to_string(), + ] + ); + + assert_eq!( + property(params, "string").type_value.as_deref(), + Some("string") + ); + + let string_enum = property(params, "string_enum"); + assert_eq!(string_enum.type_value.as_deref(), Some("string")); + assert_eq!( + string_enum.enum_value.as_ref().unwrap(), + &vec!["foo".to_string(), "bar".to_string()] + ); + + assert_eq!( + property(params, "boolean").type_value.as_deref(), + Some("boolean") + ); + assert_eq!( + property(params, "integer").type_value.as_deref(), + Some("integer") + ); + assert_eq!( + property(params, "number").type_value.as_deref(), + Some("number") + ); + + let array = property(params, "array"); + assert_eq!(array.type_value.as_deref(), Some("array")); + assert_eq!( + array.items.as_ref().unwrap().type_value.as_deref(), + Some("string") + ); + + let string_optional = property(params, "string_optional"); + assert_eq!(string_optional.type_value.as_deref(), Some("string")); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"string_optional".to_string()) + ); + + let array_optional = property(params, "array_optional"); + assert_eq!(array_optional.type_value.as_deref(), Some("array")); + assert_eq!( + array_optional.items.as_ref().unwrap().type_value.as_deref(), + Some("string") + ); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"array_optional".to_string()) + ); + } + + #[test] + fn test_tool_weather() { + let source = r#" +import os +from pathlib import Path +from typing import Optional +from urllib.parse import quote_plus +from urllib.request import urlopen + + +def run( + location: str, + llm_output: Optional[str] = None, +) -> str: + """Get the current weather in a given location + + Args: + location (str): The city and optionally the state or country (e.g., "London", "San Francisco, CA"). + + Returns: + str: A single-line formatted weather string from wttr.in (``format=4`` with metric units). + """ + url = f"https://wttr.in/{quote_plus(location)}?format=4&M" + + with urlopen(url, timeout=10) as resp: + weather = resp.read().decode("utf-8", errors="replace") + + dest = llm_output if llm_output is not None else os.environ.get("LLM_OUTPUT", "/dev/stdout") + + if dest not in {"-", "/dev/stdout"}: + path = Path(dest) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as fh: + fh.write(weather) + else: + pass + + return weather +"#; + + let declarations = parse_source(source, "get_current_weather", Path::new("tools")).unwrap(); + assert_eq!(declarations.len(), 1); + + let decl = &declarations[0]; + assert_eq!(decl.name, "get_current_weather"); + assert!(!decl.agent); + assert!( + decl.description + .starts_with("Get the current weather in a given location") + ); + + let params = &decl.parameters; + assert_eq!( + params.required.as_ref().unwrap(), + &vec!["location".to_string()] + ); + + let location = property(params, "location"); + assert_eq!(location.type_value.as_deref(), Some("string")); + assert_eq!( + location.description.as_deref(), + Some( + "The city and optionally the state or country (e.g., \"London\", \"San Francisco, CA\")." + ) + ); + + let llm_output = property(params, "llm_output"); + assert_eq!(llm_output.type_value.as_deref(), Some("string")); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"llm_output".to_string()) + ); + } + + #[test] + fn test_tool_execute_py_code() { + let source = r#" +import ast +import io +from contextlib import redirect_stdout + + +def run(code: str): + """Execute the given Python code. + Args: + code: The Python code to execute, such as `print("hello world")` + """ + output = io.StringIO() + with redirect_stdout(output): + value = exec_with_return(code, {}, {}) + + if value is not None: + output.write(str(value)) + + return output.getvalue() + + +def exec_with_return(code: str, globals: dict, locals: dict): + a = ast.parse(code) + last_expression = None + if a.body: + if isinstance(a_last := a.body[-1], ast.Expr): + last_expression = ast.unparse(a.body.pop()) + elif isinstance(a_last, ast.Assign): + last_expression = ast.unparse(a_last.targets[0]) + elif isinstance(a_last, (ast.AnnAssign, ast.AugAssign)): + last_expression = ast.unparse(a_last.target) + exec(ast.unparse(a), globals, locals) + if last_expression: + return eval(last_expression, globals, locals) +"#; + + let declarations = parse_source(source, "execute_py_code", Path::new("tools")).unwrap(); + assert_eq!(declarations.len(), 1); + + let decl = &declarations[0]; + assert_eq!(decl.name, "execute_py_code"); + assert!(!decl.agent); + + let params = &decl.parameters; + assert_eq!(properties(params).len(), 1); + let code = property(params, "code"); + assert_eq!(code.type_value.as_deref(), Some("string")); + assert_eq!( + code.description.as_deref(), + Some("The Python code to execute, such as `print(\"hello world\")`") + ); + } + + #[test] + fn test_agent_tools() { + let source = r#" +import urllib.request + +def get_ipinfo(): + """ + Get the ip info + """ + with urllib.request.urlopen("https://httpbin.org/ip") as response: + data = response.read() + return data.decode('utf-8') +"#; + + let declarations = parse_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 1); + + let decl = &declarations[0]; + assert_eq!(decl.name, "get_ipinfo"); + assert!(decl.agent); + assert_eq!(decl.description, "Get the ip info"); + assert!(properties(&decl.parameters).is_empty()); + } + + #[test] + fn test_reject_varargs() { + let source = r#" +def run(*args): + """Has docstring""" + return args +"#; + + let err = parse_source(source, "reject_varargs", Path::new("tools")).unwrap_err(); + assert!(err + .to_string() + .contains("Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions")); + } + + #[test] + fn test_reject_kwargs() { + let source = r#" +def run(**kwargs): + """Has docstring""" + return kwargs +"#; + + let err = parse_source(source, "reject_kwargs", Path::new("tools")).unwrap_err(); + assert!(err + .to_string() + .contains("Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions")); + } + + #[test] + fn test_reject_positional_only() { + let source = r#" +def run(x, /, y): + """Has docstring""" + return x + y +"#; + + let err = parse_source(source, "reject_positional_only", Path::new("tools")).unwrap_err(); + assert!(err + .to_string() + .contains("Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions")); + } + + #[test] + fn test_missing_docstring() { + let source = r#" +def run(x: str): + pass +"#; + + let err = parse_source(source, "missing_docstring", Path::new("tools")).unwrap_err(); + assert!( + err.to_string() + .contains("Missing or empty description on function: run") + ); + } + + #[test] + fn test_syntax_error() { + let source = "def run(: broken"; + let err = parse_source(source, "syntax_error", Path::new("tools")).unwrap_err(); + assert!(err.to_string().contains("failed to parse python")); + } + + #[test] + fn test_underscore_functions_skipped() { + let source = r#" +def _private(): + """Private""" + return None + +def public(): + """Public""" + return None +"#; + + let declarations = parse_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].name, "public"); + } + + #[test] + fn test_instructions_not_skipped() { + let source = r#" +def _instructions(): + """Help text""" + return None +"#; + + let declarations = parse_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].name, "instructions"); + assert_eq!(declarations[0].description, "Help text"); + assert!(declarations[0].agent); + } +} From f865892c2831d08ccddc5c42182550432158b757 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:16:35 -0600 Subject: [PATCH 02/12] refactor: Extracted common Python parser logic into a common.rs module --- src/parsers/common.rs | 269 +++++++++++++++++++++++++ src/parsers/mod.rs | 1 + src/parsers/python.rs | 443 +++++++++++++----------------------------- 3 files changed, 410 insertions(+), 303 deletions(-) create mode 100644 src/parsers/common.rs diff --git a/src/parsers/common.rs b/src/parsers/common.rs new file mode 100644 index 0000000..b0d80d6 --- /dev/null +++ b/src/parsers/common.rs @@ -0,0 +1,269 @@ +use crate::function::{FunctionDeclaration, JsonSchema}; +use anyhow::{Context, Result, anyhow, bail}; +use indexmap::IndexMap; +use serde_json::Value; +use tree_sitter::Node; + +#[derive(Debug)] +pub(crate) struct Param { + pub name: String, + pub ty_hint: String, + pub required: bool, + pub default: Option, + pub doc_type: Option, + pub doc_desc: Option, +} + +pub(crate) trait ScriptedLanguage { + fn ts_language(&self) -> tree_sitter::Language; + + fn default_runtime(&self) -> &str; + + fn lang_name(&self) -> &str; + + fn find_functions<'a>( + &self, + root: Node<'a>, + src: &str, + ) -> Vec<(Node<'a>, Node<'a>)>; + + fn function_name<'a>(&self, func_node: Node<'a>, src: &'a str) -> Result<&'a str>; + + fn extract_description( + &self, + wrapper_node: Node<'_>, + func_node: Node<'_>, + src: &str, + ) -> Option; + + fn extract_params( + &self, + func_node: Node<'_>, + src: &str, + description: &str, + ) -> Result>; +} + +pub(crate) fn build_param( + name: &str, + mut ty: String, + mut required: bool, + default: Option, +) -> Param { + if ty.ends_with('?') { + ty.pop(); + required = false; + } + + Param { + name: name.to_string(), + ty_hint: ty, + required, + default, + doc_type: None, + doc_desc: None, + } +} + +pub(crate) fn build_parameters_schema(params: &[Param], _description: &str) -> JsonSchema { + let mut props: IndexMap = IndexMap::new(); + let mut req: Vec = Vec::new(); + + for p in params { + let name = p.name.replace('-', "_"); + let mut schema = JsonSchema::default(); + + let ty = if !p.ty_hint.is_empty() { + p.ty_hint.as_str() + } else if let Some(t) = &p.doc_type { + t.as_str() + } else { + "str" + }; + + if let Some(d) = &p.doc_desc + && !d.is_empty() + { + schema.description = Some(d.clone()); + } + + apply_type_to_schema(ty, &mut schema); + + if p.default.is_none() && p.required { + req.push(name.clone()); + } + + props.insert(name, schema); + } + + JsonSchema { + type_value: Some("object".into()), + description: None, + properties: Some(props), + items: None, + any_of: None, + enum_value: None, + default: None, + required: if req.is_empty() { None } else { Some(req) }, + } +} + +pub(crate) fn apply_type_to_schema(ty: &str, s: &mut JsonSchema) { + let t = ty.trim_end_matches('?'); + if let Some(rest) = t.strip_prefix("list[") { + s.type_value = Some("array".into()); + let inner = rest.trim_end_matches(']'); + let mut item = JsonSchema::default(); + + apply_type_to_schema(inner, &mut item); + + if item.type_value.is_none() { + item.type_value = Some("string".into()); + } + s.items = Some(Box::new(item)); + + return; + } + + if let Some(rest) = t.strip_prefix("literal:") { + s.type_value = Some("string".into()); + let vals = rest + .split('|') + .map(|x| x.trim().trim_matches('"').trim_matches('\'').to_string()) + .collect::>(); + if !vals.is_empty() { + s.enum_value = Some(vals); + } + return; + } + + s.type_value = Some( + match t { + "bool" => "boolean", + "int" => "integer", + "float" => "number", + "str" | "any" | "" => "string", + _ => "string", + } + .into(), + ); +} + +pub(crate) fn underscore(s: &str) -> String { + s.chars() + .map(|c| { + if c.is_ascii_alphanumeric() { + c.to_ascii_lowercase() + } else { + '_' + } + }) + .collect::() + .split('_') + .filter(|t| !t.is_empty()) + .collect::>() + .join("_") +} + +pub(crate) fn node_text<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { + node.utf8_text(src.as_bytes()) + .map_err(|err| anyhow!("invalid utf-8 in source: {err}")) +} + +pub(crate) fn named_child(node: Node<'_>, index: usize) -> Option> { + let mut cursor = node.walk(); + node.named_children(&mut cursor).nth(index) +} + +pub(crate) fn extract_runtime(tree: &tree_sitter::Tree, src: &str, default: &str) -> String { + let root = tree.root_node(); + let mut cursor = root.walk(); + for child in root.named_children(&mut cursor) { + let text = match child.kind() { + "hash_bang_line" | "comment" => match child.utf8_text(src.as_bytes()) { + Ok(t) => t, + Err(_) => continue, + }, + _ => break, + }; + + if let Some(cmd) = text.strip_prefix("#!") { + let cmd = cmd.trim(); + if let Some(after_env) = cmd.strip_prefix("/usr/bin/env ") { + return after_env.trim().to_string(); + } + return cmd.to_string(); + } + + break; + } + default.to_string() +} + +pub(crate) fn generate_declarations( + lang: &L, + src: &str, + file_name: &str, + is_tool: bool, +) -> Result> { + let mut parser = tree_sitter::Parser::new(); + let language = lang.ts_language(); + parser.set_language(&language).with_context(|| { + format!( + "failed to initialize {} tree-sitter parser", + lang.lang_name() + ) + })?; + + let tree = parser + .parse(src.as_bytes(), None) + .ok_or_else(|| anyhow!("failed to parse {}: {file_name}", lang.lang_name()))?; + + if tree.root_node().has_error() { + bail!( + "failed to parse {}: syntax error in {file_name}", + lang.lang_name() + ); + } + + let _runtime = extract_runtime(&tree, src, lang.default_runtime()); + + let mut out = Vec::new(); + for (wrapper, func) in lang.find_functions(tree.root_node(), src) { + let func_name = lang.function_name(func, src)?; + + if func_name.starts_with('_') && func_name != "_instructions" { + continue; + } + if is_tool && func_name != "run" { + continue; + } + + let description = lang + .extract_description(wrapper, func, src) + .unwrap_or_default(); + let params = lang + .extract_params(func, src, &description) + .with_context(|| format!("in function '{func_name}' in {file_name}"))?; + let schema = build_parameters_schema(¶ms, &description); + + let name = if is_tool && func_name == "run" { + underscore(file_name) + } else { + underscore(func_name) + }; + + let desc_trim = description.trim().to_string(); + if desc_trim.is_empty() { + bail!("Missing or empty description on function: {func_name}"); + } + + out.push(FunctionDeclaration { + name, + description: desc_trim, + parameters: schema, + agent: !is_tool, + }); + } + Ok(out) +} diff --git a/src/parsers/mod.rs b/src/parsers/mod.rs index 57ae15a..5b872a9 100644 --- a/src/parsers/mod.rs +++ b/src/parsers/mod.rs @@ -1,2 +1,3 @@ pub(crate) mod bash; +pub(crate) mod common; pub(crate) mod python; diff --git a/src/parsers/python.rs b/src/parsers/python.rs index ff411b9..d8eb169 100644 --- a/src/parsers/python.rs +++ b/src/parsers/python.rs @@ -1,20 +1,124 @@ -use crate::function::{FunctionDeclaration, JsonSchema}; +use crate::function::FunctionDeclaration; +use crate::parsers::common::{self, Param, ScriptedLanguage}; use anyhow::{Context, Result, anyhow, bail}; use indexmap::IndexMap; use serde_json::Value; use std::fs::File; use std::io::Read; use std::path::Path; -use tree_sitter::{Node, Parser, Tree}; +use tree_sitter::Node; -#[derive(Debug)] -struct Param { - name: String, - ty_hint: String, - required: bool, - default: Option, - doc_type: Option, - doc_desc: Option, +pub(crate) struct PythonLanguage; + +impl ScriptedLanguage for PythonLanguage { + fn ts_language(&self) -> tree_sitter::Language { + tree_sitter_python::LANGUAGE.into() + } + + fn default_runtime(&self) -> &str { + "python" + } + + fn lang_name(&self) -> &str { + "python" + } + + fn find_functions<'a>(&self, root: Node<'a>, _src: &str) -> Vec<(Node<'a>, Node<'a>)> { + let mut cursor = root.walk(); + root.named_children(&mut cursor) + .filter_map(|stmt| unwrap_function_definition(stmt).map(|fd| (stmt, fd))) + .collect() + } + + fn function_name<'a>(&self, func_node: Node<'a>, src: &'a str) -> Result<&'a str> { + let name_node = func_node + .child_by_field_name("name") + .ok_or_else(|| anyhow!("function_definition missing name"))?; + common::node_text(name_node, src) + } + + fn extract_description( + &self, + _wrapper_node: Node<'_>, + func_node: Node<'_>, + src: &str, + ) -> Option { + get_docstring_from_function(func_node, src) + } + + fn extract_params( + &self, + func_node: Node<'_>, + src: &str, + description: &str, + ) -> Result> { + let parameters = func_node + .child_by_field_name("parameters") + .ok_or_else(|| anyhow!("function_definition missing parameters"))?; + let mut out = Vec::new(); + let mut cursor = parameters.walk(); + + for param in parameters.named_children(&mut cursor) { + match param.kind() { + "identifier" => out.push(Param { + name: common::node_text(param, src)?.to_string(), + ty_hint: String::new(), + required: true, + default: None, + doc_type: None, + doc_desc: None, + }), + "typed_parameter" => out.push(common::build_param( + parameter_name(param, src)?, + get_arg_type(param.child_by_field_name("type"), src)?, + true, + None, + )), + "default_parameter" => out.push(common::build_param( + parameter_name(param, src)?, + String::new(), + false, + Some(Value::Null), + )), + "typed_default_parameter" => out.push(common::build_param( + parameter_name(param, src)?, + get_arg_type(param.child_by_field_name("type"), src)?, + false, + Some(Value::Null), + )), + "list_splat_pattern" | "dictionary_splat_pattern" | "positional_separator" => { + let line = param.start_position().row + 1; + bail!( + "line {line}: *args/*kwargs/positional-only parameters are not supported in tool functions" + ) + } + "keyword_separator" => continue, + other => { + let line = param.start_position().row + 1; + bail!("line {line}: unsupported parameter type: {other}") + } + } + } + + let meta = parse_docstring_args(description); + for p in &mut out { + if let Some((t, d)) = meta.get(&p.name) { + if !t.is_empty() { + p.doc_type = Some(t.clone()); + } + + if !d.is_empty() { + p.doc_desc = Some(d.clone()); + } + + if t.ends_with('?') { + p.required = false; + } + } + } + + Ok(out) + } } pub fn generate_python_declarations( @@ -26,80 +130,12 @@ pub fn generate_python_declarations( tool_file .read_to_string(&mut src) .with_context(|| format!("Failed to load script at '{tool_file:?}'"))?; - let tree = parse_tree(&src, file_name)?; let is_tool = parent .and_then(|p| p.file_name()) .is_some_and(|n| n == "tools"); - python_to_function_declarations(file_name, &src, &tree, is_tool) -} - -fn parse_tree(src: &str, filename: &str) -> Result { - let mut parser = Parser::new(); - let language = tree_sitter_python::LANGUAGE.into(); - parser - .set_language(&language) - .context("failed to initialize python tree-sitter parser")?; - - let tree = parser - .parse(src.as_bytes(), None) - .ok_or_else(|| anyhow!("failed to parse python: {filename}"))?; - - if tree.root_node().has_error() { - bail!("failed to parse python: syntax error in {filename}"); - } - - Ok(tree) -} - -fn python_to_function_declarations( - file_name: &str, - src: &str, - tree: &Tree, - is_tool: bool, -) -> Result> { - let mut out = Vec::new(); - let root = tree.root_node(); - let mut cursor = root.walk(); - - for stmt in root.named_children(&mut cursor) { - let Some(fd) = unwrap_function_definition(stmt) else { - continue; - }; - - let func_name = function_name(fd, src)?.to_string(); - - if func_name.starts_with('_') && func_name != "_instructions" { - continue; - } - - if is_tool && func_name != "run" { - continue; - } - - let description = get_docstring_from_function(fd, src).unwrap_or_default(); - let params = collect_params(fd, src)?; - let schema = build_parameters_schema(¶ms, &description); - let name = if is_tool && func_name == "run" { - underscore(file_name) - } else { - underscore(&func_name) - }; - let desc_trim = description.trim().to_string(); - if desc_trim.is_empty() { - bail!("Missing or empty description on function: {func_name}"); - } - - out.push(FunctionDeclaration { - name, - description: desc_trim, - parameters: schema, - agent: !is_tool, - }); - } - - Ok(out) + common::generate_declarations(&PythonLanguage, &src, file_name, is_tool) } fn unwrap_function_definition(node: Node<'_>) -> Option> { @@ -114,13 +150,6 @@ fn unwrap_function_definition(node: Node<'_>) -> Option> { } } -fn function_name<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { - let name_node = node - .child_by_field_name("name") - .ok_or_else(|| anyhow!("function_definition missing name"))?; - node_text(name_node, src) -} - fn get_docstring_from_function(node: Node<'_>, src: &str) -> Option { let body = node.child_by_field_name("body")?; let mut cursor = body.walk(); @@ -135,7 +164,7 @@ fn get_docstring_from_function(node: Node<'_>, src: &str) -> Option { return None; } - let text = node_text(expr, src).ok()?; + let text = common::node_text(expr, src).ok()?; strip_string_quotes(text) } @@ -171,99 +200,16 @@ fn strip_string_quotes(text: &str) -> Option { Some(literal[quote.len()..literal.len() - quote.len()].to_string()) } -fn collect_params(node: Node<'_>, src: &str) -> Result> { - let parameters = node - .child_by_field_name("parameters") - .ok_or_else(|| anyhow!("function_definition missing parameters"))?; - let mut out = Vec::new(); - let mut cursor = parameters.walk(); - - for param in parameters.named_children(&mut cursor) { - match param.kind() { - "identifier" => out.push(Param { - name: node_text(param, src)?.to_string(), - ty_hint: String::new(), - required: true, - default: None, - doc_type: None, - doc_desc: None, - }), - "typed_parameter" => out.push(build_param( - parameter_name(param, src)?, - get_arg_type(param.child_by_field_name("type"), src)?, - true, - None, - )), - "default_parameter" => out.push(build_param( - parameter_name(param, src)?, - String::new(), - false, - Some(Value::Null), - )), - "typed_default_parameter" => out.push(build_param( - parameter_name(param, src)?, - get_arg_type(param.child_by_field_name("type"), src)?, - false, - Some(Value::Null), - )), - "list_splat_pattern" | "dictionary_splat_pattern" | "positional_separator" => { - bail!( - "Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions" - ) - } - "keyword_separator" => continue, - other => bail!("Unsupported parameter type: {other}"), - } - } - - if let Some(doc) = get_docstring_from_function(node, src) { - let meta = parse_docstring_args(&doc); - for p in &mut out { - if let Some((t, d)) = meta.get(&p.name) { - if !t.is_empty() { - p.doc_type = Some(t.clone()); - } - - if !d.is_empty() { - p.doc_desc = Some(d.clone()); - } - - if t.ends_with('?') { - p.required = false; - } - } - } - } - - Ok(out) -} - -fn build_param(name: &str, mut ty: String, mut required: bool, default: Option) -> Param { - if ty.ends_with('?') { - ty.pop(); - required = false; - } - - Param { - name: name.to_string(), - ty_hint: ty, - required, - default, - doc_type: None, - doc_desc: None, - } -} - fn parameter_name<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { if let Some(name) = node.child_by_field_name("name") { - return node_text(name, src); + return common::node_text(name, src); } let mut cursor = node.walk(); node.named_children(&mut cursor) .find(|child| child.kind() == "identifier") .ok_or_else(|| anyhow!("parameter missing name")) - .and_then(|name| node_text(name, src)) + .and_then(|name| common::node_text(name, src)) } fn get_arg_type(annotation: Option>, src: &str) -> Result { @@ -272,14 +218,14 @@ fn get_arg_type(annotation: Option>, src: &str) -> Result { }; match annotation.kind() { - "type" => get_arg_type(named_child(annotation, 0), src), + "type" => get_arg_type(common::named_child(annotation, 0), src), "generic_type" => { let value = annotation .child_by_field_name("type") - .or_else(|| named_child(annotation, 0)) + .or_else(|| common::named_child(annotation, 0)) .ok_or_else(|| anyhow!("generic_type missing value"))?; let value_name = if value.kind() == "identifier" { - node_text(value, src)? + common::node_text(value, src)? } else { return Ok("any".to_string()); }; @@ -287,7 +233,7 @@ fn get_arg_type(annotation: Option>, src: &str) -> Result { let inner = annotation .child_by_field_name("type_parameter") .or_else(|| annotation.child_by_field_name("parameters")) - .or_else(|| named_child(annotation, 1)) + .or_else(|| common::named_child(annotation, 1)) .ok_or_else(|| anyhow!("generic_type missing inner type"))?; match value_name { @@ -300,14 +246,14 @@ fn get_arg_type(annotation: Option>, src: &str) -> Result { _ => Ok("any".to_string()), } } - "identifier" => Ok(node_text(annotation, src)?.to_string()), + "identifier" => Ok(common::node_text(annotation, src)?.to_string()), "subscript" => { let value = annotation .child_by_field_name("value") - .or_else(|| named_child(annotation, 0)) + .or_else(|| common::named_child(annotation, 0)) .ok_or_else(|| anyhow!("subscript missing value"))?; let value_name = if value.kind() == "identifier" { - node_text(value, src)? + common::node_text(value, src)? } else { return Ok("any".to_string()); }; @@ -315,7 +261,7 @@ fn get_arg_type(annotation: Option>, src: &str) -> Result { let inner = annotation .child_by_field_name("subscript") .or_else(|| annotation.child_by_field_name("slice")) - .or_else(|| named_child(annotation, 1)) + .or_else(|| common::named_child(annotation, 1)) .ok_or_else(|| anyhow!("subscript missing inner type"))?; match value_name { "Optional" => Ok(format!("{}?", get_arg_type(Some(inner), src)?)), @@ -333,7 +279,7 @@ fn get_arg_type(annotation: Option>, src: &str) -> Result { fn generic_inner_type(node: Node<'_>, src: &str) -> Result { if node.kind() == "type_parameter" { - return get_arg_type(named_child(node, 0), src); + return get_arg_type(common::named_child(node, 0), src); } get_arg_type(Some(node), src) @@ -342,7 +288,7 @@ fn generic_inner_type(node: Node<'_>, src: &str) -> Result { fn literal_members(node: Node<'_>, src: &str) -> Result> { if node.kind() == "type" { return literal_members( - named_child(node, 0).ok_or_else(|| anyhow!("type missing inner literal"))?, + common::named_child(node, 0).ok_or_else(|| anyhow!("type missing inner literal"))?, src, ); } @@ -367,25 +313,15 @@ fn literal_members(node: Node<'_>, src: &str) -> Result> { fn expr_to_str(node: Node<'_>, src: &str) -> Result { match node.kind() { "type" => expr_to_str( - named_child(node, 0).ok_or_else(|| anyhow!("type missing expression"))?, + common::named_child(node, 0).ok_or_else(|| anyhow!("type missing expression"))?, src, ), "string" | "integer" | "float" | "true" | "false" | "none" | "identifier" - | "unary_operator" => Ok(node_text(node, src)?.trim().to_string()), + | "unary_operator" => Ok(common::node_text(node, src)?.trim().to_string()), _ => Ok("any".to_string()), } } -fn named_child(node: Node<'_>, index: usize) -> Option> { - let mut cursor = node.walk(); - node.named_children(&mut cursor).nth(index) -} - -fn node_text<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { - node.utf8_text(src.as_bytes()) - .map_err(|err| anyhow!("invalid utf-8 in python source: {err}")) -} - fn parse_docstring_args(doc: &str) -> IndexMap { let mut out = IndexMap::new(); let mut in_args = false; @@ -421,109 +357,10 @@ fn parse_docstring_args(doc: &str) -> IndexMap { out } -fn underscore(s: &str) -> String { - s.chars() - .map(|c| { - if c.is_ascii_alphanumeric() { - c.to_ascii_lowercase() - } else { - '_' - } - }) - .collect::() - .split('_') - .filter(|t| !t.is_empty()) - .collect::>() - .join("_") -} - -fn build_parameters_schema(params: &[Param], _description: &str) -> JsonSchema { - let mut props: IndexMap = IndexMap::new(); - let mut req: Vec = Vec::new(); - - for p in params { - let name = p.name.replace('-', "_"); - let mut schema = JsonSchema::default(); - - let ty = if !p.ty_hint.is_empty() { - p.ty_hint.as_str() - } else if let Some(t) = &p.doc_type { - t.as_str() - } else { - "str" - }; - - if let Some(d) = &p.doc_desc - && !d.is_empty() - { - schema.description = Some(d.clone()); - } - - apply_type_to_schema(ty, &mut schema); - - if p.default.is_none() && p.required { - req.push(name.clone()); - } - - props.insert(name, schema); - } - - JsonSchema { - type_value: Some("object".into()), - description: None, - properties: Some(props), - items: None, - any_of: None, - enum_value: None, - default: None, - required: if req.is_empty() { None } else { Some(req) }, - } -} - -fn apply_type_to_schema(ty: &str, s: &mut JsonSchema) { - let t = ty.trim_end_matches('?'); - if let Some(rest) = t.strip_prefix("list[") { - s.type_value = Some("array".into()); - let inner = rest.trim_end_matches(']'); - let mut item = JsonSchema::default(); - - apply_type_to_schema(inner, &mut item); - - if item.type_value.is_none() { - item.type_value = Some("string".into()); - } - s.items = Some(Box::new(item)); - - return; - } - - if let Some(rest) = t.strip_prefix("literal:") { - s.type_value = Some("string".into()); - let vals = rest - .split('|') - .map(|x| x.trim().trim_matches('"').trim_matches('\'').to_string()) - .collect::>(); - if !vals.is_empty() { - s.enum_value = Some(vals); - } - return; - } - - s.type_value = Some( - match t { - "bool" => "boolean", - "int" => "integer", - "float" => "number", - "str" | "any" | "" => "string", - _ => "string", - } - .into(), - ); -} - #[cfg(test)] mod tests { use super::*; + use crate::function::JsonSchema; use std::fs; use std::time::{SystemTime, UNIX_EPOCH}; @@ -844,9 +681,9 @@ def run(*args): "#; let err = parse_source(source, "reject_varargs", Path::new("tools")).unwrap_err(); - assert!(err - .to_string() - .contains("Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions")); + let msg = format!("{err:#}"); + assert!(msg.contains("*args/*kwargs/positional-only parameters are not supported")); + assert!(msg.contains("in function 'run'")); } #[test] @@ -858,9 +695,9 @@ def run(**kwargs): "#; let err = parse_source(source, "reject_kwargs", Path::new("tools")).unwrap_err(); - assert!(err - .to_string() - .contains("Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions")); + let msg = format!("{err:#}"); + assert!(msg.contains("*args/*kwargs/positional-only parameters are not supported")); + assert!(msg.contains("in function 'run'")); } #[test] @@ -872,9 +709,9 @@ def run(x, /, y): "#; let err = parse_source(source, "reject_positional_only", Path::new("tools")).unwrap_err(); - assert!(err - .to_string() - .contains("Unsupported parameter type: *args/*kwargs/positional-only parameters are not supported in tool functions")); + let msg = format!("{err:#}"); + assert!(msg.contains("*args/*kwargs/positional-only parameters are not supported")); + assert!(msg.contains("in function 'run'")); } #[test] From 4c75655f587a9728b52e2a92d60551868888b743 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:17:28 -0600 Subject: [PATCH 03/12] feat: Added TypeScript tool support using the refactored common ScriptedLanguage trait --- src/config/mod.rs | 2 +- src/function/mod.rs | 51 ++- src/parsers/mod.rs | 1 + src/parsers/typescript.rs | 793 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 839 insertions(+), 8 deletions(-) create mode 100644 src/parsers/typescript.rs diff --git a/src/config/mod.rs b/src/config/mod.rs index bd788c5..69fbdfe 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -584,7 +584,7 @@ impl Config { } pub fn agent_functions_file(name: &str) -> Result { - let allowed = ["tools.sh", "tools.py", "tools.js"]; + let allowed = ["tools.sh", "tools.py", "tools.ts", "tools.js"]; for entry in read_dir(Self::agent_data_dir(name))? { let entry = entry?; diff --git a/src/function/mod.rs b/src/function/mod.rs index 984d543..ca777c9 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -12,7 +12,7 @@ use crate::mcp::{ MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, }; -use crate::parsers::{bash, python}; +use crate::parsers::{bash, python, typescript}; use anyhow::{Context, Result, anyhow, bail}; use indexmap::IndexMap; use indoc::formatdoc; @@ -53,6 +53,7 @@ enum BinaryType<'a> { enum Language { Bash, Python, + TypeScript, Unsupported, } @@ -61,6 +62,7 @@ impl From<&String> for Language { match s.to_lowercase().as_str() { "sh" => Language::Bash, "py" => Language::Python, + "ts" => Language::TypeScript, _ => Language::Unsupported, } } @@ -72,6 +74,7 @@ impl Language { match self { Language::Bash => "bash", Language::Python => "python", + Language::TypeScript => "npx tsx", Language::Unsupported => "sh", } } @@ -80,6 +83,7 @@ impl Language { match self { Language::Bash => "sh", Language::Python => "py", + Language::TypeScript => "ts", _ => "sh", } } @@ -473,6 +477,11 @@ impl Functions { file_name, tools_file_path.parent(), ), + Language::TypeScript => typescript::generate_typescript_declarations( + tool_file, + file_name, + tools_file_path.parent(), + ), Language::Unsupported => { bail!("Unsupported tool file extension: {}", language.as_ref()) } @@ -684,6 +693,13 @@ impl Functions { let canonicalized_path = dunce::canonicalize(&executable_path)?; canonicalized_path.to_string_lossy().into_owned() } + Language::TypeScript => { + let npx_path = which::which("npx").map_err(|_| { + anyhow!("npx executable not found in PATH (required for TypeScript tools)") + })?; + let canonicalized_path = dunce::canonicalize(&npx_path)?; + format!("{} tsx", canonicalized_path.to_string_lossy()) + } _ => bail!("Unsupported language: {}", language.as_ref()), }; let bin_dir = binary_file @@ -773,13 +789,34 @@ impl Functions { "{prompt_utils_file}", &Config::bash_prompt_utils_file().to_string_lossy(), ); - if binary_file.exists() { - fs::remove_file(&binary_file)?; - } - let mut file = File::create(&binary_file)?; - file.write_all(content.as_bytes())?; - fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?; + if language == Language::TypeScript { + let bin_dir = binary_file + .parent() + .expect("Failed to get parent directory of binary file"); + let script_file = bin_dir.join(format!("run-{binary_name}.ts")); + if script_file.exists() { + fs::remove_file(&script_file)?; + } + let mut sf = File::create(&script_file)?; + sf.write_all(content.as_bytes())?; + fs::set_permissions(&script_file, fs::Permissions::from_mode(0o755))?; + + let wrapper = format!("#!/bin/sh\nexec tsx \"{}\" \"$@\"\n", script_file.display()); + if binary_file.exists() { + fs::remove_file(&binary_file)?; + } + let mut wf = File::create(&binary_file)?; + wf.write_all(wrapper.as_bytes())?; + fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?; + } else { + if binary_file.exists() { + fs::remove_file(&binary_file)?; + } + let mut file = File::create(&binary_file)?; + file.write_all(content.as_bytes())?; + fs::set_permissions(&binary_file, fs::Permissions::from_mode(0o755))?; + } Ok(()) } diff --git a/src/parsers/mod.rs b/src/parsers/mod.rs index 5b872a9..30f20c2 100644 --- a/src/parsers/mod.rs +++ b/src/parsers/mod.rs @@ -1,3 +1,4 @@ pub(crate) mod bash; pub(crate) mod common; pub(crate) mod python; +pub(crate) mod typescript; diff --git a/src/parsers/typescript.rs b/src/parsers/typescript.rs new file mode 100644 index 0000000..abaf878 --- /dev/null +++ b/src/parsers/typescript.rs @@ -0,0 +1,793 @@ +use crate::function::FunctionDeclaration; +use crate::parsers::common::{self, Param, ScriptedLanguage}; +use anyhow::{Context, Result, anyhow, bail}; +use indexmap::IndexMap; +use serde_json::Value; +use std::fs::File; +use std::io::Read; +use std::path::Path; +use tree_sitter::Node; + +pub(crate) struct TypeScriptLanguage; + +impl ScriptedLanguage for TypeScriptLanguage { + fn ts_language(&self) -> tree_sitter::Language { + tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into() + } + + fn default_runtime(&self) -> &str { + "npx tsx" + } + + fn lang_name(&self) -> &str { + "typescript" + } + + fn find_functions<'a>(&self, root: Node<'a>, _src: &str) -> Vec<(Node<'a>, Node<'a>)> { + let mut cursor = root.walk(); + root.named_children(&mut cursor) + .filter_map(|stmt| match stmt.kind() { + "export_statement" => unwrap_exported_function(stmt).map(|fd| (stmt, fd)), + _ => None, + }) + .collect() + } + + fn function_name<'a>(&self, func_node: Node<'a>, src: &'a str) -> Result<&'a str> { + let name_node = func_node + .child_by_field_name("name") + .ok_or_else(|| anyhow!("function_declaration missing name"))?; + common::node_text(name_node, src) + } + + fn extract_description( + &self, + wrapper_node: Node<'_>, + func_node: Node<'_>, + src: &str, + ) -> Option { + let text = jsdoc_text(wrapper_node, func_node, src)?; + let lines = clean_jsdoc_lines(text); + let mut description = Vec::new(); + for line in lines { + if line.starts_with('@') { + break; + } + description.push(line); + } + + let description = description.join("\n").trim().to_string(); + (!description.is_empty()).then_some(description) + } + + fn extract_params( + &self, + func_node: Node<'_>, + src: &str, + _description: &str, + ) -> Result> { + let parameters = func_node + .child_by_field_name("parameters") + .ok_or_else(|| anyhow!("function_declaration missing parameters"))?; + let mut out = Vec::new(); + let mut cursor = parameters.walk(); + + for param in parameters.named_children(&mut cursor) { + match param.kind() { + "required_parameter" | "optional_parameter" => { + let name = parameter_name(param, src)?; + let ty = get_arg_type(param.child_by_field_name("type"), src)?; + let required = param.kind() == "required_parameter" + && param.child_by_field_name("value").is_none(); + let default = param.child_by_field_name("value").map(|_| Value::Null); + out.push(common::build_param(name, ty, required, default)); + } + "rest_parameter" => { + let line = param.start_position().row + 1; + bail!("line {line}: rest parameters (...) are not supported in tool functions") + } + "object_pattern" => { + let line = param.start_position().row + 1; + bail!( + "line {line}: destructured object parameters (e.g. '{{ a, b }}: {{ a: string }}') \ + are not supported in tool functions. Use flat parameters instead (e.g. 'a: string, b: string')." + ) + } + other => { + let line = param.start_position().row + 1; + bail!("line {line}: unsupported parameter type: {other}") + } + } + } + + let wrapper = match func_node.parent() { + Some(parent) if parent.kind() == "export_statement" => parent, + _ => func_node, + }; + if let Some(doc) = jsdoc_text(wrapper, func_node, src) { + let meta = parse_jsdoc_params(doc); + for p in &mut out { + if let Some(desc) = meta.get(&p.name) + && !desc.is_empty() + { + p.doc_desc = Some(desc.clone()); + } + } + } + + Ok(out) + } +} + +pub fn generate_typescript_declarations( + mut tool_file: File, + file_name: &str, + parent: Option<&Path>, +) -> Result> { + let mut src = String::new(); + tool_file + .read_to_string(&mut src) + .with_context(|| format!("Failed to load script at '{tool_file:?}'"))?; + + let is_tool = parent + .and_then(|p| p.file_name()) + .is_some_and(|n| n == "tools"); + + common::generate_declarations(&TypeScriptLanguage, &src, file_name, is_tool) +} + +fn unwrap_exported_function(node: Node<'_>) -> Option> { + node.child_by_field_name("declaration") + .filter(|child| child.kind() == "function_declaration") + .or_else(|| { + let mut cursor = node.walk(); + node.named_children(&mut cursor) + .find(|child| child.kind() == "function_declaration") + }) +} + +fn jsdoc_text<'a>(wrapper_node: Node<'_>, func_node: Node<'_>, src: &'a str) -> Option<&'a str> { + wrapper_node + .prev_named_sibling() + .or_else(|| func_node.prev_named_sibling()) + .filter(|node| node.kind() == "comment") + .and_then(|node| common::node_text(node, src).ok()) + .filter(|text| text.trim_start().starts_with("/**")) +} + +fn clean_jsdoc_lines(doc: &str) -> Vec { + let trimmed = doc.trim(); + let inner = trimmed + .strip_prefix("/**") + .unwrap_or(trimmed) + .strip_suffix("*/") + .unwrap_or(trimmed); + + inner + .lines() + .map(|line| { + let line = line.trim(); + let line = line.strip_prefix('*').unwrap_or(line).trim_start(); + line.to_string() + }) + .collect() +} + +fn parse_jsdoc_params(doc: &str) -> IndexMap { + let mut out = IndexMap::new(); + + for line in clean_jsdoc_lines(doc) { + let Some(rest) = line.strip_prefix("@param") else { + continue; + }; + + let mut rest = rest.trim(); + if rest.starts_with('{') + && let Some(end) = rest.find('}') + { + rest = rest[end + 1..].trim_start(); + } + + if rest.is_empty() { + continue; + } + + let name_end = rest.find(char::is_whitespace).unwrap_or(rest.len()); + let mut name = rest[..name_end].trim(); + if let Some(stripped) = name.strip_suffix('?') { + name = stripped; + } + + if name.is_empty() { + continue; + } + + let mut desc = rest[name_end..].trim(); + if let Some(stripped) = desc.strip_prefix('-') { + desc = stripped.trim_start(); + } + + out.insert(name.to_string(), desc.to_string()); + } + + out +} + +fn parameter_name<'a>(node: Node<'_>, src: &'a str) -> Result<&'a str> { + if let Some(name) = node.child_by_field_name("name") { + return match name.kind() { + "identifier" => common::node_text(name, src), + "rest_pattern" => { + let line = node.start_position().row + 1; + bail!("line {line}: rest parameters (...) are not supported in tool functions") + } + "object_pattern" | "array_pattern" => { + let line = node.start_position().row + 1; + bail!( + "line {line}: destructured parameters are not supported in tool functions. \ + Use flat parameters instead (e.g. 'a: string, b: string')." + ) + } + other => { + let line = node.start_position().row + 1; + bail!("line {line}: unsupported parameter type: {other}") + } + }; + } + + let pattern = node + .child_by_field_name("pattern") + .ok_or_else(|| anyhow!("parameter missing pattern"))?; + + match pattern.kind() { + "identifier" => common::node_text(pattern, src), + "rest_pattern" => { + let line = node.start_position().row + 1; + bail!("line {line}: rest parameters (...) are not supported in tool functions") + } + "object_pattern" | "array_pattern" => { + let line = node.start_position().row + 1; + bail!( + "line {line}: destructured parameters are not supported in tool functions. \ + Use flat parameters instead (e.g. 'a: string, b: string')." + ) + } + other => { + let line = node.start_position().row + 1; + bail!("line {line}: unsupported parameter type: {other}") + } + } +} + +fn get_arg_type(annotation: Option>, src: &str) -> Result { + let Some(annotation) = annotation else { + return Ok(String::new()); + }; + + match annotation.kind() { + "type_annotation" | "type" => get_arg_type(common::named_child(annotation, 0), src), + "predefined_type" => Ok(match common::node_text(annotation, src)? { + "string" => "str", + "number" => "float", + "boolean" => "bool", + "any" | "unknown" | "void" | "undefined" => "any", + _ => "any", + } + .to_string()), + "type_identifier" | "nested_type_identifier" => Ok("any".to_string()), + "generic_type" => { + let name = annotation + .child_by_field_name("name") + .ok_or_else(|| anyhow!("generic_type missing name"))?; + let type_name = common::node_text(name, src)?; + let type_args = annotation + .child_by_field_name("type_arguments") + .ok_or_else(|| anyhow!("generic_type missing type arguments"))?; + let inner = common::named_child(type_args, 0) + .ok_or_else(|| anyhow!("generic_type missing inner type"))?; + + match type_name { + "Array" => Ok(format!("list[{}]", get_arg_type(Some(inner), src)?)), + _ => Ok("any".to_string()), + } + } + "array_type" => { + let inner = common::named_child(annotation, 0) + .ok_or_else(|| anyhow!("array_type missing inner type"))?; + Ok(format!("list[{}]", get_arg_type(Some(inner), src)?)) + } + "union_type" => resolve_union_type(annotation, src), + "literal_type" => resolve_literal_type(annotation, src), + "parenthesized_type" => get_arg_type(common::named_child(annotation, 0), src), + _ => Ok("any".to_string()), + } +} + +fn resolve_union_type(annotation: Node<'_>, src: &str) -> Result { + let members = flatten_union_members(annotation); + let has_null = members.iter().any(|member| is_nullish_type(*member, src)); + + let mut literal_values = Vec::new(); + let mut all_string_literals = true; + for member in &members { + match string_literal_member(*member, src) { + Some(value) => literal_values.push(value), + None => { + all_string_literals = false; + break; + } + } + } + + if all_string_literals && !literal_values.is_empty() { + return Ok(format!("literal:{}", literal_values.join("|"))); + } + + let mut first_non_null = None; + for member in members { + if is_nullish_type(member, src) { + continue; + } + first_non_null = Some(get_arg_type(Some(member), src)?); + break; + } + + let mut ty = first_non_null.unwrap_or_else(|| "any".to_string()); + if has_null && !ty.ends_with('?') { + ty.push('?'); + } + Ok(ty) +} + +fn flatten_union_members(node: Node<'_>) -> Vec> { + let node = if node.kind() == "type" { + match common::named_child(node, 0) { + Some(inner) => inner, + None => return vec![], + } + } else { + node + }; + + if node.kind() != "union_type" { + return vec![node]; + } + + let mut cursor = node.walk(); + let mut out = Vec::new(); + for child in node.named_children(&mut cursor) { + out.extend(flatten_union_members(child)); + } + out +} + +fn resolve_literal_type(annotation: Node<'_>, src: &str) -> Result { + let inner = common::named_child(annotation, 0) + .ok_or_else(|| anyhow!("literal_type missing inner literal"))?; + + match inner.kind() { + "string" | "number" | "true" | "false" | "unary_expression" => { + Ok(format!("literal:{}", common::node_text(inner, src)?.trim())) + } + "null" | "undefined" => Ok("any".to_string()), + _ => Ok("any".to_string()), + } +} + +fn string_literal_member(node: Node<'_>, src: &str) -> Option { + let node = if node.kind() == "type" { + common::named_child(node, 0)? + } else { + node + }; + + if node.kind() != "literal_type" { + return None; + } + + let inner = common::named_child(node, 0)?; + if inner.kind() != "string" { + return None; + } + + Some(common::node_text(inner, src).ok()?.to_string()) +} + +fn is_nullish_type(node: Node<'_>, src: &str) -> bool { + let node = if node.kind() == "type" { + match common::named_child(node, 0) { + Some(inner) => inner, + None => return false, + } + } else { + node + }; + + match node.kind() { + "literal_type" => common::named_child(node, 0) + .is_some_and(|inner| matches!(inner.kind(), "null" | "undefined")), + "predefined_type" => common::node_text(node, src) + .map(|text| matches!(text, "undefined" | "void")) + .unwrap_or(false), + _ => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::function::JsonSchema; + use std::fs; + use std::time::{SystemTime, UNIX_EPOCH}; + + fn parse_ts_source( + source: &str, + file_name: &str, + parent: &Path, + ) -> Result> { + let unique = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time") + .as_nanos(); + let path = std::env::temp_dir().join(format!("loki_ts_parser_{file_name}_{unique}.ts")); + fs::write(&path, source).expect("write"); + let file = File::open(&path).expect("open"); + let result = generate_typescript_declarations(file, file_name, Some(parent)); + let _ = fs::remove_file(&path); + result + } + + fn properties(schema: &JsonSchema) -> &IndexMap { + schema + .properties + .as_ref() + .expect("missing schema properties") + } + + fn property<'a>(schema: &'a JsonSchema, name: &str) -> &'a JsonSchema { + properties(schema) + .get(name) + .unwrap_or_else(|| panic!("missing property: {name}")) + } + + #[test] + fn test_ts_tool_demo() { + let source = r#" +/** + * Demonstrates how to create a tool using TypeScript. + * + * @param query - The search query string + * @param format - Output format + * @param count - Maximum results to return + * @param verbose - Enable verbose output + * @param tags - List of tags to filter by + * @param language - Optional language filter + * @param extra_tags - Optional extra tags + */ +export function run( + query: string, + format: "json" | "csv" | "xml", + count: number, + verbose: boolean, + tags: string[], + language?: string, + extra_tags?: Array, +): string { + return "result"; +} +"#; + + let declarations = parse_ts_source(source, "demo_ts", Path::new("tools")).unwrap(); + assert_eq!(declarations.len(), 1); + + let decl = &declarations[0]; + assert_eq!(decl.name, "demo_ts"); + assert!(!decl.agent); + + let params = &decl.parameters; + assert_eq!(params.type_value.as_deref(), Some("object")); + assert_eq!( + params.required.as_ref().unwrap(), + &vec![ + "query".to_string(), + "format".to_string(), + "count".to_string(), + "verbose".to_string(), + "tags".to_string(), + ] + ); + + assert_eq!( + property(params, "query").type_value.as_deref(), + Some("string") + ); + + let format = property(params, "format"); + assert_eq!(format.type_value.as_deref(), Some("string")); + assert_eq!( + format.enum_value.as_ref().unwrap(), + &vec!["json".to_string(), "csv".to_string(), "xml".to_string()] + ); + + assert_eq!( + property(params, "count").type_value.as_deref(), + Some("number") + ); + assert_eq!( + property(params, "verbose").type_value.as_deref(), + Some("boolean") + ); + + let tags = property(params, "tags"); + assert_eq!(tags.type_value.as_deref(), Some("array")); + assert_eq!( + tags.items.as_ref().unwrap().type_value.as_deref(), + Some("string") + ); + + let language = property(params, "language"); + assert_eq!(language.type_value.as_deref(), Some("string")); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"language".to_string()) + ); + + let extra_tags = property(params, "extra_tags"); + assert_eq!(extra_tags.type_value.as_deref(), Some("array")); + assert_eq!( + extra_tags.items.as_ref().unwrap().type_value.as_deref(), + Some("string") + ); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"extra_tags".to_string()) + ); + } + + #[test] + fn test_ts_tool_simple() { + let source = r#" +/** + * Execute the given code. + * + * @param code - The code to execute + */ +export function run(code: string): string { + return eval(code); +} +"#; + + let declarations = parse_ts_source(source, "execute_code", Path::new("tools")).unwrap(); + assert_eq!(declarations.len(), 1); + + let decl = &declarations[0]; + assert_eq!(decl.name, "execute_code"); + assert!(!decl.agent); + + let params = &decl.parameters; + assert_eq!(params.required.as_ref().unwrap(), &vec!["code".to_string()]); + assert_eq!( + property(params, "code").type_value.as_deref(), + Some("string") + ); + } + + #[test] + fn test_ts_agent_tools() { + let source = r#" +/** Get user info by ID */ +export function get_user(id: string): string { + return ""; +} + +/** List all users */ +export function list_users(): string { + return ""; +} +"#; + + let declarations = parse_ts_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].name, "get_user"); + assert_eq!(declarations[1].name, "list_users"); + assert!(declarations[0].agent); + assert!(declarations[1].agent); + } + + #[test] + fn test_ts_reject_rest_params() { + let source = r#" +/** + * Has rest params + */ +export function run(...args: string[]): string { + return ""; +} +"#; + + let err = parse_ts_source(source, "rest_params", Path::new("tools")).unwrap_err(); + let msg = format!("{err:#}"); + assert!(msg.contains("rest parameters")); + assert!(msg.contains("in function 'run'")); + } + + #[test] + fn test_ts_missing_jsdoc() { + let source = r#" +export function run(x: string): string { + return x; +} +"#; + + let err = parse_ts_source(source, "missing_jsdoc", Path::new("tools")).unwrap_err(); + assert!( + err.to_string() + .contains("Missing or empty description on function: run") + ); + } + + #[test] + fn test_ts_syntax_error() { + let source = "export function run(: broken"; + let err = parse_ts_source(source, "syntax_error", Path::new("tools")).unwrap_err(); + assert!(err.to_string().contains("failed to parse typescript")); + } + + #[test] + fn test_ts_underscore_skipped() { + let source = r#" +/** Private helper */ +function _helper(): void {} + +/** Public function */ +export function do_stuff(): string { + return ""; +} +"#; + + let declarations = parse_ts_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].name, "do_stuff"); + assert!(declarations[0].agent); + } + + #[test] + fn test_ts_non_exported_helpers_skipped() { + let source = r#" +#!/usr/bin/env tsx + +import { appendFileSync } from 'fs'; + +/** + * Get the current weather in a given location + * @param location - The city + */ +export function get_current_weather(location: string): string { + return fetchSync("https://example.com/" + location); +} + +function fetchSync(url: string): string { + return "sunny"; +} +"#; + + let declarations = parse_ts_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].name, "get_current_weather"); + } + + #[test] + fn test_ts_instructions_not_skipped() { + let source = r#" +/** Help text for the agent */ +export function _instructions(): string { + return ""; +} +"#; + + let declarations = parse_ts_source(source, "tools", Path::new("demo")).unwrap(); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].name, "instructions"); + assert!(declarations[0].agent); + } + + #[test] + fn test_ts_optional_with_null_union() { + let source = r#" +/** + * Fetch data with optional filter + * + * @param url - The URL to fetch + * @param filter - Optional filter string + */ +export function run(url: string, filter: string | null): string { + return ""; +} +"#; + + let declarations = parse_ts_source(source, "fetch_data", Path::new("tools")).unwrap(); + let params = &declarations[0].parameters; + assert!( + params + .required + .as_ref() + .unwrap() + .contains(&"url".to_string()) + ); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"filter".to_string()) + ); + assert_eq!( + property(params, "filter").type_value.as_deref(), + Some("string") + ); + } + + #[test] + fn test_ts_optional_with_default() { + let source = r#" +/** + * Search with limit + * + * @param query - Search query + * @param limit - Max results + */ +export function run(query: string, limit: number = 10): string { + return ""; +} +"#; + + let declarations = + parse_ts_source(source, "search_with_limit", Path::new("tools")).unwrap(); + let params = &declarations[0].parameters; + assert!( + params + .required + .as_ref() + .unwrap() + .contains(&"query".to_string()) + ); + assert!( + !params + .required + .as_ref() + .unwrap() + .contains(&"limit".to_string()) + ); + assert_eq!( + property(params, "limit").type_value.as_deref(), + Some("number") + ); + } + + #[test] + fn test_ts_shebang_parses() { + let source = r#"#!/usr/bin/env tsx + +/** + * Get weather + * @param location - The city + */ +export function run(location: string): string { + return location; +} +"#; + + let result = parse_ts_source(source, "get_weather", Path::new("tools")); + eprintln!("shebang parse result: {result:?}"); + assert!(result.is_ok(), "shebang should not cause parse failure"); + let declarations = result.unwrap(); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].name, "get_weather"); + } +} From f30e81af08235551a886334435fae1e33bf13fdb Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:17:53 -0600 Subject: [PATCH 04/12] fix: Added in forgotten wrapper scripts for TypeScript tools --- assets/functions/scripts/run-agent.ts | 164 ++++++++++++++++++++++++++ assets/functions/scripts/run-tool.ts | 159 +++++++++++++++++++++++++ 2 files changed, 323 insertions(+) create mode 100644 assets/functions/scripts/run-agent.ts create mode 100644 assets/functions/scripts/run-tool.ts diff --git a/assets/functions/scripts/run-agent.ts b/assets/functions/scripts/run-agent.ts new file mode 100644 index 0000000..425b5df --- /dev/null +++ b/assets/functions/scripts/run-agent.ts @@ -0,0 +1,164 @@ +#!/usr/bin/env tsx + +// Usage: ./{agent_name}.ts + +import { readFileSync, writeFileSync, existsSync } from "fs"; +import { join } from "path"; +import { pathToFileURL } from "url"; + +async function main(): Promise { + const { agentFunc, rawData } = parseArgv(); + const agentData = parseRawData(rawData); + + const configDir = "{config_dir}"; + setupEnv(configDir, agentFunc); + + const agentToolsPath = join(configDir, "agents", "{agent_name}", "tools.ts"); + await run(agentToolsPath, agentFunc, agentData); +} + +function parseRawData(data: string): Record { + if (!data) { + throw new Error("No JSON data"); + } + + try { + return JSON.parse(data); + } catch { + throw new Error("Invalid JSON data"); + } +} + +function parseArgv(): { agentFunc: string; rawData: string } { + const agentFunc = process.argv[2]; + + const toolDataFile = process.env["LLM_TOOL_DATA_FILE"]; + let agentData: string; + if (toolDataFile && existsSync(toolDataFile)) { + agentData = readFileSync(toolDataFile, "utf-8"); + } else { + agentData = process.argv[3]; + } + + if (!agentFunc || !agentData) { + process.stderr.write("Usage: ./{agent_name}.ts \n"); + process.exit(1); + } + + return { agentFunc, rawData: agentData }; +} + +function setupEnv(configDir: string, agentFunc: string): void { + loadEnv(join(configDir, ".env")); + process.env["LLM_ROOT_DIR"] = configDir; + process.env["LLM_AGENT_NAME"] = "{agent_name}"; + process.env["LLM_AGENT_FUNC"] = agentFunc; + process.env["LLM_AGENT_ROOT_DIR"] = join(configDir, "agents", "{agent_name}"); + process.env["LLM_AGENT_CACHE_DIR"] = join(configDir, "cache", "{agent_name}"); +} + +function loadEnv(filePath: string): void { + let lines: string[]; + try { + lines = readFileSync(filePath, "utf-8").split("\n"); + } catch { + return; + } + + for (const raw of lines) { + const line = raw.trim(); + if (line.startsWith("#") || !line) { + continue; + } + + const eqIdx = line.indexOf("="); + if (eqIdx === -1) { + continue; + } + + const key = line.slice(0, eqIdx).trim(); + if (key in process.env) { + continue; + } + + let value = line.slice(eqIdx + 1).trim(); + if ( + (value.startsWith('"') && value.endsWith('"')) || + (value.startsWith("'") && value.endsWith("'")) + ) { + value = value.slice(1, -1); + } + process.env[key] = value; + } +} + +async function run( + agentPath: string, + agentFunc: string, + agentData: Record, +): Promise { + const mod = await import(pathToFileURL(agentPath).href); + + if (typeof mod[agentFunc] !== "function") { + throw new Error(`No module function '${agentFunc}' at '${agentPath}'`); + } + + const value = await mod[agentFunc](agentData); + returnToLlm(value); + dumpResult(`{agent_name}:${agentFunc}`); +} + +function returnToLlm(value: unknown): void { + if (value === null || value === undefined) { + return; + } + + const output = process.env["LLM_OUTPUT"]; + const write = (s: string) => { + if (output) { + writeFileSync(output, s, "utf-8"); + } else { + process.stdout.write(s); + } + }; + + if (typeof value === "string" || typeof value === "number" || typeof value === "boolean") { + write(String(value)); + } else if (typeof value === "object") { + write(JSON.stringify(value, null, 2)); + } +} + +function dumpResult(name: string): void { + const dumpResults = process.env["LLM_DUMP_RESULTS"]; + const llmOutput = process.env["LLM_OUTPUT"]; + + if (!dumpResults || !llmOutput || !process.stdout.isTTY) { + return; + } + + try { + const pattern = new RegExp(`\\b(${dumpResults})\\b`); + if (!pattern.test(name)) { + return; + } + } catch { + return; + } + + let data: string; + try { + data = readFileSync(llmOutput, "utf-8"); + } catch { + return; + } + + process.stdout.write( + `\x1b[2m----------------------\n${data}\n----------------------\x1b[0m\n`, + ); +} + +main().catch((err) => { + process.stderr.write(`${err}\n`); + process.exit(1); +}); diff --git a/assets/functions/scripts/run-tool.ts b/assets/functions/scripts/run-tool.ts new file mode 100644 index 0000000..f3d2887 --- /dev/null +++ b/assets/functions/scripts/run-tool.ts @@ -0,0 +1,159 @@ +#!/usr/bin/env tsx + +// Usage: ./{function_name}.ts + +import { readFileSync, writeFileSync, existsSync, statSync } from "fs"; +import { join, basename } from "path"; +import { pathToFileURL } from "url"; + +async function main(): Promise { + const rawData = parseArgv(); + const toolData = parseRawData(rawData); + + const rootDir = "{root_dir}"; + setupEnv(rootDir); + + const toolPath = "{tool_path}.ts"; + await run(toolPath, "run", toolData); +} + +function parseRawData(data: string): Record { + if (!data) { + throw new Error("No JSON data"); + } + + try { + return JSON.parse(data); + } catch { + throw new Error("Invalid JSON data"); + } +} + +function parseArgv(): string { + const toolDataFile = process.env["LLM_TOOL_DATA_FILE"]; + if (toolDataFile && existsSync(toolDataFile)) { + return readFileSync(toolDataFile, "utf-8"); + } + + const toolData = process.argv[2]; + + if (!toolData) { + process.stderr.write("Usage: ./{function_name}.ts \n"); + process.exit(1); + } + + return toolData; +} + +function setupEnv(rootDir: string): void { + loadEnv(join(rootDir, ".env")); + process.env["LLM_ROOT_DIR"] = rootDir; + process.env["LLM_TOOL_NAME"] = "{function_name}"; + process.env["LLM_TOOL_CACHE_DIR"] = join(rootDir, "cache", "{function_name}"); +} + +function loadEnv(filePath: string): void { + let lines: string[]; + try { + lines = readFileSync(filePath, "utf-8").split("\n"); + } catch { + return; + } + + for (const raw of lines) { + const line = raw.trim(); + if (line.startsWith("#") || !line) { + continue; + } + + const eqIdx = line.indexOf("="); + if (eqIdx === -1) { + continue; + } + + const key = line.slice(0, eqIdx).trim(); + if (key in process.env) { + continue; + } + + let value = line.slice(eqIdx + 1).trim(); + if ( + (value.startsWith('"') && value.endsWith('"')) || + (value.startsWith("'") && value.endsWith("'")) + ) { + value = value.slice(1, -1); + } + process.env[key] = value; + } +} + +async function run( + toolPath: string, + toolFunc: string, + toolData: Record, +): Promise { + const mod = await import(pathToFileURL(toolPath).href); + + if (typeof mod[toolFunc] !== "function") { + throw new Error(`No module function '${toolFunc}' at '${toolPath}'`); + } + + const value = await mod[toolFunc](toolData); + returnToLlm(value); + dumpResult("{function_name}"); +} + +function returnToLlm(value: unknown): void { + if (value === null || value === undefined) { + return; + } + + const output = process.env["LLM_OUTPUT"]; + const write = (s: string) => { + if (output) { + writeFileSync(output, s, "utf-8"); + } else { + process.stdout.write(s); + } + }; + + if (typeof value === "string" || typeof value === "number" || typeof value === "boolean") { + write(String(value)); + } else if (typeof value === "object") { + write(JSON.stringify(value, null, 2)); + } +} + +function dumpResult(name: string): void { + const dumpResults = process.env["LLM_DUMP_RESULTS"]; + const llmOutput = process.env["LLM_OUTPUT"]; + + if (!dumpResults || !llmOutput || !process.stdout.isTTY) { + return; + } + + try { + const pattern = new RegExp(`\\b(${dumpResults})\\b`); + if (!pattern.test(name)) { + return; + } + } catch { + return; + } + + let data: string; + try { + data = readFileSync(llmOutput, "utf-8"); + } catch { + return; + } + + process.stdout.write( + `\x1b[2m----------------------\n${data}\n----------------------\x1b[0m\n`, + ); +} + +main().catch((err) => { + process.stderr.write(`${err}\n`); + process.exit(1); +}); From 4caa0355281b4f4daeca2dc0c1ab7a79d1ba511e Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:18:18 -0600 Subject: [PATCH 05/12] feat: Updated the Python demo tool to show all possible parameter types and variations --- assets/functions/tools/demo_py.py | 33 +++++++++++++++++++++---------- 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/assets/functions/tools/demo_py.py b/assets/functions/tools/demo_py.py index 8b0ca0a..d71cd52 100644 --- a/assets/functions/tools/demo_py.py +++ b/assets/functions/tools/demo_py.py @@ -1,6 +1,7 @@ import os from typing import List, Literal, Optional + def run( string: str, string_enum: Literal["foo", "bar"], @@ -9,26 +10,38 @@ def run( number: float, array: List[str], string_optional: Optional[str] = None, + integer_with_default: int = 42, + boolean_with_default: bool = True, + number_with_default: float = 3.14, + string_with_default: str = "hello", array_optional: Optional[List[str]] = None, ): - """Demonstrates how to create a tool using Python and how to use comments. + """Demonstrates all supported Python parameter types and variations. Args: - string: Define a required string property - string_enum: Define a required string property with enum - boolean: Define a required boolean property - integer: Define a required integer property - number: Define a required number property - array: Define a required string array property - string_optional: Define an optional string property - array_optional: Define an optional string array property + string: A required string property + string_enum: A required string property constrained to specific values + boolean: A required boolean property + integer: A required integer property + number: A required number (float) property + array: A required string array property + string_optional: An optional string property (Optional[str] with None default) + integer_with_default: An optional integer with a non-None default value + boolean_with_default: An optional boolean with a default value + number_with_default: An optional number with a default value + string_with_default: An optional string with a default value + array_optional: An optional string array property """ output = f"""string: {string} string_enum: {string_enum} -string_optional: {string_optional} boolean: {boolean} integer: {integer} number: {number} array: {array} +string_optional: {string_optional} +integer_with_default: {integer_with_default} +boolean_with_default: {boolean_with_default} +number_with_default: {number_with_default} +string_with_default: {string_with_default} array_optional: {array_optional}""" for key, value in os.environ.items(): From 11334149b0b28a6a33eb71f28194b1ce7f56ed85 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:18:41 -0600 Subject: [PATCH 06/12] feat: Created a demo TypeScript tool and a get_current_weather function in TypeScript --- assets/functions/tools/demo_ts.ts | 53 +++++++++++++++++++ assets/functions/tools/get_current_weather.ts | 24 +++++++++ 2 files changed, 77 insertions(+) create mode 100644 assets/functions/tools/demo_ts.ts create mode 100644 assets/functions/tools/get_current_weather.ts diff --git a/assets/functions/tools/demo_ts.ts b/assets/functions/tools/demo_ts.ts new file mode 100644 index 0000000..26dc8d1 --- /dev/null +++ b/assets/functions/tools/demo_ts.ts @@ -0,0 +1,53 @@ +/** + * Demonstrates all supported TypeScript parameter types and variations. + * + * @param string - A required string property + * @param string_enum - A required string property constrained to specific values + * @param boolean - A required boolean property + * @param number - A required number property + * @param array_bracket - A required string array using bracket syntax + * @param array_generic - A required string array using generic syntax + * @param string_optional - An optional string using the question mark syntax + * @param string_nullable - An optional string using the union-with-null syntax + * @param number_with_default - An optional number with a default value + * @param boolean_with_default - An optional boolean with a default value + * @param string_with_default - An optional string with a default value + * @param array_optional - An optional string array using the question mark syntax + */ +export function run( + string: string, + string_enum: "foo" | "bar", + boolean: boolean, + number: number, + array_bracket: string[], + array_generic: Array, + string_optional?: string, + string_nullable: string | null = null, + number_with_default: number = 42, + boolean_with_default: boolean = true, + string_with_default: string = "hello", + array_optional?: string[], +): string { + const parts = [ + `string: ${string}`, + `string_enum: ${string_enum}`, + `boolean: ${boolean}`, + `number: ${number}`, + `array_bracket: ${JSON.stringify(array_bracket)}`, + `array_generic: ${JSON.stringify(array_generic)}`, + `string_optional: ${string_optional}`, + `string_nullable: ${string_nullable}`, + `number_with_default: ${number_with_default}`, + `boolean_with_default: ${boolean_with_default}`, + `string_with_default: ${string_with_default}`, + `array_optional: ${JSON.stringify(array_optional)}`, + ]; + + for (const [key, value] of Object.entries(process.env)) { + if (key.startsWith("LLM_")) { + parts.push(`${key}: ${value}`); + } + } + + return parts.join("\n"); +} diff --git a/assets/functions/tools/get_current_weather.ts b/assets/functions/tools/get_current_weather.ts new file mode 100644 index 0000000..34bbf0d --- /dev/null +++ b/assets/functions/tools/get_current_weather.ts @@ -0,0 +1,24 @@ +#!/usr/bin/env tsx + +import { appendFileSync, mkdirSync } from "fs"; +import { dirname } from "path"; + +/** + * Get the current weather in a given location + * @param location - The city and optionally the state or country (e.g., "London", "San Francisco, CA"). + */ +export async function run(location: string): string { + const encoded = encodeURIComponent(location); + const url = `https://wttr.in/${encoded}?format=4`; + + const resp = await fetch(url); + const data = await resp.text(); + + const dest = process.env["LLM_OUTPUT"] ?? "/dev/stdout"; + if (dest !== "-" && dest !== "/dev/stdout") { + mkdirSync(dirname(dest), { recursive: true }); + appendFileSync(dest, data, "utf-8"); + } + + return data; +} From 78c3932f36c8cc7caf392d4dd0e68ab9c43fa75a Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:19:15 -0600 Subject: [PATCH 07/12] docs: Updated docs to talk about the new TypeScript-based tool support --- README.md | 1 + config.example.yaml | 2 + docs/AGENTS.md | 71 ++++++++-- docs/function-calling/CUSTOM-TOOLS.md | 185 ++++++++++++++++++++++++-- docs/function-calling/TOOLS.md | 2 + 5 files changed, 241 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index df5e59e..c3d9cf5 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g * [Function Calling](./docs/function-calling/TOOLS.md#Tools): Leverage function calling capabilities to extend Loki's functionality with custom tools * [Creating Custom Tools](./docs/function-calling/CUSTOM-TOOLS.md): You can create your own custom tools to enhance Loki's capabilities. * [Create Custom Python Tools](./docs/function-calling/CUSTOM-TOOLS.md#custom-python-based-tools) + * [Create Custom TypeScript Tools](./docs/function-calling/CUSTOM-TOOLS.md#custom-typescript-based-tools) * [Create Custom Bash Tools](./docs/function-calling/CUSTOM-BASH-TOOLS.md) * [Bash Prompt Utilities](./docs/function-calling/BASH-PROMPT-HELPERS.md) * [First-Class MCP Server Support](./docs/function-calling/MCP-SERVERS.md): Easily connect and interact with MCP servers for advanced functionality. diff --git a/config.example.yaml b/config.example.yaml index b626163..b7c1961 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -46,6 +46,7 @@ enabled_tools: null # Which tools to enable by default. (e.g. 'fs,w visible_tools: # Which tools are visible to be compiled (and are thus able to be defined in 'enabled_tools') # - demo_py.py # - demo_sh.sh +# - demo_ts.ts - execute_command.sh # - execute_py_code.py # - execute_sql_code.sh @@ -61,6 +62,7 @@ visible_tools: # Which tools are visible to be compiled (and a # - fs_write.sh - get_current_time.sh # - get_current_weather.py +# - get_current_weather.ts - get_current_weather.sh - query_jira_issues.sh # - search_arxiv.sh diff --git a/docs/AGENTS.md b/docs/AGENTS.md index d9d0aea..e9782d8 100644 --- a/docs/AGENTS.md +++ b/docs/AGENTS.md @@ -33,6 +33,7 @@ If you're looking for more example agents, refer to the [built-in agents](../ass - [.env File Support](#env-file-support) - [Python-Based Agent Tools](#python-based-agent-tools) - [Bash-Based Agent Tools](#bash-based-agent-tools) + - [TypeScript-Based Agent Tools](#typescript-based-agent-tools) - [5. Conversation Starters](#5-conversation-starters) - [6. Todo System & Auto-Continuation](#6-todo-system--auto-continuation) - [7. Sub-Agent Spawning System](#7-sub-agent-spawning-system) @@ -62,10 +63,12 @@ Agent configurations often have the following directory structure: ├── tools.sh or ├── tools.py + or + ├── tools.ts ``` This means that agent configurations often are only two files: the agent configuration file (`config.yaml`), and the -tool definitions (`agents/my-agent/tools.sh` or `tools.py`). +tool definitions (`agents/my-agent/tools.sh`, `tools.py`, or `tools.ts`). To see a full example configuration file, refer to the [example agent config file](../config.agent.example.yaml). @@ -114,10 +117,10 @@ isolated environment, so in order for an agent to use a tool or MCP server that explicitly state which tools and/or MCP servers the agent uses. Otherwise, it is assumed that the agent doesn't use any tools outside its own custom defined tools. -And if you don't define a `agents/my-agent/tools.sh` or `agents/my-agent/tools.py`, then the agent is really just a +And if you don't define a `agents/my-agent/tools.sh`, `agents/my-agent/tools.py`, or `agents/my-agent/tools.ts`, then the agent is really just a `role`. -You'll notice there's no settings for agent-specific tooling. This is because they are handled separately and +You'll notice there are no settings for agent-specific tooling. This is because they are handled separately and automatically. See the [Building Tools for Agents](#4-building-tools-for-agents) section below for more information. To see a full example configuration file, refer to the [example agent config file](../config.agent.example.yaml). @@ -205,7 +208,7 @@ variables: ### Dynamic Instructions Sometimes you may find it useful to dynamically generate instructions on startup. Whether that be via a call to Loki itself to generate them, or by some other means. Loki supports this type of behavior using a special function defined -in your `agents/my-agent/tools.py` or `agents/my-agent/tools.sh`. +in your `agents/my-agent/tools.py`, `agents/my-agent/tools.sh`, or `agents/my-agent/tools.ts`. **Example: Instructions for a JSON-reader agent that specializes on each JSON input it receives** `agents/json-reader/tools.py`: @@ -306,8 +309,8 @@ EOF } ``` -For more information on how to create custom tools for your agent and the structure of the `agent/my-agent/tools.sh` or -`agent/my-agent/tools.py` files, refer to the [Building Tools for Agents](#4-building-tools-for-agents) section below. +For more information on how to create custom tools for your agent and the structure of the `agent/my-agent/tools.sh`, +`agent/my-agent/tools.py`, or `agent/my-agent/tools.ts` files, refer to the [Building Tools for Agents](#4-building-tools-for-agents) section below. #### Variables All the same variable interpolations supported by static instructions is also supported by dynamic instructions. For @@ -337,10 +340,11 @@ defining a single function that gets executed at runtime (e.g. `main` for bash t tools define a number of *subcommands*. ### Limitations -You can only utilize either a bash-based `/agents/my-agent/tools.sh` or a Python-based -`/agents/my-agent/tools.py`. However, if it's easier to achieve a task in one language vs the other, +You can only utilize one of: a bash-based `/agents/my-agent/tools.sh`, a Python-based +`/agents/my-agent/tools.py`, or a TypeScript-based `/agents/my-agent/tools.ts`. +However, if it's easier to achieve a task in one language vs the other, you're free to define other scripts in your agent's configuration directory and reference them from the main -`tools.py/sh` file. **Any scripts *not* named `tools.{py,sh}` will not be picked up by Loki's compiler**, meaning they +tools file. **Any scripts *not* named `tools.{py,sh,ts}` will not be picked up by Loki's compiler**, meaning they can be used like any other set of scripts. It's important to keep in mind the following: @@ -428,6 +432,55 @@ the same syntax ad formatting as is used to create custom bash tools globally. For more information on how to write, [build and test](function-calling/CUSTOM-BASH-TOOLS.md#execute-and-test-your-bash-tools) tools in bash, refer to the [custom bash tools documentation](function-calling/CUSTOM-BASH-TOOLS.md). +### TypeScript-Based Agent Tools +TypeScript-based agent tools work exactly the same as TypeScript global tools. Instead of a single `run` function, +you define as many exported functions as you like. Non-exported functions are private helpers and are invisible to the +LLM. + +**Example:** +`agents/my-agent/tools.ts` +```typescript +/** + * Get your IP information + */ +export async function get_ip_info(): Promise { + const resp = await fetch("https://httpbin.org/ip"); + return await resp.text(); +} + +/** + * Find your public IP address using AWS + */ +export async function get_ip_address_from_aws(): Promise { + const resp = await fetch("https://checkip.amazonaws.com"); + return await resp.text(); +} + +// Non-exported helper — invisible to the LLM +function formatResponse(data: string): string { + return data.trim(); +} +``` + +Loki automatically compiles each exported function as a separate tool for the LLM to call. Just make sure you +follow the same JSDoc and parameter conventions as you would when creating custom TypeScript tools. + +TypeScript agent tools also support dynamic instructions via an exported `_instructions()` function: + +```typescript +import { readFileSync } from "fs"; + +/** + * Generates instructions for the agent dynamically + */ +export function _instructions(): string { + const schema = readFileSync("schema.json", "utf-8"); + return `You are an AI agent that works with the following schema:\n${schema}`; +} +``` + +For more information on how to build tools in TypeScript, refer to the [custom TypeScript tools documentation](function-calling/CUSTOM-TOOLS.md#custom-typescript-based-tools). + ## 5. Conversation Starters It's often helpful to also have some conversation starters so users know what kinds of things the agent is capable of doing. These are available in the REPL via the `.starter` command and are selectable. diff --git a/docs/function-calling/CUSTOM-TOOLS.md b/docs/function-calling/CUSTOM-TOOLS.md index 1cabcad..1de68cc 100644 --- a/docs/function-calling/CUSTOM-TOOLS.md +++ b/docs/function-calling/CUSTOM-TOOLS.md @@ -10,6 +10,8 @@ into your Loki setup. This document provides a guide on how to create and use cu - [Environment Variables](#environment-variables) - [Custom Bash-Based Tools](#custom-bash-based-tools) - [Custom Python-Based Tools](#custom-python-based-tools) + - [Custom TypeScript-Based Tools](#custom-typescript-based-tools) +- [Custom Runtime](#custom-runtime) --- @@ -19,9 +21,10 @@ Loki supports custom tools written in the following programming languages: * Python * Bash +* TypeScript ## Creating a Custom Tool -All tools are created as scripts in either Python or Bash. They should be placed in the `functions/tools` directory. +All tools are created as scripts in either Python, Bash, or TypeScript. They should be placed in the `functions/tools` directory. The location of the `functions` directory varies between systems, so you can use the following command to locate your `functions` directory: @@ -81,6 +84,7 @@ Loki and demonstrates how to create a Python-based tool: import os from typing import List, Literal, Optional + def run( string: str, string_enum: Literal["foo", "bar"], @@ -89,26 +93,38 @@ def run( number: float, array: List[str], string_optional: Optional[str] = None, + integer_with_default: int = 42, + boolean_with_default: bool = True, + number_with_default: float = 3.14, + string_with_default: str = "hello", array_optional: Optional[List[str]] = None, ): - """Demonstrates how to create a tool using Python and how to use comments. + """Demonstrates all supported Python parameter types and variations. Args: - string: Define a required string property - string_enum: Define a required string property with enum - boolean: Define a required boolean property - integer: Define a required integer property - number: Define a required number property - array: Define a required string array property - string_optional: Define an optional string property - array_optional: Define an optional string array property + string: A required string property + string_enum: A required string property constrained to specific values + boolean: A required boolean property + integer: A required integer property + number: A required number (float) property + array: A required string array property + string_optional: An optional string property (Optional[str] with None default) + integer_with_default: An optional integer with a non-None default value + boolean_with_default: An optional boolean with a default value + number_with_default: An optional number with a default value + string_with_default: An optional string with a default value + array_optional: An optional string array property """ output = f"""string: {string} string_enum: {string_enum} -string_optional: {string_optional} boolean: {boolean} integer: {integer} number: {number} array: {array} +string_optional: {string_optional} +integer_with_default: {integer_with_default} +boolean_with_default: {boolean_with_default} +number_with_default: {number_with_default} +string_with_default: {string_with_default} array_optional: {array_optional}""" for key, value in os.environ.items(): @@ -117,3 +133,150 @@ array_optional: {array_optional}""" return output ``` + +### Custom TypeScript-Based Tools +Loki supports tools written in TypeScript. TypeScript tools require [Node.js](https://nodejs.org/) and +[tsx](https://tsx.is/) (`npx tsx` is used as the default runtime). + +Each TypeScript-based tool must follow a specific structure in order for Loki to properly compile and execute it: + +* The tool must be a TypeScript file with a `.ts` file extension. +* The tool must have an `export function run(...)` that serves as the entry point for the tool. + * Non-exported functions are ignored by the compiler and can be used as private helpers. +* The `run` function must accept flat parameters that define the inputs for the tool. + * Always use type annotations to specify the data type of each parameter. + * Use `param?: type` or `type | null` to indicate optional parameters. + * Use `param: type = value` for parameters with default values. +* The `run` function must return a `string` (or `Promise` for async functions). + * For TypeScript, the return value is automatically written to the `LLM_OUTPUT` environment variable, so there's + no need to explicitly write to the environment variable within the function. +* The function must have a JSDoc comment that describes the tool and its parameters. + * Each parameter should be documented using `@param name - description` tags. + * These descriptions are passed to the LLM as the tool description, letting the LLM know what the tool does and + how to use it. +* Async functions (`export async function run(...)`) are fully supported and handled transparently. + +**Supported Parameter Types:** + +| TypeScript Type | JSON Schema | Notes | +|-------------------|--------------------------------------------------|-----------------------------| +| `string` | `{"type": "string"}` | Required string | +| `number` | `{"type": "number"}` | Required number | +| `boolean` | `{"type": "boolean"}` | Required boolean | +| `string[]` | `{"type": "array", "items": {"type": "string"}}` | Array (bracket syntax) | +| `Array` | `{"type": "array", "items": {"type": "string"}}` | Array (generic syntax) | +| `"foo" \| "bar"` | `{"type": "string", "enum": ["foo", "bar"]}` | String enum (literal union) | +| `param?: string` | `{"type": "string"}` (not required) | Optional via question mark | +| `string \| null` | `{"type": "string"}` (not required) | Optional via null union | +| `param = "value"` | `{"type": "string"}` (not required) | Optional via default value | + +**Unsupported Patterns (will produce a compile error):** + +* Rest parameters (`...args: string[]`) +* Destructured object parameters (`{ a, b }: { a: string, b: string }`) +* Arrow functions (`const run = (x: string) => ...`) +* Function expressions (`const run = function(x: string) { ... }`) + +Only `export function` declarations are recognized. Non-exported functions are invisible to the compiler. + +Below is the [`demo_ts.ts`](../../assets/functions/tools/demo_ts.ts) tool definition that comes pre-packaged with +Loki and demonstrates how to create a TypeScript-based tool: + +```typescript +/** + * Demonstrates all supported TypeScript parameter types and variations. + * + * @param string - A required string property + * @param string_enum - A required string property constrained to specific values + * @param boolean - A required boolean property + * @param number - A required number property + * @param array_bracket - A required string array using bracket syntax + * @param array_generic - A required string array using generic syntax + * @param string_optional - An optional string using the question mark syntax + * @param string_nullable - An optional string using the union-with-null syntax + * @param number_with_default - An optional number with a default value + * @param boolean_with_default - An optional boolean with a default value + * @param string_with_default - An optional string with a default value + * @param array_optional - An optional string array using the question mark syntax + */ +export function run( + string: string, + string_enum: "foo" | "bar", + boolean: boolean, + number: number, + array_bracket: string[], + array_generic: Array, + string_optional?: string, + string_nullable: string | null = null, + number_with_default: number = 42, + boolean_with_default: boolean = true, + string_with_default: string = "hello", + array_optional?: string[], +): string { + const parts = [ + `string: ${string}`, + `string_enum: ${string_enum}`, + `boolean: ${boolean}`, + `number: ${number}`, + `array_bracket: ${JSON.stringify(array_bracket)}`, + `array_generic: ${JSON.stringify(array_generic)}`, + `string_optional: ${string_optional}`, + `string_nullable: ${string_nullable}`, + `number_with_default: ${number_with_default}`, + `boolean_with_default: ${boolean_with_default}`, + `string_with_default: ${string_with_default}`, + `array_optional: ${JSON.stringify(array_optional)}`, + ]; + + for (const [key, value] of Object.entries(process.env)) { + if (key.startsWith("LLM_")) { + parts.push(`${key}: ${value}`); + } + } + + return parts.join("\n"); +} +``` + +## Custom Runtime +By default, Loki uses the following runtimes to execute tools: + +| Language | Default Runtime | Requirement | +|------------|-----------------|--------------------------------| +| Python | `python` | Python 3 on `$PATH` | +| TypeScript | `npx tsx` | Node.js + tsx (`npm i -g tsx`) | +| Bash | `bash` | Bash on `$PATH` | + +You can override the runtime for Python and TypeScript tools using a **shebang line** (`#!`) at the top of your +script. Loki reads the first line of each tool file; if it starts with `#!`, the specified interpreter is used instead +of the default. + +**Examples:** + +```python +#!/usr/bin/env python3.11 +# This Python tool will be executed with python3.11 instead of the default `python` + +def run(name: str): + """Greet someone. + Args: + name: The name to greet + """ + return f"Hello, {name}!" +``` + +```typescript +#!/usr/bin/env bun +// This TypeScript tool will be executed with Bun instead of the default `npx tsx` + +/** + * Greet someone. + * @param name - The name to greet + */ +export function run(name: string): string { + return `Hello, ${name}!`; +} +``` + +This is useful for pinning a specific Python version, using an alternative TypeScript runtime like +[Bun](https://bun.sh/) or [Deno](https://deno.com/), or working with virtual environments. diff --git a/docs/function-calling/TOOLS.md b/docs/function-calling/TOOLS.md index 7aec07b..80bf392 100644 --- a/docs/function-calling/TOOLS.md +++ b/docs/function-calling/TOOLS.md @@ -32,6 +32,7 @@ be enabled/disabled can be found in the [Configuration](#configuration) section |-------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------| | [`demo_py.py`](../../assets/functions/tools/demo_py.py) | Demonstrates how to create a tool using Python and how to use comments. | 🔴 | | [`demo_sh.sh`](../../assets/functions/tools/demo_sh.sh) | Demonstrate how to create a tool using Bash and how to use comment tags. | 🔴 | +| [`demo_ts.ts`](../../assets/functions/tools/demo_ts.ts) | Demonstrates how to create a tool using TypeScript and how to use JSDoc comments. | 🔴 | | [`execute_command.sh`](../../assets/functions/tools/execute_command.sh) | Execute the shell command. | 🟢 | | [`execute_py_code.py`](../../assets/functions/tools/execute_py_code.py) | Execute the given Python code. | 🔴 | | [`execute_sql_code.sh`](../../assets/functions/tools/execute_sql_code.sh) | Execute SQL code. | 🔴 | @@ -49,6 +50,7 @@ be enabled/disabled can be found in the [Configuration](#configuration) section | [`get_current_time.sh`](../../assets/functions/tools/get_current_time.sh) | Get the current time. | 🟢 | | [`get_current_weather.py`](../../assets/functions/tools/get_current_weather.py) | Get the current weather in a given location (Python implementation) | 🔴 | | [`get_current_weather.sh`](../../assets/functions/tools/get_current_weather.sh) | Get the current weather in a given location. | 🟢 | +| [`get_current_weather.ts`](../../assets/functions/tools/get_current_weather.ts) | Get the current weather in a given location (TypeScript implementation) | 🔴 | | [`query_jira_issues.sh`](../../assets/functions/tools/query_jira_issues.sh) | Query for jira issues using a Jira Query Language (JQL) query. | 🟢 | | [`search_arxiv.sh`](../../assets/functions/tools/search_arxiv.sh) | Search arXiv using the given search query and return the top papers. | 🔴 | | [`search_wikipedia.sh`](../../assets/functions/tools/search_wikipedia.sh) | Search Wikipedia using the given search query.
Use it to get detailed information about a public figure, interpretation of a
complex scientific concept or in-depth connectivity of a significant historical
event, etc. | 🔴 | From 7839e1dbd9063aab86daddeaece81289b4627ee3 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:28:19 -0600 Subject: [PATCH 08/12] build: upgraded dependencies to latest --- Cargo.lock | 93 ++++++++++++++++++++++++++++++++++++++---------------- Cargo.toml | 10 +++--- 2 files changed, 72 insertions(+), 31 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 162b10d..62acbb4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -205,7 +205,7 @@ checksum = "3c3610892ee6e0cbce8ae2700349fcf8f98adb0dbfbee85aec3c9179d29cc072" dependencies = [ "base64ct", "blake2", - "cpufeatures", + "cpufeatures 0.2.17", "password-hash", ] @@ -314,29 +314,17 @@ version = "1.16.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a054912289d18629dc78375ba2c3726a3afe3ff71b4edba9dedfca0e3446d1fc" dependencies = [ - "aws-lc-sys 0.39.1", + "aws-lc-sys", "zeroize", ] -[[package]] -name = "aws-lc-sys" -version = "0.37.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b092fe214090261288111db7a2b2c2118e5a7f30dc2569f1732c4069a6840549" -dependencies = [ - "bindgen", - "cc", - "cmake", - "dunce", - "fs_extra", -] - [[package]] name = "aws-lc-sys" version = "0.39.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83a25cf98105baa966497416dbd42565ce3a8cf8dbfd59803ec9ad46f3126399" dependencies = [ + "bindgen", "cc", "cmake", "dunce", @@ -1038,7 +1026,18 @@ checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" dependencies = [ "cfg-if", "cipher", - "cpufeatures", + "cpufeatures 0.2.17", +] + +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.0", ] [[package]] @@ -1048,7 +1047,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" dependencies = [ "aead", - "chacha20", + "chacha20 0.9.1", "cipher", "poly1305", "zeroize", @@ -1300,6 +1299,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32c" version = "0.6.8" @@ -2114,6 +2122,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.0", "wasip2", "wasip3", ] @@ -2132,15 +2141,15 @@ checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "gman" -version = "0.3.0" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7c3a428900217107275faf709b30c00f37e1112ec2b75742987b5ca88700eaa" +checksum = "742225eb41061a0938aa0924ce8d08a1ec48875789b72ce3f0cb02eda52ab1db" dependencies = [ "anyhow", "argon2", "async-trait", "aws-config", - "aws-lc-sys 0.37.1", + "aws-lc-sys", "aws-sdk-secretsmanager", "azure_core", "azure_identity", @@ -3086,7 +3095,7 @@ dependencies = [ "parking_lot", "path-absolutize", "pretty_assertions", - "rand 0.9.2", + "rand 0.10.0", "rayon", "reedline", "reqwest", @@ -3109,7 +3118,9 @@ dependencies = [ "tokio-graceful", "tokio-stream", "tree-sitter", + "tree-sitter-language", "tree-sitter-python", + "tree-sitter-typescript", "unicode-segmentation", "unicode-width", "url", @@ -3975,7 +3986,7 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" dependencies = [ - "cpufeatures", + "cpufeatures 0.2.17", "opaque-debug", "universal-hash", ] @@ -4235,6 +4246,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" +dependencies = [ + "chacha20 0.10.0", + "getrandom 0.4.2", + "rand_core 0.10.0", +] + [[package]] name = "rand_chacha" version = "0.9.0" @@ -4263,6 +4285,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" + [[package]] name = "rayon" version = "1.11.0" @@ -5010,7 +5038,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -5829,13 +5857,14 @@ dependencies = [ [[package]] name = "tree-sitter" -version = "0.24.7" +version = "0.26.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5387dffa7ffc7d2dae12b50c6f7aab8ff79d6210147c6613561fc3d474c6f75" +checksum = "887bd495d0582c5e3e0d8ece2233666169fa56a9644d172fc22ad179ab2d0538" dependencies = [ "cc", "regex", "regex-syntax", + "serde_json", "streaming-iterator", "tree-sitter-language", ] @@ -5848,9 +5877,19 @@ checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" [[package]] name = "tree-sitter-python" -version = "0.23.6" +version = "0.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d065aaa27f3aaceaf60c1f0e0ac09e1cb9eb8ed28e7bcdaa52129cffc7f4b04" +checksum = "6bf85fd39652e740bf60f46f4cda9492c3a9ad75880575bf14960f775cb74a1c" +dependencies = [ + "cc", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-typescript" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c5f76ed8d947a75cc446d5fccd8b602ebf0cde64ccf2ffa434d873d7a575eff" dependencies = [ "cc", "tree-sitter-language", diff --git a/Cargo.toml b/Cargo.toml index 8b2177b..29dea38 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -91,14 +91,16 @@ strum_macros = "0.27.2" indoc = "2.0.6" rmcp = { version = "0.16.0", features = ["client", "transport-child-process"] } num_cpus = "1.17.0" -tree-sitter = "0.24" -tree-sitter-python = "0.23" +tree-sitter = "0.26.8" +tree-sitter-language = "0.1" +tree-sitter-python = "0.25.0" +tree-sitter-typescript = "0.23" colored = "3.0.0" clap_complete = { version = "4.5.58", features = ["unstable-dynamic"] } -gman = "0.3.0" +gman = "0.4.1" clap_complete_nushell = "4.5.9" open = "5" -rand = "0.9.0" +rand = "0.10.0" url = "2.5.8" [dependencies.reqwest] From 6b4a45874f9cd4fa8cc393ba9cebbc37f1ee7a29 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:32:16 -0600 Subject: [PATCH 09/12] fix: TypeScript function args were being passed as objects rather than direct parameters --- assets/functions/scripts/run-agent.ts | 27 ++++++++++++++++++++++++++- assets/functions/scripts/run-tool.ts | 27 ++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/assets/functions/scripts/run-agent.ts b/assets/functions/scripts/run-agent.ts index 425b5df..a8e02a3 100644 --- a/assets/functions/scripts/run-agent.ts +++ b/assets/functions/scripts/run-agent.ts @@ -92,6 +92,29 @@ function loadEnv(filePath: string): void { } } +function extractParamNames(fn: Function): string[] { + const src = fn.toString(); + const match = src.match(/^(?:async\s+)?function\s*\w*\s*\(([^)]*)\)/); + if (!match) { + return []; + } + return match[1] + .split(",") + .map((p) => p.trim().replace(/[:=?].*/s, "").trim()) + .filter(Boolean); +} + +function spreadArgs( + fn: Function, + data: Record, +): unknown[] { + const names = extractParamNames(fn); + if (names.length === 0) { + return []; + } + return names.map((name) => data[name]); +} + async function run( agentPath: string, agentFunc: string, @@ -103,7 +126,9 @@ async function run( throw new Error(`No module function '${agentFunc}' at '${agentPath}'`); } - const value = await mod[agentFunc](agentData); + const fn = mod[agentFunc] as Function; + const args = spreadArgs(fn, agentData); + const value = await fn(...args); returnToLlm(value); dumpResult(`{agent_name}:${agentFunc}`); } diff --git a/assets/functions/scripts/run-tool.ts b/assets/functions/scripts/run-tool.ts index f3d2887..5ffd32b 100644 --- a/assets/functions/scripts/run-tool.ts +++ b/assets/functions/scripts/run-tool.ts @@ -87,6 +87,29 @@ function loadEnv(filePath: string): void { } } +function extractParamNames(fn: Function): string[] { + const src = fn.toString(); + const match = src.match(/^(?:async\s+)?function\s*\w*\s*\(([^)]*)\)/); + if (!match) { + return []; + } + return match[1] + .split(",") + .map((p) => p.trim().replace(/[:=?].*/s, "").trim()) + .filter(Boolean); +} + +function spreadArgs( + fn: Function, + data: Record, +): unknown[] { + const names = extractParamNames(fn); + if (names.length === 0) { + return []; + } + return names.map((name) => data[name]); +} + async function run( toolPath: string, toolFunc: string, @@ -98,7 +121,9 @@ async function run( throw new Error(`No module function '${toolFunc}' at '${toolPath}'`); } - const value = await mod[toolFunc](toolData); + const fn = mod[toolFunc] as Function; + const args = spreadArgs(fn, toolData); + const value = await fn(...args); returnToLlm(value); dumpResult("{function_name}"); } From d658f1d2fe477e042d23d76280e017907ea7beec Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:45:08 -0600 Subject: [PATCH 10/12] build: Pulled additional features for rand dependency --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 29dea38..5d15c52 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -100,7 +100,7 @@ clap_complete = { version = "4.5.58", features = ["unstable-dynamic"] } gman = "0.4.1" clap_complete_nushell = "4.5.9" open = "5" -rand = "0.10.0" +rand = { version = "0.10.0", features = ["default"] } url = "2.5.8" [dependencies.reqwest] From dedcef8ac53330d032646f69948ef6109c59ca60 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 13:53:52 -0600 Subject: [PATCH 11/12] test: Updated client stream tests to use the thread_rng from rand --- src/client/stream.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/client/stream.rs b/src/client/stream.rs index 3780d93..d2fe7f7 100644 --- a/src/client/stream.rs +++ b/src/client/stream.rs @@ -342,7 +342,7 @@ mod tests { use bytes::Bytes; use futures_util::stream; - use rand::Rng; + use rand::random_range; use serde_json::json; #[test] @@ -392,10 +392,9 @@ mod tests { } fn split_chunks(text: &str) -> Vec> { - let mut rng = rand::rng(); let len = text.len(); - let cut1 = rng.random_range(1..len - 1); - let cut2 = rng.random_range(cut1 + 1..len); + let cut1 = random_range(1..len - 1); + let cut2 = random_range(cut1 + 1..len); let chunk1 = text.as_bytes()[..cut1].to_vec(); let chunk2 = text.as_bytes()[cut1..cut2].to_vec(); let chunk3 = text.as_bytes()[cut2..].to_vec(); From a5899da4fbc0eb95022c3f566d0d11ec14303530 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 9 Apr 2026 14:16:02 -0600 Subject: [PATCH 12/12] feat: Automatic runtime customization using shebangs --- assets/functions/scripts/run-tool.ts | 4 +- src/function/mod.rs | 117 ++++++++++++++++++--------- src/parsers/common.rs | 35 +------- src/parsers/python.rs | 4 - src/parsers/typescript.rs | 4 - 5 files changed, 84 insertions(+), 80 deletions(-) diff --git a/assets/functions/scripts/run-tool.ts b/assets/functions/scripts/run-tool.ts index 5ffd32b..228166a 100644 --- a/assets/functions/scripts/run-tool.ts +++ b/assets/functions/scripts/run-tool.ts @@ -2,8 +2,8 @@ // Usage: ./{function_name}.ts -import { readFileSync, writeFileSync, existsSync, statSync } from "fs"; -import { join, basename } from "path"; +import { readFileSync, writeFileSync, existsSync } from "fs"; +import { join } from "path"; import { pathToFileURL } from "url"; async function main(): Promise { diff --git a/src/function/mod.rs b/src/function/mod.rs index ca777c9..96ffdb4 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -89,6 +89,26 @@ impl Language { } } +fn extract_shebang_runtime(path: &Path) -> Option { + let file = File::open(path).ok()?; + let reader = io::BufReader::new(file); + let first_line = io::BufRead::lines(reader).next()?.ok()?; + let shebang = first_line.strip_prefix("#!")?; + let cmd = shebang.trim(); + if cmd.is_empty() { + return None; + } + if let Some(after_env) = cmd.strip_prefix("/usr/bin/env ") { + let runtime = after_env.trim(); + if runtime.is_empty() { + return None; + } + Some(runtime.to_string()) + } else { + Some(cmd.to_string()) + } +} + pub async fn eval_tool_calls( config: &GlobalConfig, mut calls: Vec, @@ -522,7 +542,14 @@ impl Functions { bail!("Unsupported tool file extension: {}", language.as_ref()); } - Self::build_binaries(binary_name, language, BinaryType::Tool(agent_name))?; + let tool_path = Config::global_tools_dir().join(tool); + let custom_runtime = extract_shebang_runtime(&tool_path); + Self::build_binaries( + binary_name, + language, + BinaryType::Tool(agent_name), + custom_runtime.as_deref(), + )?; } Ok(()) @@ -563,8 +590,9 @@ impl Functions { } fn build_agent_tool_binaries(name: &str) -> Result<()> { + let tools_file = Config::agent_functions_file(name)?; let language = Language::from( - &Config::agent_functions_file(name)? + &tools_file .extension() .and_then(OsStr::to_str) .map(|s| s.to_lowercase()) @@ -577,7 +605,8 @@ impl Functions { bail!("Unsupported tool file extension: {}", language.as_ref()); } - Self::build_binaries(name, language, BinaryType::Agent) + let custom_runtime = extract_shebang_runtime(&tools_file); + Self::build_binaries(name, language, BinaryType::Agent, custom_runtime.as_deref()) } #[cfg(windows)] @@ -585,6 +614,7 @@ impl Functions { binary_name: &str, language: Language, binary_type: BinaryType, + custom_runtime: Option<&str>, ) -> Result<()> { use native::runtime; let (binary_file, binary_script_file) = match binary_type { @@ -669,38 +699,42 @@ impl Functions { binary_file.display() ); - let run = match language { - Language::Bash => { - let shell = runtime::bash_path().ok_or_else(|| anyhow!("Shell not found"))?; - format!("{shell} --noprofile --norc") + let run = if let Some(rt) = custom_runtime { + rt.to_string() + } else { + match language { + Language::Bash => { + let shell = runtime::bash_path().ok_or_else(|| anyhow!("Shell not found"))?; + format!("{shell} --noprofile --norc") + } + Language::Python if Path::new(".venv").exists() => { + let executable_path = env::current_dir()? + .join(".venv") + .join("Scripts") + .join("activate.bat"); + let canonicalized_path = dunce::canonicalize(&executable_path)?; + format!( + "call \"{}\" && {}", + canonicalized_path.to_string_lossy(), + language.to_cmd() + ) + } + Language::Python => { + let executable_path = which::which("python") + .or_else(|_| which::which("python3")) + .map_err(|_| anyhow!("Python executable not found in PATH"))?; + let canonicalized_path = dunce::canonicalize(&executable_path)?; + canonicalized_path.to_string_lossy().into_owned() + } + Language::TypeScript => { + let npx_path = which::which("npx").map_err(|_| { + anyhow!("npx executable not found in PATH (required for TypeScript tools)") + })?; + let canonicalized_path = dunce::canonicalize(&npx_path)?; + format!("{} tsx", canonicalized_path.to_string_lossy()) + } + _ => bail!("Unsupported language: {}", language.as_ref()), } - Language::Python if Path::new(".venv").exists() => { - let executable_path = env::current_dir()? - .join(".venv") - .join("Scripts") - .join("activate.bat"); - let canonicalized_path = dunce::canonicalize(&executable_path)?; - format!( - "call \"{}\" && {}", - canonicalized_path.to_string_lossy(), - language.to_cmd() - ) - } - Language::Python => { - let executable_path = which::which("python") - .or_else(|_| which::which("python3")) - .map_err(|_| anyhow!("Python executable not found in PATH"))?; - let canonicalized_path = dunce::canonicalize(&executable_path)?; - canonicalized_path.to_string_lossy().into_owned() - } - Language::TypeScript => { - let npx_path = which::which("npx").map_err(|_| { - anyhow!("npx executable not found in PATH (required for TypeScript tools)") - })?; - let canonicalized_path = dunce::canonicalize(&npx_path)?; - format!("{} tsx", canonicalized_path.to_string_lossy()) - } - _ => bail!("Unsupported language: {}", language.as_ref()), }; let bin_dir = binary_file .parent() @@ -730,6 +764,7 @@ impl Functions { binary_name: &str, language: Language, binary_type: BinaryType, + custom_runtime: Option<&str>, ) -> Result<()> { use std::os::unix::prelude::PermissionsExt; @@ -758,7 +793,7 @@ impl Functions { ) })?; let content_template = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) }; - let content = match binary_type { + let mut content = match binary_type { BinaryType::Tool(None) => { let root_dir = Config::functions_dir(); let tool_path = format!( @@ -790,6 +825,12 @@ impl Functions { &Config::bash_prompt_utils_file().to_string_lossy(), ); + if let Some(rt) = custom_runtime + && let Some(newline_pos) = content.find('\n') + { + content = format!("#!/usr/bin/env {rt}{}", &content[newline_pos..]); + } + if language == Language::TypeScript { let bin_dir = binary_file .parent() @@ -802,7 +843,11 @@ impl Functions { sf.write_all(content.as_bytes())?; fs::set_permissions(&script_file, fs::Permissions::from_mode(0o755))?; - let wrapper = format!("#!/bin/sh\nexec tsx \"{}\" \"$@\"\n", script_file.display()); + let ts_runtime = custom_runtime.unwrap_or("tsx"); + let wrapper = format!( + "#!/bin/sh\nexec {ts_runtime} \"{}\" \"$@\"\n", + script_file.display() + ); if binary_file.exists() { fs::remove_file(&binary_file)?; } diff --git a/src/parsers/common.rs b/src/parsers/common.rs index b0d80d6..58e4a9a 100644 --- a/src/parsers/common.rs +++ b/src/parsers/common.rs @@ -17,15 +17,9 @@ pub(crate) struct Param { pub(crate) trait ScriptedLanguage { fn ts_language(&self) -> tree_sitter::Language; - fn default_runtime(&self) -> &str; - fn lang_name(&self) -> &str; - fn find_functions<'a>( - &self, - root: Node<'a>, - src: &str, - ) -> Vec<(Node<'a>, Node<'a>)>; + fn find_functions<'a>(&self, root: Node<'a>, src: &str) -> Vec<(Node<'a>, Node<'a>)>; fn function_name<'a>(&self, func_node: Node<'a>, src: &'a str) -> Result<&'a str>; @@ -175,31 +169,6 @@ pub(crate) fn named_child(node: Node<'_>, index: usize) -> Option> { node.named_children(&mut cursor).nth(index) } -pub(crate) fn extract_runtime(tree: &tree_sitter::Tree, src: &str, default: &str) -> String { - let root = tree.root_node(); - let mut cursor = root.walk(); - for child in root.named_children(&mut cursor) { - let text = match child.kind() { - "hash_bang_line" | "comment" => match child.utf8_text(src.as_bytes()) { - Ok(t) => t, - Err(_) => continue, - }, - _ => break, - }; - - if let Some(cmd) = text.strip_prefix("#!") { - let cmd = cmd.trim(); - if let Some(after_env) = cmd.strip_prefix("/usr/bin/env ") { - return after_env.trim().to_string(); - } - return cmd.to_string(); - } - - break; - } - default.to_string() -} - pub(crate) fn generate_declarations( lang: &L, src: &str, @@ -226,8 +195,6 @@ pub(crate) fn generate_declarations( ); } - let _runtime = extract_runtime(&tree, src, lang.default_runtime()); - let mut out = Vec::new(); for (wrapper, func) in lang.find_functions(tree.root_node(), src) { let func_name = lang.function_name(func, src)?; diff --git a/src/parsers/python.rs b/src/parsers/python.rs index d8eb169..de9a089 100644 --- a/src/parsers/python.rs +++ b/src/parsers/python.rs @@ -15,10 +15,6 @@ impl ScriptedLanguage for PythonLanguage { tree_sitter_python::LANGUAGE.into() } - fn default_runtime(&self) -> &str { - "python" - } - fn lang_name(&self) -> &str { "python" } diff --git a/src/parsers/typescript.rs b/src/parsers/typescript.rs index abaf878..4fbea59 100644 --- a/src/parsers/typescript.rs +++ b/src/parsers/typescript.rs @@ -15,10 +15,6 @@ impl ScriptedLanguage for TypeScriptLanguage { tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into() } - fn default_runtime(&self) -> &str { - "npx tsx" - } - fn lang_name(&self) -> &str { "typescript" }