diff --git a/src/ai_response.rs b/src/ai_response.rs index c5b1d8b..55889e2 100644 --- a/src/ai_response.rs +++ b/src/ai_response.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; -use tokio::sync::RwLock; +use std::sync::RwLock; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { diff --git a/src/lib.rs b/src/lib.rs index f9c4eda..b7273f1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,7 @@ pub mod krill; pub mod project; pub mod task; pub mod todo; +pub mod tools; pub use ai_response::{Conversation, Message}; pub use coral::{Coral, CoralId}; @@ -19,3 +20,7 @@ pub use krill::{KrillConfig, KrillDescriptor, KrillId}; pub use project::{Project, ProjectFile, ProjectId, ProjectSettings}; pub use task::{Task, TaskResult}; pub use todo::ToDo; +pub use tools::{ + format_tool_error, format_tool_result, parse_tool_call_blocks, parse_tool_call_stream, + ParsedToolCall, ToolDefinition, ToolParser, +}; diff --git a/src/tools.rs b/src/tools.rs new file mode 100644 index 0000000..94880e6 --- /dev/null +++ b/src/tools.rs @@ -0,0 +1,615 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ToolDefinition { + #[serde(default)] + pub id: String, + pub name: String, + pub description: String, + #[serde(default)] + pub category: String, + pub parameters: serde_json::Value, + #[serde(default)] + pub dangerous: bool, + #[serde(default)] + pub requires_approval: bool, + #[serde(default)] + pub version: String, +} + +impl ToolDefinition { + pub fn to_openai_format(&self) -> serde_json::Value { + serde_json::json!({ + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + } + }) + } + + pub fn to_short_doc(&self) -> String { + format!( + "- {}: {} (parameters: {})\n", + self.name, self.description, self.parameters + ) + } + + /// Strip the project_id parameter from the tool's JSON schema. + /// This is used to hide the project context from agents that shouldn't + /// have to manage it manually. + pub fn strip_project_id(&mut self) { + if let Some(obj) = self.parameters.as_object_mut() { + if let Some(properties) = obj.get_mut("properties").and_then(|p| p.as_object_mut()) { + properties.remove("project_id"); + } + if let Some(required) = obj.get_mut("required").and_then(|r| r.as_array_mut()) { + required.retain(|v| v.as_str() != Some("project_id")); + } + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ParsedToolCall { + pub tool_name: String, + pub call_id: String, + pub payload: serde_json::Value, + pub raw_json: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ParserState { + Text, + Collecting, +} + +pub struct ToolParser { + state: ParserState, + current_tool_name: String, + current_accumulator: String, + leftover: String, +} + +impl ToolParser { + pub fn new() -> Self { + Self { + state: ParserState::Text, + current_tool_name: String::new(), + current_accumulator: String::new(), + leftover: String::new(), + } + } + + pub fn ingest(&mut self, chunk: &str) -> Vec { + let mut results = Vec::new(); + let text = format!("{}{}", self.leftover, chunk); + + let mut lines: Vec<&str> = text.split('\n').collect(); + + // The last element is either empty ( if text ended in \n) or the start of a new line + if let Some(last) = lines.pop() { + self.leftover = last.to_string(); + } else { + self.leftover.clear(); + } + + for line in lines { + match self.state { + ParserState::Text => { + let trimmed = line.trim(); + if let Some(tool_name) = trimmed.strip_prefix("```tool:") { + if !tool_name.trim().starts_with("-result") { + self.current_tool_name = tool_name.trim().to_string(); + self.current_accumulator.clear(); + self.state = ParserState::Collecting; + } + } + } + ParserState::Collecting => { + if line.trim() == "```" { + let parsed = parse_collected_json( + &self.current_tool_name, + &self.current_accumulator, + ); + results.push(parsed); + self.state = ParserState::Text; + } else { + if !self.current_accumulator.is_empty() { + self.current_accumulator.push('\n'); + } + self.current_accumulator.push_str(line); + } + } + } + } + results + } + + pub fn finish(mut self) -> Vec { + let mut results = Vec::new(); + if !self.leftover.is_empty() { + // Treat leftover as a final line + let line = self.leftover.clone(); + match self.state { + ParserState::Text => { + let trimmed = line.trim(); + if let Some(tool_name) = trimmed.strip_prefix("```tool:") { + if !tool_name.trim().starts_with("-result") { + self.current_tool_name = tool_name.trim().to_string(); + self.current_accumulator.clear(); + self.state = ParserState::Collecting; + } + } + } + ParserState::Collecting => { + if line.trim() == "```" { + let parsed = parse_collected_json( + &self.current_tool_name, + &self.current_accumulator, + ); + results.push(parsed); + self.state = ParserState::Text; + } else { + if !self.current_accumulator.is_empty() { + self.current_accumulator.push('\n'); + } + self.current_accumulator.push_str(&line); + } + } + } + } + + if self.state == ParserState::Collecting && !self.current_accumulator.is_empty() { + results.push(parse_collected_json( + &self.current_tool_name, + &self.current_accumulator, + )); + } + + results + } +} + +pub fn parse_tool_call_blocks(text: &str) -> Vec { + let mut parser = ToolParser::new(); + let mut results = parser.ingest(text); + results.extend(parser.finish()); + results +} + +pub fn parse_tool_call_stream<'a>(chunks: impl Iterator) -> Vec { + let mut parser = ToolParser::new(); + let mut results = Vec::new(); + for chunk in chunks { + results.extend(parser.ingest(chunk)); + } + results.extend(parser.finish()); + results +} + +fn parse_collected_json(tool_name: &str, json_str: &str) -> ParsedToolCall { + match serde_json::from_str::(json_str) { + Ok(payload) => { + let call_id = payload + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("missing_call_id") + .to_string(); + + ParsedToolCall { + tool_name: tool_name.to_string(), + call_id, + payload, + raw_json: json_str.to_string(), + } + } + Err(e) => ParsedToolCall { + tool_name: tool_name.to_string(), + call_id: "parse_error".to_string(), + payload: serde_json::json!({ + "error": "invalid_json", + "details": e.to_string() + }), + raw_json: json_str.to_string(), + }, + } +} + +pub fn format_tool_result( + tool_name: &str, + call_id: &str, + status: &str, + data: &serde_json::Value, +) -> String { + let payload = serde_json::json!({ + "call_id": call_id, + "status": status, + "data": data + }); + format!( + "```tool-result:{}\n{}\n```", + tool_name, + serde_json::to_string_pretty(&payload).unwrap_or_default() + ) +} + +pub fn format_tool_error(tool_name: &str, call_id: &str, code: &str, message: &str) -> String { + let data = serde_json::json!({ + "code": code, + "message": message + }); + format_tool_result(tool_name, call_id, "error", &data) +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCallEnvelope { + pub call_id: String, + pub tool_name: String, + pub arguments: serde_json::Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolResultEnvelope { + pub call_id: String, + pub tool_name: String, + pub status: ToolResultStatus, + pub data: Option, + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolResultStatus { + Success, + Error, + Streaming, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolErrorInfo { + pub code: String, + pub message: String, +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_parse_single_block() { + let text = r#" +Some text before. +```tool:read_file +{ + "call_id": "call_01", + "path": "test.txt" +} +``` +Some text after. +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "read_file"); + assert_eq!(results[0].call_id, "call_01"); + assert_eq!(results[0].payload["path"], "test.txt"); + } + + #[test] + fn test_parse_multiple_blocks() { + let text = r#" +```tool:tool1 +{"call_id": "c1"} +``` +Middle text. +```tool:tool2 +{"call_id": "c2"} +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 2); + assert_eq!(results[0].tool_name, "tool1"); + assert_eq!(results[0].call_id, "c1"); + assert_eq!(results[1].tool_name, "tool2"); + assert_eq!(results[1].call_id, "c2"); + } + + #[test] + fn test_parse_error_json() { + let text = r#" +```tool:bad_json +{ "invalid": +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].call_id, "parse_error"); + assert_eq!(results[0].payload["error"], "invalid_json"); + } + + #[test] + fn test_parse_missing_call_id() { + let text = r#" +```tool:no_id +{"foo": "bar"} +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].call_id, "missing_call_id"); + } + + #[test] + fn test_streaming_parser() { + let chunks = vec![ + "Some text.\n```tool:", + "my_tool\n", + "{\"call_id\": \"st", + "ream_01\"}\n", + "```\nAnd more.", + ]; + let results = parse_tool_call_stream(chunks.into_iter()); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "my_tool"); + assert_eq!(results[0].call_id, "stream_01"); + } + + #[test] + fn test_format_result() { + let data = json!({"content": "hello world"}); + let result = format_tool_result("read_file", "call_01", "success", &data); + assert!(result.contains("```tool-result:read_file")); + assert!(result.contains("call_01")); + assert!(result.contains("success")); + assert!(result.contains("hello world")); + } + + #[test] + fn test_parse_missing_closing_fence() { + let text = r#" +```tool:orphan +{"call_id": "call_01", "data": "never closed"}"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "orphan"); + assert_eq!(results[0].call_id, "call_01"); + } + + #[test] + fn test_parse_empty_input() { + let results = parse_tool_call_blocks(""); + assert!(results.is_empty()); + + let results = parse_tool_call_blocks(" \n \n "); + assert!(results.is_empty()); + } + + #[test] + fn test_parse_mixed_content() { + let text = r#"First, let me read the file. +```tool:read_file +{"call_id": "call_01", "path": "test.txt"} +``` +Now let me search. +```tool:search_files +{"call_id": "call_02", "pattern": "TODO"} +``` +Done with tools. +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 2); + assert_eq!(results[0].tool_name, "read_file"); + assert_eq!(results[0].call_id, "call_01"); + assert_eq!(results[1].tool_name, "search_files"); + assert_eq!(results[1].call_id, "call_02"); + } + + #[test] + fn test_parse_base64_content() { + let text = r#" +```tool:write_file +{"call_id": "call_01", "path": "out.bin", "content_base64": "SGVsbG8gV29ybGQ="} +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "write_file"); + assert_eq!(results[0].call_id, "call_01"); + assert_eq!(results[0].payload["content_base64"], "SGVsbG8gV29ybGQ="); + } + + #[test] + fn test_format_tool_error() { + let result = format_tool_error("read_file", "call_01", "NOT_FOUND", "File not found"); + assert!(result.contains("```tool-result:read_file")); + assert!(result.contains("NOT_FOUND")); + assert!(result.contains("File not found")); + assert!(result.contains("\"status\": \"error\"")); + } + + #[test] + fn test_parse_tool_result_block_not_parsed() { + let text = r#" +```tool-result:read_file +{"call_id": "call_01", "status": "success", "data": {"content": "hello"}} +``` +"#; + let results = parse_tool_call_blocks(text); + assert!( + results.is_empty(), + "tool-result: blocks should not be parsed by parse_tool_call_blocks" + ); + } + + #[test] + fn test_tool_definition_to_openai_format() { + let def = ToolDefinition { + name: "test_tool".to_string(), + description: "A test tool".to_string(), + parameters: json!({"type": "object", "properties": {}}), + ..Default::default() + }; + let openai = def.to_openai_format(); + assert_eq!(openai["type"], "function"); + assert_eq!(openai["function"]["name"], "test_tool"); + assert_eq!(openai["function"]["description"], "A test tool"); + } + + #[test] + fn test_tool_definition_to_short_doc() { + let def = ToolDefinition { + name: "test_tool".to_string(), + description: "A test tool".to_string(), + parameters: json!({"type": "object"}), + ..Default::default() + }; + let doc = def.to_short_doc(); + assert!(doc.contains("test_tool")); + assert!(doc.contains("A test tool")); + } + + #[test] + fn test_streaming_rejects_tool_result_block() { + let chunks = vec![ + "Some text.\n```tool-result:read_file\n", + "{\"call_id\": \"call_01\"}\n", + "```\nAnd more.", + ]; + let results = parse_tool_call_stream(chunks.into_iter()); + assert!( + results.is_empty(), + "tool-result: blocks should be rejected by streaming parser" + ); + } + + #[test] + fn test_streaming_rejects_tool_result_across_chunks() { + let chunks = vec![ + "```tool-", + "result:read_file\n", + "{\"call_id\": \"call_01\"}\n", + "```\n", + ]; + let results = parse_tool_call_stream(chunks.into_iter()); + assert!( + results.is_empty(), + "tool-result: split across chunks should be rejected" + ); + } + + #[test] + fn test_parse_tool_not_result() { + let text = r#" +```tool:result +{"call_id": "call_01", "data": "this is a tool named 'result'"} +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "result"); + assert_eq!(results[0].call_id, "call_01"); + } + + #[test] + fn test_streaming_tool_not_result() { + let chunks = vec!["```tool:result\n", "{\"call_id\": \"call_01\"}\n", "```\n"]; + let results = parse_tool_call_stream(chunks.into_iter()); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "result"); + assert_eq!(results[0].call_id, "call_01"); + } + + #[test] + fn test_finish_rejects_tool_result() { + // When finish() is called with a leftover that is a tool-result block, + // it should not parse it. + let mut parser = ToolParser::new(); + let chunks = vec!["```tool-result:read_file\n{\"call_id\": \"call_01\"}\n```"]; + let _ = parser.ingest(chunks[0]); + let results = parser.finish(); + assert!( + results.is_empty(), + "finish() should not parse tool-result blocks" + ); + } + + #[test] + fn test_parse_tool_name_with_special_chars() { + let text = r#" +```tool:my-tool_v2.with.dots +{"call_id": "call_01"} +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "my-tool_v2.with.dots"); + } + + #[test] + fn test_parse_payload_with_newlines_and_unicode() { + let text = "```tool:write_file\n{\"call_id\": \"call_01\", \"content\": \"line1\\nline2\\nunicode: 🎉\"}\n```"; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].call_id, "call_01"); + assert!(results[0].payload["content"] + .as_str() + .unwrap_or("") + .contains("🎉")); + } + + #[test] + fn test_parse_block_immediately_followed_by_text() { + let text = r#"```tool:read_file +{"call_id": "call_01", "path": "test.txt"} +```And then text right after"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "read_file"); + } + + #[test] + fn test_parse_tool_result_suffix_not_mistaken() { + // "tool-result" as a suffix of a longer tool name should not be rejected + let text = r#" +```tool:my-tool-result-processor +{"call_id": "call_01"} +``` +"#; + let results = parse_tool_call_blocks(text); + assert_eq!(results.len(), 1); + assert_eq!(results[0].tool_name, "my-tool-result-processor"); + assert_eq!(results[0].call_id, "call_01"); + } + + #[test] + fn test_tool_call_envelope_roundtrip() { + let envelope = ToolCallEnvelope { + call_id: "call_01".to_string(), + tool_name: "test_tool".to_string(), + arguments: json!({"key": "value"}), + }; + let json = serde_json::to_string(&envelope).unwrap(); + let deserialized: ToolCallEnvelope = serde_json::from_str(&json).unwrap(); + assert_eq!(deserialized.call_id, "call_01"); + assert_eq!(deserialized.tool_name, "test_tool"); + assert_eq!(deserialized.arguments["key"], "value"); + } + + #[test] + fn test_tool_result_envelope_roundtrip() { + let envelope = ToolResultEnvelope { + call_id: "call_01".to_string(), + tool_name: "test_tool".to_string(), + status: ToolResultStatus::Success, + data: Some(json!({"result": "ok"})), + error: None, + }; + let json = serde_json::to_string(&envelope).unwrap(); + let deserialized: ToolResultEnvelope = serde_json::from_str(&json).unwrap(); + assert_eq!(deserialized.call_id, "call_01"); + assert!(matches!(deserialized.status, ToolResultStatus::Success)); + assert_eq!(deserialized.data.unwrap()["result"], "ok"); + } +}