返回 CodeWhale
auto_router.rs
根目录 / crates / tui / src / tui / auto_router.rs
1 //! Auto-routing helpers: deciding when to consult the auto-route flash
2 //! model, and building the small context window it sees.
3 //!
4 //! `dispatch_user_message` calls `model_routing::resolve_auto_route_with_inventory_for_session`
5 //! directly once per user turn when `app.auto_model` is set. The remaining
6 //! helpers here build the compact recent-context summary the router sees.
7
8 use crate::models::{ContentBlock, Message};
9 use crate::tui::app::App;
10
11 /// Whether the next turn should consult the auto-route flash model.
12 pub(super) fn should_resolve_auto_model_selection(app: &App) -> bool {
13 app.auto_model
14 }
15
16 /// Build a compact recent-context summary for the auto-route prompt.
17 ///
18 /// Walks `api_messages` from the most recent turn back, skipping the
19 /// final draft (which is what the router is being asked to classify),
20 /// collects up to six non-empty rows, and reverses them so the prompt
21 /// reads oldest-first. Each row is `<role>: <truncated content>` and
22 /// is capped at 900 characters.
23 pub(super) fn recent_auto_router_context(messages: &[Message]) -> String {
24 let mut rows = Vec::new();
25 for message in messages.iter().rev().skip(1) {
26 if rows.len() >= 6 {
27 break;
28 }
29 let text = content_blocks_text(&message.content);
30 let text = text.trim();
31 if text.is_empty() {
32 continue;
33 }
34 rows.push(format!(
35 "{}: {}",
36 message.role,
37 truncate_for_auto_router(text, 900)
38 ));
39 }
40 rows.reverse();
41 if rows.is_empty() {
42 "No prior context.".to_string()
43 } else {
44 rows.join("\n")
45 }
46 }
47
48 fn content_blocks_text(blocks: &[ContentBlock]) -> String {
49 let mut out = String::new();
50 for block in blocks {
51 match block {
52 ContentBlock::Text { text, .. } => {
53 append_router_text(&mut out, text);
54 }
55 ContentBlock::Thinking { .. } => {}
56 ContentBlock::ToolUse { name, .. } => {
57 append_router_text(&mut out, &format!("[tool call: {name}]"));
58 }
59 ContentBlock::ToolResult { content, .. } => {
60 append_router_text(&mut out, &format!("[tool result] {content}"));
61 }
62 _ => {}
63 }
64 }
65 out
66 }
67
68 fn append_router_text(out: &mut String, text: &str) {
69 if !out.is_empty() {
70 out.push('\n');
71 }
72 out.push_str(text);
73 }
74
75 fn truncate_for_auto_router(text: &str, max_chars: usize) -> String {
76 let mut chars = text.chars();
77 let truncated: String = chars.by_ref().take(max_chars).collect();
78 if chars.next().is_some() {
79 format!("{truncated}...")
80 } else {
81 truncated
82 }
83 }
84
85 #[cfg(test)]
86 mod tests {
87 use super::*;
88 use crate::models::ContentBlock;
89
90 fn make_msg(role: &str, text: &str) -> Message {
91 Message {
92 role: role.to_string(),
93 content: vec![ContentBlock::Text {
94 text: text.to_string(),
95 cache_control: None,
96 }],
97 }
98 }
99
100 #[test]
101 fn truncate_for_auto_router_honors_char_budget() {
102 let s = "abcdefghij";
103 assert_eq!(truncate_for_auto_router(s, 4), "abcd...");
104 assert_eq!(truncate_for_auto_router(s, 10), "abcdefghij");
105 assert_eq!(truncate_for_auto_router(s, 100), "abcdefghij");
106 }
107
108 #[test]
109 fn recent_auto_router_context_skips_final_message_and_caps_rows() {
110 // Eight messages; final one (the draft being routed) is skipped,
111 // so we expect at most six of the remaining seven.
112 let msgs: Vec<Message> = (0..8)
113 .map(|i| {
114 make_msg(
115 if i % 2 == 0 { "user" } else { "assistant" },
116 &format!("turn {i}"),
117 )
118 })
119 .collect();
120 let context = recent_auto_router_context(&msgs);
121 assert!(!context.contains("turn 7"), "final draft must be skipped");
122 let row_count = context.lines().count();
123 assert_eq!(row_count, 6);
124 // Output is oldest-first.
125 let first = context.lines().next().unwrap();
126 assert!(first.contains("turn 1"), "got: {context}");
127 }
128
129 #[test]
130 fn recent_auto_router_context_handles_empty_history() {
131 assert_eq!(recent_auto_router_context(&[]), "No prior context.");
132 }
133
134 #[test]
135 fn recent_auto_router_context_excludes_hidden_thinking() {
136 let msgs = vec![
137 Message {
138 role: "assistant".to_string(),
139 content: vec![
140 ContentBlock::Thinking {
141 signature: None,
142 thinking: "The user seems to be asking me to classify myself.".to_string(),
143 },
144 ContentBlock::Text {
145 text: "Visible assistant answer.".to_string(),
146 cache_control: None,
147 },
148 ],
149 },
150 make_msg("user", "latest draft"),
151 ];
152
153 let context = recent_auto_router_context(&msgs);
154
155 assert!(context.contains("Visible assistant answer."));
156 assert!(!context.contains("The user seems"));
157 assert!(!context.contains("latest draft"));
158 }
159 }
160
160 lines RUST