Agents, tools

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 {

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,
+};

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<ParsedToolCall> {
+        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<ParsedToolCall> {
+        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<ParsedToolCall> {
+    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<Item = &'a str>) -> Vec<ParsedToolCall> {
+    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::<serde_json::Value>(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<serde_json::Value>,
+    pub error: Option<ToolErrorInfo>,
+}
+
+#[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");
+    }
+}