zesdex/apps/domain/src/core/tool_call.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 the
//! function name against `all_tools()` and runs it after sanitizing arguments
//! via [`sanitize_tool_arguments`] (which handles string-encoded JSON,
//! control characters, and truncation).
//!
//! # Components
//!
//! - `ToolCall` — a single tool-invocation request (id + type + function)
//! - `ToolFunction` — function name + raw arguments Value
//! - `sanitize_tool_arguments` — normalizes argument shape, repairs truncation
//! - `repair_json` — closes unclosed strings/braces/brackets in truncated JSON
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tracing;

/// A single tool-call request emitted by the model in an assistant message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// Unique identifier for this tool call (referenced by `ToolCallResult`).
    pub id: String,
    /// Discriminator, e.g. `"function"`.
    #[serde(rename = "type")]
    pub type_: String,
    /// The function to invoke (name + arguments).
    pub function: ToolFunction,
}

/// The function name and raw arguments payload for a `ToolCall`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFunction {
    /// The function/tool name to dispatch against.
    pub name: String,
    /// Arguments as a JSON Value (may be a string-encoded object before
    /// `sanitize_tool_arguments` normalises it).
    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.
///
/// Attempt to fix truncated JSON by closing open strings, braces and brackets.
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)
            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(),
    }
}

/// Repair 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.
pub 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
}