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