1use reqwest::Client;
2use serde_json::{json, Value};
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::time::Duration;
6use crate::{Result, RuntimeError, ToolRegistry};
7use std::sync::Mutex;
8use tokio::sync::{mpsc, RwLock};
9use tokio_stream::wrappers::UnboundedReceiverStream;
10use tokio_util::sync::CancellationToken;
11use futures::stream::Stream;
12use std::pin::Pin;
13
14mod types;
15mod auth;
16mod api;
17mod api_sync;
18mod request;
19mod stream;
20mod helpers;
21pub mod subagent;
22pub mod openai;
23
24pub use types::{StreamEvent, LlmEvent, SessionEvent, AgentEvent};
25use types::AuthState;
26use auth::AuthMethods;
27use api::ApiMethods;
28use stream::StreamMethods;
29use helpers::HelperMethods;
30
31pub enum BeforeToolCallDecision {
33 Continue { input: Value },
34 Block { reason: String },
35}
36
37pub async fn emit_before_tool_call(
40 hook_bus: &Arc<crate::extensions::hooks::HookBus>,
41 tool_name: &str,
42 runtime_tool_name: Option<&str>,
43 input: Value,
44) -> crate::extensions::hooks::events::HookResult {
45 let mut event = crate::extensions::hooks::events::HookEvent::before_tool_call(tool_name, input);
46 if let Some(runtime_tool_name) = runtime_tool_name {
47 event.tool_runtime_name = Some(runtime_tool_name.to_string());
48 }
49 hook_bus.emit(&event).await
50}
51
52
53pub async fn resolve_before_tool_call_result(
57 hook_result: crate::extensions::hooks::events::HookResult,
58 secret_prompt: Option<&crate::tools::SecretPromptHandle>,
59) -> crate::extensions::hooks::events::HookResult {
60 match hook_result {
61 crate::extensions::hooks::events::HookResult::Confirm { message } => {
62 let Some(prompt) = secret_prompt else {
63 return crate::extensions::hooks::events::HookResult::Block {
64 reason: format!(
65 "Tool call requires confirmation but no interactive prompt is available: {}",
66 message
67 ),
68 };
69 };
70
71 let response = prompt
72 .prompt(
73 "Confirm tool call".to_string(),
74 format!("{}\n\nType 'yes' or 'y' to allow.", message),
75 )
76 .await;
77
78 match response.as_deref().map(str::trim) {
79 Some(answer) if answer.eq_ignore_ascii_case("yes") || answer.eq_ignore_ascii_case("y") => {
80 crate::extensions::hooks::events::HookResult::Continue
81 }
82 _ => crate::extensions::hooks::events::HookResult::Block {
83 reason: format!("Tool call confirmation denied: {}", message),
84 },
85 }
86 }
87 other => other,
88 }
89}
90
91pub async fn resolve_before_tool_call_decision(
93 original_input: Value,
94 hook_result: crate::extensions::hooks::events::HookResult,
95 secret_prompt: Option<&crate::tools::SecretPromptHandle>,
96) -> BeforeToolCallDecision {
97 match resolve_before_tool_call_result(hook_result, secret_prompt).await {
98 crate::extensions::hooks::events::HookResult::Block { reason } => {
99 BeforeToolCallDecision::Block { reason }
100 }
101 crate::extensions::hooks::events::HookResult::Modify { input } => {
102 BeforeToolCallDecision::Continue { input }
103 }
104 _ => BeforeToolCallDecision::Continue { input: original_input },
105 }
106}
107
108pub async fn emit_after_tool_call(
111 hook_bus: &Arc<crate::extensions::hooks::HookBus>,
112 tool_name: &str,
113 runtime_tool_name: Option<&str>,
114 input: Value,
115 output: String,
116) -> crate::extensions::hooks::events::HookResult {
117 let mut event = crate::extensions::hooks::events::HookEvent::after_tool_call(
118 tool_name,
119 input,
120 output,
121 );
122 if let Some(runtime_tool_name) = runtime_tool_name {
123 event.tool_runtime_name = Some(runtime_tool_name.to_string());
124 }
125 hook_bus.emit(&event).await
126}
127
128pub struct Runtime {
131 client: Client,
132 auth: Arc<RwLock<AuthState>>,
133 model: String,
134 tools: Arc<RwLock<ToolRegistry>>,
135 system_prompt: Option<String>,
136 thinking_budget: u32,
137 context_window_override: Option<u64>,
142 compaction_model: Option<String>,
144 subagent_registry: Arc<Mutex<crate::runtime::subagent::SubagentRegistry>>,
146 event_queue: Arc<crate::events::EventQueue>,
148 pub watcher_exit_path: Option<PathBuf>,
150 max_tool_output: usize,
152 bash_timeout: u64,
153 bash_max_timeout: u64,
154 subagent_timeout: u64,
155 api_retries: u32,
156 session_manager: std::sync::Arc<crate::tools::shell::SessionManager>,
157 hook_bus: Arc<crate::extensions::hooks::HookBus>,
159 #[allow(dead_code)]
161 reaper_handle: Option<tokio::task::JoinHandle<()>>,
162 #[allow(dead_code)]
163 reaper_cancel: Option<tokio_util::sync::CancellationToken>,
164}
165
166impl Runtime {
167 pub async fn new() -> Result<Self> {
168 let (auth_token, auth_type, refresh_token, token_expires) = AuthMethods::get_auth_token()?;
169
170 let client = Client::builder()
171 .connect_timeout(Duration::from_secs(10))
172 .timeout(Duration::from_secs(300))
173 .build()
174 .map_err(|e| RuntimeError::Config(format!("Failed to build HTTP client: {}", e)))?;
175
176 let session_manager = {
177 let config = crate::tools::shell::ShellConfig::default();
178 crate::tools::shell::SessionManager::new(config)
179 };
180
181 let mgr = session_manager.clone();
183 let cancel = tokio_util::sync::CancellationToken::new();
184 let reaper_handle = crate::tools::shell::session::start_reaper(mgr, cancel.clone());
185
186 Ok(Runtime {
187 client,
188 auth: Arc::new(RwLock::new(AuthState {
189 auth_token,
190 auth_type,
191 refresh_token,
192 token_expires,
193 })),
194 model: crate::models::default_model().to_string(),
195 tools: Arc::new(RwLock::new(ToolRegistry::new())),
196 system_prompt: None,
197 thinking_budget: 4096,
198 context_window_override: None,
199 compaction_model: None,
200 subagent_registry: Arc::new(Mutex::new(crate::runtime::subagent::SubagentRegistry::new())),
201 event_queue: Arc::new(crate::events::EventQueue::new(1000)),
202 watcher_exit_path: None,
203 max_tool_output: 30000,
204 bash_timeout: 30,
205 bash_max_timeout: 300,
206 subagent_timeout: 300,
207 api_retries: 3,
208 session_manager,
209 hook_bus: Arc::new(crate::extensions::hooks::HookBus::new()),
210 reaper_handle: Some(reaper_handle),
211 reaper_cancel: Some(cancel),
212 })
213 }
214
215 pub fn set_system_prompt(&mut self, prompt: String) {
216 self.system_prompt = Some(prompt);
217 }
218
219 pub fn system_prompt(&self) -> Option<&str> {
220 self.system_prompt.as_deref()
221 }
222
223 pub fn set_model(&mut self, model: String) {
224 let cleaned = if let Some(pos) = model.find("claude-") {
226 model[pos..].to_string()
227 } else if let Some(pos) = model.find('/') {
228 let before = &model[..pos];
229 let key_start = before.rfind(|c: char| !c.is_ascii_alphanumeric() && c != '-' && c != '_')
230 .map(|i| i + before[i..].chars().next().map(|c| c.len_utf8()).unwrap_or(1))
231 .unwrap_or(0);
232 model[key_start..].to_string()
233 } else {
234 model
235 };
236 self.model = cleaned;
237 }
238
239 pub fn set_tools(&mut self, tools: ToolRegistry) {
240 self.tools = Arc::new(RwLock::new(tools));
241 }
242
243 pub fn subagent_registry(&self) -> &Arc<Mutex<crate::runtime::subagent::SubagentRegistry>> {
244 &self.subagent_registry
245 }
246
247 pub fn event_queue(&self) -> &Arc<crate::events::EventQueue> {
248 &self.event_queue
249 }
250
251 pub fn hook_bus(&self) -> &Arc<crate::extensions::hooks::HookBus> {
253 &self.hook_bus
254 }
255
256 pub fn tools_shared(&self) -> Arc<RwLock<ToolRegistry>> {
258 Arc::clone(&self.tools)
259 }
260
261 pub fn model(&self) -> &str {
262 &self.model
263 }
264
265 pub fn http_client(&self) -> &Client {
266 &self.client
267 }
268 pub fn set_thinking_budget(&mut self, budget: u32) {
269 self.thinking_budget = budget;
270 }
271
272 pub fn set_compaction_model(&mut self, model: Option<String>) {
273 self.compaction_model = model;
274 }
275
276 pub fn set_context_window(&mut self, window: Option<u64>) {
277 self.context_window_override = window;
278 }
279
280 pub fn compaction_model(&self) -> &str {
283 self.compaction_model.as_deref().unwrap_or("claude-sonnet-4-6")
284 }
285
286 pub fn context_window(&self) -> u64 {
287 self.context_window_override
288 .unwrap_or_else(|| crate::models::context_window_for_model(&self.model))
289 }
290
291 pub fn apply_config(&mut self, config: &crate::config::SynapsConfig) {
293 if let Some(ref model) = config.model {
294 self.set_model(model.clone());
295 }
296 if let Some(budget) = config.thinking_budget {
297 self.set_thinking_budget(budget);
298 }
299 self.context_window_override = config.context_window;
300 self.compaction_model = config.compaction_model.clone();
301 self.max_tool_output = config.max_tool_output;
302 self.bash_timeout = config.bash_timeout;
303 self.bash_max_timeout = config.bash_max_timeout;
304 self.subagent_timeout = config.subagent_timeout;
305 self.api_retries = config.api_retries;
306 }
307
308 pub fn thinking_budget(&self) -> u32 {
309 self.thinking_budget
310 }
311
312 pub fn max_tool_output(&self) -> usize {
313 self.max_tool_output
314 }
315
316 pub fn bash_timeout(&self) -> u64 {
317 self.bash_timeout
318 }
319
320 pub fn bash_max_timeout(&self) -> u64 {
321 self.bash_max_timeout
322 }
323
324 pub fn subagent_timeout(&self) -> u64 {
325 self.subagent_timeout
326 }
327
328 pub fn api_retries(&self) -> u32 {
329 self.api_retries
330 }
331
332 pub fn set_max_tool_output(&mut self, v: usize) {
333 self.max_tool_output = v;
334 }
335
336 pub fn set_bash_timeout(&mut self, v: u64) {
337 self.bash_timeout = v;
338 }
339
340 pub fn set_bash_max_timeout(&mut self, v: u64) {
341 self.bash_max_timeout = v;
342 }
343
344 pub fn set_subagent_timeout(&mut self, v: u64) {
345 self.subagent_timeout = v;
346 }
347
348 pub fn set_api_retries(&mut self, v: u32) {
349 self.api_retries = v;
350 }
351
352 pub fn thinking_level(&self) -> &str {
353 crate::core::models::thinking_level_for_budget(self.thinking_budget)
354 }
355
356 pub async fn refresh_if_needed(&self) -> Result<()> {
358 AuthMethods::refresh_if_needed(Arc::clone(&self.auth), &self.client).await
359 }
360
361 pub async fn compact_call(&self, messages: Vec<Value>) -> Result<String> {
367 self.refresh_if_needed().await?;
368
369 use crate::core::compaction::COMPACTION_SYSTEM_PROMPT;
370
371 ApiMethods::call_api_simple(
372 &self.auth,
373 &self.client,
374 self.compaction_model(),
375 COMPACTION_SYSTEM_PROMPT,
376 self.thinking_budget,
377 &messages,
378 self.api_retries,
379 ).await
380 }
381
382 pub async fn run_single(&self, prompt: &str) -> Result<String> {
385 self.refresh_if_needed().await?;
387
388 let mut messages = vec![json!({"role": "user", "content": prompt})];
389
390 loop {
391 let response = ApiMethods::call_api(
392 &self.auth,
393 &self.client,
394 &self.model,
395 &*self.tools.read().await,
396 &self.system_prompt,
397 self.thinking_budget,
398 &messages,
399 self.api_retries,
400 &api::ApiOptions {
401 use_1m_context: self.context_window_override == Some(1_000_000),
402 },
403 ).await?;
404
405 if let Some(content) = response["content"].as_array() {
407 let mut response_text = String::new();
408 let mut tool_uses = Vec::new();
409
410 for item in content {
412 match item["type"].as_str() {
413 Some("text") => {
414 if let Some(text) = item["text"].as_str() {
415 response_text.push_str(text);
416 }
417 }
418 Some("tool_use") => {
419 tool_uses.push(item.clone());
420 }
421 _ => {}
422 }
423 }
424
425 if tool_uses.is_empty() {
427 return Ok(response_text);
428 }
429
430 messages.push(json!({
432 "role": "assistant",
433 "content": content
434 }));
435
436 let mut tool_results = Vec::new();
438
439 if tool_uses.len() == 1 {
440 let tool_use = &tool_uses[0];
442 if let (Some(tool_name), Some(tool_id)) = (
443 tool_use["name"].as_str(),
444 tool_use["id"].as_str()
445 ) {
446 let input = &tool_use["input"];
447 let result = match self.tools.read().await.get(tool_name).cloned() {
448 Some(tool) => {
449 let input = self.tools.read().await.translate_input_for_api_tool(tool_name, input.clone());
450 let runtime_name = self.tools.read().await.runtime_name_for_api(tool_name).to_string();
451 let ctx = crate::ToolContext {
452 channels: crate::tools::ToolChannels {
453 tx_delta: None,
454 tx_events: None,
455 },
456 capabilities: crate::tools::ToolCapabilities {
457 watcher_exit_path: self.watcher_exit_path.clone(),
458 tool_register_tx: None,
459 session_manager: Some(self.session_manager.clone()),
460 subagent_registry: Some(self.subagent_registry.clone()),
461 event_queue: Some(self.event_queue.clone()),
462 secret_prompt: None,
463 },
464 limits: crate::tools::ToolLimits {
465 max_tool_output: self.max_tool_output,
466 bash_timeout: self.bash_timeout,
467 bash_max_timeout: self.bash_max_timeout,
468 subagent_timeout: self.subagent_timeout,
469 },
470 };
471 let decision = resolve_before_tool_call_decision(
472 input.clone(),
473 emit_before_tool_call(
474 &self.hook_bus,
475 &tool_name,
476 Some(&runtime_name),
477 input.clone(),
478 ).await,
479 None,
480 ).await;
481 if let BeforeToolCallDecision::Block { reason } = decision {
482 format!("Tool call blocked by extension: {}", reason)
483 } else {
484 let BeforeToolCallDecision::Continue { input } = decision else { unreachable!() };
485 let input_for_hook = input.clone();
486 let output = match tool.execute(input, ctx).await {
487 Ok(output) => output,
488 Err(e) => format!("Tool execution failed: {}", e),
489 };
490 let _ = emit_after_tool_call(
491 &self.hook_bus,
492 &tool_name,
493 Some(&runtime_name),
494 input_for_hook,
495 output.clone(),
496 ).await;
497 output
498 }
499 }
500 None => format!("Unknown tool: {}", tool_name),
501 };
502 tool_results.push(json!({
503 "type": "tool_result",
504 "tool_use_id": tool_id,
505 "content": HelperMethods::truncate_tool_result(&result, self.max_tool_output)
506 }));
507 }
508 } else {
509 let mut join_set = tokio::task::JoinSet::new();
511
512 let cfg_max_tool_output = self.max_tool_output;
514 let cfg_bash_timeout = self.bash_timeout;
515 let cfg_bash_max_timeout = self.bash_max_timeout;
516 let cfg_subagent_timeout = self.subagent_timeout;
517 let session_mgr = self.session_manager.clone();
518 let cfg_subagent_registry = self.subagent_registry.clone();
519 let cfg_event_queue = self.event_queue.clone();
520 let cfg_hook_bus = self.hook_bus.clone();
521
522 for tool_use in &tool_uses {
523 if let (Some(tool_name), Some(tool_id)) = (
524 tool_use["name"].as_str().map(|s| s.to_string()),
525 tool_use["id"].as_str().map(|s| s.to_string()),
526 ) {
527 let input = tool_use["input"].clone();
528 let tools_snapshot = self.tools.read().await;
529 let input = tools_snapshot.translate_input_for_api_tool(&tool_name, input);
530 let runtime_name = tools_snapshot.runtime_name_for_api(&tool_name).to_string();
531 let tool = tools_snapshot.get(&tool_name).cloned();
532 drop(tools_snapshot);
533 let exit_path = self.watcher_exit_path.clone();
534 let session_mgr_inner = session_mgr.clone();
535 let registry_inner = cfg_subagent_registry.clone();
536 let event_queue_inner = cfg_event_queue.clone();
537 let hook_bus_inner = cfg_hook_bus.clone();
538 let tool_name_for_hook = tool_name.clone();
539 let runtime_name_for_hook = runtime_name.clone();
540
541 join_set.spawn(async move {
542 let result = match tool {
543 Some(t) => {
544 let decision = crate::runtime::resolve_before_tool_call_decision(
545 input.clone(),
546 crate::runtime::emit_before_tool_call(
547 &hook_bus_inner,
548 &tool_name_for_hook,
549 Some(&runtime_name_for_hook),
550 input.clone(),
551 ).await,
552 None,
553 ).await;
554 if let crate::runtime::BeforeToolCallDecision::Block { reason } = decision {
555 format!("Tool call blocked by extension: {}", reason)
556 } else {
557 let crate::runtime::BeforeToolCallDecision::Continue { input } = decision else { unreachable!() };
558 let ctx = crate::ToolContext {
559 channels: crate::tools::ToolChannels {
560 tx_delta: None,
561 tx_events: None,
562 },
563 capabilities: crate::tools::ToolCapabilities {
564 watcher_exit_path: exit_path,
565 tool_register_tx: None,
566 session_manager: Some(session_mgr_inner),
567 subagent_registry: Some(registry_inner),
568 event_queue: Some(event_queue_inner),
569 secret_prompt: None,
570 },
571 limits: crate::tools::ToolLimits {
572 max_tool_output: cfg_max_tool_output,
573 bash_timeout: cfg_bash_timeout,
574 bash_max_timeout: cfg_bash_max_timeout,
575 subagent_timeout: cfg_subagent_timeout,
576 },
577 };
578 let input_for_hook = input.clone();
579 let output = match t.execute(input, ctx).await {
580 Ok(output) => output,
581 Err(e) => format!("Tool execution failed: {}", e),
582 };
583 let _ = crate::runtime::emit_after_tool_call(
584 &hook_bus_inner,
585 &tool_name_for_hook,
586 Some(&runtime_name_for_hook),
587 input_for_hook,
588 output.clone(),
589 ).await;
590 output
591 }
592 }
593 None => format!("Unknown tool: {}", tool_name),
594 };
595 (tool_id, result)
596 });
597 }
598 }
599
600 let mut results_map = std::collections::HashMap::new();
602 while let Some(res) = join_set.join_next().await {
603 match res {
604 Ok((tool_id, result)) => {
605 results_map.insert(tool_id, result);
606 }
607 Err(e) => {
608 tracing::error!("Parallel tool task panicked: {}", e);
610 }
611 }
612 }
613
614 for tool_use in &tool_uses {
616 if let Some(tool_id) = tool_use["id"].as_str() {
617 let result = results_map.remove(tool_id)
618 .unwrap_or_else(|| "Tool execution failed: task panicked".to_string());
619 tool_results.push(json!({
620 "type": "tool_result",
621 "tool_use_id": tool_id,
622 "content": HelperMethods::truncate_tool_result(&result, self.max_tool_output)
623 }));
624 }
625 }
626 }
627
628 messages.push(json!({
630 "role": "user",
631 "content": tool_results
632 }));
633
634 } else {
636 return Err(RuntimeError::Tool("Invalid response format".to_string()));
637 }
638 }
639 }
640
641 pub async fn run_stream(&self, prompt: String, cancel: CancellationToken) -> Pin<Box<dyn Stream<Item = StreamEvent> + Send>> {
644 self.run_stream_with_messages(vec![json!({"role": "user", "content": prompt})], cancel, None, None).await
645 }
646
647 pub async fn run_stream_with_messages(
651 &self,
652 messages: Vec<Value>,
653 cancel: CancellationToken,
654 steering_rx: Option<mpsc::UnboundedReceiver<String>>,
655 secret_prompt: Option<crate::tools::SecretPromptHandle>,
656 ) -> Pin<Box<dyn Stream<Item = StreamEvent> + Send>> {
657 let (tx, rx) = mpsc::unbounded_channel();
658
659 if let Err(e) = self.refresh_if_needed().await {
661 let _ = tx.send(StreamEvent::Session(SessionEvent::Error(e.to_string())));
662 let _ = tx.send(StreamEvent::Session(SessionEvent::Done));
663 return Box::pin(UnboundedReceiverStream::new(rx));
664 }
665
666 let auth = Arc::clone(&self.auth);
669 let client = self.client.clone();
670 let model = self.model.clone();
671 let tools = self.tools.clone();
672 let system_prompt = self.system_prompt.clone();
673 let thinking_budget = self.thinking_budget;
674 let watcher_exit_path = self.watcher_exit_path.clone();
675 let max_tool_output = self.max_tool_output;
676 let bash_timeout = self.bash_timeout;
677 let bash_max_timeout = self.bash_max_timeout;
678 let subagent_timeout = self.subagent_timeout;
679 let api_retries = self.api_retries;
680 let session_manager = self.session_manager.clone();
681 let subagent_registry = self.subagent_registry.clone();
685 let event_queue = self.event_queue.clone();
686 let options = api::ApiOptions {
687 use_1m_context: self.context_window_override == Some(1_000_000),
688 };
689
690 let session = crate::runtime::stream::StreamSession {
691 auth, client, options, api_retries,
692 model, tools, system_prompt, thinking_budget,
693 tx: tx.clone(), cancel, steering_rx,
694 watcher_exit_path, max_tool_output,
695 bash_timeout, bash_max_timeout, subagent_timeout,
696 session_manager, subagent_registry, event_queue, secret_prompt,
697 hook_bus: self.hook_bus.clone(),
698 };
699
700 tokio::spawn(async move {
701 if let Err(e) = StreamMethods::run_stream_internal(session, messages).await {
702 let _ = tx.send(StreamEvent::Session(SessionEvent::Error(e.to_string())));
703 }
704 let _ = tx.send(StreamEvent::Session(SessionEvent::Done));
705 });
706
707 Box::pin(UnboundedReceiverStream::new(rx))
708 }
709}
710
711impl Clone for Runtime {
712 fn clone(&self) -> Self {
713 Self {
714 client: self.client.clone(),
715 auth: Arc::clone(&self.auth),
716 model: self.model.clone(),
717 tools: self.tools.clone(),
718 system_prompt: self.system_prompt.clone(),
719 thinking_budget: self.thinking_budget,
720 context_window_override: self.context_window_override,
721 compaction_model: self.compaction_model.clone(),
722 subagent_registry: self.subagent_registry.clone(),
723 event_queue: self.event_queue.clone(),
724 watcher_exit_path: self.watcher_exit_path.clone(),
725 max_tool_output: self.max_tool_output,
726 bash_timeout: self.bash_timeout,
727 bash_max_timeout: self.bash_max_timeout,
728 subagent_timeout: self.subagent_timeout,
729 api_retries: self.api_retries,
730 session_manager: self.session_manager.clone(),
731 hook_bus: self.hook_bus.clone(),
732 reaper_handle: None, reaper_cancel: None, }
735 }
736}
737
738#[cfg(test)]
739mod tests {
740 use super::*;
741
742 #[tokio::test]
743 async fn confirm_without_prompt_fails_closed() {
744 let result = resolve_before_tool_call_result(
745 crate::extensions::hooks::events::HookResult::Confirm {
746 message: "Run deploy?".into(),
747 },
748 None,
749 )
750 .await;
751
752 assert!(matches!(
753 result,
754 crate::extensions::hooks::events::HookResult::Block { reason }
755 if reason.contains("requires confirmation") && reason.contains("Run deploy?")
756 ));
757 }
758
759 #[tokio::test]
760 async fn modify_result_replaces_tool_input() {
761 let result = resolve_before_tool_call_decision(
762 serde_json::json!({"command":"rm -rf /"}),
763 crate::extensions::hooks::events::HookResult::Modify {
764 input: serde_json::json!({"command":"echo safe"}),
765 },
766 None,
767 ).await;
768
769 match result {
770 BeforeToolCallDecision::Continue { input } => {
771 assert_eq!(input, serde_json::json!({"command":"echo safe"}));
772 }
773 BeforeToolCallDecision::Block { reason } => panic!("unexpected block: {reason}"),
774 }
775 }
776
777 #[tokio::test]
778 async fn confirm_prompt_yes_continues() {
779 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
780 let handle = crate::tools::SecretPromptHandle::new(tx);
781
782 let task = tokio::spawn(async move {
783 let request = rx.recv().await.expect("confirm prompt request");
784 assert_eq!(request.title, "Confirm tool call");
785 assert!(request.prompt.contains("Run deploy?"));
786 let _ = request.response_tx.send(Some("yes".to_string()));
787 });
788
789 let result = resolve_before_tool_call_result(
790 crate::extensions::hooks::events::HookResult::Confirm {
791 message: "Run deploy?".into(),
792 },
793 Some(&handle),
794 )
795 .await;
796
797 task.await.unwrap();
798 assert!(matches!(
799 result,
800 crate::extensions::hooks::events::HookResult::Continue
801 ));
802 }
803
804 #[tokio::test]
805 async fn confirm_prompt_non_yes_blocks() {
806 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
807 let handle = crate::tools::SecretPromptHandle::new(tx);
808
809 let task = tokio::spawn(async move {
810 let request = rx.recv().await.expect("confirm prompt request");
811 let _ = request.response_tx.send(Some("no".to_string()));
812 });
813
814 let result = resolve_before_tool_call_result(
815 crate::extensions::hooks::events::HookResult::Confirm {
816 message: "Run deploy?".into(),
817 },
818 Some(&handle),
819 )
820 .await;
821
822 task.await.unwrap();
823 assert!(matches!(
824 result,
825 crate::extensions::hooks::events::HookResult::Block { reason }
826 if reason.contains("confirmation denied")
827 ));
828 }
829
830 #[test]
831 fn test_max_tokens_for_model() {
832 assert_eq!(HelperMethods::max_tokens_for_model("claude-opus-4-6"), 128000);
834 assert_eq!(HelperMethods::max_tokens_for_model("opus-something"), 128000);
835
836 assert_eq!(HelperMethods::max_tokens_for_model("claude-sonnet-4-20250514"), 64000);
838 assert_eq!(HelperMethods::max_tokens_for_model("haiku"), 64000);
839 assert_eq!(HelperMethods::max_tokens_for_model("claude-3-haiku"), 64000);
840 assert_eq!(HelperMethods::max_tokens_for_model("some-other-model"), 64000);
841
842 assert_eq!(HelperMethods::max_tokens_for_model(""), 64000);
844 assert_eq!(HelperMethods::max_tokens_for_model("OPUS"), 64000); assert_eq!(HelperMethods::max_tokens_for_model("model-opus-end"), 128000); }
847
848 #[test]
849 fn test_truncate_tool_result() {
850 let default_max = 30000;
851
852 let short = "This is a short string.";
854 assert_eq!(HelperMethods::truncate_tool_result(short, default_max), short);
855
856 let exact = "x".repeat(30000);
858 assert_eq!(HelperMethods::truncate_tool_result(&exact, default_max), exact);
859
860 let too_long = "x".repeat(30001);
862 let truncated = HelperMethods::truncate_tool_result(&too_long, default_max);
863
864 assert!(truncated.starts_with(&"x".repeat(30000)));
866
867 assert!(truncated.contains("[truncated — 30001 total chars, showing first 30000]"));
869
870 assert!(truncated.len() > 30000);
872
873 let very_long = "a".repeat(50000);
875 let truncated_very_long = HelperMethods::truncate_tool_result(&very_long, default_max);
876 assert!(truncated_very_long.contains("[truncated — 50000 total chars, showing first 30000]"));
877 assert!(truncated_very_long.starts_with(&"a".repeat(30000)));
878
879 let custom_truncated = HelperMethods::truncate_tool_result(&very_long, 100);
881 assert!(custom_truncated.starts_with(&"a".repeat(100)));
882 assert!(custom_truncated.contains("[truncated — 50000 total chars, showing first 100]"));
883 }
884
885 #[test]
886 fn test_thinking_level_ranges() {
887 use crate::core::models::thinking_level_for_budget;
888
889 assert_eq!(thinking_level_for_budget(0), "adaptive");
891
892 assert_eq!(thinking_level_for_budget(1), "low");
894 assert_eq!(thinking_level_for_budget(1024), "low");
895 assert_eq!(thinking_level_for_budget(2048), "low");
896
897 assert_eq!(thinking_level_for_budget(2049), "medium");
899 assert_eq!(thinking_level_for_budget(3000), "medium");
900 assert_eq!(thinking_level_for_budget(4096), "medium");
901
902 assert_eq!(thinking_level_for_budget(4097), "high");
904 assert_eq!(thinking_level_for_budget(8192), "high");
905 assert_eq!(thinking_level_for_budget(16384), "high");
906
907 assert_eq!(thinking_level_for_budget(16385), "xhigh");
909 assert_eq!(thinking_level_for_budget(32768), "xhigh");
910 assert_eq!(thinking_level_for_budget(100000), "xhigh");
911 }
912}