1use std::collections::{HashMap, HashSet};
10use std::hash::{DefaultHasher, Hash, Hasher};
11use std::sync::LazyLock;
12
13use regex::Regex;
14
15use crate::config::UtilityScoringConfig;
16use crate::executor::ToolCall;
17
18#[must_use]
27pub fn has_explicit_tool_request(user_message: &str) -> bool {
28 static RE: LazyLock<Regex> = LazyLock::new(|| {
29 Regex::new(
30 r"(?xi)
31 using\s+a\s+tool
32 | call\s+(the\s+)?[a-z_]+\s+tool
33 | use\s+(the\s+)?[a-z_]+\s+tool
34 | run\s+(the\s+)?[a-z_]+\s+tool
35 | invoke\s+(the\s+)?[a-z_]+\s+tool
36 | execute\s+(the\s+)?[a-z_]+\s+tool
37 | show\s+me\s+the\s+result\s+of\s*:
38 | run\s*:
39 | execute\s*:
40 | what\s+(does|would|is\s+the\s+output\s+of)
41 ",
42 )
43 .expect("static regex is valid")
44 });
45 static RE_CODE: LazyLock<Regex> =
48 LazyLock::new(|| Regex::new(r"`[^`]*[|><$;&][^`]*`").expect("static regex is valid"));
49 RE.is_match(user_message) || RE_CODE.is_match(user_message)
50}
51
52fn default_gain(tool_name: &str) -> f32 {
73 if tool_name.starts_with("memory") {
74 return 0.8;
75 }
76 match tool_name {
77 "bash" | "shell" => 0.6,
78 "read" | "write" => 0.55,
79 "search_code" | "grep" | "glob" | "find_path" | "list_directory" => 0.65,
80 "diagnostics" | "edit" | "format" | "create_directory" | "delete_path" | "move_path"
81 | "copy_path" => 0.75,
82 _ => 0.5,
83 }
84}
85
86#[derive(Debug, Clone)]
88pub struct UtilityScore {
89 pub gain: f32,
91 pub cost: f32,
93 pub redundancy: f32,
95 pub uncertainty: f32,
97 pub total: f32,
99}
100
101impl UtilityScore {
102 fn is_valid(&self) -> bool {
104 self.gain.is_finite()
105 && self.cost.is_finite()
106 && self.redundancy.is_finite()
107 && self.uncertainty.is_finite()
108 && self.total.is_finite()
109 }
110}
111
112#[derive(Debug, Clone)]
114pub struct UtilityContext {
115 pub tool_calls_this_turn: usize,
117 pub tokens_consumed: usize,
119 pub token_budget: usize,
121 pub user_requested: bool,
125 pub mandated_retry: bool,
135}
136
137#[non_exhaustive]
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
140pub enum UtilityAction {
141 Respond,
143 Retrieve,
145 ToolCall,
147 Verify,
149 Stop,
151}
152
153fn call_hash(call: &ToolCall) -> u64 {
155 let mut h = DefaultHasher::new();
156 call.tool_id.hash(&mut h);
157 format!("{:?}", call.params).hash(&mut h);
161 h.finish()
162}
163
164#[derive(Debug)]
169pub struct UtilityScorer {
170 config: UtilityScoringConfig,
171 recent_calls: HashMap<u64, u32>,
173 consecutive_low: usize,
175 mandated_retries: HashSet<u64>,
180}
181
182impl UtilityScorer {
183 #[must_use]
185 pub fn new(config: UtilityScoringConfig) -> Self {
186 Self {
187 config,
188 recent_calls: HashMap::new(),
189 consecutive_low: 0,
190 mandated_retries: HashSet::new(),
191 }
192 }
193
194 #[must_use]
196 pub fn is_enabled(&self) -> bool {
197 self.config.enabled
198 }
199
200 #[must_use]
206 pub fn score(&self, call: &ToolCall, ctx: &UtilityContext) -> Option<UtilityScore> {
207 if !self.config.enabled {
208 return None;
209 }
210
211 let gain = if self.is_high_gain(call.tool_id.as_str()) {
212 0.75
213 } else {
214 default_gain(call.tool_id.as_str())
215 };
216
217 let cost = if ctx.token_budget > 0 {
218 #[allow(clippy::cast_precision_loss)]
219 (ctx.tokens_consumed as f32 / ctx.token_budget as f32).clamp(0.0, 1.0)
220 } else {
221 0.0
222 };
223
224 let hash = call_hash(call);
225 let redundancy = if self.recent_calls.contains_key(&hash) {
226 1.0_f32
227 } else {
228 0.0_f32
229 };
230
231 #[allow(clippy::cast_precision_loss)]
234 let uncertainty = (1.0_f32 - ctx.tool_calls_this_turn as f32 / 10.0).clamp(0.0, 1.0);
235
236 let total = self.config.gain_weight * gain
237 - self.config.cost_weight * cost
238 - self.config.redundancy_weight * redundancy
239 + self.config.uncertainty_bonus * uncertainty;
240
241 let score = UtilityScore {
242 gain,
243 cost,
244 redundancy,
245 uncertainty,
246 total,
247 };
248
249 if score.is_valid() { Some(score) } else { None }
250 }
251
252 #[must_use]
269 pub fn recommend_action(
270 &self,
271 score: Option<&UtilityScore>,
272 ctx: &UtilityContext,
273 ) -> UtilityAction {
274 if ctx.user_requested {
276 return UtilityAction::ToolCall;
277 }
278 if !self.config.enabled {
280 return UtilityAction::ToolCall;
281 }
282 if ctx.mandated_retry {
284 return UtilityAction::ToolCall;
285 }
286 let Some(s) = score else {
287 return UtilityAction::Stop;
289 };
290
291 if s.cost > 0.9 {
293 return UtilityAction::Stop;
294 }
295 if s.redundancy >= 1.0 {
297 return UtilityAction::Respond;
298 }
299 if s.gain >= 0.7 && s.total >= self.config.threshold {
301 return UtilityAction::ToolCall;
302 }
303 if s.gain >= 0.5 && s.uncertainty > 0.5 {
305 return UtilityAction::Retrieve;
306 }
307 if s.total < self.config.threshold && ctx.tool_calls_this_turn > 0 {
309 return UtilityAction::Verify;
310 }
311 if s.total >= self.config.threshold {
313 return UtilityAction::ToolCall;
314 }
315 UtilityAction::Respond
316 }
317
318 pub fn record_call(&mut self, call: &ToolCall) {
323 let hash = call_hash(call);
324 *self.recent_calls.entry(hash).or_insert(0) += 1;
325 }
326
327 pub fn clear(&mut self) {
329 self.recent_calls.clear();
330 self.consecutive_low = 0;
331 self.mandated_retries.clear();
332 }
333
334 pub fn mark_mandated_retry(&mut self, call: &ToolCall) {
341 self.mandated_retries.insert(call_hash(call));
342 }
343
344 pub fn take_mandated_retry(&mut self, call: &ToolCall) -> bool {
350 self.mandated_retries.remove(&call_hash(call))
351 }
352
353 pub fn note_action(&mut self, action: &UtilityAction) -> bool {
364 if *action == UtilityAction::ToolCall {
365 self.consecutive_low = 0;
366 } else {
367 self.consecutive_low = self.consecutive_low.saturating_add(1);
368 }
369 self.config.utility_window > 0 && self.consecutive_low >= self.config.utility_window
370 }
371
372 fn contains_tool_name(list: &[String], tool_name: &str) -> bool {
384 let normalize = |s: &str| s.to_lowercase().replace(':', "_");
385 let normalized = normalize(tool_name);
386 list.iter().any(|e| normalize(e) == normalized)
387 }
388
389 #[must_use]
393 pub fn is_exempt(&self, tool_name: &str) -> bool {
394 Self::contains_tool_name(&self.config.exempt_tools, tool_name)
395 }
396
397 #[must_use]
404 pub fn is_high_gain(&self, tool_name: &str) -> bool {
405 Self::contains_tool_name(&self.config.high_gain_tools, tool_name)
406 }
407
408 #[must_use]
410 pub fn threshold(&self) -> f32 {
411 self.config.threshold
412 }
413
414 #[must_use]
416 pub fn utility_window(&self) -> usize {
417 self.config.utility_window
418 }
419}
420
421#[cfg(test)]
422mod tests {
423 use super::*;
424 use crate::ToolName;
425 use serde_json::json;
426
427 fn make_call(name: &str, params: serde_json::Value) -> ToolCall {
428 ToolCall {
429 tool_id: ToolName::new(name),
430 params: if let serde_json::Value::Object(m) = params {
431 m
432 } else {
433 serde_json::Map::new()
434 },
435 caller_id: None,
436 context: None,
437
438 tool_call_id: String::new(),
439 skill_name: None,
440 }
441 }
442
443 fn default_ctx() -> UtilityContext {
444 UtilityContext {
445 tool_calls_this_turn: 0,
446 tokens_consumed: 0,
447 token_budget: 1000,
448 user_requested: false,
449 mandated_retry: false,
450 }
451 }
452
453 fn default_config() -> UtilityScoringConfig {
454 UtilityScoringConfig {
455 enabled: true,
456 ..UtilityScoringConfig::default()
457 }
458 }
459
460 #[test]
461 fn disabled_returns_none() {
462 let scorer = UtilityScorer::new(UtilityScoringConfig::default());
463 assert!(!scorer.is_enabled());
464 let call = make_call("bash", json!({}));
465 let score = scorer.score(&call, &default_ctx());
466 assert!(score.is_none());
467 assert_eq!(
469 scorer.recommend_action(score.as_ref(), &default_ctx()),
470 UtilityAction::ToolCall
471 );
472 }
473
474 #[test]
475 fn first_call_passes_default_threshold() {
476 let scorer = UtilityScorer::new(default_config());
477 let call = make_call("bash", json!({"cmd": "ls"}));
478 let score = scorer.score(&call, &default_ctx());
479 assert!(score.is_some());
480 let s = score.unwrap();
481 assert!(
482 s.total >= 0.1,
483 "first call should exceed threshold: {}",
484 s.total
485 );
486 let action = scorer.recommend_action(Some(&s), &default_ctx());
489 assert!(
490 action == UtilityAction::ToolCall || action == UtilityAction::Retrieve,
491 "first call should not be blocked, got {action:?}",
492 );
493 }
494
495 #[test]
496 fn redundant_call_penalized() {
497 let mut scorer = UtilityScorer::new(default_config());
498 let call = make_call("bash", json!({"cmd": "ls"}));
499 scorer.record_call(&call);
500 let score = scorer.score(&call, &default_ctx()).unwrap();
501 assert!((score.redundancy - 1.0).abs() < f32::EPSILON);
502 }
503
504 #[test]
505 fn clear_resets_redundancy() {
506 let mut scorer = UtilityScorer::new(default_config());
507 let call = make_call("bash", json!({"cmd": "ls"}));
508 scorer.record_call(&call);
509 scorer.clear();
510 let score = scorer.score(&call, &default_ctx()).unwrap();
511 assert!(score.redundancy.abs() < f32::EPSILON);
512 }
513
514 #[test]
515 fn user_requested_always_executes() {
516 let scorer = UtilityScorer::new(default_config());
517 let score = UtilityScore {
519 gain: 0.0,
520 cost: 1.0,
521 redundancy: 1.0,
522 uncertainty: 0.0,
523 total: -100.0,
524 };
525 let ctx = UtilityContext {
526 user_requested: true,
527 ..default_ctx()
528 };
529 assert_eq!(
530 scorer.recommend_action(Some(&score), &ctx),
531 UtilityAction::ToolCall
532 );
533 }
534
535 #[test]
536 fn none_score_fail_closed_when_enabled() {
537 let scorer = UtilityScorer::new(default_config());
538 assert_eq!(
540 scorer.recommend_action(None, &default_ctx()),
541 UtilityAction::Stop
542 );
543 }
544
545 #[test]
546 fn none_score_executes_when_disabled() {
547 let scorer = UtilityScorer::new(UtilityScoringConfig::default()); assert_eq!(
549 scorer.recommend_action(None, &default_ctx()),
550 UtilityAction::ToolCall
551 );
552 }
553
554 #[test]
555 fn cost_increases_with_token_consumption() {
556 let scorer = UtilityScorer::new(default_config());
557 let call = make_call("bash", json!({}));
558 let ctx_low = UtilityContext {
559 tokens_consumed: 100,
560 token_budget: 1000,
561 ..default_ctx()
562 };
563 let ctx_high = UtilityContext {
564 tokens_consumed: 900,
565 token_budget: 1000,
566 ..default_ctx()
567 };
568 let s_low = scorer.score(&call, &ctx_low).unwrap();
569 let s_high = scorer.score(&call, &ctx_high).unwrap();
570 assert!(s_low.cost < s_high.cost);
571 assert!(s_low.total > s_high.total);
572 }
573
574 #[test]
575 fn uncertainty_decreases_with_call_count() {
576 let scorer = UtilityScorer::new(default_config());
577 let call = make_call("bash", json!({}));
578 let ctx_early = UtilityContext {
579 tool_calls_this_turn: 0,
580 ..default_ctx()
581 };
582 let ctx_late = UtilityContext {
583 tool_calls_this_turn: 9,
584 ..default_ctx()
585 };
586 let s_early = scorer.score(&call, &ctx_early).unwrap();
587 let s_late = scorer.score(&call, &ctx_late).unwrap();
588 assert!(s_early.uncertainty > s_late.uncertainty);
589 }
590
591 #[test]
592 fn memory_tool_has_higher_gain_than_scrape() {
593 let scorer = UtilityScorer::new(default_config());
594 let mem_call = make_call("memory_search", json!({}));
595 let web_call = make_call("scrape", json!({}));
596 let s_mem = scorer.score(&mem_call, &default_ctx()).unwrap();
597 let s_web = scorer.score(&web_call, &default_ctx()).unwrap();
598 assert!(s_mem.gain > s_web.gain);
599 }
600
601 #[test]
602 fn zero_token_budget_zeroes_cost() {
603 let scorer = UtilityScorer::new(default_config());
604 let call = make_call("bash", json!({}));
605 let ctx = UtilityContext {
606 tokens_consumed: 500,
607 token_budget: 0,
608 ..default_ctx()
609 };
610 let s = scorer.score(&call, &ctx).unwrap();
611 assert!(s.cost.abs() < f32::EPSILON);
612 }
613
614 #[test]
615 fn validate_rejects_negative_weights() {
616 let cfg = UtilityScoringConfig {
617 enabled: true,
618 gain_weight: -1.0,
619 ..UtilityScoringConfig::default()
620 };
621 assert!(cfg.validate().is_err());
622 }
623
624 #[test]
625 fn validate_rejects_nan_weights() {
626 let cfg = UtilityScoringConfig {
627 enabled: true,
628 threshold: f32::NAN,
629 ..UtilityScoringConfig::default()
630 };
631 assert!(cfg.validate().is_err());
632 }
633
634 #[test]
635 fn validate_accepts_default() {
636 assert!(UtilityScoringConfig::default().validate().is_ok());
637 }
638
639 #[test]
640 fn threshold_zero_all_calls_pass() {
641 let scorer = UtilityScorer::new(UtilityScoringConfig {
643 enabled: true,
644 threshold: 0.0,
645 ..UtilityScoringConfig::default()
646 });
647 let call = make_call("bash", json!({}));
648 let score = scorer.score(&call, &default_ctx()).unwrap();
649 assert!(
651 score.total >= 0.0,
652 "total should be non-negative: {}",
653 score.total
654 );
655 let action = scorer.recommend_action(Some(&score), &default_ctx());
657 assert!(
658 action == UtilityAction::ToolCall || action == UtilityAction::Retrieve,
659 "threshold=0 should not block calls, got {action:?}",
660 );
661 }
662
663 #[test]
664 fn threshold_one_blocks_all_calls() {
665 let scorer = UtilityScorer::new(UtilityScoringConfig {
667 enabled: true,
668 threshold: 1.0,
669 ..UtilityScoringConfig::default()
670 });
671 let call = make_call("bash", json!({}));
672 let score = scorer.score(&call, &default_ctx()).unwrap();
673 assert!(
674 score.total < 1.0,
675 "realistic score should be below 1.0: {}",
676 score.total
677 );
678 assert_ne!(
680 scorer.recommend_action(Some(&score), &default_ctx()),
681 UtilityAction::ToolCall
682 );
683 }
684
685 #[test]
688 fn recommend_action_user_requested_always_tool_call() {
689 let scorer = UtilityScorer::new(default_config());
690 let score = UtilityScore {
691 gain: 0.0,
692 cost: 1.0,
693 redundancy: 1.0,
694 uncertainty: 0.0,
695 total: -100.0,
696 };
697 let ctx = UtilityContext {
698 user_requested: true,
699 ..default_ctx()
700 };
701 assert_eq!(
702 scorer.recommend_action(Some(&score), &ctx),
703 UtilityAction::ToolCall
704 );
705 }
706
707 #[test]
708 fn recommend_action_disabled_scorer_always_tool_call() {
709 let scorer = UtilityScorer::new(UtilityScoringConfig::default()); let ctx = default_ctx();
711 assert_eq!(scorer.recommend_action(None, &ctx), UtilityAction::ToolCall);
712 }
713
714 #[test]
715 fn recommend_action_none_score_enabled_stops() {
716 let scorer = UtilityScorer::new(default_config());
717 let ctx = default_ctx();
718 assert_eq!(scorer.recommend_action(None, &ctx), UtilityAction::Stop);
719 }
720
721 #[test]
722 fn recommend_action_budget_exhausted_stops() {
723 let scorer = UtilityScorer::new(default_config());
724 let score = UtilityScore {
725 gain: 0.8,
726 cost: 0.95,
727 redundancy: 0.0,
728 uncertainty: 0.5,
729 total: 0.5,
730 };
731 assert_eq!(
732 scorer.recommend_action(Some(&score), &default_ctx()),
733 UtilityAction::Stop
734 );
735 }
736
737 #[test]
738 fn recommend_action_redundant_responds() {
739 let scorer = UtilityScorer::new(default_config());
740 let score = UtilityScore {
741 gain: 0.8,
742 cost: 0.1,
743 redundancy: 1.0,
744 uncertainty: 0.5,
745 total: 0.5,
746 };
747 assert_eq!(
748 scorer.recommend_action(Some(&score), &default_ctx()),
749 UtilityAction::Respond
750 );
751 }
752
753 #[test]
754 fn recommend_action_high_gain_above_threshold_tool_call() {
755 let scorer = UtilityScorer::new(default_config());
756 let score = UtilityScore {
757 gain: 0.8,
758 cost: 0.1,
759 redundancy: 0.0,
760 uncertainty: 0.4,
761 total: 0.6,
762 };
763 assert_eq!(
764 scorer.recommend_action(Some(&score), &default_ctx()),
765 UtilityAction::ToolCall
766 );
767 }
768
769 #[test]
770 fn recommend_action_uncertain_retrieves() {
771 let scorer = UtilityScorer::new(default_config());
772 let score = UtilityScore {
774 gain: 0.6,
775 cost: 0.1,
776 redundancy: 0.0,
777 uncertainty: 0.8,
778 total: 0.4,
779 };
780 assert_eq!(
781 scorer.recommend_action(Some(&score), &default_ctx()),
782 UtilityAction::Retrieve
783 );
784 }
785
786 #[test]
787 fn recommend_action_below_threshold_with_prior_calls_verifies() {
788 let scorer = UtilityScorer::new(default_config());
789 let score = UtilityScore {
790 gain: 0.3,
791 cost: 0.1,
792 redundancy: 0.0,
793 uncertainty: 0.2,
794 total: 0.05, };
796 let ctx = UtilityContext {
797 tool_calls_this_turn: 1,
798 ..default_ctx()
799 };
800 assert_eq!(
801 scorer.recommend_action(Some(&score), &ctx),
802 UtilityAction::Verify
803 );
804 }
805
806 #[test]
807 fn recommend_action_default_responds() {
808 let scorer = UtilityScorer::new(default_config());
809 let score = UtilityScore {
810 gain: 0.3,
811 cost: 0.1,
812 redundancy: 0.0,
813 uncertainty: 0.2,
814 total: 0.05, };
816 let ctx = UtilityContext {
817 tool_calls_this_turn: 0,
818 ..default_ctx()
819 };
820 assert_eq!(
821 scorer.recommend_action(Some(&score), &ctx),
822 UtilityAction::Respond
823 );
824 }
825
826 #[test]
829 fn default_gain_direct_action_tools_reach_tool_call_threshold() {
830 for tool in [
831 "diagnostics",
832 "edit",
833 "format",
834 "create_directory",
835 "delete_path",
836 "move_path",
837 "copy_path",
838 ] {
839 let gain = default_gain(tool);
840 assert!(gain >= 0.7, "{tool} gain should be >= 0.7, got {gain}");
841 }
842 }
843
844 #[test]
845 fn default_gain_find_path_and_list_directory_match_grep_glob_tier() {
846 for tool in ["find_path", "list_directory", "grep", "glob"] {
847 let gain = default_gain(tool);
848 assert!(
849 (gain - 0.65).abs() < f32::EPSILON,
850 "{tool} gain should be 0.65, got {gain}"
851 );
852 }
853 }
854
855 #[test]
856 fn recommend_action_direct_tools_execute_on_first_call() {
857 let scorer = UtilityScorer::new(default_config());
862 let ctx = default_ctx(); for tool in [
864 "diagnostics",
865 "edit",
866 "format",
867 "create_directory",
868 "delete_path",
869 "move_path",
870 "copy_path",
871 ] {
872 let call = make_call(tool, json!({}));
873 let score = scorer.score(&call, &ctx).unwrap();
874 assert!(
875 score.gain >= 0.7,
876 "{tool} gain should be >= 0.7, got {}",
877 score.gain
878 );
879 assert_eq!(
880 scorer.recommend_action(Some(&score), &ctx),
881 UtilityAction::ToolCall,
882 "{tool} should execute immediately on first call, not stall on Retrieve"
883 );
884 }
885 }
886
887 #[test]
888 fn recommend_action_unclassified_tools_still_retrieve_on_first_call() {
889 let scorer = UtilityScorer::new(default_config());
895 let ctx = default_ctx();
896 for tool in ["fetch", "totally_unrecognized_tool_xyz"] {
897 let call = make_call(tool, json!({}));
898 let score = scorer.score(&call, &ctx).unwrap();
899 assert!((score.gain - 0.5).abs() < f32::EPSILON);
900 assert_eq!(
901 scorer.recommend_action(Some(&score), &ctx),
902 UtilityAction::Retrieve,
903 "{tool} should still be eligible for Retrieve on first call"
904 );
905 }
906 }
907
908 #[test]
909 fn recommend_action_diagnostics_never_enters_the_retrieve_redundant_respond_stall() {
910 let mut scorer = UtilityScorer::new(default_config());
917 let ctx = default_ctx();
918
919 let fetch_call = make_call("fetch", json!({"url": "https://example.com"}));
920 let fetch_score = scorer.score(&fetch_call, &ctx).unwrap();
921 assert_eq!(
922 scorer.recommend_action(Some(&fetch_score), &ctx),
923 UtilityAction::Retrieve
924 );
925 scorer.record_call(&fetch_call);
926 let fetch_retry_score = scorer.score(&fetch_call, &ctx).unwrap();
927 assert_eq!(
928 scorer.recommend_action(Some(&fetch_retry_score), &ctx),
929 UtilityAction::Respond,
930 "identical retry should be flagged as redundant, reproducing the stall"
931 );
932
933 let diagnostics_call = make_call("diagnostics", json!({}));
934 let diagnostics_score = scorer.score(&diagnostics_call, &ctx).unwrap();
935 assert_eq!(
936 scorer.recommend_action(Some(&diagnostics_score), &ctx),
937 UtilityAction::ToolCall,
938 "diagnostics must execute on the first call, bypassing the stall entirely"
939 );
940 }
941
942 #[test]
945 fn explicit_request_using_a_tool() {
946 assert!(has_explicit_tool_request(
947 "Please list the files in the current directory using a tool"
948 ));
949 }
950
951 #[test]
952 fn explicit_request_call_the_tool() {
953 assert!(has_explicit_tool_request("call the list_directory tool"));
954 }
955
956 #[test]
957 fn explicit_request_use_the_tool() {
958 assert!(has_explicit_tool_request("use the shell tool to run ls"));
959 }
960
961 #[test]
962 fn explicit_request_run_the_tool() {
963 assert!(has_explicit_tool_request("run the bash tool"));
964 }
965
966 #[test]
967 fn explicit_request_invoke_the_tool() {
968 assert!(has_explicit_tool_request("invoke the search_code tool"));
969 }
970
971 #[test]
972 fn explicit_request_execute_the_tool() {
973 assert!(has_explicit_tool_request("execute the grep tool for me"));
974 }
975
976 #[test]
977 fn explicit_request_case_insensitive() {
978 assert!(has_explicit_tool_request("USING A TOOL to find files"));
979 }
980
981 #[test]
982 fn explicit_request_no_match_plain_message() {
983 assert!(!has_explicit_tool_request("what is the weather today?"));
984 }
985
986 #[test]
987 fn explicit_request_no_match_tool_mentioned_without_invocation() {
988 assert!(!has_explicit_tool_request(
989 "the shell tool is very useful in general"
990 ));
991 }
992
993 #[test]
994 fn explicit_request_show_me_result_of() {
995 assert!(has_explicit_tool_request(
996 "show me the result of: echo hello"
997 ));
998 }
999
1000 #[test]
1001 fn explicit_request_run_colon() {
1002 assert!(has_explicit_tool_request("run: echo hello"));
1003 }
1004
1005 #[test]
1006 fn explicit_request_execute_colon() {
1007 assert!(has_explicit_tool_request("execute: ls -la"));
1008 }
1009
1010 #[test]
1011 fn explicit_request_what_does() {
1012 assert!(has_explicit_tool_request("what does echo hello output?"));
1013 }
1014
1015 #[test]
1016 fn explicit_request_what_would() {
1017 assert!(has_explicit_tool_request("what would cat /etc/hosts show?"));
1018 }
1019
1020 #[test]
1021 fn explicit_request_what_is_the_output_of() {
1022 assert!(has_explicit_tool_request(
1023 "what is the output of ls | grep foo?"
1024 ));
1025 }
1026
1027 #[test]
1028 fn explicit_request_inline_code_pipe() {
1029 assert!(has_explicit_tool_request("try running `ls | grep foo`"));
1030 }
1031
1032 #[test]
1033 fn explicit_request_inline_code_redirect() {
1034 assert!(has_explicit_tool_request("run `echo hello > /tmp/out`"));
1035 }
1036
1037 #[test]
1038 fn explicit_request_inline_code_dollar() {
1039 assert!(has_explicit_tool_request("check `$HOME/bin`"));
1040 }
1041
1042 #[test]
1043 fn explicit_request_inline_code_and() {
1044 assert!(has_explicit_tool_request("try `git fetch && git rebase`"));
1045 }
1046
1047 #[test]
1048 fn no_match_run_the_tests() {
1049 assert!(!has_explicit_tool_request("run the tests please"));
1050 }
1051
1052 #[test]
1053 fn no_match_execute_the_plan() {
1054 assert!(!has_explicit_tool_request("execute the plan we discussed"));
1055 }
1056
1057 #[test]
1058 fn no_match_inline_code_no_shell_syntax() {
1059 assert!(!has_explicit_tool_request(
1060 "the function `process_items` handles it"
1061 ));
1062 }
1063
1064 #[test]
1069 fn known_fp_what_does_function_do() {
1070 assert!(has_explicit_tool_request("what does this function do?"));
1072 }
1073
1074 #[test]
1075 fn no_match_show_me_result_without_colon() {
1076 assert!(!has_explicit_tool_request(
1078 "show me the result of running it"
1079 ));
1080 }
1081
1082 #[test]
1083 fn is_exempt_matches_case_insensitively() {
1084 let scorer = UtilityScorer::new(UtilityScoringConfig {
1085 enabled: true,
1086 exempt_tools: vec!["Read".to_owned(), "file_read".to_owned()],
1087 ..UtilityScoringConfig::default()
1088 });
1089 assert!(scorer.is_exempt("read"));
1090 assert!(scorer.is_exempt("READ"));
1091 assert!(scorer.is_exempt("FILE_READ"));
1092 assert!(!scorer.is_exempt("write"));
1093 assert!(!scorer.is_exempt("bash"));
1094 }
1095
1096 #[test]
1097 fn is_exempt_empty_list_returns_false() {
1098 let scorer = UtilityScorer::new(UtilityScoringConfig::default());
1099 assert!(!scorer.is_exempt("read"));
1100 }
1101
1102 #[test]
1105 fn is_high_gain_matches_case_insensitively() {
1106 let scorer = UtilityScorer::new(UtilityScoringConfig {
1107 enabled: true,
1108 high_gain_tools: vec!["Github_create_issue".to_owned()],
1109 ..UtilityScoringConfig::default()
1110 });
1111 assert!(scorer.is_high_gain("github_create_issue"));
1112 assert!(scorer.is_high_gain("GITHUB_CREATE_ISSUE"));
1113 assert!(!scorer.is_high_gain("bash"));
1114 }
1115
1116 #[test]
1117 fn is_high_gain_empty_list_returns_false() {
1118 let scorer = UtilityScorer::new(UtilityScoringConfig::default());
1119 assert!(!scorer.is_high_gain("github_create_issue"));
1120 }
1121
1122 #[test]
1123 fn default_gain_unconfigured_mcp_shaped_tool_id_stays_neutral() {
1124 assert!((default_gain("github_create_issue") - 0.5).abs() < f32::EPSILON);
1128 }
1129
1130 #[test]
1131 fn score_high_gain_tools_overrides_default_gain_for_mcp_shaped_tool_id() {
1132 let scorer = UtilityScorer::new(UtilityScoringConfig {
1138 enabled: true,
1139 high_gain_tools: vec!["github_create_issue".to_owned()],
1140 ..UtilityScoringConfig::default()
1141 });
1142 let ctx = default_ctx(); let call = make_call("github_create_issue", json!({}));
1144 let score = scorer.score(&call, &ctx).unwrap();
1145 assert!(
1146 (score.gain - 0.75).abs() < f32::EPSILON,
1147 "high_gain_tools entry should raise gain to 0.75, got {}",
1148 score.gain
1149 );
1150 assert_eq!(
1151 scorer.recommend_action(Some(&score), &ctx),
1152 UtilityAction::ToolCall,
1153 "high-gain MCP tool should execute immediately on first call, not stall on Retrieve"
1154 );
1155 }
1156
1157 #[test]
1158 fn score_high_gain_tools_does_not_affect_unlisted_tools() {
1159 let scorer = UtilityScorer::new(UtilityScoringConfig {
1160 enabled: true,
1161 high_gain_tools: vec!["github_create_issue".to_owned()],
1162 ..UtilityScoringConfig::default()
1163 });
1164 let ctx = default_ctx();
1165 let call = make_call("fetch", json!({}));
1166 let score = scorer.score(&call, &ctx).unwrap();
1167 assert!(
1168 (score.gain - 0.5).abs() < f32::EPSILON,
1169 "unlisted tool must keep its default_gain value, got {}",
1170 score.gain
1171 );
1172 }
1173
1174 #[test]
1177 fn is_high_gain_matches_qualified_name_config_against_sanitized_id_call() {
1178 let scorer = UtilityScorer::new(UtilityScoringConfig {
1182 enabled: true,
1183 high_gain_tools: vec!["myserver:mytool".to_owned()],
1184 ..UtilityScoringConfig::default()
1185 });
1186 assert!(scorer.is_high_gain("myserver_mytool"));
1187 }
1188
1189 #[test]
1190 fn is_high_gain_matches_sanitized_id_config_against_qualified_name_call() {
1191 let scorer = UtilityScorer::new(UtilityScoringConfig {
1193 enabled: true,
1194 high_gain_tools: vec!["myserver_mytool".to_owned()],
1195 ..UtilityScoringConfig::default()
1196 });
1197 assert!(scorer.is_high_gain("myserver:mytool"));
1198 }
1199
1200 #[test]
1201 fn is_exempt_matches_qualified_name_config_against_sanitized_id_call() {
1202 let scorer = UtilityScorer::new(UtilityScoringConfig {
1204 enabled: true,
1205 exempt_tools: vec!["myserver:mytool".to_owned()],
1206 ..UtilityScoringConfig::default()
1207 });
1208 assert!(scorer.is_exempt("myserver_mytool"));
1209 }
1210
1211 #[test]
1212 fn is_high_gain_dual_form_still_case_insensitive() {
1213 let scorer = UtilityScorer::new(UtilityScoringConfig {
1214 enabled: true,
1215 high_gain_tools: vec!["MyServer:MyTool".to_owned()],
1216 ..UtilityScoringConfig::default()
1217 });
1218 assert!(scorer.is_high_gain("myserver_mytool"));
1219 assert!(scorer.is_high_gain("MYSERVER_MYTOOL"));
1220 }
1221
1222 #[test]
1223 fn is_high_gain_dual_form_does_not_match_unrelated_tool() {
1224 let scorer = UtilityScorer::new(UtilityScoringConfig {
1225 enabled: true,
1226 high_gain_tools: vec!["myserver:mytool".to_owned()],
1227 ..UtilityScoringConfig::default()
1228 });
1229 assert!(!scorer.is_high_gain("otherserver_othertool"));
1230 }
1231
1232 #[test]
1241 fn mark_and_take_mandated_retry_is_consumed_exactly_once() {
1242 let mut scorer = UtilityScorer::new(default_config());
1243 let call = make_call("find_path", json!({"pattern": "*.rs"}));
1244
1245 assert!(
1246 !scorer.take_mandated_retry(&call),
1247 "no marker set yet — must not report a pending retry"
1248 );
1249
1250 scorer.mark_mandated_retry(&call);
1251 assert!(
1252 scorer.take_mandated_retry(&call),
1253 "marker set — first take must report the pending retry"
1254 );
1255 assert!(
1256 !scorer.take_mandated_retry(&call),
1257 "marker consumed — second take must not report a pending retry again"
1258 );
1259 }
1260
1261 #[test]
1262 fn recommend_action_mandated_retry_bypasses_redundancy_veto() {
1263 let scorer = UtilityScorer::new(default_config());
1264 let score = UtilityScore {
1267 gain: 0.65,
1268 cost: 0.1,
1269 redundancy: 1.0,
1270 uncertainty: 0.7,
1271 total: 0.5,
1272 };
1273 let ctx = UtilityContext {
1274 mandated_retry: true,
1275 ..default_ctx()
1276 };
1277 assert_eq!(
1278 scorer.recommend_action(Some(&score), &ctx),
1279 UtilityAction::ToolCall,
1280 "mandated retry must bypass the redundancy veto and execute"
1281 );
1282 }
1283
1284 #[test]
1285 fn find_path_retrieve_then_mandated_retry_executes_not_redundant_respond() {
1286 let mut scorer = UtilityScorer::new(default_config());
1287 let ctx = default_ctx(); let call = make_call("find_path", json!({"pattern": "*.rs"}));
1289
1290 let first_score = scorer.score(&call, &ctx).unwrap();
1292 assert_eq!(
1293 scorer.recommend_action(Some(&first_score), &ctx),
1294 UtilityAction::Retrieve
1295 );
1296 scorer.record_call(&call);
1300 scorer.mark_mandated_retry(&call);
1301
1302 let retry_ctx = UtilityContext {
1305 mandated_retry: scorer.take_mandated_retry(&call),
1306 ..default_ctx()
1307 };
1308 assert!(
1309 retry_ctx.mandated_retry,
1310 "retry must be recognized as mandated"
1311 );
1312 let retry_score = scorer.score(&call, &retry_ctx).unwrap();
1313 assert!(
1314 (retry_score.redundancy - 1.0).abs() < f32::EPSILON,
1315 "retry is indeed flagged redundant by the raw score — the bypass must come from \
1316 recommend_action, not from suppressing the redundancy component"
1317 );
1318 assert_eq!(
1319 scorer.recommend_action(Some(&retry_score), &retry_ctx),
1320 UtilityAction::ToolCall,
1321 "mandated retry must execute instead of being re-vetoed as a redundant duplicate"
1322 );
1323
1324 scorer.record_call(&call);
1327 let third_ctx = UtilityContext {
1328 mandated_retry: scorer.take_mandated_retry(&call),
1329 ..default_ctx()
1330 };
1331 assert!(
1332 !third_ctx.mandated_retry,
1333 "marker was consumed by the mandated retry — a third call is not exempted"
1334 );
1335 let third_score = scorer.score(&call, &third_ctx).unwrap();
1336 assert_eq!(
1337 scorer.recommend_action(Some(&third_score), &third_ctx),
1338 UtilityAction::Respond,
1339 "a genuine third identical call must be treated as a redundant duplicate"
1340 );
1341 }
1342
1343 #[test]
1344 fn clear_resets_mandated_retries() {
1345 let mut scorer = UtilityScorer::new(default_config());
1346 let call = make_call("find_path", json!({}));
1347 scorer.mark_mandated_retry(&call);
1348 scorer.clear();
1349 assert!(
1350 !scorer.take_mandated_retry(&call),
1351 "clear() must reset mandated-retry state at turn start"
1352 );
1353 }
1354
1355 #[test]
1356 fn note_action_window_zero_never_fires() {
1357 let mut scorer = UtilityScorer::new(UtilityScoringConfig {
1358 enabled: true,
1359 utility_window: 0,
1360 ..UtilityScoringConfig::default()
1361 });
1362 for _ in 0..100 {
1364 assert!(!scorer.note_action(&UtilityAction::Stop));
1365 }
1366 }
1367
1368 #[test]
1369 fn note_action_window_three_fires_on_third() {
1370 let mut scorer = UtilityScorer::new(UtilityScoringConfig {
1371 enabled: true,
1372 utility_window: 3,
1373 ..UtilityScoringConfig::default()
1374 });
1375 assert!(!scorer.note_action(&UtilityAction::Stop));
1376 assert!(!scorer.note_action(&UtilityAction::Respond));
1377 assert!(scorer.note_action(&UtilityAction::Stop));
1378 }
1379
1380 #[test]
1381 fn note_action_tool_call_resets_counter() {
1382 let mut scorer = UtilityScorer::new(UtilityScoringConfig {
1383 enabled: true,
1384 utility_window: 2,
1385 ..UtilityScoringConfig::default()
1386 });
1387 assert!(!scorer.note_action(&UtilityAction::Stop));
1388 assert!(!scorer.note_action(&UtilityAction::ToolCall));
1390 assert!(!scorer.note_action(&UtilityAction::Stop));
1392 }
1393
1394 #[test]
1395 fn note_action_clear_resets_counter() {
1396 let mut scorer = UtilityScorer::new(UtilityScoringConfig {
1397 enabled: true,
1398 utility_window: 1,
1399 ..UtilityScoringConfig::default()
1400 });
1401 assert!(scorer.note_action(&UtilityAction::Stop));
1403 scorer.clear();
1405 assert!(scorer.note_action(&UtilityAction::Stop));
1406 }
1407}