返回 CodeWhale
advisor.rs
根目录 / crates / tui / src / tools / subagent / advisor.rs
1 //! Background advisor watcher (#3982).
2 //!
3 //! When enabled, the advisor wakes on turn boundaries, reads a bounded slice
4 //! of recent tool calls from the session transcript, makes a concise LLM
5 //! advisory call on an exactly resolved provider/model client, and
6 //! emits an [`Event::AdvisoryNote`] fire-and-forget.
7 //!
8 //! Key design properties:
9 //! - **Off by default** — enabled via `[advisor] enabled = true` or `/advisor on`.
10 //! - **Bounded input** — at most `max_tool_calls` tool-call/result pairs are
11 //! included; the rest are dropped oldest-first.
12 //! - **Rate-limited** — at most one emission per `rate_limit_secs` seconds.
13 //! - **Deduplicated** — notes whose content hash matches the previous note
14 //! within `dedup_window_secs` are silently dropped.
15 //! - **Child-failure isolated** — advisor errors are logged but never surface
16 //! as parent turn failures.
17 //! - **Policy-bounded** — the advisor uses a read-only reviewer prompt and
18 //! no tool access; it cannot exceed the parent session policy.
19
20 use std::collections::hash_map::DefaultHasher;
21 use std::hash::{Hash, Hasher};
22 use std::time::{Duration, Instant};
23
24 use codewhale_config::AdvisorConfigToml;
25 use tokio::sync::mpsc;
26 use tracing::debug;
27
28 use crate::client::CodewhaleClient;
29 use crate::config::Config;
30 use crate::core::events::Event;
31 use crate::llm_client::LlmClient;
32 use crate::utils::truncate_with_ellipsis;
33 use codewhale_models::Role;
34 use codewhale_models::{ContentBlock, Message, MessageRequest, SystemPrompt};
35
36 /// Maximum tokens the advisor may generate. Kept short so the note stays
37 /// concise and does not compete with the parent turn's billing budget.
38 const ADVISOR_MAX_TOKENS: u32 = 256;
39
40 /// Maximum characters of tool input + result to include per tool-call pair.
41 const MAX_CHARS_PER_PAIR: usize = 800;
42
43 /// System prompt for the advisor LLM call. Read-only review posture — no
44 /// tool access, no code generation.
45 const ADVISOR_SYSTEM_PROMPT: &str = "You are a concise background advisor reviewing recent tool activity. \
46 Your role: identify one or two concrete concerns (correctness, risk, or \
47 missed alternatives) in the tool calls provided. \
48 If nothing notable stands out, respond with exactly the word \"ok\". \
49 Otherwise write one to three short sentences — no preamble, no markdown, \
50 no praise. Focus on signal; omit noise.";
51
52 /// A single tool-call/result pair extracted from the session transcript.
53 #[derive(Debug, Clone)]
54 pub struct ToolCallPair {
55 /// Tool name (e.g. `exec_shell`, `file_write`).
56 pub name: String,
57 /// Bounded serialization of the tool input.
58 pub input_preview: String,
59 /// Bounded serialization of the tool result.
60 pub result_preview: String,
61 }
62
63 /// Resolved advisor configuration derived from [`AdvisorConfigToml`].
64 #[derive(Debug, Clone)]
65 pub struct AdvisorConfig {
66 /// Whether the advisor is currently enabled (session-level toggle).
67 pub enabled: bool,
68 /// Max tool-call pairs to review per turn.
69 pub max_tool_calls: u32,
70 /// Min seconds between consecutive emissions.
71 pub rate_limit: Duration,
72 /// Window during which duplicate notes are suppressed.
73 pub dedup_window: Duration,
74 /// Optional model override (falls back to session model when `None`).
75 pub model: Option<String>,
76 }
77
78 impl AdvisorConfig {
79 /// Build a resolved config from the TOML schema.
80 #[must_use]
81 pub fn from_toml(toml: &AdvisorConfigToml) -> Self {
82 Self {
83 enabled: toml.enabled,
84 max_tool_calls: toml.max_tool_calls.clamp(1, 50),
85 rate_limit: Duration::from_secs(toml.rate_limit_secs.clamp(5, 3600)),
86 dedup_window: Duration::from_secs(toml.dedup_window_secs),
87 model: toml.model.clone(),
88 }
89 }
90
91 /// Default disabled config (matches `[advisor]` absent from config.toml).
92 #[must_use]
93 pub fn disabled() -> Self {
94 Self {
95 enabled: false,
96 max_tool_calls: 10,
97 rate_limit: Duration::from_secs(60),
98 dedup_window: Duration::from_secs(300),
99 model: None,
100 }
101 }
102 }
103
104 /// Runtime emission guard: tracks the last emission time and the hash of the
105 /// last advisory note to enforce rate limiting and deduplication.
106 #[derive(Debug)]
107 pub struct EmissionGuard {
108 last_emission: Option<Instant>,
109 last_note_hash: Option<u64>,
110 last_note_hash_at: Option<Instant>,
111 }
112
113 /// Accounting ownership captured while the originating turn is still live.
114 ///
115 /// Runtime turns retain their synchronous durable sink through the lease;
116 /// ordinary interactive turns fall back to the exact session cost generation
117 /// captured here. Neither path can spill into a later session.
118 #[derive(Debug)]
119 pub(crate) struct AdvisorUsageContext {
120 cost_scope: crate::cost_status::CostScopeToken,
121 runtime_usage_lease: Option<crate::cost_status::RuntimeUsageLease>,
122 }
123
124 impl AdvisorUsageContext {
125 #[must_use]
126 pub(crate) fn capture(runtime_owner: Option<&str>) -> Self {
127 Self {
128 cost_scope: crate::cost_status::scope_token(),
129 runtime_usage_lease: runtime_owner
130 .and_then(crate::cost_status::acquire_runtime_usage_lease),
131 }
132 }
133
134 fn report(
135 &self,
136 source_id: &str,
137 route: &crate::cost_status::EffectiveRouteEnvelope,
138 usage: &codewhale_models::Usage,
139 ) {
140 crate::cost_status::report_effective_route_for_runtime(
141 self.cost_scope,
142 self.runtime_usage_lease
143 .as_ref()
144 .map(crate::cost_status::RuntimeUsageLease::owner),
145 source_id,
146 route,
147 usage,
148 );
149 }
150
151 fn report_unreceipted(
152 &self,
153 source_id: &str,
154 route: &crate::cost_status::EffectiveRouteEnvelope,
155 ) {
156 crate::cost_status::report_unreceipted_provider_success(
157 self.cost_scope,
158 self.runtime_usage_lease
159 .as_ref()
160 .map(crate::cost_status::RuntimeUsageLease::owner),
161 source_id,
162 route,
163 );
164 }
165 }
166
167 impl EmissionGuard {
168 /// Create a fresh guard with no emission history.
169 #[must_use]
170 pub fn new() -> Self {
171 Self {
172 last_emission: None,
173 last_note_hash: None,
174 last_note_hash_at: None,
175 }
176 }
177
178 /// Check whether emitting `note` is allowed under `config`'s rate-limit
179 /// and dedup policy. Returns `true` when the note may be emitted.
180 #[must_use]
181 pub fn may_emit(&self, note: &str, config: &AdvisorConfig) -> bool {
182 // Suppress trivial "ok" responses from the model.
183 if note.trim().eq_ignore_ascii_case("ok") {
184 return false;
185 }
186
187 let now = Instant::now();
188
189 // Rate limit: require at least `rate_limit` since last emission.
190 if let Some(last) = self.last_emission
191 && now.duration_since(last) < config.rate_limit
192 {
193 return false;
194 }
195
196 // Dedup: suppress if the note content hash matches the previous note
197 // within the dedup window.
198 let note_hash = hash_str(note);
199 if let (Some(prev_hash), Some(prev_at)) = (self.last_note_hash, self.last_note_hash_at)
200 && prev_hash == note_hash
201 && now.duration_since(prev_at) < config.dedup_window
202 {
203 return false;
204 }
205
206 true
207 }
208
209 /// Record that `note` was emitted now. Must be called immediately after
210 /// sending the `AdvisoryNote` event.
211 pub fn record_emission(&mut self, note: &str) {
212 let now = Instant::now();
213 self.last_emission = Some(now);
214 self.last_note_hash = Some(hash_str(note));
215 self.last_note_hash_at = Some(now);
216 }
217 }
218
219 impl Default for EmissionGuard {
220 fn default() -> Self {
221 Self::new()
222 }
223 }
224
225 /// Extract bounded tool-call/result pairs from a session message slice.
226 ///
227 /// Takes the newest `max_pairs` tool calls, then returns their unambiguous
228 /// pairs oldest-first. Missing results remain pending; corrupt or ambiguous
229 /// identities are omitted rather than borrowing another execution's output.
230 #[must_use]
231 pub fn extract_tool_call_pairs(messages: &[Message], max_pairs: usize) -> Vec<ToolCallPair> {
232 let mut uses = Vec::new();
233 let mut use_counts = std::collections::HashMap::new();
234 let mut results = std::collections::HashMap::new();
235
236 for msg in messages {
237 for block in &msg.content {
238 let Some(key) = block.tool_call_key() else {
239 continue;
240 };
241 match block {
242 ContentBlock::ToolUse {
243 id, name, input, ..
244 } => {
245 let input_str = truncate_with_ellipsis(
246 &serde_json::to_string(input).unwrap_or_default(),
247 MAX_CHARS_PER_PAIR / 2,
248 "…",
249 );
250 *use_counts.entry(key).or_insert(0usize) += 1;
251 uses.push((key, id.as_str(), name.as_str(), input_str));
252 }
253 ContentBlock::ToolResult {
254 tool_use_id,
255 content,
256 ..
257 } => {
258 results
259 .entry(key)
260 .and_modify(|result| *result = None)
261 .or_insert_with(|| {
262 Some((
263 tool_use_id.as_str(),
264 truncate_with_ellipsis(content, MAX_CHARS_PER_PAIR / 2, "…"),
265 ))
266 });
267 }
268 _ => {}
269 }
270 }
271 }
272
273 let start = uses.len().saturating_sub(max_pairs);
274 uses[start..]
275 .iter()
276 .filter_map(|(key, provider_id, name, input)| {
277 if key.as_str().trim().is_empty()
278 || provider_id.trim().is_empty()
279 || use_counts.get(key) != Some(&1)
280 {
281 return None;
282 }
283 let result = match results.get(key) {
284 Some(Some((result_provider, content))) if result_provider == provider_id => {
285 content.clone()
286 }
287 Some(_) => return None,
288 None => "(pending)".to_string(),
289 };
290 Some(ToolCallPair {
291 name: (*name).to_string(),
292 input_preview: input.clone(),
293 result_preview: result,
294 })
295 })
296 .collect()
297 }
298
299 /// Build the user prompt for the advisor from a slice of tool-call pairs.
300 #[must_use]
301 pub fn build_advisor_prompt(pairs: &[ToolCallPair]) -> String {
302 let mut out = String::from("Recent tool activity to review (oldest → newest):\n\n");
303 for (i, pair) in pairs.iter().enumerate() {
304 out.push_str(&format!(
305 "{}. tool={}\n input: {}\n result: {}\n\n",
306 i + 1,
307 pair.name,
308 pair.input_preview,
309 pair.result_preview
310 ));
311 }
312 out.push_str(
313 "Provide your advisory in one to three sentences, or respond with \"ok\" if nothing notable.",
314 );
315 out
316 }
317
318 /// Run one advisor review cycle for a completed turn.
319 ///
320 /// This is the async work dispatched by `spawn_supervised` in the engine. It:
321 /// 1. Checks whether emission is allowed by `guard`.
322 /// 2. Extracts bounded tool-call pairs from `messages`.
323 /// 3. Makes a non-streaming LLM call with a short read-only prompt.
324 /// 4. Checks emission again (the LLM call may have taken time).
325 /// 5. Sends `Event::AdvisoryNote` if the note passes the guard.
326 ///
327 /// All errors are logged and swallowed — the advisor must never fail the
328 /// parent turn.
329 pub async fn run_advisor_for_turn(
330 turn_id: String,
331 messages: Vec<Message>,
332 config: AdvisorConfig,
333 client: CodewhaleClient,
334 route_config: Config,
335 session_model: String,
336 usage_context: AdvisorUsageContext,
337 guard: std::sync::Arc<tokio::sync::Mutex<EmissionGuard>>,
338 tx_event: mpsc::Sender<Event>,
339 ) {
340 // Pre-flight: skip if the guard already blocks (avoids the LLM call when
341 // rate-limited, which is the common case for rapid turn sequences).
342 {
343 let g = guard.lock().await;
344 // We don't have the note content yet, so we only check the rate limit
345 // here by testing with a placeholder. The dedup check runs after the
346 // LLM call, when we have the actual content.
347 if let Some(last) = g.last_emission
348 && std::time::Instant::now().duration_since(last) < config.rate_limit
349 {
350 debug!(target: "advisor", "rate-limited, skipping advisor run for turn {turn_id}");
351 return;
352 }
353 }
354
355 // Extract a bounded slice of tool-call pairs.
356 let pairs = extract_tool_call_pairs(&messages, config.max_tool_calls as usize);
357 if pairs.is_empty() {
358 debug!(target: "advisor", "no tool calls found; skipping advisor for turn {turn_id}");
359 return;
360 }
361
362 let tool_call_count = pairs.len() as u32;
363 let prompt = build_advisor_prompt(&pairs);
364 let model = config
365 .model
366 .clone()
367 .unwrap_or_else(|| session_model.clone());
368
369 let (client, model) = match exact_advisor_client(&route_config, client, &session_model, &model)
370 {
371 Ok(route) => route,
372 Err(error) => {
373 tracing::warn!(target: "advisor", "advisor route resolution failed for turn {turn_id}: {error}");
374 return;
375 }
376 };
377 let route = client.effective_route_envelope(&model, chrono::Utc::now());
378 let request = MessageRequest {
379 model: model.clone(),
380 messages: vec![Message {
381 role: Role::User,
382 content: vec![ContentBlock::Text {
383 text: prompt,
384 cache_control: None,
385 }],
386 }],
387 max_tokens: ADVISOR_MAX_TOKENS,
388 system: Some(SystemPrompt::Text(ADVISOR_SYSTEM_PROMPT.to_string())),
389 tools: None,
390 tool_choice: None,
391 metadata: None,
392 thinking: None,
393 // The advisor has a deliberately tiny answer contract. Hidden
394 // reasoning would spend that allowance before the note is emitted.
395 reasoning_effort: Some("off".to_string()),
396 stream: Some(false),
397 temperature: None,
398 top_p: None,
399 };
400
401 let response = match client.create_message(request).await {
402 Ok(r) => r,
403 Err(e) => {
404 tracing::warn!(target: "advisor", "advisor LLM call failed for turn {turn_id}: {e}");
405 return;
406 }
407 };
408
409 // A decoded provider response is billable even when its partial/empty
410 // content is rejected below or the emission guard suppresses a duplicate.
411 let usage_source_id = format!("advisor:{turn_id}:provider-response:0");
412 if response.usage == codewhale_models::Usage::default() {
413 usage_context.report_unreceipted(&usage_source_id, &route);
414 tracing::warn!(
415 target: "advisor",
416 "advisor provider response omitted usage for turn {turn_id}; cost coverage is unknown"
417 );
418 } else {
419 usage_context.report(&usage_source_id, &route, &response.usage);
420 }
421
422 if codewhale_models::is_incomplete_stop_reason(response.stop_reason.as_deref()) {
423 tracing::warn!(
424 target: "advisor",
425 "advisor response incomplete for turn {turn_id} (stop reason `{}`); dropping partial note",
426 codewhale_models::stop_reason_detail(response.stop_reason.as_deref())
427 );
428 return;
429 }
430
431 // Extract the text from the response.
432 let note: String = response
433 .content
434 .iter()
435 .filter_map(|block| {
436 if let ContentBlock::Text { text, .. } = block {
437 Some(text.as_str())
438 } else {
439 None
440 }
441 })
442 .collect::<Vec<_>>()
443 .join("\n")
444 .trim()
445 .to_string();
446
447 if note.is_empty() {
448 debug!(target: "advisor", "empty advisor response for turn {turn_id}; skipping");
449 return;
450 }
451
452 // Post-flight emission check (rate limit + dedup).
453 let mut guard_lock = guard.lock().await;
454 if !guard_lock.may_emit(&note, &config) {
455 debug!(target: "advisor", "emission suppressed by guard for turn {turn_id}");
456 return;
457 }
458
459 guard_lock.record_emission(&note);
460 drop(guard_lock);
461
462 let _ = tx_event
463 .send(Event::AdvisoryNote {
464 turn_id: turn_id.clone(),
465 note: note.clone(),
466 tool_call_count,
467 })
468 .await;
469
470 debug!(target: "advisor", "advisory note emitted for turn {turn_id} ({tool_call_count} tool calls reviewed)");
471 }
472
473 fn exact_advisor_client(
474 config: &Config,
475 parent_client: CodewhaleClient,
476 session_model: &str,
477 requested_model: &str,
478 ) -> anyhow::Result<(CodewhaleClient, String)> {
479 if requested_model
480 .trim()
481 .eq_ignore_ascii_case(session_model.trim())
482 {
483 return Ok((parent_client, session_model.trim().to_string()));
484 }
485
486 if config.providers.as_ref().is_some_and(|providers| {
487 providers.custom.values().any(|provider| {
488 provider
489 .model
490 .as_deref()
491 .is_some_and(|model| model.trim().eq_ignore_ascii_case(requested_model.trim()))
492 })
493 }) {
494 anyhow::bail!(
495 "advisor model `{}` belongs to a custom provider but no exact provider identity is carried",
496 requested_model.trim()
497 );
498 }
499
500 let active = config
501 .active_provider_identity()
502 .map_err(anyhow::Error::msg)?;
503 let selection =
504 crate::model_routing::resolve_explicit_route_with_inventory(config, requested_model);
505 let (identity, model) = if let Some(selection) = selection {
506 (selection.provider, selection.model)
507 } else {
508 let candidates =
509 crate::model_routing::explicit_route_candidate_providers(config, requested_model);
510 if !candidates.is_empty() && !candidates.contains(&active) {
511 anyhow::bail!(
512 "advisor model `{}` is not owned by the originating provider and has no unique exact route",
513 requested_model.trim()
514 );
515 }
516 (active, requested_model.trim().to_string())
517 };
518 let client =
519 crate::route_runtime::resolve_runtime_route_for_identity(config, &identity, Some(&model))
520 .map_err(anyhow::Error::msg)?
521 .validate()
522 .map(|route| route.client)
523 .map_err(anyhow::Error::msg)?;
524 Ok((client, model))
525 }
526
527 fn hash_str(s: &str) -> u64 {
528 let mut h = DefaultHasher::new();
529 s.hash(&mut h);
530 h.finish()
531 }
532
533 // ── Tests ──────────────────────────────────────────────────────────────────
534
535 #[cfg(test)]
536 mod tests {
537 use super::*;
538 use crate::config::{ProviderConfig, ProvidersConfig};
539 use std::time::Duration;
540 use wiremock::matchers::{method, path};
541 use wiremock::{Mock, MockServer, ResponseTemplate};
542
543 fn test_config() -> AdvisorConfig {
544 AdvisorConfig {
545 enabled: true,
546 max_tool_calls: 5,
547 rate_limit: Duration::from_secs(1),
548 dedup_window: Duration::from_secs(10),
549 model: None,
550 }
551 }
552
553 // ── enable/disable ────────────────────────────────────────────────────
554
555 #[test]
556 fn disabled_config_has_enabled_false() {
557 let cfg = AdvisorConfig::disabled();
558 assert!(!cfg.enabled);
559 }
560
561 #[test]
562 fn from_toml_clamps_max_tool_calls() {
563 let toml = AdvisorConfigToml {
564 enabled: true,
565 max_tool_calls: 999,
566 rate_limit_secs: 60,
567 dedup_window_secs: 300,
568 model: None,
569 };
570 let cfg = AdvisorConfig::from_toml(&toml);
571 assert_eq!(
572 cfg.max_tool_calls, 50,
573 "max_tool_calls must be clamped to 50"
574 );
575 }
576
577 #[test]
578 fn from_toml_clamps_rate_limit() {
579 let toml = AdvisorConfigToml {
580 enabled: true,
581 max_tool_calls: 10,
582 rate_limit_secs: 0, // below minimum of 5
583 dedup_window_secs: 300,
584 model: None,
585 };
586 let cfg = AdvisorConfig::from_toml(&toml);
587 assert!(
588 cfg.rate_limit >= Duration::from_secs(5),
589 "rate_limit must be at least 5s"
590 );
591 }
592
593 // ── bounded input ─────────────────────────────────────────────────────
594
595 fn make_messages_with_n_tool_calls(n: usize) -> Vec<Message> {
596 let mut messages = Vec::new();
597 for i in 0..n {
598 let id = format!("tool_{i}");
599 // assistant message with ToolUse
600 messages.push(Message {
601 role: Role::Assistant,
602 content: vec![ContentBlock::ToolUse {
603 execution_id: None,
604 id: id.clone(),
605 name: "exec_shell".to_string(),
606 input: serde_json::json!({"command": format!("echo {i}")}),
607 caller: None,
608 thought_signature: None,
609 }],
610 });
611 // user message with ToolResult
612 messages.push(Message {
613 role: Role::User,
614 content: vec![ContentBlock::ToolResult {
615 execution_id: None,
616 tool_use_id: id,
617 content: format!("{i}"),
618 is_error: None,
619 content_blocks: None,
620 }],
621 });
622 }
623 messages
624 }
625
626 #[test]
627 fn extract_tool_call_pairs_bounded_by_max() {
628 let messages = make_messages_with_n_tool_calls(20);
629 let pairs = extract_tool_call_pairs(&messages, 5);
630 assert_eq!(pairs.len(), 5, "must return at most max_pairs");
631 // Should be the last 5 (newest).
632 assert_eq!(pairs[0].name, "exec_shell");
633 }
634
635 #[test]
636 fn extract_tool_call_pairs_empty_when_no_tool_calls() {
637 let messages = vec![Message {
638 role: Role::User,
639 content: vec![ContentBlock::Text {
640 text: "hello".to_string(),
641 cache_control: None,
642 }],
643 }];
644 let pairs = extract_tool_call_pairs(&messages, 5);
645 assert!(pairs.is_empty());
646 }
647
648 #[test]
649 fn extract_tool_call_pairs_fewer_than_max_returns_all() {
650 let messages = make_messages_with_n_tool_calls(3);
651 let pairs = extract_tool_call_pairs(&messages, 10);
652 assert_eq!(pairs.len(), 3);
653 }
654
655 #[test]
656 fn extract_tool_call_pairs_keep_execution_and_legacy_domains_distinct() {
657 let messages: Vec<Message> = serde_json::from_value(serde_json::json!([
658 {"role":"assistant", "content":[{"type":"tool_use", "id":"wire",
659 "execution_id":"first", "name":"read_file", "input":{"path":"first.txt"}}]},
660 {"role":"user", "content":[{"type":"tool_result", "tool_use_id":"wire",
661 "execution_id":"first", "content":"first output"}]},
662 {"role":"assistant", "content":[{"type":"tool_use", "id":"wire",
663 "execution_id":"second", "name":"read_file", "input":{"path":"second.txt"}}]},
664 {"role":"user", "content":[{"type":"tool_result", "tool_use_id":"wire",
665 "execution_id":"second", "content":"second output"}]},
666 {"role":"assistant", "content":[{"type":"tool_use", "id":"first",
667 "name":"read_file", "input":{"path":"legacy.txt"}}]},
668 {"role":"user", "content":[{"type":"tool_result", "tool_use_id":"first",
669 "content":"legacy output"}]}
670 ]))
671 .expect("transcript");
672 let pairs = extract_tool_call_pairs(&messages, 3);
673 assert_eq!(pairs.len(), 3);
674 for (pair, path, output) in [
675 (&pairs[0], "first.txt", "first output"),
676 (&pairs[1], "second.txt", "second output"),
677 (&pairs[2], "legacy.txt", "legacy output"),
678 ] {
679 assert!(pair.input_preview.contains(path));
680 assert_eq!(pair.result_preview, output);
681 }
682 let bounded = extract_tool_call_pairs(&messages, 2);
683 assert_eq!(bounded.len(), 2);
684 assert_eq!(bounded[0].result_preview, "second output");
685 assert_eq!(bounded[1].result_preview, "legacy output");
686 assert!(extract_tool_call_pairs(&messages, 0).is_empty());
687 }
688
689 #[test]
690 fn extract_tool_call_pairs_refuse_corrupt_identity_without_fallback() {
691 let call = |execution_id: Option<&str>| {
692 serde_json::json!({
693 "type":"tool_use", "id":"wire", "execution_id":execution_id,
694 "name":"read_file", "input":{"path":"selected.txt"}
695 })
696 };
697 let result = |execution_id: Option<&str>, provider: &str| {
698 serde_json::json!({
699 "type":"tool_result", "tool_use_id":provider, "execution_id":execution_id,
700 "content":"must not borrow this output"
701 })
702 };
703 for (label, calls, results, pending) in [
704 (
705 "duplicate local calls",
706 vec![call(Some("exec")), call(Some("exec"))],
707 vec![result(Some("exec"), "wire")],
708 false,
709 ),
710 (
711 "duplicate legacy calls",
712 vec![call(None), call(None)],
713 vec![result(None, "wire")],
714 false,
715 ),
716 (
717 "duplicate local results",
718 vec![call(Some("exec"))],
719 vec![result(Some("exec"), "wire"), result(Some("exec"), "wire")],
720 false,
721 ),
722 (
723 "duplicate legacy results",
724 vec![call(None)],
725 vec![result(None, "wire"), result(None, "wire")],
726 false,
727 ),
728 (
729 "wrong provider",
730 vec![call(Some("exec"))],
731 vec![result(Some("exec"), "other")],
732 false,
733 ),
734 (
735 "empty local identity",
736 vec![call(Some(""))],
737 vec![result(Some(""), "wire")],
738 false,
739 ),
740 (
741 "no local fallback",
742 vec![call(Some("exec"))],
743 vec![result(Some("wrong"), "wire"), result(None, "wire")],
744 true,
745 ),
746 (
747 "no legacy fallback",
748 vec![call(None)],
749 vec![result(Some("wire"), "wire")],
750 true,
751 ),
752 ("missing result", vec![call(Some("exec"))], vec![], true),
753 ] {
754 let messages: Vec<Message> = serde_json::from_value(serde_json::json!([
755 {"role":"assistant", "content":calls}, {"role":"user", "content":results}
756 ]))
757 .expect("transcript");
758 let pairs = extract_tool_call_pairs(&messages, 5);
759 assert_eq!(pairs.len(), usize::from(pending), "{label}");
760 if pending {
761 assert_eq!(pairs[0].result_preview, "(pending)", "{label}");
762 }
763 assert!(
764 !build_advisor_prompt(&pairs).contains("must not borrow"),
765 "{label}"
766 );
767 }
768 }
769
770 // ── rate limiting ─────────────────────────────────────────────────────
771
772 #[test]
773 fn emission_guard_allows_first_emission() {
774 let guard = EmissionGuard::new();
775 let config = test_config();
776 assert!(
777 guard.may_emit("something concerning here", &config),
778 "first emission must be allowed"
779 );
780 }
781
782 #[test]
783 fn emission_guard_blocks_immediately_after_emission() {
784 let mut guard = EmissionGuard::new();
785 let config = test_config();
786 let note = "something concerning";
787 guard.record_emission(note);
788 assert!(
789 !guard.may_emit("a completely different note", &config),
790 "emission must be blocked immediately after a prior emission (rate limit)"
791 );
792 }
793
794 #[test]
795 fn emission_guard_allows_after_rate_limit_expires() {
796 let mut guard = EmissionGuard::new();
797 // Rate limit of 0ms — always expired.
798 let config = AdvisorConfig {
799 rate_limit: Duration::ZERO,
800 dedup_window: Duration::from_secs(300),
801 ..AdvisorConfig::disabled()
802 };
803 let note = "first note";
804 guard.record_emission(note);
805 assert!(
806 guard.may_emit("second different note", &config),
807 "emission must be allowed when rate limit duration is zero"
808 );
809 }
810
811 // ── deduplication ─────────────────────────────────────────────────────
812
813 #[test]
814 fn emission_guard_suppresses_ok_response() {
815 let guard = EmissionGuard::new();
816 let config = test_config();
817 assert!(!guard.may_emit("ok", &config), "\"ok\" must be suppressed");
818 assert!(!guard.may_emit("OK", &config), "\"OK\" must be suppressed");
819 assert!(
820 !guard.may_emit(" ok ", &config),
821 "\" ok \" must be suppressed"
822 );
823 }
824
825 #[test]
826 fn emission_guard_dedup_blocks_identical_note_within_window() {
827 let mut guard = EmissionGuard::new();
828 // Use a zero rate limit so only dedup is tested.
829 let config = AdvisorConfig {
830 rate_limit: Duration::ZERO,
831 dedup_window: Duration::from_secs(300),
832 ..AdvisorConfig::disabled()
833 };
834 let note = "risky shell command with no error checking";
835 guard.record_emission(note);
836 assert!(
837 !guard.may_emit(note, &config),
838 "identical note must be suppressed within the dedup window"
839 );
840 }
841
842 #[test]
843 fn emission_guard_allows_different_note_within_dedup_window() {
844 let mut guard = EmissionGuard::new();
845 let config = AdvisorConfig {
846 rate_limit: Duration::ZERO,
847 dedup_window: Duration::from_secs(300),
848 ..AdvisorConfig::disabled()
849 };
850 guard.record_emission("first note");
851 assert!(
852 guard.may_emit("entirely different note", &config),
853 "a different note must be allowed even within the dedup window"
854 );
855 }
856
857 // ── child failure isolation ───────────────────────────────────────────
858
859 #[test]
860 fn advisor_prompt_is_non_empty_for_non_empty_pairs() {
861 let pairs = vec![ToolCallPair {
862 name: "exec_shell".to_string(),
863 input_preview: r#"{"command":"ls -la"}"#.to_string(),
864 result_preview: "total 4\ndrwxr-xr-x 2 user user 4096".to_string(),
865 }];
866 let prompt = build_advisor_prompt(&pairs);
867 assert!(
868 prompt.contains("exec_shell"),
869 "prompt must include the tool name"
870 );
871 assert!(
872 prompt.contains("ls -la"),
873 "prompt must include the tool input"
874 );
875 }
876
877 #[test]
878 fn advisor_model_override_builds_the_owning_provider_client() {
879 let config = Config {
880 provider: Some("deepseek".to_string()),
881 providers: Some(ProvidersConfig {
882 deepseek: ProviderConfig {
883 api_key: Some("sk-deepseek-advisor-test".to_string()),
884 model: Some("deepseek-chat".to_string()),
885 ..ProviderConfig::default()
886 },
887 zai: ProviderConfig {
888 api_key: Some("zai-advisor-test-key".to_string()),
889 model: Some(crate::config::DEFAULT_ZAI_MODEL.to_string()),
890 ..ProviderConfig::default()
891 },
892 ..ProvidersConfig::default()
893 }),
894 ..Config::default()
895 };
896 let parent = CodewhaleClient::new(&config).expect("parent client");
897 let (advisor, resolved_model) = exact_advisor_client(
898 &config,
899 parent,
900 "deepseek-chat",
901 crate::config::DEFAULT_ZAI_MODEL,
902 )
903 .expect("cross-provider advisor route");
904 let route = advisor.effective_route_envelope(&resolved_model, chrono::Utc::now());
905
906 assert_eq!(route.provider, crate::config::ProviderKind::Zai);
907 assert_eq!(route.provider_identity, "zai");
908 assert_eq!(route.model, crate::config::DEFAULT_ZAI_MODEL);
909 }
910
911 #[test]
912 fn advisor_foreign_custom_override_fails_closed_without_exact_identity() {
913 let config = Config {
914 provider: Some("deepseek".to_string()),
915 providers: Some(ProvidersConfig {
916 deepseek: ProviderConfig {
917 api_key: Some("sk-deepseek-advisor-test".to_string()),
918 model: Some("deepseek-chat".to_string()),
919 ..ProviderConfig::default()
920 },
921 custom: [(
922 "private-route".to_string(),
923 ProviderConfig {
924 api_key: Some("custom-advisor-test-key".to_string()),
925 base_url: Some("https://custom.invalid/v1".to_string()),
926 model: Some("private-advisor-model".to_string()),
927 ..ProviderConfig::default()
928 },
929 )]
930 .into_iter()
931 .collect(),
932 ..ProvidersConfig::default()
933 }),
934 ..Config::default()
935 };
936 let parent = CodewhaleClient::new(&config).expect("parent client");
937 let error =
938 match exact_advisor_client(&config, parent, "deepseek-chat", "private-advisor-model") {
939 Ok(_) => panic!("generic custom kind cannot identify the exact foreign route"),
940 Err(error) => error,
941 };
942 assert!(
943 error.to_string().contains("exact provider identity"),
944 "{error}"
945 );
946 }
947
948 #[test]
949 fn advisor_named_custom_a_cannot_route_model_owned_by_custom_b() {
950 let config = Config {
951 provider: Some("custom-a".to_string()),
952 providers: Some(ProvidersConfig {
953 custom: [
954 (
955 "custom-a".to_string(),
956 ProviderConfig {
957 api_key: Some("custom-a-advisor-test-key".to_string()),
958 base_url: Some("https://custom-a.invalid/v1".to_string()),
959 model: Some("custom-a-model".to_string()),
960 kind: Some("openai-compatible".to_string()),
961 ..ProviderConfig::default()
962 },
963 ),
964 (
965 "custom-b".to_string(),
966 ProviderConfig {
967 api_key: Some("custom-b-advisor-test-key".to_string()),
968 base_url: Some("https://custom-b.invalid/v1".to_string()),
969 model: Some("custom-b-model".to_string()),
970 kind: Some("openai-compatible".to_string()),
971 ..ProviderConfig::default()
972 },
973 ),
974 ]
975 .into_iter()
976 .collect(),
977 ..ProvidersConfig::default()
978 }),
979 ..Config::default()
980 };
981 let parent = CodewhaleClient::new(&config).expect("active custom-a client");
982 let error = match exact_advisor_client(&config, parent, "custom-a-model", "custom-b-model")
983 {
984 Ok(_) => panic!("custom-b must not reuse custom-a's endpoint or credential"),
985 Err(error) => error,
986 };
987 assert!(
988 error.to_string().contains("exact provider identity"),
989 "{error}"
990 );
991 }
992
993 async fn run_billed_advisor_fixture(
994 note: &str,
995 stop_reason: &str,
996 suppress_as_duplicate: bool,
997 include_usage: bool,
998 ) -> (crate::cost_status::PendingBackgroundCost, Option<Event>) {
999 let _scope = crate::cost_status::test_scope();
1000 let server = MockServer::start().await;
1001 let mut provider_response = serde_json::json!({
1002 "id": "advisor-provider-response",
1003 "model": "deepseek-chat",
1004 "choices": [{
1005 "index": 0,
1006 "message": {"role": "assistant", "content": note},
1007 "finish_reason": stop_reason
1008 }]
1009 });
1010 if include_usage {
1011 provider_response["usage"] = serde_json::json!({
1012 "prompt_tokens": 9,
1013 "completion_tokens": 3,
1014 "total_tokens": 12
1015 });
1016 }
1017 Mock::given(method("POST"))
1018 .and(path("/v1/chat/completions"))
1019 .respond_with(ResponseTemplate::new(200).set_body_json(provider_response))
1020 .expect(1)
1021 .mount(&server)
1022 .await;
1023 let route_config = Config {
1024 provider: Some("deepseek".to_string()),
1025 providers: Some(ProvidersConfig {
1026 deepseek: ProviderConfig {
1027 api_key: Some("sk-deepseek-advisor-test".to_string()),
1028 model: Some("deepseek-chat".to_string()),
1029 ..ProviderConfig::default()
1030 },
1031 ..ProvidersConfig::default()
1032 }),
1033 ..Config::default()
1034 };
1035 let mut client = CodewhaleClient::new(&route_config).expect("advisor client");
1036 client.set_test_chat_transport_base_url(server.uri());
1037 let mut emission_guard = EmissionGuard::new();
1038 if suppress_as_duplicate {
1039 emission_guard.record_emission(note);
1040 }
1041 let guard = std::sync::Arc::new(tokio::sync::Mutex::new(emission_guard));
1042 let (tx, mut rx) = mpsc::channel(1);
1043 run_advisor_for_turn(
1044 "advisor-turn".to_string(),
1045 make_messages_with_n_tool_calls(1),
1046 AdvisorConfig {
1047 enabled: true,
1048 max_tool_calls: 5,
1049 rate_limit: Duration::ZERO,
1050 dedup_window: Duration::from_secs(60),
1051 model: None,
1052 },
1053 client,
1054 route_config,
1055 "deepseek-chat".to_string(),
1056 AdvisorUsageContext::capture(None),
1057 guard,
1058 tx,
1059 )
1060 .await;
1061 (crate::cost_status::drain(), rx.try_recv().ok())
1062 }
1063
1064 #[tokio::test]
1065 async fn advisor_incomplete_and_dedup_suppressed_responses_are_each_billed_once() {
1066 let (incomplete, incomplete_event) =
1067 run_billed_advisor_fixture("partial note", "max_tokens", false, true).await;
1068 assert!(incomplete_event.is_none());
1069 assert_eq!(
1070 incomplete
1071 .priced_turns
1072 .saturating_add(incomplete.unpriced_turns),
1073 1
1074 );
1075
1076 let (dedup, dedup_event) =
1077 run_billed_advisor_fixture("same advisory", "stop", true, true).await;
1078 assert!(dedup_event.is_none());
1079 assert_eq!(dedup.priced_turns.saturating_add(dedup.unpriced_turns), 1);
1080 }
1081
1082 #[tokio::test]
1083 async fn advisor_provider_success_without_usage_marks_unknown_once() {
1084 let (pending, event) =
1085 run_billed_advisor_fixture("use a smaller focused slice", "stop", false, false).await;
1086
1087 assert!(
1088 event.is_some(),
1089 "the semantic advisor response remains usable"
1090 );
1091 assert_eq!(pending.priced_turns, 0);
1092 assert_eq!(pending.unpriced_turns, 1);
1093 assert_eq!(pending.cny_unpriced_turns, 1);
1094 assert!(
1095 pending
1096 .unpriced_reasons
1097 .contains("provider_success_missing_usage")
1098 );
1099 }
1100 }
1101
1101 lines RUST