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