Files
tkmind_go/crates/goose-local-inference/src/tool_emulation.rs
T

556 lines
20 KiB
Rust

//! Shared text-based tool call emulation for local inference backends.
//!
//! Models that do not have native tool-calling support are prompted to emit shell commands
//! as `$ command` on a new line and code blocks as ```execute_typescript fenced blocks.
//! The parser converts those patterns into Goose tool-call messages.
use goose_provider_types::conversation::message::{Message, MessageContent};
use rmcp::model::{CallToolRequestParams, Tool};
use serde_json::json;
use std::borrow::Cow;
use uuid::Uuid;
pub(crate) const SHELL_TOOL: &str = "developer__shell";
pub(crate) const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript";
const HOLD_BACK_CODE_MODE: usize = " ```execute_typescript\n".len();
const HOLD_BACK_SHELL_ONLY: usize = "\n$".len();
pub(crate) fn load_tiny_model_prompt() -> String {
use std::env;
let os = if cfg!(target_os = "macos") {
"macos"
} else if cfg!(target_os = "linux") {
"linux"
} else if cfg!(target_os = "windows") {
"windows"
} else {
"unknown"
};
let working_directory = env::current_dir()
.map(|p| p.display().to_string())
.unwrap_or_else(|_| "unknown".to_string());
let shell = env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string());
let context = json!({
"os": os,
"working_directory": working_directory,
"shell": shell,
});
crate::prompt_template::render_template("tiny_model_system.md", &context).unwrap_or_else(|e| {
tracing::warn!("Failed to load tiny_model_system.md: {:?}", e);
"You are Goose, an AI assistant. You can execute shell commands by starting lines with $."
.to_string()
})
}
pub(crate) fn build_emulator_tool_description(tools: &[Tool], code_mode_enabled: bool) -> String {
let mut tool_desc = String::new();
if code_mode_enabled {
tool_desc.push_str("\n\n# Running Code\n\n");
tool_desc.push_str(
"You can call tools by writing code in a ```execute_typescript block. \
The code runs immediately — do not explain it, just run it.\n\n",
);
tool_desc.push_str("Example — counting files in /tmp:\n\n");
tool_desc.push_str("```execute_typescript\nasync function run() {\n");
tool_desc.push_str(
" const result = await Developer.shell({ command: \"ls -1 /tmp | wc -l\" });\n",
);
tool_desc.push_str(" return result;\n}\n```\n\n");
tool_desc.push_str("Rules:\n");
tool_desc.push_str("- Code MUST define async function run() and return a result\n");
tool_desc.push_str("- All function calls are async — use await\n");
tool_desc.push_str(
"- Use ```execute_typescript for tool calls, $ for simple shell one-liners\n\n",
);
tool_desc.push_str("Available functions:\n\n");
for tool in tools {
if tool.name.starts_with("code_execution__") {
continue;
}
let parts: Vec<&str> = tool.name.splitn(2, "__").collect();
if parts.len() == 2 {
let namespace = {
let mut c = parts[0].chars();
match c.next() {
None => String::new(),
Some(first) => first.to_uppercase().chain(c).collect::<String>(),
}
};
let camel_name: String = parts[1]
.split('_')
.enumerate()
.map(|(i, part)| {
if i == 0 {
part.to_string()
} else {
let mut c = part.chars();
match c.next() {
None => String::new(),
Some(first) => first.to_uppercase().chain(c).collect(),
}
}
})
.collect();
let desc = tool.description.as_ref().map(|d| d.as_ref()).unwrap_or("");
tool_desc.push_str(&format!("- {namespace}.{camel_name}(): {desc}\n"));
}
}
} else {
tool_desc.push_str("\n\n# Tools\n\nYou have access to the following tools:\n\n");
for tool in tools {
let desc = tool
.description
.as_ref()
.map(|d| d.as_ref())
.unwrap_or("No description");
tool_desc.push_str(&format!("- {}: {}\n", tool.name, desc));
}
}
tool_desc
}
pub(crate) enum EmulatorAction {
Text(String),
ShellCommand(String),
ExecuteCode(String),
}
enum ParserState {
Normal,
InCommand,
InExecuteBlock,
}
pub(crate) struct StreamingEmulatorParser {
buffer: String,
state: ParserState,
code_mode_enabled: bool,
}
impl StreamingEmulatorParser {
pub(crate) fn new(code_mode_enabled: bool) -> Self {
Self {
buffer: String::new(),
state: ParserState::Normal,
code_mode_enabled,
}
}
pub(crate) fn process_chunk(&mut self, chunk: &str) -> Vec<EmulatorAction> {
self.buffer.push_str(chunk);
let mut results = Vec::new();
loop {
match self.state {
ParserState::InCommand => {
if let Some((command_line, rest)) = self.buffer.split_once('\n') {
if let Some(command) = command_line.strip_prefix('$') {
let command = command.trim();
if !command.is_empty() {
results.push(EmulatorAction::ShellCommand(command.to_string()));
}
}
self.buffer = rest.to_string();
self.state = ParserState::Normal;
} else {
break;
}
}
ParserState::InExecuteBlock => {
if let Some(end_idx) = self.buffer.find("\n```") {
#[allow(clippy::string_slice)]
let code = self.buffer[..end_idx].to_string();
#[allow(clippy::string_slice)]
let rest = &self.buffer[end_idx + 4..];
let rest = rest.strip_prefix('\n').unwrap_or(rest);
self.buffer = rest.to_string();
self.state = ParserState::Normal;
if !code.trim().is_empty() {
results.push(EmulatorAction::ExecuteCode(code));
}
} else {
break;
}
}
ParserState::Normal => {
if self.code_mode_enabled {
if let Some((before, after)) =
self.buffer.split_once("```execute_typescript\n")
{
if !before.trim().is_empty() {
results.push(EmulatorAction::Text(before.to_string()));
}
self.buffer = after.to_string();
self.state = ParserState::InExecuteBlock;
continue;
}
if self.buffer.ends_with("```execute_typescript") {
let before = self.buffer.trim_end_matches("```execute_typescript");
if !before.trim().is_empty() {
results.push(EmulatorAction::Text(before.to_string()));
}
self.buffer.clear();
self.state = ParserState::InExecuteBlock;
continue;
}
}
if let Some((before_dollar, from_dollar)) = self.buffer.split_once("\n$") {
let text = format!("{}\n", before_dollar);
if !text.trim().is_empty() {
results.push(EmulatorAction::Text(text));
}
self.buffer = format!("${}", from_dollar);
self.state = ParserState::InCommand;
} else if self.buffer.starts_with('$') && self.buffer.len() == chunk.len() {
self.state = ParserState::InCommand;
} else {
let hold_back = if self.code_mode_enabled {
HOLD_BACK_CODE_MODE
} else {
HOLD_BACK_SHELL_ONLY
};
let char_count = self.buffer.chars().count();
if char_count > hold_back && !self.buffer.ends_with('\n') {
let mut chars = self.buffer.chars();
let emit_count = char_count - hold_back;
let emit_text: String = chars.by_ref().take(emit_count).collect();
let keep_text: String = chars.collect();
if !emit_text.is_empty() {
results.push(EmulatorAction::Text(emit_text));
}
self.buffer = keep_text;
}
break;
}
}
}
}
results
}
pub(crate) fn flush(&mut self) -> Vec<EmulatorAction> {
let mut results = Vec::new();
if !self.buffer.is_empty() {
match self.state {
ParserState::InCommand => {
let command_line = self.buffer.trim();
if let Some(command) = command_line.strip_prefix('$') {
let command = command.trim();
if !command.is_empty() {
results.push(EmulatorAction::ShellCommand(command.to_string()));
}
} else if !command_line.is_empty() {
results.push(EmulatorAction::Text(self.buffer.clone()));
}
}
ParserState::InExecuteBlock => {
let code = self.buffer.trim();
if !code.is_empty() {
results.push(EmulatorAction::ExecuteCode(code.to_string()));
}
}
ParserState::Normal => {
results.push(EmulatorAction::Text(self.buffer.clone()));
}
}
self.buffer.clear();
self.state = ParserState::Normal;
}
results
}
}
pub(crate) fn message_for_emulator_action(
action: &EmulatorAction,
message_id: &str,
) -> (Message, bool) {
match action {
EmulatorAction::Text(text) => {
let mut message = Message::assistant().with_text(text);
message.id = Some(message_id.to_string());
(message, false)
}
EmulatorAction::ShellCommand(command) => {
let tool_id = Uuid::new_v4().to_string();
let mut args = serde_json::Map::new();
args.insert("command".to_string(), json!(command));
let tool_call =
CallToolRequestParams::new(Cow::Borrowed(SHELL_TOOL)).with_arguments(args);
let mut message = Message::assistant();
message
.content
.push(MessageContent::tool_request(tool_id, Ok(tool_call)));
message.id = Some(message_id.to_string());
(message, true)
}
EmulatorAction::ExecuteCode(code) => {
let tool_id = Uuid::new_v4().to_string();
let wrapped = if code.contains("async function run()") {
code.clone()
} else {
format!("async function run() {{\n{}\n}}", code)
};
let mut args = serde_json::Map::new();
args.insert("code".to_string(), json!(wrapped));
let tool_call =
CallToolRequestParams::new(Cow::Borrowed(CODE_EXECUTION_TOOL)).with_arguments(args);
let mut message = Message::assistant();
message
.content
.push(MessageContent::tool_request(tool_id, Ok(tool_call)));
message.id = Some(message_id.to_string());
(message, true)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn parse_chunks(chunks: &[&str], code_mode: bool) -> Vec<EmulatorAction> {
let mut parser = StreamingEmulatorParser::new(code_mode);
let mut actions = Vec::new();
for chunk in chunks {
actions.extend(parser.process_chunk(chunk));
}
actions.extend(parser.flush());
actions
}
fn parse_all(input: &str, code_mode: bool) -> Vec<EmulatorAction> {
parse_chunks(&[input], code_mode)
}
fn assert_text(action: &EmulatorAction, expected: &str) {
match action {
EmulatorAction::Text(t) => assert_eq!(t.trim(), expected.trim(), "text mismatch"),
other => panic!("expected Text, got {:?}", action_label(other)),
}
}
fn assert_shell(action: &EmulatorAction, expected: &str) {
match action {
EmulatorAction::ShellCommand(cmd) => {
assert_eq!(cmd, expected, "shell command mismatch")
}
other => panic!("expected ShellCommand, got {:?}", action_label(other)),
}
}
fn assert_execute(action: &EmulatorAction, expected: &str) {
match action {
EmulatorAction::ExecuteCode(code) => {
assert_eq!(code.trim(), expected.trim(), "execute code mismatch")
}
other => panic!("expected ExecuteCode, got {:?}", action_label(other)),
}
}
fn action_label(a: &EmulatorAction) -> &'static str {
match a {
EmulatorAction::Text(_) => "Text",
EmulatorAction::ShellCommand(_) => "ShellCommand",
EmulatorAction::ExecuteCode(_) => "ExecuteCode",
}
}
#[test]
fn plain_text_no_tools() {
let actions = parse_all("Hello, world!", false);
let all_text: String = actions
.iter()
.map(|a| match a {
EmulatorAction::Text(t) => t.as_str(),
_ => panic!("expected only Text actions"),
})
.collect();
assert_eq!(all_text.trim(), "Hello, world!");
}
#[test]
fn single_shell_command() {
let actions = parse_all("$ ls -la\n", false);
assert_eq!(actions.len(), 1);
assert_shell(&actions[0], "ls -la");
}
#[test]
fn text_then_shell_command() {
let actions = parse_all("Let me check:\n$ ls -la\n", false);
assert!(actions.len() >= 2);
assert_text(&actions[0], "Let me check:");
assert_shell(&actions[actions.len() - 1], "ls -la");
}
#[test]
fn shell_command_at_start_of_output() {
let actions = parse_all("$ whoami\n", false);
assert_eq!(actions.len(), 1);
assert_shell(&actions[0], "whoami");
}
#[test]
fn shell_command_without_trailing_newline() {
let actions = parse_all("$ whoami", false);
assert_eq!(actions.len(), 1);
assert_shell(&actions[0], "whoami");
}
#[test]
fn dollar_sign_mid_sentence_is_not_command() {
let actions = parse_all("It costs $50 per month", false);
for action in &actions {
assert!(matches!(action, EmulatorAction::Text(_)));
}
let all_text: String = actions
.iter()
.filter_map(|a| match a {
EmulatorAction::Text(t) => Some(t.as_str()),
_ => None,
})
.collect();
assert_eq!(all_text.trim(), "It costs $50 per month");
}
#[test]
fn execute_block() {
let input = "Here's the code:\n```execute_typescript\nconsole.log('hi');\n```\n";
let actions = parse_all(input, true);
assert!(actions.len() >= 2);
assert_text(&actions[0], "Here's the code:");
assert_execute(&actions[actions.len() - 1], "console.log('hi');");
}
#[test]
fn tool_description_uses_parser_execute_fence() {
let description = build_emulator_tool_description(&[], true);
assert!(description.contains("```execute_typescript"));
assert!(!description.contains("```execute block"));
assert!(!description.contains("Use ```execute for tool calls"));
}
#[test]
fn execute_block_not_detected_without_code_mode() {
let input = "```execute_typescript\nconsole.log('hi');\n```\n";
let actions = parse_all(input, false);
for action in &actions {
assert!(matches!(action, EmulatorAction::Text(_)));
}
}
#[test]
fn dollar_split_across_chunks() {
let actions = parse_chunks(&["Let me check\n", "$ ls -la\n"], false);
let shells: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
.collect();
assert_eq!(shells.len(), 1);
assert_shell(shells[0], "ls -la");
}
#[test]
fn execute_fence_split_across_chunks() {
let actions = parse_chunks(
&["Here:\n```ex", "ecute_typescript\nlet x = 1;\n", "```\n"],
true,
);
let executes: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ExecuteCode(_)))
.collect();
assert_eq!(executes.len(), 1);
assert_execute(executes[0], "let x = 1;");
}
#[test]
fn multiple_commands_on_separate_lines() {
let actions = parse_chunks(&["Here:\n$ cd /tmp\n", "Done.\n$ ls\n"], false);
let shells: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
.collect();
assert_eq!(shells.len(), 2);
assert_shell(shells[0], "cd /tmp");
assert_shell(shells[1], "ls");
}
#[test]
fn regular_code_fence_not_treated_as_execute() {
let input = "```python\nprint('hi')\n```\n";
let actions = parse_all(input, true);
for action in &actions {
assert!(matches!(action, EmulatorAction::Text(_)));
}
}
#[test]
fn empty_command_ignored() {
let actions = parse_all("$\n", false);
let shells: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
.collect();
assert_eq!(shells.len(), 0);
}
#[test]
fn token_by_token_streaming() {
let input = "$ echo hello\n";
let chars: Vec<String> = input.chars().map(|c| c.to_string()).collect();
let chunks: Vec<&str> = chars.iter().map(|s| s.as_str()).collect();
let actions = parse_chunks(&chunks, false);
let shells: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ShellCommand(_)))
.collect();
assert_eq!(shells.len(), 1);
assert_shell(shells[0], "echo hello");
}
#[test]
fn execute_block_with_multiline_code() {
let input = "```execute_typescript\nasync function run() {\n const r = await Developer.shell({ command: \"ls\" });\n return r;\n}\n```\n";
let actions = parse_all(input, true);
let executes: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ExecuteCode(_)))
.collect();
assert_eq!(executes.len(), 1);
match executes[0] {
EmulatorAction::ExecuteCode(code) => {
assert!(code.contains("async function run()"));
assert!(code.contains("Developer.shell"));
}
_ => unreachable!(),
}
}
#[test]
fn unclosed_execute_block_flushed() {
let input = "```execute_typescript\nlet x = 1;";
let actions = parse_all(input, true);
let executes: Vec<_> = actions
.iter()
.filter(|a| matches!(a, EmulatorAction::ExecuteCode(_)))
.collect();
assert_eq!(executes.len(), 1);
assert_execute(executes[0], "let x = 1;");
}
}