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