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