Skip to main content

agentic_core/tool/
custom.rs

1use std::collections::HashMap;
2
3use serde_json::{Map, Value};
4
5use crate::events::WireEvent;
6use crate::types::io::{CustomToolCall, FunctionTool, FunctionToolCall, OutputItem, ToolChoice};
7use crate::types::tools::{CustomToolParam, ResponsesTool};
8use crate::utils::common::serialize_to_value_or_custom_default;
9
10use super::{ToolEntry, ToolError, ToolHandler, ToolType};
11
12/// Request-scoped mapping from normalized function names to their original
13/// public custom-tool declarations.
14#[derive(Debug, Default)]
15pub(crate) struct CustomToolMap {
16    declarations: HashMap<String, CustomToolParam>,
17}
18
19impl CustomToolMap {
20    fn from_tools(tools: &[ResponsesTool]) -> Option<Self> {
21        let declarations = tools
22            .iter()
23            .filter_map(|tool| match tool {
24                ResponsesTool::Custom(param) => Some((param.name.as_str().to_owned(), param.clone())),
25                _ => None,
26            })
27            .collect::<HashMap<_, _>>();
28        (!declarations.is_empty()).then_some(Self { declarations })
29    }
30
31    fn declaration(&self, name: &str) -> Option<&CustomToolParam> {
32        self.declarations.get(name)
33    }
34}
35
36/// Handler for client-owned `type: "custom"` tools.
37///
38/// Custom tools are normalized for the model but are executed by the client,
39/// so this intentionally implements [`ToolHandler`] without
40/// [`super::GatewayExecutor`].
41#[derive(Debug)]
42pub struct CustomHandler;
43
44impl CustomHandler {
45    #[must_use]
46    pub(crate) fn build_tool_map(tools: &[ResponsesTool]) -> Option<CustomToolMap> {
47        CustomToolMap::from_tools(tools)
48    }
49
50    pub(crate) fn validate_tool_choice(
51        tools: Option<&[ResponsesTool]>,
52        tool_choice: &ToolChoice,
53    ) -> Result<(), ToolError> {
54        let map = tools.and_then(CustomToolMap::from_tools);
55        match tool_choice {
56            ToolChoice::Custom { name } => validate_custom_selector(map.as_ref(), name.as_str()),
57            ToolChoice::AllowedTools { tools, .. } => {
58                for tool in tools {
59                    if tool.type_.as_str() == "custom" {
60                        validate_custom_selector(map.as_ref(), tool.name.as_str())?;
61                    }
62                }
63                Ok(())
64            }
65            _ => Ok(()),
66        }
67    }
68
69    #[must_use]
70    pub fn to_function_call(param: &CustomToolParam) -> FunctionTool {
71        FunctionTool {
72            type_: "function".to_owned(),
73            name: param.name.as_str().to_owned(),
74            description: Some(model_visible_description(param)),
75            parameters: Some(serde_json::json!({
76                "type": "object",
77                "properties": {
78                    "input": {
79                        "type": "string",
80                        "description": "Raw custom tool input. Follow the tool description and declared format exactly."
81                    }
82                },
83                "required": ["input"],
84                "additionalProperties": false
85            })),
86            strict: Some(true),
87        }
88    }
89
90    #[must_use]
91    pub(crate) fn output_item(call: &FunctionToolCall) -> OutputItem {
92        OutputItem::CustomToolCall(CustomToolCall {
93            id: public_item_id(&call.id),
94            status: Some(call.status),
95            call_id: call.call_id.clone(),
96            name: call.name.clone(),
97            input: input_from_arguments(&call.arguments),
98        })
99    }
100
101    /// Restores normalized custom-tool declarations in response lifecycle
102    /// metadata before the event is emitted to the client.
103    pub(crate) fn restore_response_wire(wire: &mut WireEvent, map: Option<&CustomToolMap>) -> bool {
104        let Some(map) = map else {
105            return false;
106        };
107        restore_response_map(&mut wire.rest, map)
108    }
109}
110
111fn validate_custom_selector(map: Option<&CustomToolMap>, name: &str) -> Result<(), ToolError> {
112    if map.and_then(|map| map.declaration(name)).is_some() {
113        return Ok(());
114    }
115    Err(ToolError::Config(format!(
116        "tool_choice selects custom tool '{name}', but no matching custom tool is declared"
117    )))
118}
119
120fn restore_response_map(object: &mut Map<String, Value>, map: &CustomToolMap) -> bool {
121    let mut changed = restore_response_metadata(object, map);
122    for key in ["response", "payload"] {
123        if let Some(nested) = object.get_mut(key).and_then(Value::as_object_mut) {
124            changed |= restore_response_map(nested, map);
125        }
126    }
127    changed
128}
129
130fn restore_response_metadata(object: &mut Map<String, Value>, map: &CustomToolMap) -> bool {
131    let mut changed = false;
132    if let Some(tools) = object.get_mut("tools").and_then(Value::as_array_mut) {
133        for tool in tools {
134            changed |= restore_custom_declaration(tool, map);
135        }
136    }
137    if let Some(tool_choice) = object.get_mut("tool_choice") {
138        changed |= restore_custom_tool_choice(tool_choice, map);
139    }
140    changed
141}
142
143fn restore_custom_declaration(tool: &mut Value, map: &CustomToolMap) -> bool {
144    let Some(name) = normalized_custom_name(tool, map) else {
145        return false;
146    };
147    let Some(param) = map.declaration(&name) else {
148        return false;
149    };
150    let Some(mut declaration) =
151        serialize_to_value_or_custom_default(param, "custom tool metadata serialization failed", Some, None)
152    else {
153        return false;
154    };
155    let Some(object) = declaration.as_object_mut() else {
156        return false;
157    };
158    object.insert("type".to_owned(), Value::String("custom".to_owned()));
159    *tool = declaration;
160    true
161}
162
163fn restore_custom_tool_choice(choice: &mut Value, map: &CustomToolMap) -> bool {
164    let Some(object) = choice.as_object_mut() else {
165        return false;
166    };
167    if object.get("type").and_then(Value::as_str) == Some("allowed_tools") {
168        let Some(tools) = object.get_mut("tools").and_then(Value::as_array_mut) else {
169            return false;
170        };
171        return tools
172            .iter_mut()
173            .map(|tool| restore_custom_choice_type(tool, map))
174            .fold(false, |changed, restored| changed | restored);
175    }
176    restore_custom_choice_type(choice, map)
177}
178
179fn restore_custom_choice_type(choice: &mut Value, map: &CustomToolMap) -> bool {
180    if normalized_custom_name(choice, map).is_none() {
181        return false;
182    }
183    let Some(object) = choice.as_object_mut() else {
184        return false;
185    };
186    object.insert("type".to_owned(), Value::String("custom".to_owned()));
187    object.remove("namespace");
188    true
189}
190
191fn normalized_custom_name(value: &Value, map: &CustomToolMap) -> Option<String> {
192    let object = value.as_object()?;
193    if object.get("type").and_then(Value::as_str) != Some("function") {
194        return None;
195    }
196    let name = object.get("name")?.as_str()?;
197    map.declaration(name).map(|_| name.to_owned())
198}
199
200fn model_visible_description(param: &CustomToolParam) -> String {
201    let mut fragments = Vec::new();
202    if let Some(description) = param
203        .description
204        .as_deref()
205        .map(str::trim)
206        .filter(|value| !value.is_empty())
207    {
208        fragments.push(description.to_owned());
209    }
210
211    fragments.push("Provide the raw tool input in the `input` string field.".to_owned());
212
213    if !param.extra.is_empty()
214        && let Ok(extra) = serde_json::to_string(&param.extra)
215    {
216        fragments.push(format!(
217            "Additional custom tool declaration fields that must be respected:\n{extra}"
218        ));
219    }
220
221    fragments.join("\n\n")
222}
223
224impl ToolHandler for CustomHandler {
225    fn tool_type(&self) -> ToolType {
226        ToolType::Custom
227    }
228
229    fn validate(&self, param: &serde_json::Value) -> Result<(), ToolError> {
230        let param = serde_json::from_value::<CustomToolParam>(param.clone())
231            .map_err(|error| ToolError::Config(format!("invalid custom tool config: {error}")))?;
232        if param
233            .format
234            .as_ref()
235            .is_some_and(|format| format.get("type").and_then(Value::as_str) != Some("text"))
236        {
237            return Err(ToolError::Config(format!(
238                "custom tool '{}' uses an unsupported format; gateway normalization cannot preserve constrained decoding",
239                param.name
240            )));
241        }
242        Ok(())
243    }
244
245    fn normalize(&self, param: &serde_json::Value) -> Vec<FunctionTool> {
246        match serde_json::from_value::<CustomToolParam>(param.clone()) {
247            Ok(param) => vec![Self::to_function_call(&param)],
248            Err(error) => {
249                tracing::warn!(%error, "invalid custom tool param");
250                Vec::new()
251            }
252        }
253    }
254}
255
256pub(crate) fn insert_custom_entry(entries: &mut HashMap<String, ToolEntry>, param: &CustomToolParam) {
257    serialize_to_value_or_custom_default(
258        param,
259        "custom tool config serialization failed",
260        |config| {
261            entries.insert(
262                param.name.as_str().to_owned(),
263                ToolEntry {
264                    tool_type: ToolType::Custom,
265                    config,
266                    server_label: None,
267                    handler: None,
268                },
269            );
270        },
271        (),
272    );
273}
274
275pub(crate) fn public_item_id(item_id: &str) -> String {
276    if item_id.starts_with("ctc_") {
277        return item_id.to_owned();
278    }
279    if let Some(suffix) = item_id.strip_prefix("fc_").filter(|suffix| !suffix.is_empty()) {
280        return format!("ctc_{suffix}");
281    }
282    format!("ctc_{:016x}", stable_name_hash(item_id))
283}
284
285fn stable_name_hash(value: &str) -> u64 {
286    value.as_bytes().iter().fold(0xcbf2_9ce4_8422_2325_u64, |hash, byte| {
287        (hash ^ u64::from(*byte)).wrapping_mul(0x0000_0100_0000_01b3)
288    })
289}
290
291pub(crate) fn input_from_arguments(arguments: &str) -> String {
292    try_input_from_arguments(arguments).unwrap_or_else(|| {
293        tracing::debug!(
294            argument_bytes = arguments.len(),
295            "custom tool arguments did not match the normalized input envelope; forwarding raw arguments"
296        );
297        arguments.to_owned()
298    })
299}
300
301pub(crate) fn try_input_from_arguments(arguments: &str) -> Option<String> {
302    match serde_json::from_str::<serde_json::Value>(arguments).ok()? {
303        serde_json::Value::String(input) => Some(input),
304        serde_json::Value::Object(fields) => fields
305            .get("input")
306            .and_then(serde_json::Value::as_str)
307            .map(str::to_owned),
308        _ => None,
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use crate::types::event::MessageStatus;
316
317    #[test]
318    fn function_fallback_uses_public_custom_tool_shape() {
319        let call = FunctionToolCall {
320            id: "fc_1".to_owned(),
321            call_id: "call_1".to_owned(),
322            name: "raw_echo".to_owned(),
323            namespace: None,
324            arguments: r#"{"input":"hello"}"#.to_owned(),
325            status: MessageStatus::Completed,
326        };
327
328        let OutputItem::CustomToolCall(completed) = CustomHandler::output_item(&call) else {
329            panic!("expected custom output item");
330        };
331        assert_eq!(completed.id, "ctc_1");
332        assert_eq!(completed.input, "hello");
333        assert_eq!(completed.status, Some(MessageStatus::Completed));
334    }
335
336    #[test]
337    fn custom_call_id_is_stable_for_every_source_item_id() {
338        assert_eq!(public_item_id("fc_item"), "ctc_item");
339        assert_eq!(public_item_id("ctc_item"), "ctc_item");
340        assert_eq!(public_item_id("provider_item"), public_item_id("provider_item"));
341    }
342
343    #[test]
344    fn custom_declaration_normalizes_to_function_with_raw_input() {
345        let param = serde_json::from_value::<CustomToolParam>(serde_json::json!({
346            "name": "raw_echo",
347            "description": "Echo raw input.",
348            "x-provider-field": {"mode": "strict"}
349        }))
350        .expect("custom tool");
351
352        let value = serde_json::to_value(param).expect("custom tool value");
353        let mut tools = CustomHandler.normalize(&value);
354        let tool = tools.pop().expect("normalized custom tool");
355
356        assert_eq!(tool.type_, "function");
357        assert_eq!(tool.name, "raw_echo");
358        assert_eq!(
359            tool.parameters.as_ref().unwrap()["properties"]["input"]["type"],
360            "string"
361        );
362        assert_eq!(tool.parameters.as_ref().unwrap()["required"][0], "input");
363        let description = tool.description.as_deref().expect("model-visible description");
364        assert!(description.contains("Echo raw input."));
365        assert!(description.contains("raw tool input in the `input` string field"));
366        assert!(description.contains("x-provider-field"));
367        assert!(description.contains("strict"));
368    }
369
370    #[test]
371    fn grammar_formats_are_rejected() {
372        for syntax in ["lark", "regex"] {
373            let param = serde_json::from_value::<CustomToolParam>(serde_json::json!({
374                "name": "constrained_input",
375                "format": {
376                    "type": "grammar",
377                    "syntax": syntax,
378                    "definition": "start: value"
379                }
380            }))
381            .expect("custom tool");
382
383            let value = serde_json::to_value(param).expect("custom tool value");
384            let error = CustomHandler.validate(&value).expect_err("grammar must be rejected");
385            assert!(error.to_string().contains("cannot preserve constrained decoding"));
386        }
387    }
388
389    #[test]
390    fn explicit_text_format_is_supported() {
391        let param = serde_json::json!({
392            "name": "freeform",
393            "format": {"type": "text"}
394        });
395
396        CustomHandler
397            .validate(&param)
398            .expect("unconstrained text is representable");
399    }
400
401    #[test]
402    fn response_lifecycle_metadata_restores_public_custom_tool_shape() {
403        let param = serde_json::from_value::<CustomToolParam>(serde_json::json!({
404            "name": "raw_echo",
405            "description": "Echo raw input."
406        }))
407        .expect("custom tool");
408        let tools = vec![ResponsesTool::Custom(param)];
409        let map = CustomHandler::build_tool_map(&tools);
410        let mut wire = WireEvent::new("response.created");
411        wire.rest.insert(
412            "response".to_owned(),
413            serde_json::json!({
414                "tools": [{
415                    "type": "function",
416                    "name": "raw_echo",
417                    "description": "normalized description",
418                    "parameters": {"type": "object"}
419                }],
420                "tool_choice": {"type": "function", "name": "raw_echo"}
421            }),
422        );
423
424        assert!(CustomHandler::restore_response_wire(&mut wire, map.as_ref()));
425        let response = &wire.rest["response"];
426        assert_eq!(response["tools"][0]["type"], "custom");
427        assert_eq!(response["tools"][0]["description"], "Echo raw input.");
428        assert!(response["tools"][0].get("parameters").is_none());
429        assert_eq!(response["tool_choice"]["type"], "custom");
430        assert_eq!(response["tool_choice"]["name"], "raw_echo");
431    }
432
433    #[test]
434    fn allowed_tools_metadata_restores_custom_selector_type() {
435        let param = serde_json::from_value::<CustomToolParam>(serde_json::json!({
436            "name": "raw_echo"
437        }))
438        .expect("custom tool");
439        let tools = vec![ResponsesTool::Custom(param)];
440        let map = CustomHandler::build_tool_map(&tools);
441        let mut wire = WireEvent::new("response.in_progress");
442        wire.rest.insert(
443            "response".to_owned(),
444            serde_json::json!({
445                "tool_choice": {
446                    "type": "allowed_tools",
447                    "mode": "required",
448                    "tools": [
449                        {"type": "function", "name": "ordinary"},
450                        {"type": "function", "name": "raw_echo"}
451                    ]
452                }
453            }),
454        );
455
456        assert!(CustomHandler::restore_response_wire(&mut wire, map.as_ref()));
457        let tools = &wire.rest["response"]["tool_choice"]["tools"];
458        assert_eq!(tools[0]["type"], "function");
459        assert_eq!(tools[1]["type"], "custom");
460    }
461}