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