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