diff --git a/src/ai_response.rs b/src/ai_response.rs index 55889e2..c5b1d8b 100644 --- a/src/ai_response.rs +++ b/src/ai_response.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; -use std::sync::RwLock; +use tokio::sync::RwLock; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { diff --git a/src/lib.rs b/src/lib.rs index b7273f1..f9c4eda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,7 +10,6 @@ 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}; @@ -20,7 +19,3 @@ 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 deleted file mode 100644 index da20468..0000000 --- a/src/tools.rs +++ /dev/null @@ -1,619 +0,0 @@ -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, - /// Whether this tool requires a project context to function. - /// Project-scoped tools are filtered out when no project is selected. - #[serde(default)] - pub project_scoped: bool, -} - -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!( - "- {}: {}\n", - self.name, self.description - ) - } - - /// 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"); - } -}