1#[cfg(feature = "derive")]
154pub use aither_derive::tool;
155use alloc::borrow::Cow;
156use serde_json::Value;
157
158use crate::Result;
159use alloc::format;
160use alloc::string::{String, ToString};
161use alloc::vec::Vec;
162use alloc::{boxed::Box, collections::BTreeMap};
163use core::any::Any;
164use core::fmt::{Debug, Display};
165use core::{future::Future, pin::Pin};
166pub use mime::Mime;
167use schemars::{JsonSchema, Schema, schema_for};
168use serde::{Serialize, de::DeserializeOwned};
169
170#[derive(Debug, Clone, PartialEq, Eq)]
189#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
190#[cfg_attr(feature = "serde", serde(tag = "kind", rename_all = "snake_case"))]
191pub enum ToolResult {
192 Done,
194
195 Text {
197 text: String,
199 },
200
201 Tsv {
203 text: String,
205 },
206
207 Json {
209 value: Value,
211 },
212
213 Binary {
215 mime: String,
217 content: Vec<u8>,
219 },
220
221 Error {
223 message: String,
225 },
226
227 Parts {
234 parts: Vec<ToolResultPart>,
236 },
237}
238
239#[derive(Debug, Clone, PartialEq, Eq)]
244#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
245#[cfg_attr(feature = "serde", serde(tag = "kind", rename_all = "snake_case"))]
246pub enum ToolResultPart {
247 Text {
249 text: String,
251 },
252
253 Tsv {
255 text: String,
257 },
258
259 Json {
261 value: Value,
263 },
264
265 Binary {
267 mime: String,
269 content: Vec<u8>,
271 },
272}
273
274impl ToolResult {
275 #[must_use]
277 pub fn text(s: impl Into<String>) -> Self {
278 Self::Text { text: s.into() }
279 }
280
281 #[must_use]
283 pub fn tsv(s: impl Into<String>) -> Self {
284 Self::Tsv { text: s.into() }
285 }
286
287 pub fn json<T: Serialize>(value: &T) -> Result<Self> {
293 Ok(Self::Json {
294 value: serde_json::to_value(value)?,
295 })
296 }
297
298 #[must_use]
300 pub const fn json_value(value: Value) -> Self {
301 Self::Json { value }
302 }
303
304 #[must_use]
306 pub fn image(data: Vec<u8>, media_type: &str) -> Self {
307 Self::Binary {
308 mime: parse_media_type_or_octet_stream(media_type),
309 content: data,
310 }
311 }
312
313 #[must_use]
315 pub fn binary(data: Vec<u8>) -> Self {
316 Self::Binary {
317 mime: mime::APPLICATION_OCTET_STREAM.essence_str().to_string(),
318 content: data,
319 }
320 }
321
322 #[must_use]
324 pub fn error(message: impl Into<String>) -> Self {
325 Self::Error {
326 message: message.into(),
327 }
328 }
329
330 #[must_use]
336 pub fn parts(parts: Vec<ToolResultPart>) -> Self {
337 match <[ToolResultPart; 1]>::try_from(parts) {
338 Ok([part]) => part.into_result(),
339 Err(parts) if parts.is_empty() => Self::Done,
340 Err(parts) => Self::Parts { parts },
341 }
342 }
343
344 #[must_use]
346 pub const fn is_done(&self) -> bool {
347 matches!(self, Self::Done)
348 }
349
350 #[must_use]
352 pub const fn is_error(&self) -> bool {
353 matches!(self, Self::Error { .. })
354 }
355
356 #[must_use]
358 pub fn as_text(&self) -> Option<&str> {
359 match self {
360 Self::Text { text } | Self::Tsv { text } => Some(text),
361 Self::Error { message } => Some(message),
362 Self::Done | Self::Json { .. } | Self::Binary { .. } | Self::Parts { .. } => None,
363 }
364 }
365
366 #[must_use]
368 pub fn error_message(&self) -> Option<&str> {
369 match self {
370 Self::Error { message } => Some(message),
371 Self::Done
372 | Self::Text { .. }
373 | Self::Tsv { .. }
374 | Self::Json { .. }
375 | Self::Binary { .. }
376 | Self::Parts { .. } => None,
377 }
378 }
379
380 pub fn render_for_model(&self) -> Result<String> {
386 match self {
387 Self::Done => Ok(String::new()),
388 Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
389 Self::Json { value } => Ok(serde_json::to_string(value)?),
390 Self::Binary { mime, content } => Ok(render_binary_placeholder(mime, content)),
391 Self::Error { message } => Ok(message.clone()),
392 Self::Parts { parts } => join_part_renders(parts, false),
393 }
394 }
395
396 pub fn render_for_cli(&self) -> Result<String> {
402 match self {
403 Self::Done => Ok(String::new()),
404 Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
405 Self::Json { value } => Ok(serde_json::to_string_pretty(value)?),
406 Self::Binary { mime, content } => Ok(render_binary_placeholder(mime, content)),
407 Self::Error { message } => Ok(message.clone()),
408 Self::Parts { parts } => join_part_renders(parts, true),
409 }
410 }
411
412 #[must_use]
417 pub fn mime(&self) -> Option<Mime> {
418 match self {
419 Self::Binary { mime, .. } => mime.parse().ok(),
420 Self::Parts { parts } => match parts.as_slice() {
421 [ToolResultPart::Binary { mime, .. }] => mime.parse().ok(),
422 _ => None,
423 },
424 Self::Done
425 | Self::Text { .. }
426 | Self::Tsv { .. }
427 | Self::Json { .. }
428 | Self::Error { .. } => None,
429 }
430 }
431
432 #[must_use]
436 pub fn content(&self) -> Option<&[u8]> {
437 match self {
438 Self::Binary { content, .. } => Some(content),
439 Self::Parts { parts } => match parts.as_slice() {
440 [ToolResultPart::Binary { content, .. }] => Some(content),
441 _ => None,
442 },
443 Self::Done
444 | Self::Text { .. }
445 | Self::Tsv { .. }
446 | Self::Json { .. }
447 | Self::Error { .. } => None,
448 }
449 }
450}
451
452impl ToolResultPart {
453 #[must_use]
455 pub fn text(s: impl Into<String>) -> Self {
456 Self::Text { text: s.into() }
457 }
458
459 #[must_use]
461 pub fn tsv(s: impl Into<String>) -> Self {
462 Self::Tsv { text: s.into() }
463 }
464
465 pub fn json<T: Serialize>(value: &T) -> Result<Self> {
471 Ok(Self::Json {
472 value: serde_json::to_value(value)?,
473 })
474 }
475
476 #[must_use]
478 pub const fn json_value(value: Value) -> Self {
479 Self::Json { value }
480 }
481
482 #[must_use]
484 pub fn image(data: Vec<u8>, media_type: &str) -> Self {
485 Self::Binary {
486 mime: parse_media_type_or_octet_stream(media_type),
487 content: data,
488 }
489 }
490
491 #[must_use]
493 pub fn binary(data: Vec<u8>) -> Self {
494 Self::Binary {
495 mime: mime::APPLICATION_OCTET_STREAM.essence_str().to_string(),
496 content: data,
497 }
498 }
499
500 #[must_use]
502 pub fn as_text(&self) -> Option<&str> {
503 match self {
504 Self::Text { text } | Self::Tsv { text } => Some(text),
505 Self::Json { .. } | Self::Binary { .. } => None,
506 }
507 }
508
509 pub fn render_for_model(&self) -> Result<String> {
515 self.render(false)
516 }
517
518 pub fn render_for_cli(&self) -> Result<String> {
524 self.render(true)
525 }
526
527 fn render(&self, pretty: bool) -> Result<String> {
528 match self {
529 Self::Text { text } | Self::Tsv { text } => Ok(text.clone()),
530 Self::Json { value } => Ok(if pretty {
531 serde_json::to_string_pretty(value)?
532 } else {
533 serde_json::to_string(value)?
534 }),
535 Self::Binary { mime, content } => Ok(render_binary_placeholder(mime, content)),
536 }
537 }
538
539 #[must_use]
541 pub fn mime(&self) -> Option<Mime> {
542 match self {
543 Self::Binary { mime, .. } => mime.parse().ok(),
544 Self::Text { .. } | Self::Tsv { .. } | Self::Json { .. } => None,
545 }
546 }
547
548 #[must_use]
550 pub fn content(&self) -> Option<&[u8]> {
551 match self {
552 Self::Binary { content, .. } => Some(content),
553 Self::Text { .. } | Self::Tsv { .. } | Self::Json { .. } => None,
554 }
555 }
556
557 #[must_use]
559 pub fn into_result(self) -> ToolResult {
560 match self {
561 Self::Text { text } | Self::Tsv { text } => ToolResult::Text { text },
562 Self::Json { value } => ToolResult::Json { value },
563 Self::Binary { mime, content } => ToolResult::Binary { mime, content },
564 }
565 }
566}
567
568fn render_binary_placeholder(mime: &str, content: &[u8]) -> String {
569 let mut rendered = String::new();
570 rendered.push_str("[binary tool result: ");
571 rendered.push_str(mime);
572 rendered.push_str(", ");
573 rendered.push_str(content.len().to_string().as_str());
574 rendered.push_str(" bytes]");
575 rendered
576}
577
578fn join_part_renders(parts: &[ToolResultPart], pretty: bool) -> Result<String> {
579 let mut rendered = String::new();
580 for part in parts {
581 let text = if pretty {
582 part.render_for_cli()?
583 } else {
584 part.render_for_model()?
585 };
586 if text.is_empty() {
587 continue;
588 }
589 if !rendered.is_empty() {
590 rendered.push('\n');
591 }
592 rendered.push_str(&text);
593 }
594 Ok(rendered)
595}
596
597pub trait IntoToolResult {
603 fn into_tool_result(self) -> Result<ToolResult>;
609}
610
611impl IntoToolResult for ToolResult {
612 fn into_tool_result(self) -> Result<ToolResult> {
613 Ok(self)
614 }
615}
616
617impl IntoToolResult for () {
618 fn into_tool_result(self) -> Result<ToolResult> {
619 Ok(ToolResult::Done)
620 }
621}
622
623impl IntoToolResult for String {
624 fn into_tool_result(self) -> Result<ToolResult> {
625 Ok(ToolResult::text(self))
626 }
627}
628
629impl IntoToolResult for &str {
630 fn into_tool_result(self) -> Result<ToolResult> {
631 Ok(ToolResult::text(self))
632 }
633}
634
635impl IntoToolResult for Cow<'_, str> {
636 fn into_tool_result(self) -> Result<ToolResult> {
637 Ok(ToolResult::text(self.into_owned()))
638 }
639}
640
641impl IntoToolResult for Value {
642 fn into_tool_result(self) -> Result<ToolResult> {
643 Ok(ToolResult::json_value(self))
644 }
645}
646
647impl<T> IntoToolResult for Option<T>
648where
649 T: IntoToolResult,
650{
651 fn into_tool_result(self) -> Result<ToolResult> {
652 self.map_or_else(|| Ok(ToolResult::Done), IntoToolResult::into_tool_result)
653 }
654}
655
656impl IntoToolResult for Vec<ToolResultPart> {
657 fn into_tool_result(self) -> Result<ToolResult> {
658 Ok(ToolResult::parts(self))
659 }
660}
661
662impl<T, E> IntoToolResult for core::result::Result<T, E>
663where
664 T: Serialize,
665 E: Display,
666{
667 fn into_tool_result(self) -> Result<ToolResult> {
668 match self {
669 Ok(value) => serialize_success_value(&value),
670 Err(error) => Ok(ToolResult::error(error.to_string())),
671 }
672 }
673}
674
675fn parse_media_type_or_octet_stream(media_type: &str) -> String {
676 media_type
677 .parse::<Mime>()
678 .unwrap_or(mime::APPLICATION_OCTET_STREAM)
679 .essence_str()
680 .to_string()
681}
682
683fn serialize_success_value<T: Serialize>(value: &T) -> Result<ToolResult> {
684 let value = serde_json::to_value(value)?;
685 if let Some(tsv) = json_value_to_tsv(&value) {
686 return Ok(ToolResult::tsv(tsv));
687 }
688
689 match value {
690 Value::String(text) => Ok(ToolResult::text(text)),
691 other => Ok(ToolResult::json_value(other)),
692 }
693}
694
695#[must_use]
697pub fn json_value_to_tsv(value: &Value) -> Option<String> {
698 let rows = match value {
699 Value::Array(arr) if !arr.is_empty() => arr
700 .iter()
701 .map(|value| flatten_json_value(value, ""))
702 .collect::<Vec<_>>(),
703 Value::Object(_) => alloc::vec![flatten_json_value(value, "")],
704 Value::Array(_) | Value::String(_) | Value::Number(_) | Value::Bool(_) | Value::Null => {
705 return None;
706 }
707 };
708
709 if rows.is_empty() {
710 return None;
711 }
712
713 let mut columns: Vec<String> = Vec::new();
714 let mut seen: alloc::collections::BTreeSet<String> = alloc::collections::BTreeSet::new();
715 for row in &rows {
716 for (key, _) in row {
717 if seen.insert(key.clone()) {
718 columns.push(key.clone());
719 }
720 }
721 }
722
723 if columns.is_empty() {
724 return None;
725 }
726
727 let mut tsv = String::new();
728 for (index, column) in columns.iter().enumerate() {
729 if index > 0 {
730 tsv.push('\t');
731 }
732 tsv.push_str(&escape_tsv_field(column));
733 }
734 tsv.push('\n');
735
736 for row in &rows {
737 let row_map: alloc::collections::BTreeMap<&str, &str> = row
738 .iter()
739 .map(|(key, value)| (key.as_str(), value.as_str()))
740 .collect::<alloc::collections::BTreeMap<&str, &str>>();
741 for (index, column) in columns.iter().enumerate() {
742 if index > 0 {
743 tsv.push('\t');
744 }
745 if let Some(value) = row_map.get(column.as_str()) {
746 tsv.push_str(&escape_tsv_field(value));
747 }
748 }
749 tsv.push('\n');
750 }
751
752 Some(tsv)
753}
754
755fn flatten_json_value(value: &Value, prefix: &str) -> Vec<(String, String)> {
756 let mut flattened = Vec::new();
757 match value {
758 Value::Object(map) => {
759 for (key, child) in map {
760 let full_key = if prefix.is_empty() {
761 key.clone()
762 } else {
763 format!("{prefix}.{key}")
764 };
765 flattened.extend(flatten_json_value(child, &full_key));
766 }
767 }
768 Value::Array(_) => {
769 let serialized = serde_json::to_string(value).unwrap_or_default();
770 flattened.push((prefix.to_string(), serialized));
771 }
772 Value::String(text) => {
773 flattened.push((prefix.to_string(), text.clone()));
774 }
775 Value::Number(number) => {
776 flattened.push((prefix.to_string(), number.to_string()));
777 }
778 Value::Bool(boolean) => {
779 flattened.push((prefix.to_string(), boolean.to_string()));
780 }
781 Value::Null => {
782 flattened.push((prefix.to_string(), String::new()));
783 }
784 }
785 flattened
786}
787
788fn escape_tsv_field(value: &str) -> String {
789 value.replace(['\t', '\n', '\r'], " ")
790}
791
792pub trait Tool: Send + Sync {
832 fn name(&self) -> Cow<'static, str>;
834
835 fn description(&self) -> Cow<'static, str> {
845 description_from_schema::<Self::Arguments>().unwrap_or_default()
846 }
847
848 type Arguments: Send + JsonSchema + DeserializeOwned;
851
852 type Res: IntoToolResult + Send;
854
855 fn call(&self, arguments: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send;
861}
862
863pub fn json<T: Serialize>(value: &T) -> Result<String> {
887 let value = serde_json::to_value(value)?;
888
889 Ok(value
890 .as_str()
891 .map_or_else(|| format!("{value:#}"), ToString::to_string))
892}
893
894trait ToolImpl: Send + Sync + Any {
895 fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>>;
896
897 fn definition(&self) -> &ToolDefinition;
900
901 fn as_any(&self) -> &dyn Any;
906
907 fn as_any_mut(&mut self) -> &mut dyn Any;
909}
910
911struct DynToolImpl<F>
913where
914 F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>> + Send + Sync,
915{
916 definition: ToolDefinition,
917 handler: F,
918}
919
920impl<F> ToolImpl for DynToolImpl<F>
921where
922 F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>> + Send + Sync + 'static,
923{
924 fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>> {
925 (self.handler)(args)
926 }
927
928 fn definition(&self) -> &ToolDefinition {
929 &self.definition
930 }
931
932 fn as_any(&self) -> &dyn Any {
933 self
934 }
935
936 fn as_any_mut(&mut self) -> &mut dyn Any {
937 self
938 }
939}
940
941fn schema_is_object(value: &Value) -> bool {
944 matches!(value.get("type").and_then(Value::as_str), Some("object"))
945 || value.get("properties").is_some()
946 || value.get("oneOf").is_some()
947 || value.get("anyOf").is_some()
948 || value.get("$defs").is_some()
949}
950
951fn is_object<T: JsonSchema>() -> bool {
952 schema_is_object(&schema_for!(T).to_value())
953}
954
955fn arguments_schema<T: JsonSchema>() -> Schema {
958 if is_object::<T>() {
959 schema_for!(T)
960 } else {
961 schema_for!(ToolArgument<T>)
962 }
963}
964
965fn description_from_schema<T: JsonSchema>() -> Option<Cow<'static, str>> {
967 schema_for!(T)
968 .to_value()
969 .get("description")
970 .and_then(Value::as_str)
971 .filter(|text| !text.trim().is_empty())
972 .map(|text| Cow::Owned(text.to_string()))
973}
974
975struct RegisteredTool<T: Tool> {
981 tool: T,
982 definition: ToolDefinition,
983 args_are_object: bool,
984}
985
986impl<T: Tool> RegisteredTool<T> {
987 fn new(tool: T) -> Self {
988 let definition = ToolDefinition::new(&tool);
989 let args_are_object = is_object::<T::Arguments>();
990 Self {
991 tool,
992 definition,
993 args_are_object,
994 }
995 }
996}
997
998impl<T: Tool + 'static> ToolImpl for RegisteredTool<T> {
999 fn call(&self, args: &str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send + '_>> {
1000 let result = if self.args_are_object {
1001 serde_json::from_str::<T::Arguments>(args)
1002 } else {
1003 serde_json::from_str::<ToolArgument<T::Arguments>>(args).map(|wrapper| wrapper.value)
1004 };
1005
1006 let Ok(arguments) = result else {
1007 let name = self.definition.name().to_string();
1010 let schema_str =
1011 serde_json::to_string_pretty(&self.definition.arguments_openai_schema())
1012 .unwrap_or_else(|_| "{}".to_string());
1013 return Box::pin(async move {
1014 Err(anyhow::Error::msg(format!(
1015 "Invalid arguments for tool '{name}'. Expected schema:\n{schema_str}"
1016 )))
1017 });
1018 };
1019
1020 Box::pin(async move { Tool::call(&self.tool, arguments).await?.into_tool_result() })
1021 }
1022
1023 fn definition(&self) -> &ToolDefinition {
1024 &self.definition
1025 }
1026
1027 fn as_any(&self) -> &dyn Any {
1028 &self.tool
1029 }
1030
1031 fn as_any_mut(&mut self) -> &mut dyn Any {
1032 &mut self.tool
1033 }
1034}
1035
1036#[derive(Debug, Clone, PartialEq, Eq)]
1041pub struct InvalidSchema {
1042 name: Cow<'static, str>,
1044}
1045
1046impl InvalidSchema {
1047 #[must_use]
1049 pub fn name(&self) -> &str {
1050 &self.name
1051 }
1052}
1053
1054impl Display for InvalidSchema {
1055 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1056 write!(
1057 f,
1058 "tool '{}' has an argument schema that is neither an object nor a boolean",
1059 self.name
1060 )
1061 }
1062}
1063
1064impl core::error::Error for InvalidSchema {}
1065
1066#[derive(Debug, Clone, PartialEq, Eq)]
1068pub enum RegisterError {
1069 DuplicateName(Cow<'static, str>),
1073
1074 EmptyDescription(Cow<'static, str>),
1079}
1080
1081impl Display for RegisterError {
1082 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1083 match self {
1084 Self::DuplicateName(name) => {
1085 write!(f, "a tool named '{name}' is already registered")
1086 }
1087 Self::EmptyDescription(name) => write!(
1088 f,
1089 "tool '{name}' has an empty description; add a rustdoc comment to its \
1090 Arguments type or implement Tool::description"
1091 ),
1092 }
1093 }
1094}
1095
1096impl core::error::Error for RegisterError {}
1097
1098pub struct Tools {
1112 tools: BTreeMap<Cow<'static, str>, Box<dyn ToolImpl>>,
1113}
1114
1115impl Debug for Tools {
1116 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1117 f.debug_struct("Tools")
1118 .field("tools", &self.tools.keys().collect::<Vec<_>>())
1119 .finish()
1120 }
1121}
1122
1123#[derive(Debug, Clone, PartialEq)]
1127#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
1128pub struct ToolDefinition {
1129 name: Cow<'static, str>,
1131 description: Cow<'static, str>,
1133 arguments: Schema,
1135}
1136
1137impl ToolDefinition {
1138 #[must_use]
1143 pub fn new<T: Tool>(tool: &T) -> Self {
1144 Self {
1145 name: tool.name(),
1146 description: tool.description(),
1147 arguments: arguments_schema::<T::Arguments>(),
1148 }
1149 }
1150
1151 pub fn from_parts(
1162 name: Cow<'static, str>,
1163 description: Cow<'static, str>,
1164 schema: Value,
1165 ) -> core::result::Result<Self, InvalidSchema> {
1166 let arguments: Schema = schema
1167 .try_into()
1168 .map_err(|_| InvalidSchema { name: name.clone() })?;
1169
1170 Ok(Self {
1171 name,
1172 description,
1173 arguments,
1174 })
1175 }
1176
1177 #[must_use]
1179 pub fn name(&self) -> &str {
1180 &self.name
1181 }
1182
1183 #[must_use]
1185 pub fn description(&self) -> &str {
1186 &self.description
1187 }
1188
1189 #[must_use]
1193 pub fn arguments_openai_schema(&self) -> serde_json::Value {
1194 let mut inner = self.arguments.clone().to_value();
1195 clean_schema(&mut inner);
1196
1197 inner
1198 }
1199}
1200
1201#[derive(Debug, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
1202struct ToolArgument<T> {
1203 value: T,
1204}
1205
1206fn clean_schema(value: &mut Value) {
1207 let defs = extract_defs(value);
1209
1210 resolve_and_clean(value, &defs);
1212
1213 if let Value::Object(map) = value {
1215 map.remove("description");
1218
1219 if map.contains_key("properties") && !map.contains_key("type") {
1221 map.insert("type".to_string(), Value::String("object".to_string()));
1222 }
1223 }
1224}
1225
1226fn extract_defs(value: &Value) -> serde_json::Map<String, Value> {
1228 if let Value::Object(map) = value
1229 && let Some(Value::Object(defs)) = map.get("$defs").or_else(|| map.get("definitions"))
1230 {
1231 return defs.clone();
1232 }
1233 serde_json::Map::new()
1234}
1235
1236#[allow(clippy::too_many_lines)]
1238fn resolve_and_clean(value: &mut Value, defs: &serde_json::Map<String, Value>) {
1239 resolve_and_clean_inner(value, defs, false);
1240}
1241
1242#[allow(clippy::too_many_lines)]
1244fn resolve_and_clean_inner(
1245 value: &mut Value,
1246 defs: &serde_json::Map<String, Value>,
1247 inside_properties: bool,
1248) {
1249 match value {
1250 Value::Object(map) => {
1251 if let Some(Value::String(ref_path)) = map.remove("$ref")
1253 && let Some(Value::Object(resolved_map)) = resolve_ref(&ref_path, defs)
1254 {
1255 let existing_description = map.remove("description");
1258 for (k, v) in resolved_map {
1259 map.entry(k).or_insert(v);
1260 }
1261 if let Some(desc) = existing_description {
1263 map.insert("description".to_string(), desc);
1264 }
1265 }
1266
1267 if let Some(const_val) = map.remove("const") {
1269 map.insert("enum".to_string(), Value::Array(alloc::vec![const_val]));
1270 }
1271
1272 if let Some(Value::Array(variants)) =
1274 map.remove("oneOf").or_else(|| map.remove("anyOf"))
1275 {
1276 let is_simple_enum = variants.iter().all(|v| {
1278 if let Value::Object(vm) = v {
1279 (vm.contains_key("const") || vm.contains_key("enum"))
1280 && !vm.contains_key("properties")
1281 } else {
1282 false
1283 }
1284 });
1285
1286 if is_simple_enum {
1287 let mut enum_values: alloc::vec::Vec<Value> = alloc::vec::Vec::new();
1289 let mut variant_type: Option<String> = None;
1290
1291 for variant in &variants {
1292 if let Value::Object(vm) = variant {
1293 if let Some(const_val) = vm.get("const")
1294 && !enum_values.contains(const_val)
1295 {
1296 enum_values.push(const_val.clone());
1297 }
1298 if let Some(Value::Array(arr)) = vm.get("enum") {
1299 for val in arr {
1300 if !enum_values.contains(val) {
1301 enum_values.push(val.clone());
1302 }
1303 }
1304 }
1305 if variant_type.is_none()
1306 && let Some(Value::String(t)) = vm.get("type")
1307 {
1308 variant_type = Some(t.clone());
1309 }
1310 }
1311 }
1312
1313 if !enum_values.is_empty() {
1314 map.insert("enum".to_string(), Value::Array(enum_values));
1315 if let Some(t) = variant_type {
1316 map.insert("type".to_string(), Value::String(t));
1317 }
1318 }
1319 } else {
1320 let mut all_properties = serde_json::Map::new();
1322
1323 for variant in variants {
1324 if let Value::Object(variant_map) = variant
1325 && let Some(Value::Object(props)) = variant_map.get("properties")
1326 {
1327 for (key, val) in props {
1328 let new_values: Option<alloc::vec::Vec<Value>> =
1330 if let Value::Object(val_obj) = val {
1331 if let Some(Value::Array(arr)) = val_obj.get("enum") {
1332 Some(arr.clone())
1333 } else {
1334 val_obj
1335 .get("const")
1336 .map(|const_val| alloc::vec![const_val.clone()])
1337 }
1338 } else {
1339 None
1340 };
1341
1342 if all_properties.contains_key(key) {
1343 if let Some(values) = new_values
1345 && let Some(Value::Object(existing_obj)) =
1346 all_properties.get_mut(key)
1347 && let Some(Value::Array(existing_enum)) =
1348 existing_obj.get_mut("enum")
1349 {
1350 for e in values {
1351 if !existing_enum.contains(&e) {
1352 existing_enum.push(e);
1353 }
1354 }
1355 }
1356 } else {
1357 let mut val_clone = val.clone();
1359 if let Value::Object(obj) = &mut val_clone
1360 && let Some(const_val) = obj.remove("const")
1361 {
1362 obj.insert(
1363 "enum".to_string(),
1364 Value::Array(alloc::vec![const_val]),
1365 );
1366 }
1367 all_properties.insert(key.clone(), val_clone);
1368 }
1369 }
1370 }
1371 }
1372
1373 if !all_properties.is_empty() {
1375 map.insert("type".to_string(), Value::String("object".to_string()));
1376 map.insert("properties".to_string(), Value::Object(all_properties));
1377 }
1378 }
1379 }
1380
1381 if !inside_properties {
1384 let allowed = [
1385 "type",
1386 "description",
1387 "properties",
1388 "required",
1389 "items",
1390 "enum",
1391 "nullable",
1392 ];
1393 map.retain(|k, _| allowed.contains(&k.as_str()));
1394 }
1395
1396 if let Some(Value::Array(types)) = map.get("type") {
1398 let non_null: Vec<&Value> = types
1400 .iter()
1401 .filter(|t| !matches!(t, Value::String(s) if s == "null"))
1402 .collect();
1403 if non_null.len() == 1 {
1404 map.insert("type".to_string(), non_null[0].clone());
1405 }
1406 }
1407
1408 for (key, v) in map.iter_mut() {
1410 let child_inside_props = key == "properties";
1412 resolve_and_clean_inner(v, defs, child_inside_props);
1413 }
1414 }
1415 Value::Array(arr) => {
1416 for v in arr {
1417 resolve_and_clean_inner(v, defs, false);
1418 }
1419 }
1420 _ => {}
1421 }
1422}
1423
1424fn resolve_ref(ref_path: &str, defs: &serde_json::Map<String, Value>) -> Option<Value> {
1426 let name = ref_path
1428 .strip_prefix("#/$defs/")
1429 .or_else(|| ref_path.strip_prefix("#/definitions/"))?;
1430
1431 defs.get(name).cloned()
1432}
1433
1434impl Default for Tools {
1435 fn default() -> Self {
1436 Self::new()
1437 }
1438}
1439
1440impl Tools {
1441 #[must_use]
1443 pub const fn new() -> Self {
1444 Self {
1445 tools: BTreeMap::new(),
1446 }
1447 }
1448
1449 #[must_use]
1453 pub fn get<T>(&self) -> Option<&T>
1454 where
1455 T: Tool + 'static,
1456 {
1457 self.tools
1458 .values()
1459 .find_map(|tool| tool.as_any().downcast_ref::<T>())
1460 }
1461
1462 #[must_use]
1466 pub fn get_mut<T>(&mut self) -> Option<&mut T>
1467 where
1468 T: Tool + 'static,
1469 {
1470 self.tools
1471 .values_mut()
1472 .find_map(|tool| tool.as_any_mut().downcast_mut::<T>())
1473 }
1474
1475 #[must_use]
1477 pub fn definitions(&self) -> Vec<ToolDefinition> {
1478 self.tools
1479 .values()
1480 .map(|tool| tool.definition().clone())
1481 .collect()
1482 }
1483
1484 pub fn register<T: Tool + 'static>(
1495 &mut self,
1496 tool: T,
1497 ) -> core::result::Result<(), RegisterError> {
1498 self.insert(Box::new(RegisteredTool::new(tool)))
1499 }
1500
1501 pub fn register_dyn<F>(
1510 &mut self,
1511 definition: ToolDefinition,
1512 handler: F,
1513 ) -> core::result::Result<(), RegisterError>
1514 where
1515 F: Fn(&str) -> Pin<Box<dyn Future<Output = Result<ToolResult>> + Send>>
1516 + Send
1517 + Sync
1518 + 'static,
1519 {
1520 self.insert(Box::new(DynToolImpl {
1521 definition,
1522 handler,
1523 }))
1524 }
1525
1526 fn insert(&mut self, tool: Box<dyn ToolImpl>) -> core::result::Result<(), RegisterError> {
1527 let name = tool.definition().name.clone();
1528 if self.tools.contains_key(&name) {
1529 return Err(RegisterError::DuplicateName(name));
1530 }
1531 if tool.definition().description().trim().is_empty() {
1532 return Err(RegisterError::EmptyDescription(name));
1533 }
1534 self.tools.insert(name, tool);
1535 Ok(())
1536 }
1537
1538 pub fn unregister(&mut self, name: &str) {
1540 self.tools.remove(name);
1541 }
1542
1543 pub async fn call(&self, name: &str, args: &str) -> Result<ToolResult> {
1550 if let Some(tool) = self.tools.get(name) {
1551 tool.call(args).await
1552 } else {
1553 Err(anyhow::Error::msg(format!("Tool '{name}' not found")))
1554 }
1555 }
1556}
1557
1558#[cfg(test)]
1559mod tests {
1560 use super::*;
1561 use alloc::{format, string::ToString, vec};
1562 use schemars::JsonSchema;
1563 use serde::{Deserialize, Serialize};
1564
1565 #[derive(JsonSchema, Deserialize, Debug, PartialEq)]
1567 struct CalculatorArgs {
1568 operation: String,
1569 a: f64,
1570 b: f64,
1571 }
1572
1573 struct Calculator;
1574
1575 impl Tool for Calculator {
1576 fn name(&self) -> Cow<'static, str> {
1577 "calculator".into()
1578 }
1579 type Arguments = CalculatorArgs;
1580 type Res = ToolResult;
1581
1582 fn call(&self, args: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send {
1583 core::future::ready(match args.operation.as_str() {
1584 "add" => Ok(ToolResult::text((args.a + args.b).to_string())),
1585 "subtract" => Ok(ToolResult::text((args.a - args.b).to_string())),
1586 "multiply" => Ok(ToolResult::text((args.a * args.b).to_string())),
1587 "divide" => {
1588 if args.b == 0.0 {
1589 Err(anyhow::Error::msg("Division by zero"))
1590 } else {
1591 Ok(ToolResult::text((args.a / args.b).to_string()))
1592 }
1593 }
1594 _ => Err(anyhow::Error::msg(format!(
1595 "Unknown operation: {}",
1596 args.operation
1597 ))),
1598 })
1599 }
1600 }
1601
1602 #[derive(JsonSchema, Deserialize)]
1604 struct GreetArgs {
1605 name: String,
1606 }
1607
1608 struct Greeter;
1609
1610 impl Tool for Greeter {
1611 fn name(&self) -> Cow<'static, str> {
1612 "greeter".into()
1613 }
1614 type Arguments = GreetArgs;
1615 type Res = ToolResult;
1616
1617 fn call(&self, args: Self::Arguments) -> impl Future<Output = Result<Self::Res>> + Send {
1618 core::future::ready(Ok(ToolResult::text(format!("Hello, {}!", args.name))))
1619 }
1620 }
1621
1622 #[test]
1623 fn from_parts_accepts_object_and_boolean_schemas() {
1624 for schema in [
1626 serde_json::json!({"type": "object"}),
1627 serde_json::json!(true),
1628 ] {
1629 assert!(
1630 ToolDefinition::from_parts("t".into(), "does a thing".into(), schema.clone())
1631 .is_ok(),
1632 "{schema} should be accepted"
1633 );
1634 }
1635 }
1636
1637 #[test]
1638 fn from_parts_rejects_non_schema_values() {
1639 for schema in [
1642 serde_json::json!("a string"),
1643 serde_json::json!([1, 2, 3]),
1644 serde_json::json!(7),
1645 serde_json::json!(null),
1646 ] {
1647 let result = ToolDefinition::from_parts("weird".into(), "d".into(), schema.clone());
1648 let Err(err) = result else {
1649 panic!("{schema} should be rejected");
1650 };
1651 assert_eq!(err.name(), "weird");
1652 }
1653 }
1654
1655 #[test]
1656 fn json_utility() {
1657 let value = serde_json::json!({
1658 "name": "test",
1659 "value": 42
1660 });
1661
1662 let json_str = json(&value).expect("a JSON value always serializes");
1663 assert!(json_str.contains("\"name\": \"test\""));
1664 assert!(json_str.contains("\"value\": 42"));
1665 }
1666
1667 #[test]
1668 fn tool_definition_creation() {
1669 let calculator = Calculator;
1670 let definition = ToolDefinition::new(&calculator);
1671
1672 assert_eq!(definition.name, "calculator");
1673 assert_eq!(
1674 definition.description,
1675 "Performs basic mathematical operations."
1676 );
1677 }
1680
1681 #[test]
1682 fn tools_creation() {
1683 let tools = Tools::new();
1684 assert_eq!(tools.definitions().len(), 0);
1685 }
1686
1687 #[test]
1688 fn tools_default() {
1689 let tools = Tools::default();
1690 assert_eq!(tools.definitions().len(), 0);
1691 }
1692
1693 #[tokio::test]
1694 async fn tools_register_and_call() {
1695 let mut tools = Tools::new();
1696 tools.register(Calculator).expect("calculator registers");
1697
1698 let definitions = tools.definitions();
1699 assert_eq!(definitions.len(), 1);
1700 assert_eq!(definitions[0].name, "calculator");
1701
1702 let result = tools
1703 .call("calculator", r#"{"operation": "add", "a": 5, "b": 3}"#)
1704 .await;
1705 assert!(result.is_ok());
1706 assert_eq!(result.unwrap().as_text(), Some("8"));
1707 }
1708
1709 #[tokio::test]
1710 async fn calculator_operations() {
1711 let mut tools = Tools::new();
1712 tools.register(Calculator).expect("calculator registers");
1713
1714 let result = tools
1716 .call("calculator", r#"{"operation": "add", "a": 10, "b": 5}"#)
1717 .await;
1718 assert_eq!(result.unwrap().as_text(), Some("15"));
1719
1720 let result = tools
1722 .call(
1723 "calculator",
1724 r#"{"operation": "subtract", "a": 10, "b": 3}"#,
1725 )
1726 .await;
1727 assert_eq!(result.unwrap().as_text(), Some("7"));
1728
1729 let result = tools
1731 .call("calculator", r#"{"operation": "multiply", "a": 4, "b": 3}"#)
1732 .await;
1733 assert_eq!(result.unwrap().as_text(), Some("12"));
1734
1735 let result = tools
1737 .call("calculator", r#"{"operation": "divide", "a": 15, "b": 3}"#)
1738 .await;
1739 assert_eq!(result.unwrap().as_text(), Some("5"));
1740 }
1741
1742 #[tokio::test]
1743 async fn calculator_division_by_zero() {
1744 let mut tools = Tools::new();
1745 tools.register(Calculator).expect("calculator registers");
1746
1747 let result = tools
1748 .call("calculator", r#"{"operation": "divide", "a": 10, "b": 0}"#)
1749 .await;
1750 assert!(result.is_err());
1751 assert!(result.unwrap_err().to_string().contains("Division by zero"));
1752 }
1753
1754 #[tokio::test]
1755 async fn calculator_unknown_operation() {
1756 let mut tools = Tools::new();
1757 tools.register(Calculator).expect("calculator registers");
1758
1759 let result = tools
1760 .call("calculator", r#"{"operation": "modulo", "a": 10, "b": 3}"#)
1761 .await;
1762 assert!(result.is_err());
1763 assert!(
1764 result
1765 .unwrap_err()
1766 .to_string()
1767 .contains("Unknown operation")
1768 );
1769 }
1770
1771 #[tokio::test]
1772 async fn multiple_tools() {
1773 let mut tools = Tools::new();
1774 tools.register(Calculator).expect("calculator registers");
1775 tools.register(Greeter).expect("greeter registers");
1776
1777 let definitions = tools.definitions();
1778 assert_eq!(definitions.len(), 2);
1779
1780 let calc_def = definitions.iter().find(|d| d.name == "calculator").unwrap();
1782 let greet_def = definitions.iter().find(|d| d.name == "greeter").unwrap();
1783
1784 assert_eq!(
1785 calc_def.description,
1786 "Performs basic mathematical operations."
1787 );
1788 assert_eq!(greet_def.description, "Greets a person by name.");
1789
1790 let calc_result = tools
1792 .call("calculator", r#"{"operation": "add", "a": 2, "b": 3}"#)
1793 .await;
1794 assert_eq!(calc_result.unwrap().as_text(), Some("5"));
1795
1796 let greet_result = tools.call("greeter", r#"{"name": "Alice"}"#).await;
1797 assert_eq!(greet_result.unwrap().as_text(), Some("Hello, Alice!"));
1798 }
1799
1800 #[derive(Debug, Serialize)]
1801 struct TableRow {
1802 name: &'static str,
1803 count: u32,
1804 }
1805
1806 #[derive(Debug, Serialize)]
1807 struct NestedTableRow {
1808 user: TableRow,
1809 ok: bool,
1810 }
1811
1812 #[derive(Debug)]
1813 struct ToolFailure(&'static str);
1814
1815 impl core::fmt::Display for ToolFailure {
1816 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1817 f.write_str(self.0)
1818 }
1819 }
1820
1821 impl core::error::Error for ToolFailure {}
1822
1823 #[test]
1824 fn into_tool_result_string_is_plain_text() {
1825 assert_eq!(
1826 String::from("hello").into_tool_result().unwrap(),
1827 ToolResult::text("hello")
1828 );
1829 }
1830
1831 #[test]
1832 fn into_tool_result_str_is_plain_text() {
1833 assert_eq!(
1834 "hello".into_tool_result().unwrap(),
1835 ToolResult::text("hello")
1836 );
1837 }
1838
1839 #[test]
1840 fn into_tool_result_option_none_is_done() {
1841 let result = Option::<String>::None.into_tool_result().unwrap();
1842 assert_eq!(result, ToolResult::Done);
1843 }
1844
1845 #[test]
1846 fn into_tool_result_option_some_delegates() {
1847 let result = Some("hello").into_tool_result().unwrap();
1848 assert_eq!(result, ToolResult::text("hello"));
1849 }
1850
1851 #[test]
1852 fn into_tool_result_result_ok_string_is_text() {
1853 let result = core::result::Result::<String, ToolFailure>::Ok(String::from("hello"))
1854 .into_tool_result()
1855 .unwrap();
1856 assert_eq!(result, ToolResult::text("hello"));
1857 }
1858
1859 #[test]
1860 fn into_tool_result_result_ok_object_is_tsv() {
1861 let result = core::result::Result::<TableRow, ToolFailure>::Ok(TableRow {
1862 name: "alpha",
1863 count: 3,
1864 })
1865 .into_tool_result()
1866 .unwrap();
1867
1868 assert_eq!(result, ToolResult::tsv("count\tname\n3\talpha\n"));
1869 }
1870
1871 #[test]
1872 fn into_tool_result_result_ok_array_of_objects_is_tsv() {
1873 let result = core::result::Result::<Vec<TableRow>, ToolFailure>::Ok(vec![
1874 TableRow {
1875 name: "alpha",
1876 count: 3,
1877 },
1878 TableRow {
1879 name: "beta",
1880 count: 5,
1881 },
1882 ])
1883 .into_tool_result()
1884 .unwrap();
1885
1886 assert_eq!(result, ToolResult::tsv("count\tname\n3\talpha\n5\tbeta\n"));
1887 }
1888
1889 #[test]
1890 fn into_tool_result_result_ok_scalar_is_json() {
1891 let result = core::result::Result::<bool, ToolFailure>::Ok(true)
1892 .into_tool_result()
1893 .unwrap();
1894 assert_eq!(result, ToolResult::json_value(Value::Bool(true)));
1895 }
1896
1897 #[test]
1898 fn into_tool_result_result_err_is_typed_error() {
1899 let result = core::result::Result::<TableRow, ToolFailure>::Err(ToolFailure("boom"))
1900 .into_tool_result()
1901 .unwrap();
1902 assert_eq!(result, ToolResult::error("boom"));
1903 assert!(result.is_error());
1904 assert_eq!(result.error_message(), Some("boom"));
1905 }
1906
1907 #[test]
1908 fn json_value_to_tsv_flattens_nested_objects() {
1909 let value = serde_json::to_value(NestedTableRow {
1910 user: TableRow {
1911 name: "alpha",
1912 count: 3,
1913 },
1914 ok: true,
1915 })
1916 .unwrap();
1917
1918 assert_eq!(
1919 json_value_to_tsv(&value),
1920 Some("ok\tuser.count\tuser.name\ntrue\t3\talpha\n".to_string())
1921 );
1922 }
1923
1924 #[tokio::test]
1925 async fn tool_not_found() {
1926 let tools = Tools::new();
1927
1928 let result = tools.call("nonexistent", "{}").await;
1929 assert!(result.is_err());
1930 assert!(
1931 result
1932 .unwrap_err()
1933 .to_string()
1934 .contains("Tool 'nonexistent' not found")
1935 );
1936 }
1937
1938 #[tokio::test]
1939 async fn invalid_json() {
1940 let mut tools = Tools::new();
1941 tools.register(Calculator).expect("calculator registers");
1942
1943 let result = tools.call("calculator", "invalid json").await;
1944 assert!(result.is_err());
1945 }
1946
1947 #[test]
1948 fn tools_unregister() {
1949 let mut tools = Tools::new();
1950 tools.register(Calculator).expect("calculator registers");
1951 tools.register(Greeter).expect("greeter registers");
1952
1953 assert_eq!(tools.definitions().len(), 2);
1954
1955 tools.unregister("calculator");
1956 assert_eq!(tools.definitions().len(), 1);
1957
1958 let remaining = &tools.definitions()[0];
1959 assert_eq!(remaining.name, "greeter");
1960
1961 tools.unregister("greeter");
1962 assert_eq!(tools.definitions().len(), 0);
1963 }
1964
1965 #[test]
1966 fn tools_debug() {
1967 let mut tools = Tools::new();
1968 tools.register(Calculator).expect("calculator registers");
1969 tools.register(Greeter).expect("greeter registers");
1970
1971 let debug_str = format!("{tools:?}");
1972 assert!(debug_str.contains("Tools"));
1973 assert!(debug_str.contains("calculator"));
1974 assert!(debug_str.contains("greeter"));
1975 }
1976
1977 #[test]
1978 fn tool_definition_debug() {
1979 let calculator = Calculator;
1980 let definition = ToolDefinition::new(&calculator);
1981 let debug_str = format!("{definition:?}");
1982
1983 assert!(debug_str.contains("ToolDefinition"));
1984 assert!(debug_str.contains("calculator"));
1985 assert!(debug_str.contains("Performs basic mathematical operations"));
1986 }
1987
1988 #[test]
1989 fn tool_definition_clone() {
1990 let calculator = Calculator;
1991 let original = ToolDefinition::new(&calculator);
1992 let cloned = original.clone();
1993
1994 assert_eq!(original.name, cloned.name);
1995 assert_eq!(original.description, cloned.description);
1996 }
1997
1998 #[test]
1999 fn schema_preserves_enum() {
2000 #[derive(JsonSchema, Deserialize)]
2001 #[serde(rename_all = "snake_case")]
2002 enum Status {
2003 Pending,
2004 InProgress,
2005 Completed,
2006 }
2007
2008 #[allow(dead_code)]
2009 #[derive(JsonSchema, Deserialize)]
2010 struct Item {
2011 status: Status,
2012 }
2013
2014 #[allow(dead_code)]
2015 #[derive(JsonSchema, Deserialize)]
2016 struct Args {
2017 items: Vec<Item>,
2018 }
2019
2020 struct TestTool;
2021
2022 impl Tool for TestTool {
2023 fn name(&self) -> Cow<'static, str> {
2024 "test".into()
2025 }
2026 type Arguments = Args;
2027 type Res = ToolResult;
2028
2029 fn call(
2030 &self,
2031 _args: Self::Arguments,
2032 ) -> impl Future<Output = Result<Self::Res>> + Send {
2033 core::future::ready(Ok(ToolResult::text("ok")))
2034 }
2035 }
2036
2037 let tool = TestTool;
2038 let def = ToolDefinition::new(&tool);
2039 let schema = def.arguments_openai_schema();
2040
2041 let schema_obj = schema.as_object().expect("schema should be object");
2043 let properties = schema_obj
2044 .get("properties")
2045 .expect("should have properties")
2046 .as_object()
2047 .unwrap();
2048 let items = properties
2049 .get("items")
2050 .expect("should have items")
2051 .as_object()
2052 .unwrap();
2053 let item_props = items
2054 .get("items")
2055 .expect("items should have items schema")
2056 .as_object()
2057 .unwrap();
2058 let item_properties = item_props
2059 .get("properties")
2060 .expect("item should have properties")
2061 .as_object()
2062 .unwrap();
2063 let status = item_properties
2064 .get("status")
2065 .expect("should have status")
2066 .as_object()
2067 .unwrap();
2068
2069 assert!(
2071 status.contains_key("enum"),
2072 "Status should have enum field. Full schema: {}",
2073 serde_json::to_string_pretty(&schema).unwrap()
2074 );
2075 }
2076
2077 #[test]
2078 fn schema_ref_resolution() {
2079 let raw_schema = serde_json::json!({
2081 "type": "object",
2082 "properties": {
2083 "status": {
2084 "$ref": "#/$defs/Status"
2085 }
2086 },
2087 "$defs": {
2088 "Status": {
2089 "type": "string",
2090 "enum": ["pending", "in_progress", "completed"]
2091 }
2092 }
2093 });
2094
2095 let mut schema = raw_schema;
2096 clean_schema(&mut schema);
2097
2098 let props = schema.get("properties").unwrap().as_object().unwrap();
2099 let status = props.get("status").unwrap().as_object().unwrap();
2100
2101 assert!(
2102 status.contains_key("enum"),
2103 "Status should have enum after ref resolution. Got: {}",
2104 serde_json::to_string_pretty(&schema).unwrap()
2105 );
2106 }
2107
2108 #[test]
2109 fn schema_nested_ref_in_array() {
2110 let raw_schema = serde_json::json!({
2112 "type": "object",
2113 "properties": {
2114 "todos": {
2115 "type": "array",
2116 "items": {
2117 "$ref": "#/$defs/TodoItem"
2118 }
2119 }
2120 },
2121 "$defs": {
2122 "TodoItem": {
2123 "type": "object",
2124 "properties": {
2125 "content": { "type": "string" },
2126 "status": { "$ref": "#/$defs/TodoStatus" }
2127 },
2128 "required": ["content", "status"]
2129 },
2130 "TodoStatus": {
2131 "type": "string",
2132 "enum": ["pending", "in_progress", "completed"]
2133 }
2134 }
2135 });
2136
2137 let mut schema = raw_schema;
2138 clean_schema(&mut schema);
2139
2140 let props = schema.get("properties").unwrap().as_object().unwrap();
2142 let todos = props.get("todos").unwrap().as_object().unwrap();
2143 let items = todos.get("items").unwrap().as_object().unwrap();
2144 let item_props = items.get("properties").unwrap().as_object().unwrap();
2145 let status = item_props.get("status").unwrap().as_object().unwrap();
2146
2147 assert!(
2148 status.contains_key("enum"),
2149 "Nested status should have enum. Full schema: {}",
2150 serde_json::to_string_pretty(&schema).unwrap()
2151 );
2152 }
2153
2154 #[test]
2155 fn parts_collapses_degenerate_inputs() {
2156 assert_eq!(ToolResult::parts(vec![]), ToolResult::Done);
2157 assert_eq!(
2158 ToolResult::parts(vec![ToolResultPart::text("only")]),
2159 ToolResult::text("only")
2160 );
2161 let multi = ToolResult::parts(vec![
2162 ToolResultPart::text("a"),
2163 ToolResultPart::image(vec![1, 2], "image/png"),
2164 ]);
2165 let ToolResult::Parts { parts } = &multi else {
2166 panic!("two parts must stay a Parts result");
2167 };
2168 assert_eq!(parts.len(), 2);
2169 }
2170
2171 #[test]
2172 fn parts_render_joins_text_and_marks_binary() {
2173 let result = ToolResult::parts(vec![
2174 ToolResultPart::text("header"),
2175 ToolResultPart::json_value(serde_json::json!({"n": 1})),
2176 ToolResultPart::image(vec![0x89], "image/png"),
2177 ]);
2178 assert_eq!(
2179 result.render_for_model().unwrap(),
2180 "header\n{\"n\":1}\n[binary tool result: image/png, 1 bytes]"
2181 );
2182 assert!(result.mime().is_none() && result.content().is_none());
2183 }
2184
2185 #[test]
2186 fn parts_with_one_binary_expose_mime_and_content() {
2187 let result = ToolResult::Parts {
2188 parts: vec![ToolResultPart::image(vec![9, 9], "image/png")],
2189 };
2190 assert_eq!(result.mime().unwrap().essence_str(), "image/png");
2191 assert_eq!(result.content().unwrap(), &[9, 9]);
2192 }
2193
2194 #[cfg(feature = "serde")]
2195 #[test]
2196 fn parts_serde_roundtrip_uses_kind_tag() {
2197 let result = ToolResult::parts(vec![
2198 ToolResultPart::text("note"),
2199 ToolResultPart::json_value(serde_json::json!({"k": true})),
2200 ]);
2201 let value = serde_json::to_value(&result).unwrap();
2202 assert_eq!(value["kind"], "parts");
2203 assert_eq!(value["parts"][0]["kind"], "text");
2204 let back: ToolResult = serde_json::from_value(value).unwrap();
2205 assert_eq!(back, result);
2206 }
2207}