asepharyana
1c08b8e4e9
- Added functionality to generate lessons based on code reviews, including prompts and instructions for the reviewer. - Updated the `.gitignore` to exclude lesson documentation files. - Removed the `/lesson` command from the help menu and command parsing. - Enhanced the status bar to display a message when a lesson is being generated. - Introduced a new `lesson_running` state to track lesson generation progress. - Updated various components to handle the new lesson generation workflow, including subagent events and token usage tracking.
251 lines
8.2 KiB
Rust
251 lines
8.2 KiB
Rust
//! 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(),
|
||
}
|
||
}
|