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#[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#[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 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(¶m.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(¶m)],
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(¶m)
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}