path: root/crates/atuin-ai/src/fsm/mod.rs

//! Agent conversation FSM.
//!
//! Pure state machine that returns effects as data.
//! The driver is responsible for executing effects and feeding events back.
//!
//! The FSM owns the conversation event log and tool lifecycle state.
//! It never performs IO directly.

pub(crate) mod effects;
pub(crate) mod events;
pub(crate) mod tools;

#[cfg(test)]
mod tests;

use std::collections::HashMap;

use serde_json::Value;

use crate::context_window::ContextWindowBuilder;
use crate::tui::state::ConversationEvent;

use effects::{Effect, ExitAction, PermissionTarget, TimeoutKind};
use events::{Event, PermissionChoice, PermissionResponse};
use tools::{ToolManager, ToolState};

// ============================================================================
// State
// ============================================================================

/// The discrete states of the agent FSM.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum AgentState {
    /// Waiting for user input.
    Idle {
        confirmation: Option<PendingConfirmation>,
    },

    /// A conversation turn is in progress.
    Turn { stream: StreamPhase },

    /// Unrecoverable error. User can retry or exit.
    Error(String),
}

/// Stream connection lifecycle within a Turn.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum StreamPhase {
    /// Request sent, awaiting first stream frame.
    Connecting,
    /// Actively receiving streamed response.
    Streaming { status: Option<StreamingStatus> },
    /// Stream connection has ended (Done received).
    Done,
}

/// Streaming status indicators from server.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum StreamingStatus {
    Processing,
    Searching,
    Thinking,
    WaitingForTools,
}

impl StreamingStatus {
    pub(crate) fn from_str(s: &str) -> Self {
        match s {
            "processing" => Self::Processing,
            "searching" => Self::Searching,
            "waiting_for_tools" => Self::WaitingForTools,
            _ => Self::Thinking,
        }
    }
}

/// Pending dangerous command confirmation state.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct PendingConfirmation {
    pub command: String,
    pub timeout_id: u64,
}

// ============================================================================
// Context
// ============================================================================

/// Shared context owned by the FSM.
#[derive(Debug, Clone)]
pub(crate) struct AgentContext {
    /// The full conversation event log (source of truth for API + persistence).
    pub events: Vec<ConversationEvent>,
    /// Server-assigned session ID.
    pub session_id: Option<String>,
    /// Accumulated text from current stream (committed to events on tool call or stream end).
    pub current_response: String,
    /// Per-tool lifecycle state and cached render data.
    /// Tools persist across turns for rendering history.
    pub tools: ToolManager,
    /// Tool IDs that belong to the current turn. Cleared on continuation start.
    /// Used to determine whether a turn needs continuation (has unprocessed results).
    current_turn_tool_ids: Vec<String>,
    /// Maps timeout_id → tool_id for active tool execution timeouts.
    /// Cleaned up when a tool completes naturally, so stale timeouts are ignored.
    tool_timeout_ids: HashMap<u64, String>,
    /// Counter for generating unique timeout IDs.
    next_timeout_id: u64,
    /// Capabilities advertised to the server.
    pub capabilities: Vec<String>,
    /// Unique invocation ID for this CLI invocation.
    pub invocation_id: String,

    // ─── View state (owned by FSM for atomic transitions) ───────
    /// Index into events where the current TUI invocation starts.
    /// Events before this are context for the API but not rendered.
    pub view_start_index: usize,
    /// Whether this session was resumed from a prior invocation.
    pub is_resumed: bool,
    /// Time of the last event from a previous invocation.
    pub last_event_time: Option<chrono::DateTime<chrono::Utc>>,
    /// Events from archived sessions (/new) still rendered on screen.
    pub archived_events: Vec<ConversationEvent>,
}

impl AgentContext {
    fn next_timeout_id(&mut self) -> u64 {
        let id = self.next_timeout_id;
        self.next_timeout_id += 1;
        id
    }
}

// ============================================================================
// The Agent FSM
// ============================================================================

/// The agent finite state machine.
///
/// Pure state machine — `handle()` takes an event, mutates internal state,
/// and returns effects as data for the driver to execute.
#[derive(Debug, Clone)]
pub(crate) struct AgentFsm {
    pub state: AgentState,
    pub ctx: AgentContext,
}

impl AgentFsm {
    /// Create a new FSM in Idle state.
    pub fn new(capabilities: Vec<String>, invocation_id: String) -> Self {
        Self {
            state: AgentState::Idle { confirmation: None },
            ctx: AgentContext {
                events: Vec::new(),
                session_id: None,
                current_response: String::new(),
                tools: ToolManager::new(),
                current_turn_tool_ids: Vec::new(),
                tool_timeout_ids: HashMap::new(),
                next_timeout_id: 0,
                capabilities,
                invocation_id,
                view_start_index: 0,
                is_resumed: false,
                last_event_time: None,
                archived_events: Vec::new(),
            },
        }
    }

    /// Create an FSM from saved session state (for resume).
    pub fn from_session(
        events: Vec<ConversationEvent>,
        session_id: Option<String>,
        capabilities: Vec<String>,
        invocation_id: String,
        view_start_index: usize,
        is_resumed: bool,
        last_event_time: Option<chrono::DateTime<chrono::Utc>>,
    ) -> Self {
        Self {
            state: AgentState::Idle { confirmation: None },
            ctx: AgentContext {
                events,
                session_id,
                current_response: String::new(),
                tools: ToolManager::new(),
                current_turn_tool_ids: Vec::new(),
                tool_timeout_ids: HashMap::new(),
                next_timeout_id: 0,
                capabilities,
                invocation_id,
                view_start_index,
                is_resumed,
                last_event_time,
                archived_events: Vec::new(),
            },
        }
    }

    /// Handle an event, returning effects to execute.
    pub fn handle(&mut self, event: Event) -> Vec<Effect> {
        match (&self.state, event) {
            // ================================================================
            // Idle state
            // ================================================================
            (AgentState::Idle { confirmation: None }, Event::UserSubmit(msg)) => {
                self.start_turn(msg)
            }

            (
                AgentState::Idle {
                    confirmation: Some(_),
                },
                Event::UserSubmit(msg),
            ) => self.start_turn(msg),

            (AgentState::Idle { confirmation: None }, Event::ExecuteCommand) => {
                let cmd = self.current_command();
                let Some(cmd) = cmd else {
                    // No command suggested — exit
                    return vec![Effect::ExitApp(ExitAction::Cancel)];
                };
                if self.is_current_command_dangerous() {
                    let timeout_id = self.ctx.next_timeout_id();
                    self.state = AgentState::Idle {
                        confirmation: Some(PendingConfirmation {
                            command: cmd,
                            timeout_id,
                        }),
                    };
                    vec![Effect::ScheduleTimeout {
                        timeout_id,
                        duration: std::time::Duration::from_secs(5),
                        kind: TimeoutKind::Confirmation,
                    }]
                } else {
                    vec![Effect::ExitApp(ExitAction::Execute(cmd))]
                }
            }

            (
                AgentState::Idle {
                    confirmation: Some(_),
                },
                Event::ExecuteCommand,
            ) => {
                let confirm = self.state_confirmation().unwrap().clone();
                self.state = AgentState::Idle { confirmation: None };
                vec![Effect::ExitApp(ExitAction::Execute(confirm.command))]
            }

            (AgentState::Idle { .. }, Event::InsertCommand) => {
                let cmd = self.current_command();
                match cmd {
                    Some(cmd) => vec![Effect::ExitApp(ExitAction::Insert(cmd))],
                    None => vec![],
                }
            }

            (
                AgentState::Idle {
                    confirmation: Some(_),
                },
                Event::Cancel,
            ) => {
                self.state = AgentState::Idle { confirmation: None };
                vec![]
            }

            (AgentState::Idle { confirmation: None }, Event::Cancel) => {
                vec![Effect::ExitApp(ExitAction::Cancel)]
            }

            (AgentState::Idle { .. }, Event::ConfirmationTimeout { timeout_id }) => {
                if self
                    .state_confirmation()
                    .is_some_and(|c| c.timeout_id == timeout_id)
                {
                    self.state = AgentState::Idle { confirmation: None };
                }
                vec![]
            }

            (AgentState::Idle { .. }, Event::NewSession) => {
                // Archive visible events so they remain on screen but aren't
                // sent to the API. Tools persist for rendering.
                let visible = self.ctx.events[self.ctx.view_start_index..].to_vec();
                self.ctx.archived_events.extend(visible);

                self.ctx.events.clear();
                self.ctx.session_id = None;
                self.ctx.current_turn_tool_ids.clear();
                self.ctx.view_start_index = 0;
                self.ctx.is_resumed = false;

                // Add OOB indicator for the new session
                self.ctx.events.push(ConversationEvent::OutOfBandOutput {
                    name: "System".to_string(),
                    command: Some("/new".to_string()),
                    content: "Started a new session.".to_string(),
                });

                self.state = AgentState::Idle { confirmation: None };
                vec![Effect::ArchiveSession, Effect::Persist]
            }

            (AgentState::Idle { .. }, Event::SlashCommand { command, content }) => {
                self.handle_slash_command(&command, &content);
                vec![]
            }

            (
                AgentState::Idle { .. },
                Event::SkillLoaded {
                    name,
                    arguments,
                    content,
                },
            ) => {
                self.ctx.events.push(ConversationEvent::SkillInvocation {
                    name,
                    arguments,
                    content,
                });
                self.ctx.current_response.clear();
                self.ctx.current_turn_tool_ids.clear();

                let messages = self.build_messages();
                let session_id = self.ctx.session_id.clone();
                self.state = AgentState::Turn {
                    stream: StreamPhase::Connecting,
                };
                vec![Effect::StartStream {
                    messages,
                    session_id,
                }]
            }

            // ================================================================
            // Turn — stream lifecycle
            // ================================================================
            (
                AgentState::Turn {
                    stream: StreamPhase::Connecting,
                },
                Event::StreamStarted,
            ) => {
                self.state = AgentState::Turn {
                    stream: StreamPhase::Streaming { status: None },
                };
                vec![]
            }

            (
                AgentState::Turn {
                    stream: StreamPhase::Connecting,
                },
                Event::StreamError(e),
            ) => {
                self.state = AgentState::Error(e);
                vec![]
            }

            (
                AgentState::Turn {
                    stream: StreamPhase::Streaming { .. },
                },
                Event::StreamChunk(text),
            ) => {
                self.ctx.current_response.push_str(&text);
                vec![]
            }

            (
                AgentState::Turn {
                    stream: StreamPhase::Streaming { .. },
                },
                Event::StreamStatusChanged(status),
            ) => {
                self.state = AgentState::Turn {
                    stream: StreamPhase::Streaming {
                        status: Some(StreamingStatus::from_str(&status)),
                    },
                };
                vec![]
            }

            (AgentState::Turn { .. }, Event::StreamToolCall { id, name, input }) => {
                self.commit_streaming_text();
                self.handle_stream_tool_call(id, name, input)
            }

            (AgentState::Turn { .. }, Event::SuggestCommand { id, input }) => {
                self.commit_streaming_text();
                // Push the suggest_command as a ToolCall event (protocol requirement)
                self.ctx.events.push(ConversationEvent::ToolCall {
                    id,
                    name: "suggest_command".to_string(),
                    input,
                });
                self.state = AgentState::Idle { confirmation: None };
                vec![Effect::Persist]
            }

            (
                AgentState::Turn {
                    stream: StreamPhase::Streaming { .. },
                },
                Event::StreamServerToolResult {
                    tool_use_id,
                    content,
                    is_error,
                    remote,
                    content_length,
                },
            ) => {
                self.ctx.events.push(ConversationEvent::ToolResult {
                    tool_use_id,
                    content,
                    is_error,
                    remote,
                    content_length,
                });
                vec![]
            }

            (AgentState::Turn { .. }, Event::StreamDone { session_id }) => {
                self.commit_streaming_text();
                if !session_id.is_empty() {
                    self.ctx.session_id = Some(session_id);
                }
                self.state = AgentState::Turn {
                    stream: StreamPhase::Done,
                };
                self.check_turn_completion()
            }

            (
                AgentState::Turn {
                    stream: StreamPhase::Streaming { .. },
                },
                Event::StreamError(e),
            ) => {
                // Abort any executing tools on stream error
                let abort_effects: Vec<_> = self
                    .ctx
                    .tools
                    .executing_ids()
                    .into_iter()
                    .map(|tool_id| Effect::AbortTool { tool_id })
                    .collect();
                self.ctx.tool_timeout_ids.clear();
                self.state = AgentState::Error(e);
                abort_effects
            }

            // ================================================================
            // Turn — tool lifecycle (any stream phase)
            // ================================================================
            (AgentState::Turn { .. }, Event::PermissionResolved { tool_id, response }) => {
                self.handle_permission_resolved(tool_id, response)
            }

            (AgentState::Turn { .. }, Event::PermissionUserChoice { tool_id, choice }) => {
                self.handle_permission_choice(tool_id, choice)
            }

            (
                AgentState::Turn { .. },
                Event::ToolExecutionDone {
                    tool_id,
                    outcome,
                    preview,
                },
            ) => self.handle_tool_done(tool_id, outcome, preview),

            (
                AgentState::Turn { .. },
                Event::ToolPreviewUpdate {
                    tool_id,
                    lines,
                    exit_code,
                },
            ) => {
                if let Some(tracked) = self.ctx.tools.get_mut(&tool_id) {
                    if tracked.is_resolved() {
                        // Tool already completed — a late preview update raced with
                        // ToolExecutionDone. Update lines (they may carry the final
                        // screen) but preserve the finalized exit_code/interrupted.
                        if let Some(tools::ToolPreviewData::Shell {
                            lines: existing_lines,
                            ..
                        }) = &mut tracked.preview
                        {
                            *existing_lines = lines;
                        }
                    } else {
                        tracked.preview = Some(tools::ToolPreviewData::Shell {
                            lines,
                            exit_code,
                            interrupted: None,
                        });
                    }
                }
                vec![]
            }

            (AgentState::Turn { .. }, Event::InterruptTools) => {
                let ids = self.ctx.tools.executing_ids();
                for id in &ids {
                    if let Some(tracked) = self.ctx.tools.get_mut(id) {
                        tracked.interrupt_reason = Some(tools::InterruptReason::User);
                    }
                    // Clear any pending execution timeout for this tool
                    self.ctx.tool_timeout_ids.retain(|_, tid| tid != id);
                }
                ids.into_iter()
                    .map(|tool_id| Effect::AbortTool { tool_id })
                    .collect()
            }

            (
                AgentState::Turn { .. },
                Event::ToolExecutionTimeout {
                    timeout_id,
                    tool_id,
                },
            ) => self.handle_tool_execution_timeout(timeout_id, tool_id),

            // ─── Cancel during Turn ─────────────────────────────────────
            (AgentState::Turn { stream }, Event::Cancel) => {
                let mut effects = Vec::new();

                // Abort stream if still active
                if !matches!(stream, StreamPhase::Done) {
                    effects.push(Effect::AbortStream);
                }

                // Cancel all pending tools
                let pending = self.ctx.tools.pending_ids();
                for id in &pending {
                    if let Some(tracked) = self.ctx.tools.get_mut(id) {
                        if tracked.state == ToolState::Executing {
                            effects.push(Effect::AbortTool {
                                tool_id: id.clone(),
                            });
                        }
                        tracked.state = ToolState::Completed;
                    }
                    self.ctx.events.push(ConversationEvent::ToolResult {
                        tool_use_id: id.clone(),
                        content: "Error: user cancelled this operation".to_string(),
                        is_error: true,
                        remote: false,
                        content_length: None,
                    });
                }

                // Commit any partial streaming text
                self.commit_streaming_text_as_cancelled();

                // Add context so the LLM knows what happened
                if !pending.is_empty() {
                    self.ctx.events.push(ConversationEvent::SystemContext {
                        content: "The user cancelled the previous generation. Tool calls that were in progress have been aborted.".to_string(),
                    });
                }

                // Clear timeout mappings — stale timeouts will be ignored by the guard
                self.ctx.tool_timeout_ids.clear();

                self.state = AgentState::Idle { confirmation: None };
                effects.push(Effect::Persist);
                effects
            }

            // ================================================================
            // Error state
            // ================================================================
            (AgentState::Error(_), Event::Retry) => {
                let messages = self.build_messages();
                let session_id = self.ctx.session_id.clone();
                self.state = AgentState::Turn {
                    stream: StreamPhase::Connecting,
                };
                vec![Effect::StartStream {
                    messages,
                    session_id,
                }]
            }

            (AgentState::Error(_), Event::Cancel) => {
                vec![Effect::ExitApp(ExitAction::Cancel)]
            }

            // ================================================================
            // Fallthrough — ignore events with no valid transition
            // ================================================================

            // StreamDone can arrive after SuggestCommand (which already moved to Idle).
            // We still need to capture the session_id from it.
            (_, Event::StreamDone { session_id }) => {
                if !session_id.is_empty() {
                    self.ctx.session_id = Some(session_id);
                }
                vec![Effect::Persist]
            }

            (_, Event::SlashCommand { command, content }) => {
                self.handle_slash_command(&command, &content);
                vec![]
            }

            // RequestSkillLoad during non-idle: still emit the effect
            (_, Event::RequestSkillLoad { name, arguments }) => {
                vec![Effect::LoadSkill { name, arguments }]
            }

            // SkillLoaded during non-idle: queue so it's visible
            // in context for the next turn.
            (
                _,
                Event::SkillLoaded {
                    name,
                    arguments,
                    content,
                },
            ) => {
                self.ctx.events.push(ConversationEvent::SkillInvocation {
                    name,
                    arguments,
                    content,
                });
                vec![]
            }

            _ => vec![],
        }
    }

    // ────────────────────────────────────────────────────────────────────
    // Private helpers
    // ────────────────────────────────────────────────────────────────────

    /// Start a new turn: push user message, build messages, emit StartStream.
    fn start_turn(&mut self, msg: String) -> Vec<Effect> {
        self.ctx
            .events
            .push(ConversationEvent::UserMessage { content: msg });
        // Don't clear tools — completed tools persist for rendering history.
        // Tools are only cleared on /new (session reset).
        self.ctx.current_response.clear();
        self.ctx.current_turn_tool_ids.clear();

        let messages = self.build_messages();
        let session_id = self.ctx.session_id.clone();
        self.state = AgentState::Turn {
            stream: StreamPhase::Connecting,
        };
        vec![Effect::StartStream {
            messages,
            session_id,
        }]
    }

    /// Build API messages from the conversation event log.
    fn build_messages(&self) -> Vec<Value> {
        ContextWindowBuilder::with_default_budget().build(&self.ctx.events)
    }

    /// Commit accumulated streaming text to the event log.
    fn commit_streaming_text(&mut self) {
        let text = std::mem::take(&mut self.ctx.current_response);
        let trimmed = text.trim_start().to_string();
        if !trimmed.is_empty() {
            self.ctx
                .events
                .push(ConversationEvent::Text { content: trimmed });
        }
    }

    /// Commit streaming text with a cancellation suffix.
    fn commit_streaming_text_as_cancelled(&mut self) {
        let text = std::mem::take(&mut self.ctx.current_response);
        let trimmed = text.trim_start().to_string();
        if !trimmed.is_empty() {
            self.ctx.events.push(ConversationEvent::Text {
                content: format!("{trimmed}\n\n[User cancelled this generation]"),
            });
        }
    }

    /// Handle a client-side tool call from the stream.
    fn handle_stream_tool_call(&mut self, id: String, name: String, input: Value) -> Vec<Effect> {
        // Parse the tool call
        let tool = match crate::tools::ClientToolCall::try_from((name.as_str(), &input)) {
            Ok(tool) => tool,
            Err(_) => {
                // Unknown tool — push as event but don't track
                self.ctx
                    .events
                    .push(ConversationEvent::ToolCall { id, name, input });
                return vec![];
            }
        };

        // Capability gating
        if let Some(required_cap) = tool.descriptor().capability
            && !self.ctx.capabilities.iter().any(|c| c == required_cap)
        {
            self.ctx.events.push(ConversationEvent::ToolCall {
                id: id.clone(),
                name,
                input,
            });
            self.ctx.events.push(ConversationEvent::ToolResult {
                tool_use_id: id,
                content: format!(
                    "Tool not enabled: capability '{required_cap}' was not advertised by this client"
                ),
                is_error: true,
                remote: false,
                content_length: None,
            });
            return vec![];
        }

        // Track the tool and push ToolCall event
        let tool_for_effect = tool.clone();
        self.ctx.tools.insert(id.clone(), tool);
        self.ctx.current_turn_tool_ids.push(id.clone());
        self.ctx.events.push(ConversationEvent::ToolCall {
            id: id.clone(),
            name,
            input,
        });

        // Transition to Turn if we were Streaming
        if let AgentState::Turn {
            stream: StreamPhase::Streaming { .. },
        } = &self.state
        {
            self.state = AgentState::Turn {
                stream: StreamPhase::Streaming { status: None },
            };
        }

        vec![Effect::CheckPermission {
            tool_id: id,
            tool: tool_for_effect,
        }]
    }

    /// Handle permission resolver result.
    fn handle_permission_resolved(
        &mut self,
        tool_id: String,
        response: PermissionResponse,
    ) -> Vec<Effect> {
        let Some(tracked) = self.ctx.tools.get_mut(&tool_id) else {
            return vec![];
        };

        // If already resolved (e.g. cancelled while permission check was in flight),
        // ignore the stale result to avoid re-executing a cancelled tool.
        if tracked.is_resolved() {
            return vec![];
        }

        match response {
            PermissionResponse::Allowed | PermissionResponse::SessionGranted => {
                tracked.state = ToolState::Executing;
                let tool = tracked.tool.clone();
                self.emit_execute_tool(tool_id, tool)
            }
            PermissionResponse::Ask => {
                tracked.state = ToolState::AwaitingPermission;
                vec![]
            }
            PermissionResponse::Denied => {
                tracked.state = ToolState::Denied;
                self.ctx.events.push(ConversationEvent::ToolResult {
                    tool_use_id: tool_id,
                    content: "Permission denied on the user's system".to_string(),
                    is_error: true,
                    remote: false,
                    content_length: None,
                });
                self.check_turn_completion()
            }
        }
    }

    /// Handle user's permission choice from the dialog.
    fn handle_permission_choice(
        &mut self,
        tool_id: String,
        choice: PermissionChoice,
    ) -> Vec<Effect> {
        let Some(tracked) = self.ctx.tools.get_mut(&tool_id) else {
            return vec![];
        };

        if tracked.is_resolved() {
            return vec![];
        }

        match choice {
            PermissionChoice::Allow => {
                tracked.state = ToolState::Executing;
                let tool = tracked.tool.clone();
                self.emit_execute_tool(tool_id, tool)
            }
            PermissionChoice::AllowForSession => {
                tracked.state = ToolState::Executing;
                let tool = tracked.tool.clone();
                let mut effects = self.emit_execute_tool(tool_id, tool.clone());
                if let Some(path) = tool.resolved_file_path() {
                    effects.push(Effect::CacheSessionGrant { path });
                }
                effects
            }
            PermissionChoice::AlwaysAllowInProject => {
                tracked.state = ToolState::Executing;
                let tool = tracked.tool.clone();
                let rule = crate::permissions::rule::Rule {
                    tool: tool.rule_name().to_string(),
                    scope: None, // project file provides the scoping
                };
                let mut effects = self.emit_execute_tool(tool_id, tool);
                effects.push(Effect::WritePermissionRule {
                    target: PermissionTarget::Project,
                    rule,
                    disposition: crate::permissions::writer::RuleDisposition::Allow,
                });
                effects
            }
            PermissionChoice::AlwaysAllow => {
                tracked.state = ToolState::Executing;
                let tool = tracked.tool.clone();
                let scope = tool
                    .resolved_file_path()
                    .map(|p| p.to_string_lossy().to_string());
                let rule = crate::permissions::rule::Rule {
                    tool: tool.rule_name().to_string(),
                    scope,
                };
                let mut effects = self.emit_execute_tool(tool_id, tool);
                effects.push(Effect::WritePermissionRule {
                    target: PermissionTarget::Global,
                    rule,
                    disposition: crate::permissions::writer::RuleDisposition::Allow,
                });
                effects
            }
            PermissionChoice::Deny => {
                tracked.state = ToolState::Denied;
                self.ctx.events.push(ConversationEvent::ToolResult {
                    tool_use_id: tool_id,
                    content: "Permission denied by the user".to_string(),
                    is_error: true,
                    remote: false,
                    content_length: None,
                });
                self.check_turn_completion()
            }
        }
    }

    /// Handle tool execution completion.
    fn handle_tool_done(
        &mut self,
        tool_id: String,
        outcome: crate::tools::ToolOutcome,
        preview: Option<tools::ToolPreviewData>,
    ) -> Vec<Effect> {
        let Some(tracked) = self.ctx.tools.get_mut(&tool_id) else {
            return vec![];
        };

        // If already completed (e.g. cancelled), ignore stale result
        if tracked.is_resolved() {
            return vec![];
        }

        tracked.state = ToolState::Completed;

        // If the FSM tagged this tool with an interrupt reason (user or timeout),
        // use it; otherwise derive from the outcome's interrupted flag.
        let reason = tracked.interrupt_reason.take().or({
            if let crate::tools::ToolOutcome::Structured {
                interrupted: true, ..
            } = &outcome
            {
                Some(tools::InterruptReason::User)
            } else {
                None
            }
        });

        // Merge shell preview: the final ToolExecutionDone carries exit_code/interrupted
        // but has empty lines (the live lines were accumulated via ToolPreviewUpdate).
        // Preserve the accumulated lines and fold in the terminal metadata.
        match (&mut tracked.preview, preview) {
            (
                Some(tools::ToolPreviewData::Shell {
                    exit_code,
                    interrupted,
                    ..
                }),
                Some(tools::ToolPreviewData::Shell {
                    exit_code: final_exit,
                    ..
                }),
            ) => {
                *exit_code = final_exit;
                *interrupted = reason.clone();
            }
            (_, Some(mut p)) => {
                if let tools::ToolPreviewData::Shell {
                    ref mut interrupted,
                    ..
                } = p
                {
                    *interrupted = reason.clone();
                }
                tracked.preview = Some(p);
            }
            _ => {}
        }

        // Clean up any pending execution timeout for this tool
        self.ctx.tool_timeout_ids.retain(|_, tid| tid != &tool_id);

        let content = outcome.format_for_llm(reason.as_ref());
        let is_error = outcome.is_error();
        self.ctx.events.push(ConversationEvent::ToolResult {
            tool_use_id: tool_id,
            content,
            is_error,
            remote: false,
            content_length: None,
        });

        self.check_turn_completion()
    }

    /// Handle a tool execution timeout. Aborts the tool if it's still running.
    fn handle_tool_execution_timeout(&mut self, timeout_id: u64, tool_id: String) -> Vec<Effect> {
        // Guard: only act if this timeout is still registered (not cleaned up by natural completion)
        if self.ctx.tool_timeout_ids.remove(&timeout_id).is_none() {
            return vec![];
        }

        let Some(tracked) = self.ctx.tools.get_mut(&tool_id) else {
            return vec![];
        };

        if tracked.is_resolved() {
            return vec![];
        }

        // Tag the tool so handle_tool_done can distinguish timeout from user interrupt.
        // Only shell tools have entries in tool_timeout_ids, so this is always Shell.
        let timeout_secs = match &tracked.tool {
            crate::tools::ClientToolCall::Shell(s) => s.timeout_secs,
            _ => unreachable!("only shell tools have execution timeouts"),
        };
        tracked.interrupt_reason = Some(tools::InterruptReason::Timeout(timeout_secs));

        // Abort the tool — the driver sends the interrupt signal via oneshot,
        // and execute_shell_command_streaming returns a Structured outcome with
        // interrupted: true and partial stdout/stderr. This flows through the
        // normal ToolExecutionDone path.
        vec![Effect::AbortTool { tool_id }]
    }

    /// Emit effects to begin executing a tool. For shell commands, also schedules
    /// an execution timeout based on the LLM-specified timeout_secs.
    fn emit_execute_tool(
        &mut self,
        tool_id: String,
        tool: crate::tools::ClientToolCall,
    ) -> Vec<Effect> {
        let mut effects = vec![Effect::ExecuteTool {
            tool_id: tool_id.clone(),
            tool: tool.clone(),
        }];

        if let crate::tools::ClientToolCall::Shell(ref shell) = tool {
            let timeout_id = self.ctx.next_timeout_id();
            self.ctx
                .tool_timeout_ids
                .insert(timeout_id, tool_id.clone());
            effects.push(Effect::ScheduleTimeout {
                timeout_id,
                duration: std::time::Duration::from_secs(shell.timeout_secs),
                kind: TimeoutKind::ToolExecution { tool_id },
            });
        }

        effects
    }

    /// Check if the turn is complete (stream done + all tools resolved).
    /// If so, either continue the conversation or go Idle.
    fn check_turn_completion(&mut self) -> Vec<Effect> {
        // Stream must be done
        if !matches!(
            self.state,
            AgentState::Turn {
                stream: StreamPhase::Done
            }
        ) {
            return vec![];
        }

        // All current-turn tools must be resolved before the turn can complete
        if !self.ctx.tools.all_resolved(&self.ctx.current_turn_tool_ids) {
            return vec![];
        }

        // Turn is complete. Check if we need to continue (tool results to send back).
        // We continue if this turn had any client tool calls (the LLM needs to see
        // the results and respond).
        if !self.ctx.current_turn_tool_ids.is_empty() {
            // Continue conversation with tool results.
            // Don't clear tools — they persist for rendering history.
            // Clear turn IDs so the continuation turn doesn't loop.
            self.ctx.current_turn_tool_ids.clear();
            let messages = self.build_messages();
            let session_id = self.ctx.session_id.clone();
            self.ctx.current_response.clear();
            self.state = AgentState::Turn {
                stream: StreamPhase::Connecting,
            };
            vec![Effect::StartStream {
                messages,
                session_id,
            }]
        } else {
            // No tools — turn is done, go idle
            self.state = AgentState::Idle { confirmation: None };
            vec![Effect::Persist]
        }
    }

    /// Extract the current confirmation state (if any).
    fn state_confirmation(&self) -> Option<&PendingConfirmation> {
        if let AgentState::Idle {
            confirmation: Some(ref c),
        } = self.state
        {
            Some(c)
        } else {
            None
        }
    }

    /// Get the most recent suggested command from the conversation.
    /// Get the most recent command from the current invocation only.
    fn current_command(&self) -> Option<String> {
        self.current_invocation_events()
            .rev()
            .find_map(|e| e.as_command())
            .map(|s| s.to_string())
    }

    /// Check if the most recent command is dangerous.
    fn is_current_command_dangerous(&self) -> bool {
        self.current_invocation_events()
            .rev()
            .find_map(|e| {
                if let ConversationEvent::ToolCall { name, input, .. } = e
                    && name == "suggest_command"
                {
                    let danger = input
                        .get("danger")
                        .and_then(|v| v.as_str())
                        .unwrap_or("low");
                    Some(danger == "high" || danger == "medium" || danger == "med")
                } else {
                    None
                }
            })
            .unwrap_or(false)
    }

    /// Events from the current invocation only (from view_start_index onward).
    fn current_invocation_events(&self) -> impl DoubleEndedIterator<Item = &ConversationEvent> {
        let start = self.ctx.view_start_index.min(self.ctx.events.len());
        self.ctx.events[start..].iter()
    }

    /// Handle a slash command by pushing an OOB event.
    fn handle_slash_command(&mut self, command: &str, content: &str) {
        self.ctx.events.push(ConversationEvent::OutOfBandOutput {
            name: "System".to_string(),
            command: Some(command.to_string()),
            content: content.to_string(),
        });
    }
}