Skip to main content

albert_runtime/
conversation.rs

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, // Ternary Intelligence Stack: +1 Success, 0 Neutral/Halt, -1 Failure
50}
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    /// Wire up an external cancellation flag. When `true`, `run_turn` exits at
180    /// the next safe checkpoint (between API calls), so `handle.join()` returns
181    /// promptly and the main submit loop can continue.
182    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        // Bail out at every safe checkpoint so handle.join() returns promptly on ESC.
218        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        // Check again immediately after the blocking HTTP call returns.
235        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                                    // Neurosymbolic Gap Recovery: Try to resolve state 0 autonomously via local graph
314                                    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                                        // BET VM: @sparseskip - drop neutral paths and hit memory matrix
321                                        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                                        // Allow one more iteration to try and hit +1 state
339                                        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; // Actually halt if recovery failed
352                                }
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                // Request disambiguation from the user
393                self.session.messages.push(ConversationMessage::user_text(
394                    "Could you please clarify your request?".to_string(),
395                ));
396                break;
397            }
398            -1 => {
399                // Generate an alternative plan
400                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    // For now, we'll just return +1 (Proceed) by default.
491    // This can be replaced with a more sophisticated evaluation logic later.
492    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                // Task events are handled by the TUI in real-time and don't affect 
512                // the static ConversationMessage structure.
513            }
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}