zesdex/src/dto/chat/tool.rs

//! Tool-call DTOs embedded in assistant chat messages.
//!
//! Flow: provider response/stream carries `tool_calls` on an assistant
//! message → deserialized into `ToolCall`/`ToolFunction` → harness resolves
//! `function.name` against `all_tools()` and runs it with
//! `sanitize_tool_arguments(function.arguments)`.
//!
//! Why: kept separate from `dto::provider` because tool calls are a property
//! of a chat *message*, not of the request/response envelope.

use serde::{Deserialize, Serialize};
use serde_json::Value;

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

    #[test]
    fn repair_json_closes_string() {
        assert_eq!(repair_json("{\"a\": \"bc"), "{\"a\": \"bc\"}");
    }

    #[test]
    fn repair_json_closes_brace() {
        assert_eq!(repair_json("{\"a\": 1"), "{\"a\": 1}");
    }

    #[test]
    fn repair_json_closes_bracket() {
        assert_eq!(repair_json("{\"a\": [1, 2"), "{\"a\": [1, 2]}");
    }

    #[test]
    fn repair_json_nested() {
        assert_eq!(
            repair_json("{\"a\": {\"b\": [1, 2"),
            "{\"a\": {\"b\": [1, 2]}}"
        );
    }

    #[test]
    fn repair_json_bracket_then_brace() {
        // `[` opened first → `]` must close first, then `}`
        assert_eq!(
            repair_json("[[1, 2, {\"a\": 3"),
            "[[1, 2, {\"a\": 3}]]"
        );
    }

    #[test]
    fn repair_json_handles_escape() {
        assert_eq!(repair_json("{\"a\": \"hello\\"), "{\"a\": \"hello\"}");
    }

    #[test]
    fn repair_json_handles_escaped_quote() {
        assert_eq!(
            repair_json("{\"a\": \"he said \\\"hi\\\""),
            "{\"a\": \"he said \\\"hi\\\"\"}"
        );
    }

    #[test]
    fn repair_json_handles_nested_brackets_and_braces() {
        assert_eq!(
            repair_json("{\"a\": [1, {\"b\": 2"),
            "{\"a\": [1, {\"b\": 2}]}"
        );
    }

    #[test]
    fn repair_json_unchanged_for_valid() {
        let v = "{\"a\": 1, \"b\": [2, 3]}";
        assert_eq!(repair_json(v), v);
    }

    #[test]
    fn sanitize_repairs_truncated_string() {
        let args = Value::String("{\"path\": \"a.txt\", \"content\": \"short\"}".to_string());
        let result = sanitize_tool_arguments(&args);
        assert!(result.is_object());
        assert_eq!(result.get("path").and_then(|v| v.as_str()), Some("a.txt"));
    }

    #[test]
    fn sanitize_passes_object_through() {
        let args = serde_json::json!({"path": "a.txt"});
        let result = sanitize_tool_arguments(&args);
        assert_eq!(result, args);
    }

    #[test]
    fn sanitize_falls_back_to_raw_on_unrepairable() {
        // Completely garbage — not even close to JSON
        let args = Value::String("not even close".to_string());
        let result = sanitize_tool_arguments(&args);
        assert!(result.is_object());
        assert!(result.get("_raw").is_some());
        assert!(result.get("_parse_error").is_some());
    }
}

/// A single tool-call request emitted by the model in an assistant message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    #[serde(rename = "type")]
    pub type_: String,
    pub function: ToolFunction,
}

/// The function name and raw arguments payload for a `ToolCall`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFunction {
    pub name: String,
    pub arguments: Value,
}

/// Normalize tool-call arguments into a JSON object/value.
///
/// Flow: some providers send `arguments` as a JSON-encoded string rather
/// than a nested object; if `args` is a string, attempt to parse it as
/// JSON. Objects and other value types pass through unchanged.
///
/// Security: on parse failure we wrap the raw string in `{ "_raw": "..." }`
/// instead of passing it through as a raw string, so tools that expect a
/// JSON object (via `args.get("key")`) get `None` rather than unexpectedly
/// receiving a plain string value.
///
/// Return: the parsed `Value`, or a wrapper object on parse failure.
/// Attempt to fix truncated JSON by closing open strings, braces and brackets.
///
/// Flow: single-pass character scan tracking string/escape state with a
/// LIFO stack for `{`/`[` → append missing `"`, `]`, `}` in the right
/// (reverse nesting) order.
///
/// Why: LLM output can be cut off mid‑JSON (`max_tokens` hit, connection
/// drop).  This gives tools a chance to act on whatever was emitted.
///
/// Why LIFO vs. depth counters: `{` inside `[` must close with `}` before
/// `]`. Simple depth counters get the nesting order wrong.
fn repair_json(s: &str) -> String {
    let mut stack: Vec<char> = Vec::new();
    let mut in_string = false;
    let mut prev_was_backslash = false;
    let mut ends_with_unclosed_escape = false;

    for c in s.chars() {
        if prev_was_backslash {
            prev_was_backslash = false;
            ends_with_unclosed_escape = false;
            continue;
        }
        if c == '\\' && in_string {
            prev_was_backslash = true;
            ends_with_unclosed_escape = true;
            continue;
        }
        ends_with_unclosed_escape = false;
        if c == '"' {
            in_string = !in_string;
            continue;
        }
        if in_string {
            continue;
        }
        match c {
            '{' | '[' => stack.push(c),
            '}' | ']' => {
                stack.pop();
            }
            _ => {}
        }
    }

    let mut result = s.to_string();
    if ends_with_unclosed_escape {
        result.pop();
    }
    if in_string {
        result.push('"');
    }
    for &opener in stack.iter().rev() {
        match opener {
            '{' => result.push('}'),
            '[' => result.push(']'),
            _ => {}
        }
    }
    result
}

/// Normalize tool-call arguments into a JSON object/value.
///
/// Flow: some providers send `arguments` as a JSON-encoded string rather
/// than a nested object; if `args` is a string, attempt to parse it as
/// JSON. Objects and other value types pass through unchanged.
///
/// Security: on parse failure we wrap the raw string in `{ "_raw": "..." }`
/// instead of passing it through as a raw string, so tools that expect a
/// JSON object (via `args.get("key")`) get `None` rather than unexpectedly
/// receiving a plain string value.
///
/// Return: the parsed `Value`, or a wrapper object on parse failure.
pub fn sanitize_tool_arguments(args: &Value) -> Value {
    match args {
        Value::String(s) => {
            // Attempt 1: direct parse.
            if let Ok(v) = serde_json::from_str::<Value>(s) {
                return v;
            }
            // Attempt 2: strip control chars (0x00-0x1F except \t, \n)
            // that some LLM providers emit as literal bytes in JSON strings
            // (e.g. multi-line commit messages), then retry.
            let cleaned: String = s.chars()
                .filter(|&c| !c.is_control() || c == '\t' || c == '\n' || c == '\r')
                .collect();
            if cleaned.len() != s.len() {
                if let Ok(v) = serde_json::from_str::<Value>(&cleaned) {
                    tracing::warn!(
                        "tool argument contained control characters — stripped \
                         and reparsed successfully",
                    );
                    return v;
                }
            }
            // Attempt 3: repair truncated JSON and retry.
            let input = if cleaned.len() == s.len() { s } else { &cleaned };
            let repaired = repair_json(input);
            match serde_json::from_str::<Value>(&repaired) {
                Ok(v) => {
                    tracing::warn!(
                        "tool argument string was truncated — repaired successfully",
                    );
                    v
                }
                Err(e2) => {
                    tracing::error!(
                        "tool argument is a JSON string but failed to parse. \
                         Wrapping in object. Error: {}. Raw (first 200): {}",
                        e2, s.chars().take(200).collect::<String>(),
                    );
                    serde_json::json!({"_raw": s, "_parse_error": e2.to_string()})
                }
            }
        }
        obj @ Value::Object(_) => obj.clone(),
        other => other.clone(),
    }
}