aimux_core/
parse_tool_call.rs

1//! Parse and validate provider tool calls at the Core boundary.
2//!
3//! Rust port of the AI SDK's `parse-tool-call.ts`
4//! (`packages/ai/src/generate-text/parse-tool-call.ts`). Providers deliver the
5//! model's raw argument text; Core owns JSON parsing, prototype-pollution
6//! rejection, schema validation, and the one-shot repair callback.
7
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14use ts_rs::TS;
15
16use crate::error::AiMuxError;
17use crate::tool::{Tool, ToolCall};
18use crate::types::ProviderMetadata;
19
20/// Provider-facing tool call before Core parses and validates its input.
21///
22/// The wire shape matches [`ToolCall`] field for field, except that `input`
23/// is the provider's raw argument *text* rather than a parsed value.
24#[derive(Debug, Clone, Serialize, Deserialize, TS)]
25#[ts(export)]
26pub struct RawToolCall {
27    pub tool_call_id: String,
28    pub tool_name: String,
29    pub input: String,
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub provider_executed: Option<bool>,
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub dynamic: Option<bool>,
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub thought_signature: Option<String>,
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub provider_metadata: Option<ProviderMetadata>,
38}
39
40/// Context supplied to a one-shot tool-call repair callback.
41#[derive(Debug, Clone)]
42pub struct ToolCallRepairContext {
43    pub instructions: Option<String>,
44    /// Deprecated AI SDK-compatible alias for `instructions`.
45    pub system: Option<String>,
46    pub messages: Vec<crate::message::ModelMessage>,
47    pub tool_call: RawToolCall,
48    pub tools: Vec<Tool>,
49    pub error: AiMuxError,
50}
51
52impl ToolCallRepairContext {
53    /// Return the JSON Schema for a named function tool in this repair step.
54    ///
55    /// Never fails, matching the AI SDK's `inputSchema` repair argument: a
56    /// name that does not resolve to a function tool (unknown — the NoSuchTool
57    /// repair scenario — or a provider tool, which carries no schema at this
58    /// layer) yields the AI SDK's default empty-object schema.
59    #[must_use]
60    pub fn input_schema(&self, tool_name: &str) -> Value {
61        self.tools
62            .iter()
63            .find_map(|tool| match tool {
64                Tool::Function(tool) if tool.name == tool_name => Some(tool.input_schema.clone()),
65                _ => None,
66            })
67            .unwrap_or_else(|| {
68                serde_json::json!({
69                    "type": "object",
70                    "properties": {},
71                    "additionalProperties": false
72                })
73            })
74    }
75}
76
77type ToolCallRepairFuture =
78    Pin<Box<dyn Future<Output = Result<Option<RawToolCall>, AiMuxError>> + Send>>;
79
80/// Async callback that may replace one invalid tool call.
81///
82/// Core invokes this at most once, and parses and validates the returned call
83/// from scratch. Returning `None` keeps the original validation error.
84#[derive(Clone)]
85pub struct ToolCallRepair(Arc<dyn Fn(ToolCallRepairContext) -> ToolCallRepairFuture + Send + Sync>);
86
87impl ToolCallRepair {
88    #[must_use]
89    pub fn new<F, Fut>(repair: F) -> Self
90    where
91        F: Fn(ToolCallRepairContext) -> Fut + Send + Sync + 'static,
92        Fut: Future<Output = Result<Option<RawToolCall>, AiMuxError>> + Send + 'static,
93    {
94        Self(Arc::new(move |context| Box::pin(repair(context))))
95    }
96
97    async fn repair(
98        &self,
99        context: ToolCallRepairContext,
100    ) -> Result<Option<RawToolCall>, AiMuxError> {
101        (self.0)(context).await
102    }
103}
104
105impl std::fmt::Debug for ToolCallRepair {
106    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        formatter.write_str("ToolCallRepair(<callback>)")
108    }
109}
110
111/// Parse and validate a provider tool call using the AI SDK operation contract.
112///
113/// JSON is parsed exactly; partial-JSON repair is deliberately not automatic.
114/// When a tool set was supplied, lookup, parsing, or schema validation failure
115/// gives the optional repair callback one attempt. As in AI SDK, calls made
116/// without a tool set bypass repair. A remaining failure is represented on the
117/// returned call so callers retain both the model output and its typed error.
118pub async fn parse_tool_call(
119    tool_call: RawToolCall,
120    tools: Option<&[Tool]>,
121    repair_tool_call: Option<&ToolCallRepair>,
122    messages: &[crate::message::ModelMessage],
123    instructions: Option<&str>,
124) -> ToolCall {
125    let Some(tools) = tools else {
126        let parsed = if tool_call.provider_executed == Some(true) && tool_call.dynamic == Some(true)
127        {
128            parse_json_input(&tool_call)
129        } else {
130            Err(AiMuxError::NoSuchTool {
131                tool_name: tool_call.tool_name.clone(),
132                available_tools: None,
133                tool_input: Some(tool_call.input.clone()),
134            })
135        };
136        return match parsed {
137            Ok(input) => valid_tool_call(tool_call, input, Some(true)),
138            Err(error) => invalid_tool_call(tool_call, error),
139        };
140    };
141
142    match parse_and_validate_tool_call(&tool_call, tools) {
143        Ok((input, dynamic)) => valid_tool_call(tool_call, input, dynamic),
144        Err(original_error) => {
145            let Some(repair_tool_call) = repair_tool_call else {
146                return invalid_tool_call(tool_call, original_error);
147            };
148            let context = ToolCallRepairContext {
149                instructions: instructions.map(str::to_owned),
150                system: instructions.map(str::to_owned),
151                messages: messages.to_vec(),
152                tool_call: tool_call.clone(),
153                tools: tools.to_vec(),
154                error: original_error.clone(),
155            };
156            let outcome = match repair_tool_call.repair(context).await {
157                Ok(Some(repaired)) => RepairOutcome::Repaired(repaired),
158                Ok(None) => RepairOutcome::Unchanged,
159                Err(repair_error) => RepairOutcome::Failed(repair_error),
160            };
161            apply_repair_outcome(tool_call, original_error, tools, outcome)
162        }
163    }
164}
165
166/// What a repair attempt produced, independent of how it was obtained (the
167/// in-process [`ToolCallRepair`] closure, or a host reply that crossed the
168/// bindings as JSON).
169enum RepairOutcome {
170    Repaired(RawToolCall),
171    Unchanged,
172    Failed(AiMuxError),
173}
174
175/// The AI SDK `parseToolCall` post-repair contract, shared by both paths.
176fn apply_repair_outcome(
177    tool_call: RawToolCall,
178    original_error: AiMuxError,
179    tools: &[Tool],
180    outcome: RepairOutcome,
181) -> ToolCall {
182    let cause = match outcome {
183        RepairOutcome::Repaired(repaired) => match parse_and_validate_tool_call(&repaired, tools) {
184            Ok((input, dynamic)) => return valid_tool_call(repaired, input, dynamic),
185            Err(repaired_error) => repaired_error,
186        },
187        RepairOutcome::Unchanged => return invalid_tool_call(tool_call, original_error),
188        RepairOutcome::Failed(repair_error) => repair_error,
189    };
190    invalid_tool_call(
191        tool_call,
192        AiMuxError::ToolCallRepair {
193            original_error: Box::new(original_error),
194            cause: Box::new(cause),
195        },
196    )
197}
198
199fn parse_and_validate_tool_call(
200    tool_call: &RawToolCall,
201    tools: &[Tool],
202) -> Result<(Value, Option<bool>), AiMuxError> {
203    let tool = tools.iter().find(|tool| match tool {
204        Tool::Function(tool) => tool.name == tool_call.tool_name,
205        Tool::Provider(tool) => tool.name == tool_call.tool_name,
206    });
207
208    let provider_dynamic =
209        tool_call.provider_executed == Some(true) && tool_call.dynamic == Some(true);
210    let Some(tool) = tool else {
211        if provider_dynamic {
212            return parse_json_input(tool_call).map(|input| (input, Some(true)));
213        }
214        return Err(AiMuxError::NoSuchTool {
215            tool_name: tool_call.tool_name.clone(),
216            tool_input: Some(tool_call.input.clone()),
217            available_tools: Some(
218                tools
219                    .iter()
220                    .map(|tool| match tool {
221                        Tool::Function(tool) => tool.name.clone(),
222                        Tool::Provider(tool) => tool.name.clone(),
223                    })
224                    .collect(),
225            ),
226        });
227    };
228
229    let input = parse_json_input(tool_call)?;
230    let Tool::Function(function_tool) = tool else {
231        return Ok((input, None));
232    };
233    let validator = jsonschema::validator_for(&function_tool.input_schema).map_err(|error| {
234        AiMuxError::InvalidToolInput {
235            tool_name: tool_call.tool_name.clone(),
236            tool_input: tool_call.input.clone(),
237            cause: format!("input schema is invalid: {error}"),
238        }
239    })?;
240    // Cause wording matches the AI SDK's `TypeValidationError` template;
241    // `Value` Display is compact JSON, the `JSON.stringify` equivalent.
242    validator
243        .validate(&input)
244        .map_err(|error| AiMuxError::InvalidToolInput {
245            tool_name: tool_call.tool_name.clone(),
246            tool_input: tool_call.input.clone(),
247            cause: format!("Type validation failed: Value: {input}.\nError message: {error}"),
248        })?;
249    Ok((input, None))
250}
251
252fn parse_json_input(tool_call: &RawToolCall) -> Result<Value, AiMuxError> {
253    if tool_call.input.trim().is_empty() {
254        return Ok(Value::Object(serde_json::Map::new()));
255    }
256    // Cause wording matches the AI SDK's `JSONParseError` template.
257    let value: Value =
258        serde_json::from_str(&tool_call.input).map_err(|error| AiMuxError::InvalidToolInput {
259            tool_name: tool_call.tool_name.clone(),
260            tool_input: tool_call.input.clone(),
261            cause: format!(
262                "JSON parsing failed: Text: {}.\nError message: {error}",
263                tool_call.input
264            ),
265        })?;
266    if contains_forbidden_prototype(&value) {
267        return Err(AiMuxError::InvalidToolInput {
268            tool_name: tool_call.tool_name.clone(),
269            tool_input: tool_call.input.clone(),
270            cause: format!(
271                "JSON parsing failed: Text: {}.\nError message: Object contains forbidden prototype property",
272                tool_call.input
273            ),
274        });
275    }
276    Ok(value)
277}
278
279// Port of the AI SDK's secure JSON parse (fastify/secure-json-parse): a
280// `__proto__` key, or a `constructor` object carrying a `prototype` key,
281// anywhere in the tree marks the input invalid. Rust has no prototype
282// pollution, but the parsed value crosses the FFI into JS and Python, and
283// the valid/invalid classification must match upstream.
284fn contains_forbidden_prototype(value: &Value) -> bool {
285    match value {
286        Value::Object(map) => {
287            let constructor_prototype = map
288                .get("constructor")
289                .and_then(Value::as_object)
290                .is_some_and(|constructor| constructor.contains_key("prototype"));
291            if map.contains_key("__proto__") || constructor_prototype {
292                return true;
293            }
294            map.values().any(contains_forbidden_prototype)
295        }
296        Value::Array(items) => items.iter().any(contains_forbidden_prototype),
297        _ => false,
298    }
299}
300
301/// Recover the provider's raw argument text: string inputs pass through
302/// verbatim (possibly malformed JSON awaiting parse/repair); anything already
303/// structured re-serializes.
304pub(crate) fn raw_tool_input(input: &Value) -> String {
305    match input {
306        Value::String(input) => input.clone(),
307        input => serde_json::to_string(input).expect("serializing serde_json::Value cannot fail"),
308    }
309}
310
311fn valid_tool_call(tool_call: RawToolCall, input: Value, dynamic: Option<bool>) -> ToolCall {
312    ToolCall {
313        tool_call_id: tool_call.tool_call_id,
314        tool_name: tool_call.tool_name,
315        input,
316        provider_executed: tool_call.provider_executed,
317        dynamic,
318        thought_signature: tool_call.thought_signature,
319        provider_metadata: tool_call.provider_metadata,
320        invalid: None,
321        error: None,
322    }
323}
324
325fn invalid_tool_call(tool_call: RawToolCall, error: AiMuxError) -> ToolCall {
326    // Best effort, for every failure including `NoSuchTool`, exactly as the
327    // AI SDK's `parseToolCall` catch-all does: the parsed value when the text
328    // is valid JSON, the verbatim text otherwise. Callers replaying an
329    // invalid call rely on this — `response_messages` only carries a
330    // structured input into the next turn's transcript, so an unknown tool
331    // called with perfectly good arguments must still parse here.
332    let input = serde_json::from_str(&tool_call.input)
333        .unwrap_or_else(|_| Value::String(tool_call.input.clone()));
334    ToolCall {
335        tool_call_id: tool_call.tool_call_id,
336        tool_name: tool_call.tool_name,
337        input,
338        provider_executed: tool_call.provider_executed,
339        dynamic: Some(true),
340        thought_signature: tool_call.thought_signature,
341        provider_metadata: tool_call.provider_metadata,
342        invalid: Some(true),
343        error: Some(error),
344    }
345}
346
347// ─────────────────────────────────────────────────────────────────────────────
348// Host-side tool-call repair (RFC-0035)
349// ─────────────────────────────────────────────────────────────────────────────
350
351/// A host's answer to one tool-call repair request.
352///
353/// The bindings cannot carry a [`ToolCallRepair`] closure, so a host instead
354/// reads the invalid call off the result, repairs it in its own language, and
355/// hands the answer back as this value. Wire shape:
356/// `{"type":"repaired","tool_call":{…}}`, `{"type":"unchanged"}`, or
357/// `{"type":"failed","message":"…"}`.
358#[derive(Debug, Clone, Serialize, Deserialize, TS)]
359#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
360#[ts(export)]
361pub enum ToolCallRepairReply {
362    /// Re-validate this replacement call, exactly as the closure path does
363    /// for a returned `Some(call)`.
364    Repaired { tool_call: RawToolCall },
365    /// Keep the original failure — the closure path's `None`.
366    Unchanged,
367    /// The host's repair attempt itself failed. Reported as
368    /// `ToolCallRepair { cause: Other(message) }`, since a host error has no
369    /// typed counterpart on this side.
370    Failed { message: String },
371}
372
373impl From<ToolCallRepairReply> for RepairOutcome {
374    fn from(reply: ToolCallRepairReply) -> Self {
375        match reply {
376            ToolCallRepairReply::Repaired { tool_call } => RepairOutcome::Repaired(tool_call),
377            ToolCallRepairReply::Unchanged => RepairOutcome::Unchanged,
378            ToolCallRepairReply::Failed { message } => {
379                RepairOutcome::Failed(AiMuxError::Other(message))
380            }
381        }
382    }
383}
384
385/// The provider's verbatim argument text, when the error kept it.
386pub(crate) fn raw_tool_call_text(error: &AiMuxError) -> Option<String> {
387    match error {
388        AiMuxError::InvalidToolInput { tool_input, .. } => Some(tool_input.clone()),
389        AiMuxError::NoSuchTool { tool_input, .. } => tool_input.clone(),
390        AiMuxError::ToolCallRepair { original_error, .. } => raw_tool_call_text(original_error),
391        _ => None,
392    }
393}
394
395/// Reconstruct the provider-facing call an invalid [`ToolCall`] was built
396/// from.
397///
398/// The argument text comes from the error, which keeps it verbatim.
399/// [`invalid_tool_call`] parses the text into `input` when it is valid JSON,
400/// and that loses the JSON string quoting of a bare literal (`"Tokyo"` would
401/// come back as `Tokyo`), duplicate keys, number spelling and whitespace, all
402/// of which matter when a host re-submits the text unchanged under a new tool
403/// name. Only errors from before the text was recorded fall back to
404/// [`raw_tool_input`].
405///
406/// `dynamic` is not recovered: an invalid call always reports `Some(true)`.
407/// That only matters for re-deriving the original error, which is read off the
408/// call instead.
409fn raw_from_invalid(tool_call: &ToolCall) -> Result<(RawToolCall, AiMuxError), AiMuxError> {
410    if tool_call.invalid != Some(true) {
411        return Err(AiMuxError::InvalidArgument(format!(
412            "tool call '{}' is not invalid; nothing to repair",
413            tool_call.tool_call_id
414        )));
415    }
416    let error = tool_call.error.clone().ok_or_else(|| {
417        AiMuxError::InvalidArgument(format!(
418            "invalid tool call '{}' carries no error",
419            tool_call.tool_call_id
420        ))
421    })?;
422    Ok((
423        RawToolCall {
424            tool_call_id: tool_call.tool_call_id.clone(),
425            tool_name: tool_call.tool_name.clone(),
426            input: raw_tool_call_text(&error).unwrap_or_else(|| raw_tool_input(&tool_call.input)),
427            provider_executed: tool_call.provider_executed,
428            dynamic: tool_call.dynamic,
429            thought_signature: tool_call.thought_signature.clone(),
430            provider_metadata: tool_call.provider_metadata.clone(),
431        },
432        error,
433    ))
434}
435
436/// The AI SDK's "no tool set, no repair" rule, shared by every entry point
437/// that resolves a reply: a call made without tools was never a candidate for
438/// repair, so reaching this point with `None` is a caller mistake.
439fn required_tools(tools: Option<&[Tool]>) -> Result<&[Tool], AiMuxError> {
440    tools.ok_or_else(|| {
441        AiMuxError::InvalidArgument(
442            "tool calls made without a tool set are not repairable".to_string(),
443        )
444    })
445}
446
447/// Recover the three repair inputs from the same `prompt` and `options` a
448/// caller already passed to `generate_text` / `stream_text`.
449///
450/// Exists so no binding has to re-derive the transcript the model saw: the
451/// prompt-to-messages rule and the "instructions travel beside the messages"
452/// rule are `generate_text`'s, and stay in one place.
453#[must_use]
454pub fn tool_call_repair_inputs(
455    prompt: impl Into<crate::message::ModelPrompt>,
456    options: &crate::generate::GenerateTextOptions,
457) -> (
458    Vec<crate::message::ModelMessage>,
459    Option<&str>,
460    Option<&[Tool]>,
461) {
462    let (messages, instructions) =
463        crate::generate::split_prompt(prompt.into(), options.instructions.as_deref());
464    (messages, instructions, options.tools.as_deref())
465}
466
467/// Build the repair argument a host needs, as JSON, or `None` when the call
468/// was made without a tool set.
469///
470/// Mirrors the AI SDK `repairToolCall` argument: `{tool_call, error,
471/// input_schema, tools, messages, instructions}`. `tool_call` is the
472/// provider-facing call (raw argument *text*), `input_schema` the schema of
473/// the named tool (an empty-object schema when the name does not resolve),
474/// and `error` the serialized failure the call already carries.
475///
476/// `tools` is `Option` for the same reason [`parse_tool_call`]'s is: the AI
477/// SDK never repairs a call made without a tool set. That is not a failure —
478/// it is the answer "this call is not repairable", and the caller skips it.
479///
480/// # Errors
481///
482/// [`AiMuxError::InvalidArgument`] when `tool_call` is not an invalid call, or
483/// carries no error.
484pub fn tool_call_repair_context(
485    tool_call: &ToolCall,
486    tools: Option<&[Tool]>,
487    messages: &[crate::message::ModelMessage],
488    instructions: Option<&str>,
489) -> Result<Option<Value>, AiMuxError> {
490    let Some(tools) = tools else {
491        return Ok(None);
492    };
493    let (raw, error) = raw_from_invalid(tool_call)?;
494    let context = ToolCallRepairContext {
495        instructions: instructions.map(str::to_owned),
496        system: instructions.map(str::to_owned),
497        messages: messages.to_vec(),
498        tool_call: raw,
499        tools: tools.to_vec(),
500        error,
501    };
502    Ok(Some(serde_json::json!({
503        "tool_call": context.tool_call,
504        "error": context.error,
505        "input_schema": context.input_schema(&context.tool_call.tool_name),
506        "tools": context.tools,
507        "messages": context.messages,
508        "instructions": context.instructions,
509    })))
510}
511
512/// Resolve one invalid tool call against a host's repair reply.
513///
514/// Produces the same [`ToolCall`] the in-process [`ToolCallRepair`] closure
515/// would have produced for the equivalent return value; both paths share
516/// the same post-repair logic.
517///
518/// # Errors
519///
520/// [`AiMuxError::InvalidArgument`] when `tools` is `None` (the call was never
521/// repairable — [`tool_call_repair_context`] already answered `None` for it),
522/// or when `tool_call` is not an invalid call, or carries no error.
523pub fn apply_tool_call_repair(
524    tool_call: &ToolCall,
525    tools: Option<&[Tool]>,
526    reply: ToolCallRepairReply,
527) -> Result<ToolCall, AiMuxError> {
528    let tools = required_tools(tools)?;
529    let (raw, error) = raw_from_invalid(tool_call)?;
530    Ok(apply_repair_outcome(raw, error, tools, reply.into()))
531}
532
533/// Apply a repair reply to a serialized `GenerateTextResult` or
534/// `GenerateObjectResult`, or `StreamTextResultAggregated`, returning the
535/// patched result.
536///
537/// Both `tool_calls[]` and the matching `response_messages[]` tool-call part
538/// are rewritten, the latter under the same rule `to_response_messages` uses
539/// (malformed primitive input is not replayed as a prompt input). A
540/// `GenerateObjectResult` is recognised by its nested `raw` result; `object`
541/// is untouched, because it is parsed from the model's text, not from a tool
542/// call.
543///
544/// The OpenAI-shaped `ChatCompletion` is deliberately not supported: it drops
545/// `invalid` and `error`, so a host cannot tell from it that a call needs
546/// repairing in the first place. Repair the native result and convert.
547///
548/// # Errors
549///
550/// [`AiMuxError::InvalidArgument`] when `tools` is `None`, when the document
551/// has no `tool_calls` array, when `tool_call_id` is missing or ambiguous,
552/// when the matched call is not invalid, or when a replacement id conflicts
553/// with another call.
554pub fn apply_tool_call_repair_to_result(
555    result: &Value,
556    tools: Option<&[Tool]>,
557    tool_call_id: &str,
558    reply: ToolCallRepairReply,
559) -> Result<Value, AiMuxError> {
560    let mut patched = result.clone();
561    // `GenerateObjectResult` carries the whole text result under `raw`.
562    let target = if patched.get("tool_calls").is_some() {
563        &mut patched
564    } else {
565        patched.get_mut("raw").ok_or_else(|| {
566            AiMuxError::InvalidArgument("result: neither tool_calls nor raw".to_string())
567        })?
568    };
569
570    let calls = target
571        .get_mut("tool_calls")
572        .and_then(Value::as_array_mut)
573        .ok_or_else(|| AiMuxError::InvalidArgument("result: no tool_calls array".to_string()))?;
574    // Both `tool_calls` and the transcript are keyed by id here, so a
575    // duplicated id cannot be patched unambiguously; refuse rather than
576    // rewrite the first entry twice.
577    let target_index = calls
578        .iter()
579        .position(|call| call.get("tool_call_id").and_then(Value::as_str) == Some(tool_call_id))
580        .ok_or_else(|| {
581            AiMuxError::InvalidArgument(format!("result: no tool call '{tool_call_id}'"))
582        })?;
583    if calls[target_index + 1..]
584        .iter()
585        .any(|call| call.get("tool_call_id").and_then(Value::as_str) == Some(tool_call_id))
586    {
587        return Err(AiMuxError::InvalidArgument(format!(
588            "result: tool call id '{tool_call_id}' is not unique"
589        )));
590    }
591    let original: ToolCall = serde_json::from_value(calls[target_index].clone())
592        .map_err(|error| AiMuxError::InvalidArgument(format!("result.tool_calls: {error}")))?;
593
594    let repaired = apply_tool_call_repair(&original, tools, reply)?;
595    if repaired.tool_call_id != tool_call_id
596        && calls.iter().enumerate().any(|(index, call)| {
597            index != target_index
598                && call.get("tool_call_id").and_then(Value::as_str)
599                    == Some(repaired.tool_call_id.as_str())
600        })
601    {
602        return Err(AiMuxError::InvalidArgument(format!(
603            "result: repaired tool call id '{}' already exists",
604            repaired.tool_call_id
605        )));
606    }
607    calls[target_index] = serde_json::to_value(&repaired)
608        .map_err(|error| AiMuxError::InvalidArgument(format!("tool call: {error}")))?;
609
610    // The replayed transcript must agree with `tool_calls`, or the next turn
611    // sends the model the unrepaired arguments.
612    let replay_input =
613        crate::response_messages::response_tool_call_input(&repaired.input, repaired.invalid);
614    for part in target
615        .get_mut("response_messages")
616        .and_then(Value::as_array_mut)
617        .into_iter()
618        .flatten()
619        .filter_map(|message| message.get_mut("content")?.as_array_mut())
620        .flatten()
621    {
622        if part.get("type").and_then(Value::as_str) != Some("tool_call")
623            || part.get("tool_call_id").and_then(Value::as_str) != Some(tool_call_id)
624        {
625            continue;
626        }
627        let Some(part) = part.as_object_mut() else {
628            continue;
629        };
630        // A repair may rename the call; the transcript must keep pointing at
631        // the same entry as `tool_calls`.
632        part.insert(
633            "tool_call_id".to_string(),
634            repaired.tool_call_id.clone().into(),
635        );
636        part.insert("tool_name".to_string(), repaired.tool_name.clone().into());
637        part.insert("input".to_string(), replay_input.clone());
638    }
639
640    Ok(patched)
641}
aimux_core/parse_tool_call.rs

aimux_core/
parse_tool_call.rs
