1use std::collections::BTreeMap;
2use std::fmt::{Display, Formatter};
3use std::sync::{Arc, atomic::{AtomicBool, Ordering}};
4
5use serde::{Deserialize, Serialize};
6
7use crate::compact::{
8 compact_session, estimate_session_tokens, CompactionConfig, CompactionResult,
9};
10use crate::config::RuntimeFeatureConfig;
11use crate::hooks::{HookRunResult, HookRunner};
12use crate::permissions::{PermissionOutcome, PermissionPolicy, PermissionPrompter};
13use crate::session::{ContentBlock, ConversationMessage, Session};
14use crate::usage::{TokenUsage, UsageTracker};
15
16const DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD: u32 = 200_000;
17const AUTO_COMPACTION_THRESHOLD_ENV_VAR: &str = "CLAUDE_CODE_AUTO_COMPACT_INPUT_TOKENS";
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct ApiRequest {
21 pub system_prompt: Vec<String>,
22 pub messages: Vec<ConversationMessage>,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum AssistantEvent {
27 TextDelta(String),
28 ToolUse {
29 id: String,
30 name: String,
31 input: String,
32 },
33 TaskStarted {
34 id: String,
35 label: String,
36 },
37 TaskCompleted {
38 id: String,
39 success: bool,
40 },
41 Usage(TokenUsage),
42 MessageStop,
43 ToolTelemetry { text: String },
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
47pub struct ToolResult {
48 pub output: String,
49 pub state: i8, }
51
52pub trait ApiClient {
53 fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError>;
54}
55
56pub trait ToolExecutor {
57 fn execute(&mut self, tool_name: &str, input: &str) -> Result<ToolResult, ToolError>;
58 fn query_memory(&mut self, query: &str) -> Result<String, ToolError>;
59}
60
61#[derive(Debug, Clone, PartialEq, Eq)]
62pub struct ToolError {
63 message: String,
64}
65
66impl ToolError {
67 #[must_use]
68 pub fn new(message: impl Into<String>) -> Self {
69 Self {
70 message: message.into(),
71 }
72 }
73}
74
75impl Display for ToolError {
76 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
77 write!(f, "{}", self.message)
78 }
79}
80
81impl std::error::Error for ToolError {}
82
83#[derive(Debug, Clone, PartialEq, Eq)]
84pub struct RuntimeError {
85 message: String,
86}
87
88impl RuntimeError {
89 #[must_use]
90 pub fn new(message: impl Into<String>) -> Self {
91 Self {
92 message: message.into(),
93 }
94 }
95}
96
97impl Display for RuntimeError {
98 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
99 write!(f, "{}", self.message)
100 }
101}
102
103impl std::error::Error for RuntimeError {}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct TurnSummary {
107 pub assistant_messages: Vec<ConversationMessage>,
108 pub tool_results: Vec<ConversationMessage>,
109 pub iterations: usize,
110 pub usage: TokenUsage,
111 pub auto_compaction: Option<AutoCompactionEvent>,
112}
113
114#[derive(Debug, Clone, Copy, PartialEq, Eq)]
115pub struct AutoCompactionEvent {
116 pub removed_message_count: usize,
117}
118
119pub struct ConversationRuntime<C, T> {
120 session: Session,
121 api_client: C,
122 tool_executor: T,
123 permission_policy: PermissionPolicy,
124 system_prompt: Vec<String>,
125 max_iterations: usize,
126 usage_tracker: UsageTracker,
127 hook_runner: HookRunner,
128 auto_compaction_input_tokens_threshold: u32,
129 cancel_token: Option<Arc<AtomicBool>>,
130}
131
132impl<C, T> ConversationRuntime<C, T>
133where
134 C: ApiClient,
135 T: ToolExecutor,
136{
137 #[must_use]
138 pub fn new(
139 session: Session,
140 api_client: C,
141 tool_executor: T,
142 permission_policy: PermissionPolicy,
143 system_prompt: Vec<String>,
144 ) -> Self {
145 Self::new_with_features(
146 session,
147 api_client,
148 tool_executor,
149 permission_policy,
150 system_prompt,
151 RuntimeFeatureConfig::default(),
152 )
153 }
154
155 #[must_use]
156 pub fn new_with_features(
157 session: Session,
158 api_client: C,
159 tool_executor: T,
160 permission_policy: PermissionPolicy,
161 system_prompt: Vec<String>,
162 feature_config: RuntimeFeatureConfig,
163 ) -> Self {
164 let usage_tracker = UsageTracker::from_session(&session);
165 Self {
166 session,
167 api_client,
168 tool_executor,
169 permission_policy,
170 system_prompt,
171 max_iterations: usize::MAX,
172 usage_tracker,
173 hook_runner: HookRunner::from_feature_config(&feature_config),
174 auto_compaction_input_tokens_threshold: auto_compaction_threshold_from_env(),
175 cancel_token: None,
176 }
177 }
178
179 pub fn set_cancel_token(&mut self, token: Arc<AtomicBool>) {
183 self.cancel_token = Some(token);
184 }
185
186 #[inline]
187 fn is_cancelled(&self) -> bool {
188 self.cancel_token.as_ref().map_or(false, |t| t.load(Ordering::Relaxed))
189 }
190
191 #[must_use]
192 pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
193 self.max_iterations = max_iterations;
194 self
195 }
196
197 #[must_use]
198 pub fn with_auto_compaction_input_tokens_threshold(mut self, threshold: u32) -> Self {
199 self.auto_compaction_input_tokens_threshold = threshold;
200 self
201 }
202
203 pub fn run_turn(
204 &mut self,
205 user_input: impl Into<String>,
206 prompter: Option<&mut dyn PermissionPrompter>,
207 ) -> Result<TurnSummary, RuntimeError> {
208 self.session.messages.push(ConversationMessage::user_text(user_input.into()));
209 self.run_conversation_loop(prompter)
210 }
211
212 pub fn run_turn_with_images(
213 &mut self,
214 user_input: impl Into<String>,
215 images: Vec<(String, String)>,
216 prompter: Option<&mut dyn PermissionPrompter>,
217 ) -> Result<TurnSummary, RuntimeError> {
218 use crate::session::MessageRole;
219 let mut blocks: Vec<ContentBlock> = images
220 .into_iter()
221 .map(|(media_type, data)| ContentBlock::Image { media_type, data })
222 .collect();
223 blocks.push(ContentBlock::Text { text: user_input.into() });
224 self.session.messages.push(ConversationMessage {
225 role: MessageRole::User,
226 blocks,
227 usage: None,
228 });
229 self.run_conversation_loop(prompter)
230 }
231
232 fn run_conversation_loop(
233 &mut self,
234 mut prompter: Option<&mut dyn PermissionPrompter>,
235) -> Result<TurnSummary, RuntimeError> {
236 let mut assistant_messages = Vec::new();
237 let mut tool_results = Vec::new();
238 let mut iterations = 0;
239
240 loop {
241 if self.is_cancelled() {
243 return Err(RuntimeError::new("cancelled"));
244 }
245
246 iterations += 1;
247 if iterations > self.max_iterations {
248 return Err(RuntimeError::new(
249 "conversation loop exceeded the maximum number of iterations",
250 ));
251 }
252
253 let request = ApiRequest {
254 system_prompt: self.system_prompt.clone(),
255 messages: self.session.messages.clone(),
256 };
257 let events = self.api_client.stream(request)?;
258 if self.is_cancelled() {
260 return Err(RuntimeError::new("cancelled"));
261 }
262 let (assistant_message, usage) = build_assistant_message(events)?;
263 if let Some(usage) = usage {
264 self.usage_tracker.record(usage);
265 }
266
267 let evaluation = get_consensus_evaluation(&assistant_message);
268
269 match evaluation {
270 1 => {
271 let pending_tool_uses = assistant_message
272 .blocks
273 .iter()
274 .filter_map(|block| match block {
275 ContentBlock::ToolUse { id, name, input } => {
276 Some((id.clone(), name.clone(), input.clone()))
277 }
278 _ => None,
279 })
280 .collect::<Vec<_>>();
281
282 self.session.messages.push(assistant_message.clone());
283 assistant_messages.push(assistant_message);
284
285 if pending_tool_uses.is_empty() {
286 break;
287 }
288
289 for (tool_use_id, tool_name, input) in pending_tool_uses {
290 let permission_outcome = if let Some(prompt) = prompter.as_mut() {
291 self.permission_policy
292 .authorize(&tool_name, &input, Some(*prompt))
293 } else {
294 self.permission_policy.authorize(&tool_name, &input, None)
295 };
296
297 let result_message = match permission_outcome {
298 PermissionOutcome::Allow | PermissionOutcome::AllowWithEdits { .. } => {
299 let (effective_input, _is_human_edit) = match permission_outcome {
300 PermissionOutcome::AllowWithEdits { new_input } => (new_input, true),
301 _ => (input, false),
302 };
303
304 let pre_hook_result =
305 self.hook_runner.run_pre_tool_use(&tool_name, &effective_input);
306 if pre_hook_result.is_denied() {
307 let deny_message =
308 format!("PreToolUse hook denied tool `{tool_name}`");
309 ConversationMessage::tool_result(
310 tool_use_id,
311 tool_name,
312 format_hook_message(&pre_hook_result, &deny_message),
313 true,
314 )
315 } else {
316 let (output, mut is_error, validation_state) =
317 match self.tool_executor.execute(&tool_name, &effective_input) {
318 Ok(res) => (res.output, res.state == -1, res.state),
319 Err(error) => {
320 let err_msg = error.to_string();
321 let reflection_prompt = format!(
322 "The tool '{}' failed with the following error: {}. Please analyze the error and provide a corrected tool call.",
323 tool_name, err_msg
324 );
325 self.session.messages.push(ConversationMessage::user_text(reflection_prompt));
326 return Ok(TurnSummary {
327 assistant_messages: assistant_messages.clone(),
328 tool_results: tool_results.clone(),
329 iterations,
330 usage: self.usage_tracker.cumulative_usage(),
331 auto_compaction: None,
332 });
333 },
334 };
335
336 if validation_state == 0 {
337 let mut recovered = false;
339 let query_terms: Vec<&str> = effective_input.split(|c: char| !c.is_alphanumeric())
340 .filter(|s| s.len() > 3)
341 .collect();
342
343 for term in query_terms {
344 if let Ok(memory_context) = self.tool_executor.query_memory(term) {
346 if !memory_context.contains("[]") && memory_context.len() > 10 {
347 let recovery_prompt = format!(
348 "AUTONOMOUS RECOVERY (State 0 -> +1):\n\
349 Tool `{tool_name}` halted on ambiguous input. Found matching context in local knowledge graph for `{term}`:\n\
350 {}\n\
351 Please rewrite your tool call using this context to resolve the ambiguity.",
352 memory_context
353 );
354 self.session.messages.push(ConversationMessage::user_text(recovery_prompt));
355 recovered = true;
356 break;
357 }
358 }
359 }
360
361 if recovered {
362 continue;
364 }
365
366 let halt_msg = format!("Tool `{tool_name}` requested manual authorization or clarification (State 0).");
367 let result_msg = ConversationMessage::tool_result(
368 tool_use_id,
369 tool_name,
370 halt_msg,
371 true,
372 );
373 self.session.messages.push(result_msg.clone());
374 tool_results.push(result_msg);
375 break; }
377
378 let mut final_output = merge_hook_feedback(
379 pre_hook_result.messages(),
380 output,
381 false,
382 );
383
384 let post_hook_result = self.hook_runner.run_post_tool_use(
385 &tool_name,
386 &effective_input,
387 &final_output,
388 is_error,
389 );
390 if post_hook_result.is_denied() {
391 is_error = true;
392 }
393 final_output = merge_hook_feedback(
394 post_hook_result.messages(),
395 final_output,
396 post_hook_result.is_denied(),
397 );
398
399 ConversationMessage::tool_result(
400 tool_use_id,
401 tool_name,
402 final_output,
403 is_error,
404 )
405 }
406 }
407 PermissionOutcome::Deny { reason } => {
408 ConversationMessage::tool_result(tool_use_id, tool_name, reason, true)
409 }
410 };
411 self.session.messages.push(result_message.clone());
412 tool_results.push(result_message);
413 }
414 }
415 0 => {
416 self.session.messages.push(ConversationMessage::user_text(
418 "Could you please clarify your request?".to_string(),
419 ));
420 break;
421 }
422 -1 => {
423 self.session.messages.push(ConversationMessage::user_text(
425 "Let me try a different approach.".to_string(),
426 ));
427 }
428 _ => {
429 return Err(RuntimeError::new("invalid consensus evaluation"));
430 }
431 }
432 }
433
434 let auto_compaction = self.maybe_auto_compact();
435
436 Ok(TurnSummary {
437 assistant_messages,
438 tool_results,
439 iterations,
440 usage: self.usage_tracker.cumulative_usage(),
441 auto_compaction,
442 })
443}
444
445 #[must_use]
446 pub fn compact(&self, config: CompactionConfig) -> CompactionResult {
447 compact_session(&self.session, config)
448 }
449
450 #[must_use]
451 pub fn estimated_tokens(&self) -> usize {
452 estimate_session_tokens(&self.session)
453 }
454
455 #[must_use]
456 pub fn usage(&self) -> &UsageTracker {
457 &self.usage_tracker
458 }
459
460 #[must_use]
461 pub fn session(&self) -> &Session {
462 &self.session
463 }
464
465 #[must_use]
466 pub fn into_session(self) -> Session {
467 self.session
468 }
469
470 fn maybe_auto_compact(&mut self) -> Option<AutoCompactionEvent> {
471 if self.usage_tracker.cumulative_usage().input_tokens
472 < self.auto_compaction_input_tokens_threshold
473 {
474 return None;
475 }
476
477 let result = compact_session(
478 &self.session,
479 CompactionConfig {
480 max_estimated_tokens: 0,
481 ..CompactionConfig::default()
482 },
483 );
484
485 if result.removed_message_count == 0 {
486 return None;
487 }
488
489 self.session = result.compacted_session;
490 Some(AutoCompactionEvent {
491 removed_message_count: result.removed_message_count,
492 })
493 }
494}
495
496#[must_use]
497pub fn auto_compaction_threshold_from_env() -> u32 {
498 parse_auto_compaction_threshold(
499 std::env::var(AUTO_COMPACTION_THRESHOLD_ENV_VAR)
500 .ok()
501 .as_deref(),
502 )
503}
504
505#[must_use]
506fn parse_auto_compaction_threshold(value: Option<&str>) -> u32 {
507 value
508 .and_then(|raw| raw.trim().parse::<u32>().ok())
509 .filter(|threshold| *threshold > 0)
510 .unwrap_or(DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD)
511}
512
513fn get_consensus_evaluation(_reasoning: &ConversationMessage) -> i8 {
514 1
517}
518
519fn build_assistant_message(
520 events: Vec<AssistantEvent>,
521) -> Result<(ConversationMessage, Option<TokenUsage>), RuntimeError> {
522 let mut text = String::new();
523 let mut blocks = Vec::new();
524 let mut finished = false;
525 let mut usage = None;
526
527 for event in events {
528 match event {
529 AssistantEvent::TextDelta(delta) => text.push_str(&delta),
530 AssistantEvent::ToolUse { id, name, input } => {
531 flush_text_block(&mut text, &mut blocks);
532 blocks.push(ContentBlock::ToolUse { id, name, input });
533 }
534 AssistantEvent::TaskStarted { .. } | AssistantEvent::TaskCompleted { .. } => {
535 }
538 AssistantEvent::Usage(u) => usage = Some(u),
539
540 AssistantEvent::MessageStop => {
541 finished = true;
542 }
543 AssistantEvent::ToolTelemetry { .. } => {}
544 }
545 }
546
547 flush_text_block(&mut text, &mut blocks);
548
549 if !finished {
550 return Err(RuntimeError::new(
551 "assistant stream ended without a message stop event",
552 ));
553 }
554
555 Ok((
556 ConversationMessage::assistant_with_usage(blocks, usage),
557 usage,
558 ))
559}
560
561fn flush_text_block(text: &mut String, blocks: &mut Vec<ContentBlock>) {
562 if !text.is_empty() {
563 blocks.push(ContentBlock::Text {
564 text: std::mem::take(text),
565 });
566 }
567}
568
569fn format_hook_message(result: &HookRunResult, fallback: &str) -> String {
570 if result.messages().is_empty() {
571 fallback.to_string()
572 } else {
573 result.messages().join("\n")
574 }
575}
576
577fn merge_hook_feedback(messages: &[String], output: String, denied: bool) -> String {
578 if messages.is_empty() {
579 return output;
580 }
581
582 let mut sections = Vec::new();
583 if !output.trim().is_empty() {
584 sections.push(output);
585 }
586 let label = if denied {
587 "Hook feedback (denied)"
588 } else {
589 "Hook feedback"
590 };
591 sections.push(format!("{label}:\n{}", messages.join("\n")));
592 sections.join("\n\n")
593}
594
595type ToolHandler = Box<dyn FnMut(&str) -> Result<String, ToolError>>;
596
597#[derive(Default)]
598pub struct StaticToolExecutor {
599 handlers: BTreeMap<String, ToolHandler>,
600}
601
602impl StaticToolExecutor {
603 #[must_use]
604 pub fn new() -> Self {
605 Self::default()
606 }
607
608 #[must_use]
609 pub fn register(
610 mut self,
611 tool_name: impl Into<String>,
612 handler: impl FnMut(&str) -> Result<String, ToolError> + 'static,
613 ) -> Self {
614 self.handlers.insert(tool_name.into(), Box::new(handler));
615 self
616 }
617}
618
619impl ToolExecutor for StaticToolExecutor {
620 fn execute(&mut self, tool_name: &str, input: &str) -> Result<ToolResult, ToolError> {
621 self.handlers
622 .get_mut(tool_name)
623 .ok_or_else(|| ToolError::new(format!("unknown tool: {tool_name}")))?(input)
624 .map(|output| ToolResult { output, state: 1 })
625 }
626
627 fn query_memory(&mut self, _query: &str) -> Result<String, ToolError> {
628 Ok("[]".to_string())
629 }
630}
631
632#[cfg(test)]
633mod tests {
634 use super::{
635 parse_auto_compaction_threshold, ApiClient, ApiRequest, AssistantEvent,
636 AutoCompactionEvent, ConversationRuntime, RuntimeError, StaticToolExecutor,
637 DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD,
638 };
639 use crate::compact::CompactionConfig;
640 use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
641 use crate::permissions::{
642 PermissionMode, PermissionPolicy, PermissionPromptDecision, PermissionPrompter,
643 PermissionRequest,
644 };
645 use crate::prompt::{ProjectContext, SystemPromptBuilder};
646 use crate::session::{ContentBlock, MessageRole, Session};
647 use crate::usage::TokenUsage;
648 use std::path::PathBuf;
649
650 struct ScriptedApiClient {
651 call_count: usize,
652 }
653
654 impl ApiClient for ScriptedApiClient {
655 fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
656 self.call_count += 1;
657 match self.call_count {
658 1 => {
659 assert!(request
660 .messages
661 .iter()
662 .any(|message| message.role == MessageRole::User));
663 Ok(vec![
664 AssistantEvent::TextDelta("Let me calculate that.".to_string()),
665 AssistantEvent::ToolUse {
666 id: "tool-1".to_string(),
667 name: "add".to_string(),
668 input: "2,2".to_string(),
669 },
670 AssistantEvent::Usage(TokenUsage {
671 input_tokens: 20,
672 output_tokens: 6,
673 cache_creation_input_tokens: 1,
674 cache_read_input_tokens: 2,
675 }),
676 AssistantEvent::MessageStop,
677 ])
678 }
679 2 => {
680 let last_message = request
681 .messages
682 .last()
683 .expect("tool result should be present");
684 assert_eq!(last_message.role, MessageRole::Tool);
685 Ok(vec![
686 AssistantEvent::TextDelta("The answer is 4.".to_string()),
687 AssistantEvent::Usage(TokenUsage {
688 input_tokens: 24,
689 output_tokens: 4,
690 cache_creation_input_tokens: 1,
691 cache_read_input_tokens: 3,
692 }),
693 AssistantEvent::MessageStop,
694 ])
695 }
696 _ => Err(RuntimeError::new("unexpected extra API call")),
697 }
698 }
699 }
700
701 struct PromptAllowOnce;
702
703 impl PermissionPrompter for PromptAllowOnce {
704 fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
705 assert_eq!(request.tool_name, "add");
706 PermissionPromptDecision::Allow
707 }
708 }
709
710 #[test]
711 fn runs_user_to_tool_to_result_loop_end_to_end_and_tracks_usage() {
712 let api_client = ScriptedApiClient { call_count: 0 };
713 let tool_executor = StaticToolExecutor::new().register("add", |input| {
714 let total = input
715 .split(',')
716 .map(|part| part.parse::<i32>().expect("input must be valid integer"))
717 .sum::<i32>();
718 Ok(total.to_string())
719 });
720 let permission_policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite);
721 let system_prompt = SystemPromptBuilder::new()
722 .with_project_context(ProjectContext {
723 cwd: PathBuf::from("/tmp/project"),
724 current_date: "2026-03-31".to_string(),
725 git_status: None,
726 git_diff: None,
727 instruction_files: Vec::new(),
728 })
729 .with_os("linux", "6.8")
730 .build();
731 let mut runtime = ConversationRuntime::new(
732 Session::new(),
733 api_client,
734 tool_executor,
735 permission_policy,
736 system_prompt,
737 );
738
739 let summary = runtime
740 .run_turn("what is 2 + 2?", Some(&mut PromptAllowOnce))
741 .expect("conversation loop should succeed");
742
743 assert_eq!(summary.iterations, 2);
744 assert_eq!(summary.assistant_messages.len(), 2);
745 assert_eq!(summary.tool_results.len(), 1);
746 assert_eq!(runtime.session().messages.len(), 4);
747 assert_eq!(summary.usage.output_tokens, 10);
748 assert_eq!(summary.auto_compaction, None);
749 assert!(matches!(
750 runtime.session().messages[1].blocks[1],
751 ContentBlock::ToolUse { .. }
752 ));
753 assert!(matches!(
754 runtime.session().messages[2].blocks[0],
755 ContentBlock::ToolResult {
756 is_error: false,
757 ..
758 }
759 ));
760 }
761
762 #[test]
763 fn records_denied_tool_results_when_prompt_rejects() {
764 struct RejectPrompter;
765 impl PermissionPrompter for RejectPrompter {
766 fn decide(&mut self, _request: &PermissionRequest) -> PermissionPromptDecision {
767 PermissionPromptDecision::Deny {
768 reason: "not now".to_string(),
769 }
770 }
771 }
772
773 struct SingleCallApiClient;
774 impl ApiClient for SingleCallApiClient {
775 fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
776 if request
777 .messages
778 .iter()
779 .any(|message| message.role == MessageRole::Tool)
780 {
781 return Ok(vec![
782 AssistantEvent::TextDelta("I could not use the tool.".to_string()),
783 AssistantEvent::MessageStop,
784 ]);
785 }
786 Ok(vec![
787 AssistantEvent::ToolUse {
788 id: "tool-1".to_string(),
789 name: "blocked".to_string(),
790 input: "secret".to_string(),
791 },
792 AssistantEvent::MessageStop,
793 ])
794 }
795 }
796
797 let mut runtime = ConversationRuntime::new(
798 Session::new(),
799 SingleCallApiClient,
800 StaticToolExecutor::new(),
801 PermissionPolicy::new(PermissionMode::WorkspaceWrite),
802 vec!["system".to_string()],
803 );
804
805 let summary = runtime
806 .run_turn("use the tool", Some(&mut RejectPrompter))
807 .expect("conversation should continue after denied tool");
808
809 assert_eq!(summary.tool_results.len(), 1);
810 assert!(matches!(
811 &summary.tool_results[0].blocks[0],
812 ContentBlock::ToolResult { is_error: true, output, .. } if output == "not now"
813 ));
814 }
815
816 #[test]
817 fn denies_tool_use_when_pre_tool_hook_blocks() {
818 struct SingleCallApiClient;
819 impl ApiClient for SingleCallApiClient {
820 fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
821 if request
822 .messages
823 .iter()
824 .any(|message| message.role == MessageRole::Tool)
825 {
826 return Ok(vec![
827 AssistantEvent::TextDelta("blocked".to_string()),
828 AssistantEvent::MessageStop,
829 ]);
830 }
831 Ok(vec![
832 AssistantEvent::ToolUse {
833 id: "tool-1".to_string(),
834 name: "blocked".to_string(),
835 input: r#"{"path":"secret.txt"}"#.to_string(),
836 },
837 AssistantEvent::MessageStop,
838 ])
839 }
840 }
841
842 let mut runtime = ConversationRuntime::new_with_features(
843 Session::new(),
844 SingleCallApiClient,
845 StaticToolExecutor::new().register("blocked", |_input| {
846 panic!("tool should not execute when hook denies")
847 }),
848 PermissionPolicy::new(PermissionMode::DangerFullAccess),
849 vec!["system".to_string()],
850 RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
851 vec![shell_snippet("printf 'blocked by hook'; exit 2")],
852 Vec::new(),
853 )),
854 );
855
856 let summary = runtime
857 .run_turn("use the tool", None)
858 .expect("conversation should continue after hook denial");
859
860 assert_eq!(summary.tool_results.len(), 1);
861 let ContentBlock::ToolResult {
862 is_error, output, ..
863 } = &summary.tool_results[0].blocks[0]
864 else {
865 panic!("expected tool result block");
866 };
867 assert!(
868 *is_error,
869 "hook denial should produce an error result: {output}"
870 );
871 assert!(
872 output.contains("denied tool") || output.contains("blocked by hook"),
873 "unexpected hook denial output: {output:?}"
874 );
875 }
876
877 #[test]
878 fn appends_post_tool_hook_feedback_to_tool_result() {
879 struct TwoCallApiClient {
880 calls: usize,
881 }
882
883 impl ApiClient for TwoCallApiClient {
884 fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
885 self.calls += 1;
886 match self.calls {
887 1 => Ok(vec![
888 AssistantEvent::ToolUse {
889 id: "tool-1".to_string(),
890 name: "add".to_string(),
891 input: r#"{"lhs":2,"rhs":2}"#.to_string(),
892 },
893 AssistantEvent::MessageStop,
894 ]),
895 2 => {
896 assert!(request
897 .messages
898 .iter()
899 .any(|message| message.role == MessageRole::Tool));
900 Ok(vec![
901 AssistantEvent::TextDelta("done".to_string()),
902 AssistantEvent::MessageStop,
903 ])
904 }
905 _ => Err(RuntimeError::new("unexpected extra API call")),
906 }
907 }
908 }
909
910 let mut runtime = ConversationRuntime::new_with_features(
911 Session::new(),
912 TwoCallApiClient { calls: 0 },
913 StaticToolExecutor::new().register("add", |_input| Ok("4".to_string())),
914 PermissionPolicy::new(PermissionMode::DangerFullAccess),
915 vec!["system".to_string()],
916 RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
917 vec![shell_snippet("printf 'pre hook ran'")],
918 vec![shell_snippet("printf 'post hook ran'")],
919 )),
920 );
921
922 let summary = runtime
923 .run_turn("use add", None)
924 .expect("tool loop succeeds");
925
926 assert_eq!(summary.tool_results.len(), 1);
927 let ContentBlock::ToolResult {
928 is_error, output, ..
929 } = &summary.tool_results[0].blocks[0]
930 else {
931 panic!("expected tool result block");
932 };
933 assert!(
934 !*is_error,
935 "post hook should preserve non-error result: {output:?}"
936 );
937 assert!(
938 output.contains("4"),
939 "tool output missing value: {output:?}"
940 );
941 assert!(
942 output.contains("pre hook ran"),
943 "tool output missing pre hook feedback: {output:?}"
944 );
945 assert!(
946 output.contains("post hook ran"),
947 "tool output missing post hook feedback: {output:?}"
948 );
949 }
950
951 #[test]
952 fn reconstructs_usage_tracker_from_restored_session() {
953 struct SimpleApi;
954 impl ApiClient for SimpleApi {
955 fn stream(
956 &mut self,
957 _request: ApiRequest,
958 ) -> Result<Vec<AssistantEvent>, RuntimeError> {
959 Ok(vec![
960 AssistantEvent::TextDelta("done".to_string()),
961 AssistantEvent::MessageStop,
962 ])
963 }
964 }
965
966 let mut session = Session::new();
967 session
968 .messages
969 .push(crate::session::ConversationMessage::assistant_with_usage(
970 vec![ContentBlock::Text {
971 text: "earlier".to_string(),
972 }],
973 Some(TokenUsage {
974 input_tokens: 11,
975 output_tokens: 7,
976 cache_creation_input_tokens: 2,
977 cache_read_input_tokens: 1,
978 }),
979 ));
980
981 let runtime = ConversationRuntime::new(
982 session,
983 SimpleApi,
984 StaticToolExecutor::new(),
985 PermissionPolicy::new(PermissionMode::DangerFullAccess),
986 vec!["system".to_string()],
987 );
988
989 assert_eq!(runtime.usage().turns(), 1);
990 assert_eq!(runtime.usage().cumulative_usage().total_tokens(), 21);
991 }
992
993 #[test]
994 fn compacts_session_after_turns() {
995 struct SimpleApi;
996 impl ApiClient for SimpleApi {
997 fn stream(
998 &mut self,
999 _request: ApiRequest,
1000 ) -> Result<Vec<AssistantEvent>, RuntimeError> {
1001 Ok(vec![
1002 AssistantEvent::TextDelta("done".to_string()),
1003 AssistantEvent::MessageStop,
1004 ])
1005 }
1006 }
1007
1008 let mut runtime = ConversationRuntime::new(
1009 Session::new(),
1010 SimpleApi,
1011 StaticToolExecutor::new(),
1012 PermissionPolicy::new(PermissionMode::DangerFullAccess),
1013 vec!["system".to_string()],
1014 );
1015 runtime.run_turn("a", None).expect("turn a");
1016 runtime.run_turn("b", None).expect("turn b");
1017 runtime.run_turn("c", None).expect("turn c");
1018
1019 let result = runtime.compact(CompactionConfig {
1020 preserve_recent_messages: 2,
1021 max_estimated_tokens: 1,
1022 });
1023 assert!(result.summary.contains("Conversation summary"));
1024 assert_eq!(
1025 result.compacted_session.messages[0].role,
1026 MessageRole::System
1027 );
1028 }
1029
1030 #[cfg(windows)]
1031 fn shell_snippet(script: &str) -> String {
1032 script.replace('\'', "\"")
1033 }
1034
1035 #[cfg(not(windows))]
1036 fn shell_snippet(script: &str) -> String {
1037 script.to_string()
1038 }
1039
1040 #[test]
1041 fn auto_compacts_when_cumulative_input_threshold_is_crossed() {
1042 struct SimpleApi;
1043 impl ApiClient for SimpleApi {
1044 fn stream(
1045 &mut self,
1046 _request: ApiRequest,
1047 ) -> Result<Vec<AssistantEvent>, RuntimeError> {
1048 Ok(vec![
1049 AssistantEvent::TextDelta("done".to_string()),
1050 AssistantEvent::Usage(TokenUsage {
1051 input_tokens: 120_000,
1052 output_tokens: 4,
1053 cache_creation_input_tokens: 0,
1054 cache_read_input_tokens: 0,
1055 }),
1056 AssistantEvent::MessageStop,
1057 ])
1058 }
1059 }
1060
1061 let session = Session {
1062 version: 1,
1063 messages: vec![
1064 crate::session::ConversationMessage::user_text("one"),
1065 crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
1066 text: "two".to_string(),
1067 }]),
1068 crate::session::ConversationMessage::user_text("three"),
1069 crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
1070 text: "four".to_string(),
1071 }]),
1072 ],
1073 };
1074
1075 let mut runtime = ConversationRuntime::new(
1076 session,
1077 SimpleApi,
1078 StaticToolExecutor::new(),
1079 PermissionPolicy::new(PermissionMode::DangerFullAccess),
1080 vec!["system".to_string()],
1081 )
1082 .with_auto_compaction_input_tokens_threshold(100_000);
1083
1084 let summary = runtime
1085 .run_turn("trigger", None)
1086 .expect("turn should succeed");
1087
1088 assert_eq!(
1089 summary.auto_compaction,
1090 Some(AutoCompactionEvent {
1091 removed_message_count: 2,
1092 })
1093 );
1094 assert_eq!(runtime.session().messages[0].role, MessageRole::System);
1095 }
1096
1097 #[test]
1098 fn skips_auto_compaction_below_threshold() {
1099 struct SimpleApi;
1100 impl ApiClient for SimpleApi {
1101 fn stream(
1102 &mut self,
1103 _request: ApiRequest,
1104 ) -> Result<Vec<AssistantEvent>, RuntimeError> {
1105 Ok(vec![
1106 AssistantEvent::TextDelta("done".to_string()),
1107 AssistantEvent::Usage(TokenUsage {
1108 input_tokens: 99_999,
1109 output_tokens: 4,
1110 cache_creation_input_tokens: 0,
1111 cache_read_input_tokens: 0,
1112 }),
1113 AssistantEvent::MessageStop,
1114 ])
1115 }
1116 }
1117
1118 let mut runtime = ConversationRuntime::new(
1119 Session::new(),
1120 SimpleApi,
1121 StaticToolExecutor::new(),
1122 PermissionPolicy::new(PermissionMode::DangerFullAccess),
1123 vec!["system".to_string()],
1124 )
1125 .with_auto_compaction_input_tokens_threshold(100_000);
1126
1127 let summary = runtime
1128 .run_turn("trigger", None)
1129 .expect("turn should succeed");
1130 assert_eq!(summary.auto_compaction, None);
1131 assert_eq!(runtime.session().messages.len(), 2);
1132 }
1133
1134 #[test]
1135 fn auto_compaction_threshold_defaults_and_parses_values() {
1136 assert_eq!(
1137 parse_auto_compaction_threshold(None),
1138 DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD
1139 );
1140 assert_eq!(parse_auto_compaction_threshold(Some("4321")), 4321);
1141 assert_eq!(
1142 parse_auto_compaction_threshold(Some("not-a-number")),
1143 DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD
1144 );
1145 }
1146}