diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f49fef70..c7067860a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1213,6 +1213,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 registers the pre-evaluated `inputSchema` and Swift's delayed abort as differences from upstream. +### Changed + +- Tool-call repair is a trait. New `RepairToolCall` + (`aimux_core::parse_tool_call`, also in the prelude) with one async method, + `repair(&self, ToolCallRepairContext) -> Result, AiMuxError>`; + every matching closure implements it. `ToolCallRepair` holds an + `Arc`: `ToolCallRepair::new(closure)` is unchanged and + `ToolCallRepair::from(hook)` takes any implementation. Repair behaviour, + error variants, JSON, the C ABI and the bindings are unchanged. + ## [0.5.0] - 2026-09-27 **Breaking release.** 13 PRs since 0.3.0: the cross-language error model diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index 50738b44b..794766538 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -121,8 +121,9 @@ pub mod prelude { }; pub use crate::output::Output; pub use crate::parse_tool_call::{ - ToolCallRepair, ToolCallRepairContext, ToolCallRepairReply, apply_tool_call_repair, - apply_tool_call_repair_to_result, tool_call_repair_context, tool_call_repair_inputs, + RepairToolCall, ToolCallRepair, ToolCallRepairContext, ToolCallRepairReply, + apply_tool_call_repair, apply_tool_call_repair_to_result, tool_call_repair_context, + tool_call_repair_inputs, }; pub use crate::provider::Provider; pub use crate::reranking_model::{ diff --git a/aimux-core/src/parse_tool_call.rs b/aimux-core/src/parse_tool_call.rs index 42bc55a33..30347fe02 100644 --- a/aimux-core/src/parse_tool_call.rs +++ b/aimux-core/src/parse_tool_call.rs @@ -6,9 +6,9 @@ //! rejection, schema validation, and the one-shot repair callback. use std::future::Future; -use std::pin::Pin; use std::sync::Arc; +use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::Value; use ts_rs::TS; @@ -55,31 +55,92 @@ impl ToolCallRepairContext { } } -type ToolCallRepairFuture = - Pin, AiMuxError>> + Send>>; +/// A one-shot repair hook for invalid tool calls. +/// +/// The Rust form of the AI SDK's `experimental_repairToolCall`. Core calls +/// [`repair`](Self::repair) at most once per invalid call, and parses and +/// validates the returned call from scratch. Returning `Ok(None)` keeps the +/// original validation error; an `Err` becomes a `ToolCallRepair` error that +/// keeps both the original failure and the cause. +/// +/// Any `Fn(ToolCallRepairContext) -> impl Future, AiMuxError>>` +/// implements this trait, so a closure works wherever a hook does. +/// +/// ``` +/// use aimux_core::error::AiMuxError; +/// use aimux_core::parse_tool_call::{RepairToolCall, ToolCallRepair, ToolCallRepairContext}; +/// use aimux_core::tool::RawToolCall; +/// +/// struct RenameTool(&'static str); +/// +/// #[async_trait::async_trait] +/// impl RepairToolCall for RenameTool { +/// async fn repair( +/// &self, +/// context: ToolCallRepairContext, +/// ) -> Result, AiMuxError> { +/// Ok(Some(RawToolCall { +/// tool_name: self.0.into(), +/// ..context.tool_call +/// })) +/// } +/// } +/// +/// let repair = ToolCallRepair::from(RenameTool("weather")); +/// ``` +#[async_trait] +pub trait RepairToolCall: Send + Sync { + /// Return a replacement call, `None` to keep the original error, or an + /// error to report the repair as failed. + async fn repair( + &self, + context: ToolCallRepairContext, + ) -> Result, AiMuxError>; +} -/// Async callback that may replace one invalid tool call. +#[async_trait] +impl RepairToolCall for F +where + F: Fn(ToolCallRepairContext) -> Fut + Send + Sync, + Fut: Future, AiMuxError>> + Send, +{ + async fn repair( + &self, + context: ToolCallRepairContext, + ) -> Result, AiMuxError> { + self(context).await + } +} + +/// A shared [`RepairToolCall`] hook, as carried by `GenerateTextOptions`. /// -/// Core invokes this at most once, and parses and validates the returned call -/// from scratch. Returning `None` keeps the original validation error. +/// Build one from a closure with [`ToolCallRepair::new`], or from any +/// [`RepairToolCall`] implementation with `From`. #[derive(Clone)] -pub struct ToolCallRepair(Arc ToolCallRepairFuture + Send + Sync>); +pub struct ToolCallRepair(Arc); impl ToolCallRepair { + /// Wrap a repair closure. #[must_use] pub fn new(repair: F) -> Self where F: Fn(ToolCallRepairContext) -> Fut + Send + Sync + 'static, Fut: Future, AiMuxError>> + Send + 'static, { - Self(Arc::new(move |context| Box::pin(repair(context)))) + Self(Arc::new(repair)) } async fn repair( &self, context: ToolCallRepairContext, ) -> Result, AiMuxError> { - (self.0)(context).await + self.0.repair(context).await + } +} + +impl From for ToolCallRepair { + fn from(repair: T) -> Self { + Self(Arc::new(repair)) } } diff --git a/aimux-core/tests/tool_input_test.rs b/aimux-core/tests/tool_input_test.rs index 25c843ed6..3aa7c9e5f 100644 --- a/aimux-core/tests/tool_input_test.rs +++ b/aimux-core/tests/tool_input_test.rs @@ -5,7 +5,8 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use aimux_core::error::AiMuxError; use aimux_core::parse_tool_call::{ - ToolCallRepair, ToolCallRepairReply, apply_tool_call_repair, parse_tool_call, + RepairToolCall, ToolCallRepair, ToolCallRepairContext, ToolCallRepairReply, + apply_tool_call_repair, parse_tool_call, }; use aimux_core::tool::{FunctionTool, RawToolCall, Tool}; use serde_json::json; @@ -263,6 +264,46 @@ async fn a_repaired_reply_that_fails_validation_is_not_wrapped() { )); } +/// A struct implementing `RepairToolCall` is used as the hook, like a closure. +#[tokio::test] +async fn repair_hook_can_be_a_trait_implementation() { + struct FixInput { + calls: Arc, + } + + #[async_trait::async_trait] + impl RepairToolCall for FixInput { + async fn repair( + &self, + context: ToolCallRepairContext, + ) -> Result, AiMuxError> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(Some(RawToolCall { + input: r#"{"city":"Singapore","days":3}"#.into(), + ..context.tool_call + })) + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let repair = ToolCallRepair::from(FixInput { + calls: Arc::clone(&calls), + }); + + let call = parse_tool_call( + raw("weather", r#"{"city":"Singapore"#), + Some(&[weather_tool()]), + Some(&repair), + &[], + None, + ) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(call.input, json!({ "city": "Singapore", "days": 3 })); + assert_eq!(call.invalid, None); +} + /// TS: should throw NoSuchToolError when tools is null #[tokio::test] async fn missing_tools_bypasses_repair_like_ai_sdk() {