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