heartbit_core/tool/

mod.rs

//! Tool trait and built-in tool implementations (filesystem, web, MCP, A2A, etc.).

#[cfg(feature = "a2a")]
pub mod a2a;
pub mod advisor;
pub mod builtins;
pub mod handoff;
pub mod mcp;
pub mod mcp_presets;
pub mod mcp_server;
pub mod set_goal;
pub mod set_scope;

use std::future::Future;
use std::pin::Pin;

use crate::error::Error;
use crate::llm::types::ToolDefinition;

/// Output of a tool execution.
#[derive(Debug, Clone)]
pub struct ToolOutput {
    /// Stringified tool output (may be JSON, plain text, etc.).
    pub content: String,
    /// `true` when the tool returned an error.
    pub is_error: bool,
}

impl ToolOutput {
    /// Construct a successful tool output.
    pub fn success(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
            is_error: false,
        }
    }

    /// Construct a failed tool output (`is_error = true`).
    pub fn error(content: impl Into<String>) -> Self {
        Self {
            content: content.into(),
            is_error: true,
        }
    }

    /// Truncate content if it exceeds `max_bytes`, preserving UTF-8 validity.
    ///
    /// When truncated, appends a `[truncated: N bytes omitted]` suffix so the
    /// LLM knows data was cut. Content within the limit is returned unchanged.
    /// A `max_bytes` of 0 is treated as no-op (returns content unchanged).
    ///
    /// Note: the suffix itself is not counted toward `max_bytes`, so the
    /// result may slightly exceed the limit.
    pub fn truncated(mut self, max_bytes: usize) -> Self {
        if max_bytes == 0 {
            return self;
        }
        if self.content.len() > max_bytes {
            let mut cut = max_bytes;
            while cut > 0 && !self.content.is_char_boundary(cut) {
                cut -= 1;
            }
            let omitted = self.content.len() - cut;
            self.content.truncate(cut);
            self.content
                .push_str(&format!("\n\n[truncated: {omitted} bytes omitted]"));
        }
        self
    }
}

/// Trait for tools that agents can invoke.
///
/// Uses `Pin<Box<dyn Future>>` return type for dyn-compatibility,
/// allowing tools to be stored as `Arc<dyn Tool>`.
///
/// # Example
///
/// Implementing a simple synchronous tool that echoes its input:
///
/// ```rust
/// use std::future::Future;
/// use std::pin::Pin;
/// use heartbit_core::{Tool, ToolOutput};
/// use heartbit_core::llm::types::ToolDefinition;
///
/// struct EchoTool;
///
/// impl Tool for EchoTool {
///     fn definition(&self) -> ToolDefinition {
///         ToolDefinition {
///             name: "echo".into(),
///             description: "Echo back the input string.".into(),
///             input_schema: serde_json::json!({
///                 "type": "object",
///                 "properties": { "text": { "type": "string" } },
///                 "required": ["text"]
///             }),
///         }
///     }
///
///     fn execute(
///         &self,
///         _ctx: &heartbit_core::ExecutionContext,
///         input: serde_json::Value,
///     ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, heartbit_core::Error>> + Send + '_>> {
///         Box::pin(async move {
///             let text = input.get("text").and_then(|v| v.as_str()).unwrap_or("");
///             Ok(ToolOutput::success(text.to_string()))
///         })
///     }
/// }
/// ```
pub trait Tool: Send + Sync {
    /// Return the tool's definition (name, description, input schema) for the LLM.
    fn definition(&self) -> ToolDefinition;

    /// Execute the tool with the LLM-supplied `input` JSON.
    ///
    /// Input is validated against the tool's declared schema before reaching this method.
    fn execute(
        &self,
        ctx: &crate::ExecutionContext,
        input: serde_json::Value,
    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, Error>> + Send + '_>>;

    /// Optional override: produce a compact version of the tool output
    /// suitable for inclusion in conversation history sent to subsequent
    /// LLM turns. The default implementation returns the input verbatim.
    ///
    /// Tools that emit large blobs (e.g., base64 image data) should
    /// override this to strip or replace the blob with a small
    /// placeholder. This prevents conversation-history bloat that can
    /// trip per-model context limits (e.g., Anthropic's 200k token cap).
    ///
    /// The full, unredacted output is preserved on
    /// `AgentOutput.tool_call_results` for the caller's use.
    fn redact_for_history(&self, output: &str) -> String {
        output.to_string()
    }
}

/// Validate tool input against the tool's declared JSON Schema.
///
/// Returns `Ok(())` if valid, `Err(error_message)` if the input
/// does not conform. The error message is suitable for sending back
/// to the LLM so it can self-correct.
pub fn validate_tool_input(
    schema: &serde_json::Value,
    input: &serde_json::Value,
) -> Result<(), String> {
    let validator = match jsonschema::validator_for(schema) {
        Ok(v) => v,
        Err(e) => {
            // If the schema itself is invalid, skip validation rather than
            // rejecting every call. Log a warning for the operator.
            tracing::warn!(error = %e, "invalid tool schema, skipping validation");
            return Ok(());
        }
    };

    let errors: Vec<String> = validator
        .iter_errors(input)
        .map(|e| e.to_string())
        .collect();
    if errors.is_empty() {
        Ok(())
    } else {
        let mut msg = format!("Input validation failed: {}", errors.join("; "));
        // An EMPTY input object with required properties is the signature of a
        // model that put its payload in the text channel instead of the tool
        // call (live /analyze finding: write({}) alongside a 4000-token text
        // response, repeated). Name the actual mistake so it can self-correct.
        if input.as_object().is_some_and(|o| o.is_empty()) {
            msg.push_str(
                ". You sent NO arguments — parameters must be passed in the \
                 tool call's arguments (e.g. file content goes in the tool's \
                 input fields), not in your text response.",
            );
        }
        Err(msg)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn tool_output_success() {
        let output = ToolOutput::success("result data");
        assert_eq!(output.content, "result data");
        assert!(!output.is_error);
    }

    // Live /analyze finding: a model generated 4000-token file bodies in its
    // TEXT response and called write with input={} — twice, identically. The
    // bare "is a required property" error never told it the load-bearing
    // fact: arguments go in the tool call, not the text channel.
    #[test]
    fn empty_input_validation_error_names_the_channel() {
        let schema = json!({
            "type": "object",
            "properties": {
                "file_path": {"type": "string"},
                "content": {"type": "string"}
            },
            "required": ["file_path", "content"]
        });
        let err = validate_tool_input(&schema, &json!({})).unwrap_err();
        assert!(err.contains("required property"), "got: {err}");
        assert!(
            err.contains("tool call's arguments") && err.contains("not in your text"),
            "empty input must carry the channel hint, got: {err}"
        );
        // Non-empty-but-invalid input keeps the plain error (no hint spam).
        let err = validate_tool_input(&schema, &json!({"file_path": "x"})).unwrap_err();
        assert!(
            !err.contains("text response"),
            "partial input must not get the empty-args hint: {err}"
        );
    }

    #[test]
    fn tool_output_error() {
        let output = ToolOutput::error("something failed");
        assert_eq!(output.content, "something failed");
        assert!(output.is_error);
    }

    #[test]
    fn tool_output_truncated_noop_when_within_limit() {
        let output = ToolOutput::success("short text");
        let truncated = output.truncated(100);
        assert_eq!(truncated.content, "short text");
        assert!(!truncated.is_error);
    }

    #[test]
    fn tool_output_truncated_cuts_long_content() {
        let output = ToolOutput::success("a".repeat(1000));
        let truncated = output.truncated(100);
        assert!(truncated.content.len() < 1000);
        assert!(truncated.content.starts_with("aaaa"));
        assert!(truncated.content.contains("[truncated:"));
        assert!(truncated.content.contains("bytes omitted]"));
        assert!(!truncated.is_error); // preserves is_error flag
    }

    #[test]
    fn tool_output_truncated_preserves_utf8() {
        // "é" is 2 bytes in UTF-8. A cut at byte 5 would split a char boundary.
        let output = ToolOutput::success("ééééé"); // 10 bytes
        let truncated = output.truncated(5);
        // Should cut at char boundary (4 bytes = 2 chars), not mid-char
        assert!(truncated.content.starts_with("éé"));
        assert!(truncated.content.contains("[truncated:"));
    }

    #[test]
    fn tool_output_truncated_exact_boundary_noop() {
        let output = ToolOutput::success("hello"); // 5 bytes
        let truncated = output.truncated(5);
        assert_eq!(truncated.content, "hello");
    }

    #[test]
    fn tool_output_truncated_zero_is_noop() {
        let output = ToolOutput::success("some content");
        let truncated = output.truncated(0);
        assert_eq!(truncated.content, "some content"); // unchanged
    }

    #[test]
    fn tool_output_truncated_error_also_truncates() {
        let output = ToolOutput::error("e".repeat(200));
        let truncated = output.truncated(50);
        assert!(truncated.content.contains("[truncated:"));
        assert!(truncated.is_error); // preserves error flag
    }

    #[test]
    fn validate_accepts_valid_input() {
        let schema = json!({
            "type": "object",
            "properties": {
                "query": {"type": "string"}
            },
            "required": ["query"]
        });
        let input = json!({"query": "test"});
        assert!(validate_tool_input(&schema, &input).is_ok());
    }

    #[test]
    fn validate_rejects_missing_required() {
        let schema = json!({
            "type": "object",
            "properties": {
                "query": {"type": "string"}
            },
            "required": ["query"]
        });
        let input = json!({});
        let err = validate_tool_input(&schema, &input).unwrap_err();
        assert!(err.contains("validation failed"), "got: {err}");
    }

    #[test]
    fn validate_rejects_wrong_type() {
        let schema = json!({
            "type": "object",
            "properties": {
                "query": {"type": "string"}
            },
            "required": ["query"]
        });
        let input = json!({"query": 42});
        let err = validate_tool_input(&schema, &input).unwrap_err();
        assert!(err.contains("validation failed"), "got: {err}");
    }

    #[test]
    fn validate_accepts_any_for_minimal_schema() {
        let schema = json!({"type": "object"});
        let input = json!({});
        assert!(validate_tool_input(&schema, &input).is_ok());
    }

    #[test]
    fn validate_skips_on_invalid_schema() {
        // An invalid schema should not block tool execution
        let schema = json!({"type": "not-a-real-type"});
        let input = json!({"anything": true});
        // Should not fail even though schema is invalid — skips validation
        assert!(validate_tool_input(&schema, &input).is_ok());
    }

    #[test]
    fn redact_for_history_default_is_identity() {
        // A tool with no override should return the output verbatim.
        struct NoopTool;
        impl Tool for NoopTool {
            fn definition(&self) -> ToolDefinition {
                ToolDefinition {
                    name: "noop".into(),
                    description: "noop".into(),
                    input_schema: json!({"type": "object"}),
                }
            }

            fn execute(
                &self,
                _ctx: &crate::ExecutionContext,
                _input: serde_json::Value,
            ) -> std::pin::Pin<
                Box<
                    dyn std::future::Future<Output = Result<ToolOutput, crate::error::Error>>
                        + Send
                        + '_,
                >,
            > {
                Box::pin(async { Ok(ToolOutput::success("ok")) })
            }
        }

        let tool = NoopTool;
        let raw = "anything goes here, including [IMAGE:base64:abc]";
        assert_eq!(tool.redact_for_history(raw), raw);
    }

    #[test]
    fn validate_accepts_extra_properties() {
        let schema = json!({
            "type": "object",
            "properties": {
                "query": {"type": "string"}
            },
            "required": ["query"]
        });
        // Extra properties are allowed by default in JSON Schema
        let input = json!({"query": "test", "extra": true});
        assert!(validate_tool_input(&schema, &input).is_ok());
    }
}