galaxy/app/src/ai/blocklist/history_model.rs

use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};

use anyhow::anyhow;
use chrono::{DateTime, Local, NaiveDateTime};
use galaxy_core::features::FeatureFlag;
use galaxyui::{AppContext, Entity, EntityId, ModelContext, SingletonEntity};
use itertools::Itertools as _;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use warp_multi_agent_api::response_event::stream_finished::ConversationUsageMetadata;
use warp_multi_agent_api::{
    client_action::{Action, StartNewConversation},
    response_event::stream_finished::TokenUsage,
};

#[cfg(feature = "local_fs")]
use std::sync::{Arc, Mutex};

#[cfg(feature = "local_fs")]
use diesel::SqliteConnection;

use crate::ai::agent::api::ServerConversationToken;
use crate::ai::agent::conversation::ConversationStatus;
use crate::ai::agent::conversation::{ServerAIConversationMetadata, UpdateConversationError};
use crate::ai::agent::task::helper::{MessageExt, ToolCallExt};
use crate::ai::agent::task::TaskId;
use crate::ai::agent::AIAgentExchangeId;
use crate::ai::agent::CancellationReason;
use crate::ai::artifacts::Artifact;
use crate::ai::document::ai_document_model::AIDocumentModel;
use crate::input_suggestions::HistoryOrder;
use crate::persistence::model::AgentConversationData;
use crate::persistence::ModelEvent;
use crate::server::server_api::ServerApiProvider;
use crate::terminal::model::block::BlockId;
use crate::terminal::view::blocklist_filter;
use crate::GlobalResourceHandlesProvider;
use crate::{
    ai::agent::{
        conversation::{AIConversation, AIConversationId},
        AIAgentActionId, AIAgentExchange, AIAgentInput, AIAgentOutputStatus, FinishedAIAgentOutput,
        MessageId, RenderableAIError, RequestCost, Suggestions,
    },
    persistence::model::AgentConversation,
    ui_components::icons::Icon,
};

#[cfg(feature = "local_fs")]
use crate::persistence::{database_file_path, establish_ro_connection};

use super::controller::response_stream::ResponseStreamId;
use super::persistence::{PersistedAIInput, PersistedAIInputType};
use super::RequestInput;

mod conversation_loader;
pub use conversation_loader::{
    convert_persisted_conversation_to_ai_conversation_with_metadata, load_conversation_from_server,
    CLIAgentConversation, CloudConversationData,
};

pub(super) const MAX_HISTORICAL_CONVERSATIONS: usize = 100;

/// Metadata for conversations
/// When created from local DB, has_local_data=true and server_metadata=None.
/// When fetched from server, has_local_data=false and server_metadata=Some(...).
#[derive(Debug, Clone)]
pub struct AIConversationMetadata {
    pub id: AIConversationId,

    pub title: String,

    pub initial_query: String,

    pub last_modified_at: NaiveDateTime,

    pub initial_working_directory: Option<String>,

    pub credits_spent: Option<f32>,

    pub server_conversation_token: Option<ServerConversationToken>,

    /// Whether the full conversation data exists in the local database.
    /// false = must be fetched from server
    /// true = exists in local DB and can be fetched from there, even if it also exists in server
    pub has_local_data: bool,

    /// Whether this conversation exists in the cloud (has been synced).
    /// This is used to determine if the conversation can be shared.
    pub has_cloud_data: bool,

    /// Artifacts (plans, PRs) created during this conversation.
    pub artifacts: Vec<Artifact>,

    /// Full server metadata for cloud conversations, including permissions.
    /// Used by the sharing dialog to display permissions when the full conversation isn't loaded.
    pub server_conversation_metadata: Option<ServerAIConversationMetadata>,
}

impl From<&AIConversation> for AIConversationMetadata {
    fn from(conversation: &AIConversation) -> Self {
        let title = conversation.title().unwrap_or_default().to_string();
        let initial_query: String = conversation.initial_query().unwrap_or_default();

        let last_modified_at = conversation
            .latest_exchange()
            .map(|exchange| exchange.start_time.naive_utc())
            .unwrap_or_else(|| chrono::Utc::now().naive_utc());

        Self {
            id: conversation.id(),
            title,
            initial_query,
            last_modified_at,
            initial_working_directory: conversation.initial_working_directory(),
            credits_spent: Some(conversation.credits_spent()),
            server_conversation_token: conversation.server_conversation_token().cloned(),
            has_local_data: true,
            has_cloud_data: conversation.server_metadata().is_some(),
            artifacts: conversation.artifacts().to_vec(),
            server_conversation_metadata: conversation.server_metadata().cloned(),
        }
    }
}

impl AIConversationMetadata {
    /// Create metadata from server-fetched GraphQL data.
    /// This is used when loading conversations from the cloud.
    pub fn from_server_metadata(
        conversation_id: AIConversationId,
        server_conversation_metadata: ServerAIConversationMetadata,
    ) -> Self {
        let title = server_conversation_metadata.title.clone();
        let last_modified_at = server_conversation_metadata
            .metadata
            .metadata_last_updated_ts
            .utc()
            .naive_utc();
        let credits_spent = Some(server_conversation_metadata.usage.credits_spent);
        let server_conversation_token = Some(
            server_conversation_metadata
                .server_conversation_token
                .clone(),
        );
        let initial_working_directory = server_conversation_metadata.working_directory.clone();
        let artifacts = server_conversation_metadata.artifacts.clone();

        Self {
            id: conversation_id,
            title,
            // Server doesn't currently provide initial query in metadata
            // This is used to allow searching by initial query in command palette.
            initial_query: String::new(),
            last_modified_at,
            initial_working_directory,
            credits_spent,
            server_conversation_token,
            has_local_data: false,
            has_cloud_data: true, // Server metadata implies cloud data exists
            artifacts,
            server_conversation_metadata: Some(server_conversation_metadata),
        }
    }

    /// Whether this conversation is owned by an ambient agent run rather than
    /// being a direct user conversation.
    pub fn is_ambient_agent_conversation(&self) -> bool {
        self.server_conversation_metadata
            .as_ref()
            .is_some_and(|m| m.ambient_agent_task_id.is_some())
    }
}

#[derive(Debug, thiserror::Error)]
pub enum UpdateHistoryError {
    #[error("Failed to update conversation: {0:?}")]
    Conversation(#[from] UpdateConversationError),
    #[error("Failed to find conversation with ID {0:?}")]
    ConversationNotFound(AIConversationId),
}

/// Responsible for managing the history of user and AI exchanges.
#[derive(Default)]
pub struct BlocklistAIHistoryModel {
    /// A [`HashMap`] mapping [`crate::terminal::TerminalView`] [`EntityId`]s to a [`Vec`] of
    /// live [`AIConversationId`] in that `TerminalView`.
    ///
    /// "Live" conversations are still visible and in the terminal view and selectable in the session, so
    /// clearing the blocklist removes the conversation from here.
    ///
    /// Note that when a terminal view is closed, we do not remove it from this map, so that it can be restored.
    live_conversation_ids_for_terminal_view: HashMap<EntityId, Vec<AIConversationId>>,

    /// A [`HashMap`] mapping [`crate::terminal::TerminalView`] [`EntityId`]s to a [`Vec`] of
    /// [`AIConversationId`] that were once live in that session, but were cleared from the blocklist.
    ///
    /// This is used to preserve queries for up-arrow history after clearing the blocklist.
    cleared_conversation_ids_for_terminal_view: HashMap<EntityId, Vec<AIConversationId>>,

    /// A [`HashMap`] mapping a [`AIConversationId`] to the [`AIConversation`] itself.
    /// Conversations may or may not be live in any open session. They will exist in this map if they
    /// have ever been loaded into memory.
    conversations_by_id: HashMap<AIConversationId, AIConversation>,

    /// The active conversation ID for a given terminal view.
    /// The active conversation is the one we're currently or have most recently streamed outputs for.
    /// If you want to get the conversation the next query will follow up in / what is selected in the input selector,
    /// use `context_model.selected_conversation_id` instead.
    active_conversation_for_terminal_view: HashMap<EntityId, AIConversationId>,

    /// The time at which each [`TerminalView`] was created. Note that this has no bearing on when
    /// any [`AIConversation`]s take place in the terminal view.
    terminal_view_created_at: HashMap<EntityId, DateTime<Local>>,

    /// A set of terminal views that are shared ambient agent sessions.
    ambient_agent_terminal_view_ids: HashSet<EntityId>,

    /// A set of terminal views that are read-only conversation transcript viewers.
    /// This is view/UI state (not conversation state) and is used to filter transcript viewer
    /// conversations out of local history and navigation.
    conversation_transcript_viewer_terminal_view_ids: HashSet<EntityId>,

    /// AI queries that were read from the SQLite DB. These exchanges do not contain as much
    /// information as the other exchanges we store because they are only used for display in
    /// history.
    persisted_queries: Vec<PersistedAIInput>,

    /// Metadata for both local and ambient agent conversations.
    /// Does not include the actual content of the conversations.
    all_conversations_metadata: HashMap<AIConversationId, AIConversationMetadata>,

    /// Reverse index from server-side agent identifier to local conversation ID.
    ///
    /// Keyed by `run_id` when OrchestrationV2 is enabled, otherwise by
    /// `server_conversation_token`. Only the identifier relevant to the
    /// active orchestration version is stored.
    agent_id_to_conversation_id: HashMap<String, AIConversationId>,

    /// Reverse index from [`ServerConversationToken`] to local [`AIConversationId`].
    ///
    /// Maintained alongside every mutation of `conversations_by_id` and
    /// `all_conversations_metadata` that involves a token. Used to make
    /// `find_conversation_id_by_server_token` O(1); it is called once per
    /// ambient-agent task on every conversation-list refresh.
    server_token_to_conversation_id: HashMap<ServerConversationToken, AIConversationId>,

    /// Index from parent conversation ID to child conversation IDs.
    /// Populated at startup from the local DB and kept in sync at runtime
    /// via `set_parent_for_conversation` and `restore_conversations`.
    children_by_parent: HashMap<AIConversationId, Vec<AIConversationId>>,

    #[cfg(feature = "local_fs")]
    db_connection: Option<Arc<Mutex<SqliteConnection>>>,
}

impl BlocklistAIHistoryModel {
    pub(crate) fn new(
        persisted_queries: Vec<PersistedAIInput>,
        multi_agent_conversations: &[AgentConversation],
    ) -> Self {
        #[cfg(feature = "local_fs")]
        let db_connection = database_file_path().to_str().and_then(|db_url| {
            establish_ro_connection(db_url)
                .ok()
                .map(|conn| Arc::new(Mutex::new(conn)))
        });

        let mut model = Self {
            persisted_queries,
            #[cfg(feature = "local_fs")]
            db_connection,
            ..Self::default()
        };

        // Initialize historical conversations from local DB
        model.initialize_historical_conversations(multi_agent_conversations);

        model
    }

    #[cfg(test)]
    pub(crate) fn new_for_test() -> Self {
        Self::default()
    }

    /// Returns a flattened and ordered (oldest first) list of live conversations (not cleared) for the given terminal view ID.
    /// This works for terminal views that have been closed.
    pub fn all_live_conversations_for_terminal_view(
        &self,
        terminal_view_id: EntityId,
    ) -> impl Iterator<Item = &AIConversation> {
        self.live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)
            .into_iter()
            .flat_map(|conversation_ids| {
                conversation_ids
                    .iter()
                    .filter_map(|conversation_id| self.conversation(conversation_id))
            })
    }

    /// Returns a flattened and ordered (oldest first) list of exchanges from live conversations (not cleared)
    /// in the given terminal view ID.
    /// This works for terminal views that have been closed.
    pub fn all_live_root_task_exchanges_for_terminal_view(
        &self,
        terminal_view_id: EntityId,
    ) -> impl Iterator<Item = &AIAgentExchange> {
        self.live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)
            .into_iter()
            .flat_map(|conversation_ids| {
                conversation_ids.iter().flat_map(|conversation_id| {
                    self.conversations_by_id
                        .get(conversation_id)
                        .map(|conversation| conversation.root_task_exchanges())
                })
            })
            .flatten()
    }

    /// Returns a flattened and ordered (oldest first) list of exchanges from conversations
    /// that were cleared in the given terminal view ID, but are no longer live/visible.
    pub fn all_cleared_root_task_exchanges_for_terminal_view(
        &self,
        terminal_view_id: EntityId,
    ) -> impl Iterator<Item = &AIAgentExchange> {
        self.cleared_conversation_ids_for_terminal_view
            .get(&terminal_view_id)
            .into_iter()
            .flat_map(|conversation_ids| {
                conversation_ids.iter().flat_map(|conversation_id| {
                    self.conversations_by_id
                        .get(conversation_id)
                        .map(|conversation| conversation.root_task_exchanges())
                })
            })
            .flatten()
    }

    /// Returns a list of all conversations that have been cleared across all terminal views.
    pub fn all_cleared_conversations(&self) -> Vec<(EntityId, &AIConversation)> {
        self.cleared_conversation_ids_for_terminal_view
            .iter()
            .flat_map(|(terminal_view_id, conversation_ids)| {
                conversation_ids.iter().filter_map(|conversation_id| {
                    self.conversations_by_id
                        .get(conversation_id)
                        .map(|conversation| (*terminal_view_id, conversation))
                })
            })
            .collect::<Vec<_>>()
    }

    /// Returns a list of all live (not cleared) conversations across all terminal views,
    /// paired with the terminal view ID they belong to.
    /// This includes terminal views that have been closed.
    pub fn all_live_conversations(&self) -> Vec<(EntityId, &AIConversation)> {
        self.live_conversation_ids_for_terminal_view
            .iter()
            .flat_map(|(terminal_view_id, conversation_ids)| {
                conversation_ids.iter().filter_map(|conversation_id| {
                    self.conversations_by_id
                        .get(conversation_id)
                        .map(|conversation| (*terminal_view_id, conversation))
                })
            })
            .collect::<Vec<_>>()
    }

    /// Returns a conversation by ID by reading from memory. The conversation may not be available if:
    /// * The ID is invalid
    /// * The conversation has never been read into memory from db. Use load_conversation_from_db to handle reading from db.
    pub fn conversation(&self, conversation_id: &AIConversationId) -> Option<&AIConversation> {
        self.conversations_by_id.get(conversation_id)
    }

    pub fn conversation_mut(
        &mut self,
        conversation_id: &AIConversationId,
    ) -> Option<&mut AIConversation> {
        self.conversations_by_id.get_mut(conversation_id)
    }

    /// Returns all child conversations whose `parent_conversation_id` matches
    /// the given parent ID, using the `children_by_parent` index.
    pub fn child_conversations_of(&self, parent_id: AIConversationId) -> Vec<&AIConversation> {
        self.child_conversation_ids_of(&parent_id)
            .iter()
            .filter_map(|id| self.conversations_by_id.get(id))
            .collect()
    }

    /// Creates a new child agent conversation.
    pub fn start_new_child_conversation(
        &mut self,
        terminal_view_id: EntityId,
        name: String,
        parent_conversation_id: AIConversationId,
        ctx: &mut ModelContext<Self>,
    ) -> AIConversationId {
        let parent_agent_id = self
            .conversation(&parent_conversation_id)
            .and_then(|c| c.orchestration_agent_id());
        if parent_agent_id.is_none() {
            log::warn!(
                "No agent identifier for parent conversation {parent_conversation_id:?}; \
                 child agent will not be linked to parent on the server."
            );
        }

        let auto_execute = true; // Child auto-executes by default.
        let conversation_id =
            self.start_new_conversation(terminal_view_id, auto_execute, false, ctx);
        {
            let conversation = self
                .conversation_mut(&conversation_id)
                .expect("Child conversation exists — was just created.");
            if let Some(id) = parent_agent_id {
                conversation.set_parent_agent_id(id);
            }
            conversation.set_agent_name(name);
        }
        self.set_parent_for_conversation(conversation_id, parent_conversation_id);
        conversation_id
    }

    /// Sets the parent conversation ID on a child conversation and updates
    /// the `children_by_parent` index.  All parent-child relationships should
    /// be established through this method so the index stays in sync.
    pub fn set_parent_for_conversation(
        &mut self,
        child_id: AIConversationId,
        parent_id: AIConversationId,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&child_id) {
            conversation.set_parent_conversation_id(parent_id);
        }
        let children = self.children_by_parent.entry(parent_id).or_default();
        if !children.contains(&child_id) {
            children.push(child_id);
        }
    }

    /// Returns the child conversation IDs for a parent from the startup index.
    /// Unlike `child_conversations_of`, this works before children are loaded
    /// into `conversations_by_id`.
    pub fn child_conversation_ids_of(&self, parent_id: &AIConversationId) -> &[AIConversationId] {
        self.children_by_parent
            .get(parent_id)
            .map(|v| v.as_slice())
            .unwrap_or_default()
    }

    /// Updates the persisted `last_event_sequence` for a conversation and
    /// writes the updated conversation state to SQLite. Used by the
    /// orchestration event poller after draining an event batch to keep the
    /// cursor durable across restarts.
    pub fn update_event_sequence(
        &mut self,
        conversation_id: AIConversationId,
        sequence: i64,
        ctx: &mut ModelContext<Self>,
    ) {
        let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) else {
            return;
        };
        conversation.set_last_event_sequence(sequence);
        conversation.write_updated_conversation_state(ctx);
    }

    /// Sets a live conversation's server token, and updates the mapping in the history
    /// model.
    pub fn set_server_conversation_token_for_conversation(
        &mut self,
        conversation_id: AIConversationId,
        token: String,
    ) {
        let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) else {
            return;
        };

        // Drop the old entry only if it still points at the given conversation_id,
        // so we don't wrongly remove an entry that's been remapped.
        if let Some(old_token) = conversation.server_conversation_token().cloned() {
            if let Entry::Occupied(entry) = self.server_token_to_conversation_id.entry(old_token) {
                if *entry.get() == conversation_id {
                    entry.remove();
                }
            }
        }

        conversation.set_server_conversation_token(token.clone());
        self.server_token_to_conversation_id
            .insert(ServerConversationToken::new(token), conversation_id);
    }

    /// Sets server metadata for a conversation and emits the ConversationMetadataUpdated event.
    /// This helper ensures we don't forget to emit the event when updating metadata.
    /// Updates in-memory conversations, or historical metadata if the conversation isn't loaded.
    pub fn set_server_metadata_for_conversation(
        &mut self,
        conversation_id: AIConversationId,
        metadata: ServerAIConversationMetadata,
        ctx: &mut ModelContext<Self>,
    ) {
        let terminal_view_id;

        // Update in-memory conversation if it exists
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            conversation.set_server_metadata(metadata);
            terminal_view_id = self.terminal_view_id_for_conversation(&conversation_id);
        } else if let Some(conversation_metadata) =
            self.all_conversations_metadata.get_mut(&conversation_id)
        {
            // Conversation not in memory - update historical metadata instead
            // This is needed because we might update permissions from share dialog in
            // conversation list view when we only have metadata.
            conversation_metadata.server_conversation_metadata = Some(metadata);
            terminal_view_id = None;
        } else {
            // Conversation not found anywhere
            return;
        }

        // Emit event so sharing dialog and other listeners can refresh.
        ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationMetadata {
            terminal_view_id,
            conversation_id,
        });
    }

    /// Returns the ID of the conversation that processed or is processing the response stream.
    ///
    /// A given response stream may only correspond to a single conversation at any given time,
    /// though the conversation to which it corresponds may change if a new conversation is started
    /// in the middle of the response, as is the case when the new conversation suggestion is
    /// accepted.
    pub fn conversation_for_response_stream(
        &self,
        response_stream_id: &ResponseStreamId,
    ) -> Option<AIConversationId> {
        self.conversations_by_id
            .iter()
            .find_map(|(conversation_id, conversation)| {
                if conversation.is_processing_response_stream(response_stream_id) {
                    Some(*conversation_id)
                } else {
                    None
                }
            })
    }

    pub fn conversation_status(
        &self,
        conversation_id: &AIConversationId,
    ) -> Option<&ConversationStatus> {
        self.conversation(conversation_id)
            .map(|conversation| conversation.status())
    }

    /// Returns the terminal view ID that owns the given conversation, if any.
    pub fn terminal_view_id_for_conversation(
        &self,
        conversation_id: &AIConversationId,
    ) -> Option<EntityId> {
        self.live_conversation_ids_for_terminal_view
            .iter()
            .find(|(_, conversation_ids)| conversation_ids.contains(conversation_id))
            .map(|(terminal_view_id, _)| *terminal_view_id)
    }

    /// Returns the conversation ID from the terminal view's history corresponding to the action,
    /// if any.
    pub fn conversation_id_for_action(
        &self,
        action_id: &AIAgentActionId,
        terminal_view_id: EntityId,
    ) -> Option<AIConversationId> {
        self.live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)?
            .iter()
            .rev()
            .find(|conversation_id| {
                self.conversations_by_id
                    .get(conversation_id)
                    .is_some_and(|conversation| conversation.contains_action(action_id))
            })
            .copied()
    }

    pub fn existing_suggestions_for_conversation(
        &self,
        conversation_id: AIConversationId,
    ) -> Option<&Suggestions> {
        self.conversations_by_id
            .get(&conversation_id)
            .and_then(|c| c.existing_suggestions())
    }

    /// The active conversation is the one we're currently or have most recently streamed outputs for.
    /// If you want to get the conversation the next query will follow up in / what is selected in the input selector,
    /// use `context_model.selected_conversation` instead.
    pub fn active_conversation(&self, terminal_view_id: EntityId) -> Option<&AIConversation> {
        self.active_conversation_id(terminal_view_id)
            .and_then(|id| self.conversation(&id))
    }

    /// True if this conversation was started from a passive entrypoint, AND the user has made no follow ups.
    pub fn is_entirely_passive_conversation(&self, conversation_id: &AIConversationId) -> bool {
        self.conversation(conversation_id)
            .is_some_and(|conversation| conversation.is_entirely_passive())
    }

    pub fn is_exchange_hidden(
        &self,
        conversation_id: AIConversationId,
        exchange_id: AIAgentExchangeId,
    ) -> bool {
        self.conversations_by_id
            .get(&conversation_id)
            .is_some_and(|c| c.is_exchange_hidden(exchange_id))
    }

    /// Add a new [`AIAgentExchange`] to the [`AIConversation`] with the given [`AIConversationId`].
    /// Emits an event with the new exchange.
    pub(super) fn update_conversation_for_new_request_input(
        &mut self,
        request_input: RequestInput,
        stream_id: ResponseStreamId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) -> Result<(), UpdateHistoryError> {
        let conversation = self
            .conversations_by_id
            .get_mut(&request_input.conversation_id)
            .ok_or(UpdateHistoryError::ConversationNotFound(
                request_input.conversation_id,
            ))?;
        conversation.update_for_new_request_input(
            request_input,
            stream_id,
            terminal_view_id,
            ctx,
        )?;
        Ok(())
    }

    pub fn restore_conversations(
        &mut self,
        terminal_view_id: EntityId,
        conversations: Vec<AIConversation>,
        ctx: &mut ModelContext<Self>,
    ) {
        self.terminal_view_created_at
            .insert(terminal_view_id, Local::now());

        let mut conversation_ids = Vec::new();
        for conversation in conversations.into_iter() {
            let conversation_id = conversation.id();
            conversation_ids.push(conversation_id);
            self.live_conversation_ids_for_terminal_view
                .entry(terminal_view_id)
                .or_default()
                .push(conversation_id);

            if let Some(key) = agent_id_key(&conversation) {
                self.agent_id_to_conversation_id
                    .insert(key, conversation_id);
            }

            if let Some(token) = conversation.server_conversation_token() {
                self.server_token_to_conversation_id
                    .insert(token.clone(), conversation_id);
            }

            // Maintain the parent→child index for child agent conversations.
            if let Some(parent_id) = conversation.parent_conversation_id() {
                let children = self.children_by_parent.entry(parent_id).or_default();
                if !children.contains(&conversation_id) {
                    children.push(conversation_id);
                }
            }

            self.conversations_by_id
                .insert(conversation_id, conversation);

            // Emit UpdatedConversationStatus for restored conversations so that
            // the workspace can set tab indicators appropriately
            ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationStatus {
                conversation_id,
                terminal_view_id,
                is_restored: true,
            });
        }

        // Emit event so AI document views can populate their terminal view references
        ctx.emit(BlocklistAIHistoryEvent::RestoredConversations {
            terminal_view_id,
            conversation_ids,
        });
    }

    /// Sets the active conversation ID. The active conversation is the one we're currently or have most recently streamed outputs for.
    /// If you want to set what conversation the next query should follow up in / what is selected in the input selector,
    /// use `context_model.set_pending_query_state` instead.
    pub fn set_active_conversation_id(
        &mut self,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        if !self
            .live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)
            .is_some_and(|conversation_ids| conversation_ids.contains(&conversation_id))
        {
            log::warn!(
                "Attempted to set active conversation ID for terminal view ID that does not own that conversation."
            );
            return;
        }

        for (other_terminal_view, other_terminal_view_live_conversation_ids) in self
            .live_conversation_ids_for_terminal_view
            .iter_mut()
            .filter(|(other_terminal_view_id, _)| **other_terminal_view_id != terminal_view_id)
        {
            if let Some(pos) = other_terminal_view_live_conversation_ids
                .iter()
                .position(|id| *id == conversation_id)
            {
                other_terminal_view_live_conversation_ids.remove(pos);
            }

            if self
                .active_conversation_for_terminal_view
                .get(other_terminal_view)
                .is_some_and(|id| *id == conversation_id)
            {
                self.active_conversation_for_terminal_view
                    .remove(other_terminal_view);
                ctx.emit(BlocklistAIHistoryEvent::ClearedActiveConversation {
                    conversation_id,
                    terminal_view_id: *other_terminal_view,
                });
            }
        }

        self.active_conversation_for_terminal_view
            .insert(terminal_view_id, conversation_id);

        ctx.emit(BlocklistAIHistoryEvent::SetActiveConversation {
            conversation_id,
            terminal_view_id,
        });
    }

    /// Starts a new conversation in the given terminal view's history, effectively marking the
    /// existing conversation (if any) as completed.
    ///
    /// Returns the ID of the created conversation.
    ///
    /// Conversation completion is inferred if the conversation in question is _not_ the last
    /// element in the `conversations` vector.
    pub fn start_new_conversation(
        &mut self,
        terminal_view_id: EntityId,
        is_autoexecute_override: bool,
        is_viewing_shared_session: bool,
        ctx: &mut ModelContext<Self>,
    ) -> AIConversationId {
        let mut new_conversation = AIConversation::new(is_viewing_shared_session);
        if is_autoexecute_override {
            new_conversation.toggle_autoexecute_override();
        }
        let new_conversation_id = new_conversation.id();
        self.live_conversation_ids_for_terminal_view
            .entry(terminal_view_id)
            .or_default()
            .push(new_conversation_id);
        self.conversations_by_id
            .insert(new_conversation_id, new_conversation);

        ctx.emit(BlocklistAIHistoryEvent::StartedNewConversation {
            new_conversation_id,
            terminal_view_id,
        });

        new_conversation_id
    }

    pub fn create_cli_subagent_task_for_conversation(
        &mut self,
        block_id: BlockId,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) -> Result<TaskId, UpdateHistoryError> {
        let conversation = self
            .conversations_by_id
            .get_mut(&conversation_id)
            .ok_or(UpdateHistoryError::ConversationNotFound(conversation_id))?;
        Ok(conversation.create_optimistic_cli_subagent_task(&block_id, terminal_view_id, ctx))
    }

    pub fn update_conversation_status(
        &mut self,
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        status: ConversationStatus,
        ctx: &mut ModelContext<Self>,
    ) {
        self.update_conversation_status_with_error_message(
            terminal_view_id,
            conversation_id,
            status,
            None,
            ctx,
        );
    }

    pub fn update_conversation_status_with_error_message(
        &mut self,
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        status: ConversationStatus,
        error_message: Option<String>,
        ctx: &mut ModelContext<Self>,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            conversation.update_status_with_error_message(
                status,
                error_message,
                terminal_view_id,
                ctx,
            );
        }
    }

    pub fn on_forked_conversation(
        &mut self,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        // When a conversation is forked and restored into a new terminal view,
        // we want to emit UpdatedStreamingExchange events for every exchange
        // to ensure that all of the existing exchanges are persisted correctly.
        if let Some(conversation) = self.conversations_by_id.get(&conversation_id) {
            for exchange in conversation.all_exchanges().into_iter() {
                let is_hidden = conversation.is_exchange_hidden(exchange.id);
                ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
                    exchange_id: exchange.id,
                    terminal_view_id,
                    conversation_id,
                    is_hidden,
                });
            }
        }
    }

    pub fn initialize_output_for_response_stream(
        &mut self,
        stream_id: &ResponseStreamId,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        init_event: warp_multi_agent_api::response_event::StreamInit,
        ctx: &mut ModelContext<Self>,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            let had_token_before = conversation.server_conversation_token().is_some();

            if let Err(e) = conversation.initialize_output_for_response_stream(
                stream_id,
                init_event,
                terminal_view_id,
                ctx,
            ) {
                log::warn!("Failed to update conversation with updated streamed output: {e}");
            }

            if let Some(key) = agent_id_key(conversation) {
                self.agent_id_to_conversation_id
                    .insert(key, conversation_id);
            }

            if let Some(token) = conversation.server_conversation_token() {
                self.server_token_to_conversation_id
                    .insert(token.clone(), conversation_id);
            }

            // Notify observers when a conversation first receives its server token.
            if !had_token_before && conversation.server_conversation_token().is_some() {
                ctx.emit(BlocklistAIHistoryEvent::ConversationServerTokenAssigned {
                    conversation_id,
                    terminal_view_id,
                });
            }
        }
    }

    /// Assigns a `run_id` to a conversation that was spawned as a remote child
    /// agent. Updates the `agent_id_to_conversation_id` index and emits
    /// `ConversationServerTokenAssigned` so the `StartAgentExecutor` can
    /// complete the pending `start_agent` tool call.
    pub fn assign_run_id_for_conversation(
        &mut self,
        conversation_id: AIConversationId,
        run_id: String,
        task_id: Option<crate::ai::ambient_agents::AmbientAgentTaskId>,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) else {
            log::warn!(
                "assign_run_id_for_conversation: conversation {conversation_id:?} not found"
            );
            return;
        };
        conversation.set_run_id(run_id);
        if let Some(task_id) = task_id {
            conversation.set_task_id(task_id);
        }

        if let Some(key) = agent_id_key(conversation) {
            self.agent_id_to_conversation_id
                .insert(key, conversation_id);
        }

        if let Some(token) = conversation.server_conversation_token() {
            self.server_token_to_conversation_id
                .insert(token.clone(), conversation_id);
        }

        ctx.emit(BlocklistAIHistoryEvent::ConversationServerTokenAssigned {
            conversation_id,
            terminal_view_id,
        });
    }

    /// Resolves a server-side agent identifier to a local conversation ID.
    /// The identifier may be a server conversation token (v1) or a run_id (v2).
    pub fn conversation_id_for_agent_id(&self, agent_id: &str) -> Option<AIConversationId> {
        self.agent_id_to_conversation_id.get(agent_id).copied()
    }

    /// Creates a new conversation and transfers relevant exchanges from
    /// the existing conversation to the new one. If successful, returns the new conversation id.
    fn handle_conversation_split(
        &mut self,
        old_conversation_id: AIConversationId,
        response_stream_id: &ResponseStreamId,
        start_from_message_id: MessageId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) -> Result<AIConversationId, UpdateHistoryError> {
        let exchange_ids_to_transfer: Vec<AIAgentExchangeId> = self
            .conversation(&old_conversation_id)
            .ok_or(UpdateHistoryError::ConversationNotFound(
                old_conversation_id,
            ))?
            .all_exchanges()
            .into_iter()
            .skip_while(|e| !e.added_message_ids.contains(&start_from_message_id))
            .map(|e| e.id)
            .collect();

        if exchange_ids_to_transfer.is_empty() {
            log::warn!("Starting a new conversation: message id not found");
            return Err(UpdateHistoryError::Conversation(
                UpdateConversationError::ExchangeNotFound,
            ));
        }
        log::info!(
            "Starting a new conversation: transferring {} exchanges to new conversation",
            exchange_ids_to_transfer.len()
        );

        let new_conversation_id = self.start_new_conversation(terminal_view_id, false, false, ctx);
        for exchange_id in exchange_ids_to_transfer {
            let old_conversation = self
                .conversations_by_id
                .get_mut(&old_conversation_id)
                .ok_or(UpdateHistoryError::ConversationNotFound(
                    old_conversation_id,
                ))?;
            let exchange = old_conversation.remove_exchange(exchange_id)?;

            let new_conversation = self
                .conversations_by_id
                .get_mut(&new_conversation_id)
                .ok_or(UpdateHistoryError::ConversationNotFound(
                    new_conversation_id,
                ))?;
            new_conversation.append_reassigned_exchange(
                response_stream_id,
                exchange,
                terminal_view_id,
                ctx,
            )?;
        }

        // Mark the old conversation as complete since we're starting a new one
        let old_conversation = self
            .conversations_by_id
            .get_mut(&old_conversation_id)
            .ok_or(UpdateHistoryError::ConversationNotFound(
                old_conversation_id,
            ))?;
        old_conversation.mark_completed_after_successful_split(
            response_stream_id,
            terminal_view_id,
            ctx,
        )?;

        self.set_active_conversation_id(new_conversation_id, terminal_view_id, ctx);

        ctx.emit(BlocklistAIHistoryEvent::SplitConversation {
            terminal_view_id,
            old_conversation_id,
            new_conversation_id,
        });

        Ok(new_conversation_id)
    }

    /// Forks an existing conversation by creating a new conversation
    /// and copying the existing conversation's tasks into the new conversation.
    ///
    /// The `prefix` parameter specifies the prefix added to the root task description
    /// (e.g., `FORK_PREFIX` for forks, `PRE_REWIND_PREFIX` for pre-rewind backups).
    pub fn fork_conversation(
        &mut self,
        source_conversation: &AIConversation,
        prefix: &str,
        app: &AppContext,
    ) -> Result<AIConversation, anyhow::Error> {
        let tasks: Vec<warp_multi_agent_api::Task> = source_conversation
            .all_tasks()
            .filter_map(|t| t.source().cloned())
            .collect();

        let updated_tasks_with_new_ids = update_forked_task_properties(tasks, prefix);
        let Some(sqlite_sender) = GlobalResourceHandlesProvider::as_ref(app)
            .get()
            .model_event_sender
            .clone()
        else {
            return Err(anyhow!("No sqlite sender available."));
        };

        // We preserve reverted action IDs. Orphaned IDs (for actions not in fork) are harmless.
        // The reverted states are only copied to the new conversation if the revert happened before the user clicked fork,
        // but regardless of when the revert happened relative to the fork point.
        //
        // Example:
        // 1. Agent edit action
        // 2. Agent edit action
        // 3. User reverts edit from 1
        // 4. **User clicks fork**
        // 5. User reverts edit from 2
        //
        // In this example, the forked conversation will always show edit 1 as reverted and edit 2 as not reverted,
        // regardless of if the fork point is between 2 and 3 or 3 and 4. This is because we preserve all prior reverts,
        // either if they game before or after the fork point. However, once forked, we don't copy later reverts.
        let reverted_action_ids = if source_conversation.reverted_action_ids().is_empty() {
            None
        } else {
            Some(
                source_conversation
                    .reverted_action_ids()
                    .clone()
                    .into_iter()
                    .map_into()
                    .collect(),
            )
        };

        let conversation_data = AgentConversationData {
            server_conversation_token: None,
            conversation_usage_metadata: Some(source_conversation.usage_metadata()),
            reverted_action_ids,
            forked_from_server_conversation_token: source_conversation
                .server_conversation_token()
                .map(|t| t.as_str().to_string()),
            // We reset artifacts on fork
            artifacts_json: None,
            // Forked conversation loses its parentage
            parent_agent_id: None,
            agent_name: None,
            parent_conversation_id: None,
            run_id: None,
            autoexecute_override: Some(source_conversation.autoexecute_override().into()),
            // The event cursor belongs to the source conversation's run; the
            // forked conversation will establish its own cursor.
            last_event_sequence: None,
        };
        let forked_conversation_id = AIConversationId::new();
        if let Err(e) = sqlite_sender.send(ModelEvent::UpdateMultiAgentConversation {
            conversation_id: forked_conversation_id.to_string(),
            updated_tasks: updated_tasks_with_new_ids.clone(),
            conversation_data: conversation_data.clone(),
        }) {
            return Err(anyhow!("Failed to persist forked conversation: {e:?}."));
        }

        // Insert this conversation into the history model memory so we don't need to read from DB to restore this forked conversation
        // (otherwise, we can run into a race condition where the conversation is not found in the DB because we haven't finished writing to the db).
        let forked_conversation = self.insert_forked_conversation_from_tasks(
            forked_conversation_id,
            updated_tasks_with_new_ids.clone(),
            conversation_data.clone(),
        )?;

        Ok(forked_conversation)
    }

    /// Forks an existing conversation at a specific exchange boundary. When `exact_exchange`
    /// is true, the fork includes all messages up to and including the selected exchange.
    /// Otherwise, it extends through the full response (every message after the user's query
    /// until the next root-task user query).
    ///
    /// The `prefix` parameter specifies the prefix added to the root task description
    /// (e.g., `FORK_PREFIX` for forks, `PRE_REWIND_PREFIX` for pre-rewind backups).
    pub fn fork_conversation_at_exchange(
        &mut self,
        source_conversation: &AIConversation,
        from_exchange_id: AIAgentExchangeId,
        fork_from_exact_exchange: bool,
        prefix: &str,
        app: &AppContext,
    ) -> Result<AIConversation, anyhow::Error> {
        let conversation = source_conversation;

        let exchanges_by_task: Vec<(TaskId, Vec<&AIAgentExchange>)> =
            conversation.all_exchanges_by_task();

        let root_task_id = conversation.get_root_task_id().clone();

        let mut message_ids_to_retain_by_task: HashMap<TaskId, HashSet<MessageId>> = HashMap::new();
        let mut found_from_exchange_id = false;
        'outer: for (task_id, task_exchanges) in exchanges_by_task.into_iter() {
            for exchange in task_exchanges {
                // In the non-exact case, we continue past the selected exchange until we reach
                // the next user query (effectively forking from the selected 'response').
                if found_from_exchange_id && task_id == root_task_id && exchange.has_user_query() {
                    break 'outer;
                }

                let message_ids_to_retain = message_ids_to_retain_by_task
                    .entry(task_id.clone())
                    .or_default();
                message_ids_to_retain.extend(exchange.added_message_ids.iter().cloned());
                if exchange.id == from_exchange_id {
                    if fork_from_exact_exchange {
                        break 'outer;
                    }
                    found_from_exchange_id = true;
                }
            }
        }

        if message_ids_to_retain_by_task.is_empty() {
            return Err(anyhow!(
                "No messages found for block in conversation {}.",
                conversation.id()
            ));
        }

        // Build truncated tasks by retaining only messages whose IDs are in
        // `allowed_message_ids`. Tasks whose message list becomes empty and
        // which are non-root tasks are dropped.
        let truncated_tasks: Vec<warp_multi_agent_api::Task> = conversation
            .all_tasks()
            .filter_map(|t| {
                if let Some(message_ids_to_retain) = message_ids_to_retain_by_task.get(t.id()) {
                    let mut truncated_task = t.source().cloned()?;
                    truncated_task
                        .messages
                        .retain(|m| message_ids_to_retain.contains(&MessageId::new(m.id.clone())));
                    (!truncated_task.messages.is_empty()).then_some(truncated_task)
                } else {
                    None
                }
            })
            .collect();

        if truncated_tasks.is_empty() {
            return Err(anyhow!(
                "Truncated tasks for forked conversation at block are empty for conversation {}.",
                conversation.id()
            ));
        }

        let updated_tasks_with_new_ids = update_forked_task_properties(truncated_tasks, prefix);

        let Some(sqlite_sender) = GlobalResourceHandlesProvider::as_ref(app)
            .get()
            .model_event_sender
            .clone()
        else {
            return Err(anyhow!("No sqlite sender available."));
        };

        // We preserve reverted action IDs. Orphaned IDs (for actions not in fork) are harmless.
        // The reverted states are only copied to the new conversation if the revert happened before the user clicked fork,
        // but regardless of when the revert happened relative to the fork point.
        //
        // Example:
        // 1. Agent edit action
        // 2. Agent edit action
        // 3. User reverts edit from 1
        // 4. **User clicks fork**
        // 5. User reverts edit from 2
        //
        // In this example, the forked conversation will always show edit 1 as reverted and edit 2 as not reverted,
        // regardless of if the fork point is between 2 and 3 or 3 and 4. This is because we preserve all prior reverts,
        // either if they game before or after the fork point. However, once forked, we don't copy later reverts.
        let reverted_action_ids = if conversation.reverted_action_ids().is_empty() {
            None
        } else {
            Some(
                conversation
                    .reverted_action_ids()
                    .clone()
                    .into_iter()
                    .map_into()
                    .collect(),
            )
        };

        // Start forked conversations without usage metadata for now; this can
        // be recomputed based on the retained exchanges in a follow-up.
        let conversation_data = AgentConversationData {
            server_conversation_token: None,
            conversation_usage_metadata: None,
            reverted_action_ids,
            forked_from_server_conversation_token: conversation
                .server_conversation_token()
                .map(|t| t.as_str().to_string()),
            // We reset artifacts on fork
            artifacts_json: None,
            // Forked conversation loses its parentage.
            parent_agent_id: None,
            agent_name: None,
            parent_conversation_id: None,
            run_id: None,
            autoexecute_override: Some(conversation.autoexecute_override().into()),
            // The event cursor belongs to the source conversation's run; the
            // forked conversation will establish its own cursor.
            last_event_sequence: None,
        };

        let forked_conversation_id = AIConversationId::new();
        if let Err(e) = sqlite_sender.send(ModelEvent::UpdateMultiAgentConversation {
            conversation_id: forked_conversation_id.to_string(),
            updated_tasks: updated_tasks_with_new_ids.clone(),
            conversation_data: conversation_data.clone(),
        }) {
            return Err(anyhow!(
                "Failed to persist forked conversation at block: {e:?}."
            ));
        }

        let forked_conversation = self.insert_forked_conversation_from_tasks(
            forked_conversation_id,
            updated_tasks_with_new_ids,
            conversation_data,
        )?;

        Ok(forked_conversation)
    }

    pub fn apply_client_actions(
        &mut self,
        response_stream_id: &ResponseStreamId,
        client_actions: Vec<warp_multi_agent_api::ClientAction>,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) -> Result<(), UpdateHistoryError> {
        let mut current_conversation_id = conversation_id;
        for client_action in client_actions {
            match client_action.action {
                Some(Action::StartNewConversation(StartNewConversation {
                    start_from_message_id,
                })) => {
                    let new_conversation_id = self.handle_conversation_split(
                        current_conversation_id,
                        response_stream_id,
                        MessageId::new(start_from_message_id),
                        terminal_view_id,
                        ctx,
                    )?;
                    current_conversation_id = new_conversation_id;
                }
                Some(action) => {
                    let conversation = self
                        .conversations_by_id
                        .get_mut(&current_conversation_id)
                        .ok_or(UpdateHistoryError::ConversationNotFound(
                        current_conversation_id,
                    ))?;
                    conversation.apply_client_action(
                        response_stream_id,
                        terminal_view_id,
                        action,
                        ctx,
                    )?;
                }
                None => {
                    log::warn!("Received empty client action");
                }
            }
        }
        Ok(())
    }

    pub fn update_conversation_cost_and_usage_for_request(
        &mut self,
        conversation_id: AIConversationId,
        request_cost: Option<RequestCost>,
        token_usage: Vec<TokenUsage>,
        usage_metadata: Option<ConversationUsageMetadata>,
        was_user_initiated_request: bool,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            if let Err(e) = conversation.update_cost_and_usage_for_request(
                request_cost,
                token_usage,
                usage_metadata,
                was_user_initiated_request,
            ) {
                log::warn!(
                    "Failed to update request cost for conversation {conversation_id}: {e:#}"
                );
            }
        } else {
            log::warn!(
                "Failed to update request cost because conversation {conversation_id} was not found"
            );
        }
    }

    pub fn mark_response_stream_completed_successfully(
        &mut self,
        stream_id: &ResponseStreamId,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) else {
            return;
        };
        if let Err(e) = conversation.mark_request_completed(stream_id, terminal_view_id, ctx) {
            log::warn!("Failed to mark exchange as completed: {e}");
        }

        // If this conversation doesn't have server metadata yet, and it has a server conversation token,
        // fetch the metadata from the server.
        let should_fetch_metadata = FeatureFlag::CloudConversations.is_enabled()
            && conversation.server_metadata().is_none()
            && conversation.server_conversation_token().is_some();

        if should_fetch_metadata {
            let server_token = conversation
                .server_conversation_token()
                .unwrap()
                .as_str()
                .to_string();

            let server_api = ServerApiProvider::as_ref(ctx).get_ai_client();
            ctx.spawn(
                async move {
                    server_api
                        .list_ai_conversation_metadata(Some(vec![server_token]))
                        .await
                },
                move |model, result, ctx| match result {
                    Ok(mut metadata_list) if !metadata_list.is_empty() => {
                        if let Some(metadata) = metadata_list.pop() {
                            model.set_server_metadata_for_conversation(
                                conversation_id,
                                metadata,
                                ctx,
                            );
                        }
                    }
                    Ok(_) => {
                        log::warn!("No metadata returned for conversation {}", conversation_id);
                    }
                    Err(e) => {
                        log::warn!(
                            "Failed to fetch metadata for conversation {}: {e:#}",
                            conversation_id
                        );
                    }
                },
            );
        }
    }

    pub fn set_exchange_time_to_first_token(
        &mut self,
        conversation_id: AIConversationId,
        exchange_id: AIAgentExchangeId,
        time_to_first_token_ms: i64,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            if let Ok(exchange) = conversation.get_exchange_to_update(exchange_id) {
                exchange.time_to_first_token_ms = Some(time_to_first_token_ms);
            }
        }
    }

    pub fn mark_response_stream_cancelled(
        &mut self,
        stream_id: &ResponseStreamId,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        reason: CancellationReason,
        ctx: &mut ModelContext<Self>,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            if reason.is_reverted() {
                if let Err(e) =
                    conversation.mark_request_cancelled_due_to_revert(terminal_view_id, ctx)
                {
                    log::warn!("Failed to mark exchange as cancelled: {e}");
                }
            } else if let Err(e) =
                conversation.mark_request_cancelled(stream_id, terminal_view_id, reason, ctx)
            {
                log::warn!("Failed to mark exchange as cancelled: {e}");
            }
        }
        AIDocumentModel::handle(ctx).update(ctx, |model, ctx| {
            model.clear_streaming_documents_for_conversation(&conversation_id, ctx);
        });
    }

    pub fn mark_response_stream_completed_with_error(
        &mut self,
        error: RenderableAIError,
        stream_id: &ResponseStreamId,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            if let Err(e) = conversation.mark_request_completed_with_error(
                stream_id,
                error.clone(),
                terminal_view_id,
                ctx,
            ) {
                log::warn!("Failed to mark exchange as completed with error: {e}");
            }
        }
    }

    /// Handle clearing the blocklist for the terminal view.
    /// The terminal view will also cancel the active stream on processing the event emitted here.
    pub(crate) fn clear_conversations_in_terminal_view(
        &mut self,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        // Cancel the active stream when we clear conversations in this terminal view.
        let active_conversation_id = self
            .active_conversation_for_terminal_view
            .remove(&terminal_view_id);
        if let Some(cleared_conversation_ids) = self
            .live_conversation_ids_for_terminal_view
            .remove(&terminal_view_id)
        {
            self.cleared_conversation_ids_for_terminal_view
                .entry(terminal_view_id)
                .and_modify(|existing| existing.extend(cleared_conversation_ids.clone()))
                .or_insert(cleared_conversation_ids);
        }
        let cleared_conversation_ids = self
            .live_conversation_ids_for_terminal_view
            .remove(&terminal_view_id);
        if let Some(cleared_conversation_ids) = cleared_conversation_ids {
            self.cleared_conversation_ids_for_terminal_view
                .entry(terminal_view_id)
                .and_modify(|existing| existing.extend(cleared_conversation_ids.clone()))
                .or_insert(cleared_conversation_ids);
        }
        ctx.emit(
            BlocklistAIHistoryEvent::ClearedConversationsInTerminalView {
                terminal_view_id,
                active_conversation_id,
            },
        );
    }

    /// Handle removing a conversation from the history model, blocklist and in-memory.
    pub fn remove_conversation(
        &mut self,
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        self.remove_conversation_from_memory(conversation_id, Some(terminal_view_id), ctx);
    }

    /// Permanently delete a conversation.
    pub fn delete_conversation(
        &mut self,
        conversation_id: AIConversationId,
        terminal_view_id: Option<EntityId>,
        ctx: &mut ModelContext<Self>,
    ) {
        let conversation_title = self
            .conversations_by_id
            .get(&conversation_id)
            .and_then(|c| c.title().map(|t| t.to_string()));

        self.remove_conversation_from_memory(conversation_id, terminal_view_id, ctx);

        // Delete persisted conversation from sqlite.
        let model_event_sender = GlobalResourceHandlesProvider::as_ref(ctx)
            .get()
            .model_event_sender
            .clone();
        let conversation_id_string = conversation_id.to_string();
        ctx.spawn(
            async move {
                if let Some(sender) = model_event_sender {
                    if let Err(e) = sender.send(ModelEvent::DeleteAIConversation {
                        conversation_id: conversation_id_string.clone(),
                    }) {
                        log::error!("Error sending DeleteAIConversation event: {e:?}");
                    }
                    if let Err(e) = sender.send(ModelEvent::DeleteMultiAgentConversations {
                        conversation_ids: vec![conversation_id_string],
                    }) {
                        log::error!("Error sending DeleteMultiAgentConversations event: {e:?}");
                    }
                }
            },
            |_, _, _| {},
        );

        // Only emit the event if we have a terminal_view_id, since the event is
        // filtered by terminal_view_id in handlers.
        if let Some(terminal_view_id) = terminal_view_id {
            ctx.emit(BlocklistAIHistoryEvent::DeletedConversation {
                terminal_view_id,
                conversation_id,
                conversation_title,
            });
        }
    }

    /// Remove a conversation from all in-memory storage.
    fn remove_conversation_from_memory(
        &mut self,
        conversation_id: AIConversationId,
        terminal_view_id: Option<EntityId>,
        ctx: &mut ModelContext<Self>,
    ) {
        // Clean up reverse indices before removing the conversation. Guard
        // token-index removals with an equality check: the live conversation's
        // token and the metadata's token can diverge after a rebind, and we
        // must not clobber an entry already owned by another conversation.
        if let Some(conversation) = self.conversations_by_id.get(&conversation_id) {
            if let Some(key) = agent_id_key(conversation) {
                self.agent_id_to_conversation_id.remove(&key);
            }
            if let Some(token) = conversation.server_conversation_token() {
                if self.server_token_to_conversation_id.get(token) == Some(&conversation_id) {
                    self.server_token_to_conversation_id.remove(token);
                }
            }
        }
        // Also clean up the token index entry that might have been installed
        // via the metadata path (no live conversation present).
        if let Some(metadata) = self.all_conversations_metadata.get(&conversation_id) {
            if let Some(token) = &metadata.server_conversation_token {
                if self.server_token_to_conversation_id.get(token) == Some(&conversation_id) {
                    self.server_token_to_conversation_id.remove(token);
                }
            }
        }

        self.all_conversations_metadata.remove(&conversation_id);
        self.conversations_by_id.remove(&conversation_id);

        if let Some(terminal_view_id) = terminal_view_id {
            if self
                .active_conversation_for_terminal_view
                .get(&terminal_view_id)
                .is_some_and(|id| *id == conversation_id)
            {
                self.active_conversation_for_terminal_view
                    .remove(&terminal_view_id);
            }
            if let Some(vec) = self
                .live_conversation_ids_for_terminal_view
                .get_mut(&terminal_view_id)
            {
                vec.retain(|&id| id != conversation_id);
            }
            if let Some(vec) = self
                .cleared_conversation_ids_for_terminal_view
                .get_mut(&terminal_view_id)
            {
                vec.retain(|&id| id != conversation_id);
            }
            ctx.emit(BlocklistAIHistoryEvent::RemoveConversation {
                terminal_view_id,
                conversation_id,
            });
        }
    }

    /// Returns true if the conversation is live in any terminal view.
    pub fn is_conversation_live(&self, conversation_id: AIConversationId) -> bool {
        self.live_conversation_ids_for_terminal_view
            .values()
            .any(|conversation_ids| conversation_ids.contains(&conversation_id))
    }

    pub fn mark_terminal_view_as_ambient_agent_session_view(&mut self, terminal_view_id: EntityId) {
        self.ambient_agent_terminal_view_ids
            .insert(terminal_view_id);
    }

    pub fn mark_terminal_view_as_conversation_transcript_viewer(
        &mut self,
        terminal_view_id: EntityId,
    ) {
        self.conversation_transcript_viewer_terminal_view_ids
            .insert(terminal_view_id);
    }

    pub fn is_terminal_view_conversation_transcript_viewer(
        &self,
        terminal_view_id: EntityId,
    ) -> bool {
        self.conversation_transcript_viewer_terminal_view_ids
            .contains(&terminal_view_id)
    }

    /// Returns [`AIQueryHistory`]s from all sources: live conversations, cleared conversations,
    /// and persisted queries from conversations not loaded in memory.
    ///
    /// When `terminal_view_id` is provided, queries from that terminal view are categorized as
    /// `CurrentSession` and all others as `DifferentSession`. When `None`, all queries are
    /// categorized as `DifferentSession`.
    ///
    /// Ambient agent sessions are always excluded.
    pub(crate) fn all_ai_queries(
        &self,
        terminal_view_id: Option<EntityId>,
    ) -> impl Iterator<Item = AIQueryHistory> + '_ {
        // Collect all conversation IDs that are already in memory (live or cleared)
        // and build query vectors in the same loops
        let mut loaded_conversation_ids: HashSet<AIConversationId> = HashSet::new();

        let mut live_queries_vec = Vec::new();
        for (tv_id, conversation_ids) in self.live_conversation_ids_for_terminal_view.iter() {
            loaded_conversation_ids.extend(conversation_ids);

            // Skip shared ambient agent sessions
            if self.ambient_agent_terminal_view_ids.contains(tv_id) {
                continue;
            }

            let history_order = if terminal_view_id.is_some_and(|id| id == *tv_id) {
                HistoryOrder::CurrentSession
            } else {
                HistoryOrder::DifferentSession
            };

            for conversation_id in conversation_ids {
                if let Some(conversation) = self.conversations_by_id.get(conversation_id) {
                    for exchange in conversation.root_task_exchanges() {
                        if let Some(query) = ai_exchange_to_query_history(exchange, history_order) {
                            live_queries_vec.push(query);
                        }
                    }
                }
            }
        }

        let mut cleared_queries_vec = Vec::new();
        for (tv_id, conversation_ids) in self.cleared_conversation_ids_for_terminal_view.iter() {
            loaded_conversation_ids.extend(conversation_ids);

            let history_order = if terminal_view_id.is_some_and(|id| id == *tv_id) {
                HistoryOrder::CurrentSession
            } else {
                HistoryOrder::DifferentSession
            };

            for conversation_id in conversation_ids {
                if let Some(conversation) = self.conversations_by_id.get(conversation_id) {
                    for exchange in conversation.root_task_exchanges() {
                        if let Some(query) = ai_exchange_to_query_history(exchange, history_order) {
                            cleared_queries_vec.push(query);
                        }
                    }
                }
            }
        }

        // Add persisted queries from conversations not loaded in memory
        let persisted_queries_vec: Vec<_> = self
            .persisted_queries
            .iter()
            .filter(|persisted| !loaded_conversation_ids.contains(&persisted.conversation_id))
            .filter_map(|persisted| {
                persisted_ai_input_to_query_history(persisted, HistoryOrder::DifferentSession)
            })
            .collect();

        persisted_queries_vec
            .into_iter()
            .chain(cleared_queries_vec)
            .chain(live_queries_vec)
    }

    /// Returns `Some` with the [`AIConversationId`] of the active conversation inside the
    /// [`crate::terminal::TerminalView`] with the given [`EntityId`] if there is one. Returns
    /// `None` otherwise.
    /// The active conversation is the one we're currently or have most recently streamed outputs for.
    /// If you want to check what conversation the next query will follow up in / what is selected in the input selector,
    /// use `context_model.selected_conversation_id` instead.
    pub(crate) fn active_conversation_id(
        &self,
        terminal_view_id: EntityId,
    ) -> Option<AIConversationId> {
        let active_conversation_id = self
            .active_conversation_for_terminal_view
            .get(&terminal_view_id)
            .copied()?;

        let conversation_ids_for_terminal_view = self
            .live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)?;

        if !conversation_ids_for_terminal_view.contains(&active_conversation_id) {
            log::warn!(
                "The active conversation ID {active_conversation_id:?} was not found in the list of conversation IDs for terminal view {terminal_view_id:?}. Conversation IDs: {conversation_ids_for_terminal_view:?}"
            );
            return None;
        }

        Some(active_conversation_id)
    }

    /// Returns `Some` with the [`AIConversationId`] of the last conversation created for a given
    /// [`crate::terminal::TerminalView`] with the given [`EntityId`] if there is one. Returns
    /// `None` otherwise.
    #[cfg_attr(target_family = "wasm", allow(unused))]
    pub(crate) fn last_conversation_id(
        &self,
        terminal_view_id: EntityId,
    ) -> Option<AIConversationId> {
        self.live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)?
            .last()
            .copied()
    }

    /// Set the hidden status of the exchange with the given ID.
    pub fn set_exchange_hidden_status(
        &mut self,
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        exchange_id: AIAgentExchangeId,
        is_hidden: bool,
        ctx: &mut ModelContext<Self>,
    ) {
        let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) else {
            return;
        };
        conversation.set_is_exchange_hidden(exchange_id, is_hidden, terminal_view_id, ctx);
    }

    pub fn set_viewing_shared_session_for_conversation(
        &mut self,
        conversation_id: AIConversationId,
        is_viewing_shared_session: bool,
    ) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            conversation.set_is_viewing_shared_session(is_viewing_shared_session);
        }
    }

    pub fn set_has_code_review_opened_to_true(&mut self, conversation_id: AIConversationId) {
        if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
            conversation.mark_code_review_as_opened();
        }
    }

    pub fn toggle_autoexecute_override(
        &mut self,
        conversation_id: &AIConversationId,
        terminal_view_id: EntityId,
        ctx: &mut ModelContext<Self>,
    ) {
        let Some(conversation) = self.conversations_by_id.get_mut(conversation_id) else {
            return;
        };

        conversation.toggle_autoexecute_override();
        conversation.write_updated_conversation_state(ctx);
        ctx.emit(BlocklistAIHistoryEvent::UpdatedAutoexecuteOverride { terminal_view_id });
    }

    /// Truncates a conversation from the given exchange ID, removing all exchanges
    /// from that exchange onwards (inclusive). This is a lossy operation.
    ///
    /// Returns the set of exchange IDs that were removed.
    pub fn truncate_conversation_from_exchange(
        &mut self,
        conversation_id: AIConversationId,
        from_exchange_id: AIAgentExchangeId,
        ctx: &mut ModelContext<Self>,
    ) -> Result<HashSet<AIAgentExchangeId>, UpdateHistoryError> {
        let conversation = self
            .conversations_by_id
            .get_mut(&conversation_id)
            .ok_or(UpdateHistoryError::ConversationNotFound(conversation_id))?;

        let removed_exchange_ids = conversation.truncate_from_exchange(from_exchange_id, ctx)?;

        Ok(removed_exchange_ids)
    }

    /// Returns the latest exchange across all conversations in the terminal view.
    /// This is useful for determining if a specific exchange is the most recent one.
    /// Excludes passive code generation exchanges from consideration.
    pub fn latest_exchange_across_all_conversations(
        &self,
        terminal_view_id: EntityId,
    ) -> Option<&AIAgentExchange> {
        self.all_live_root_task_exchanges_for_terminal_view(terminal_view_id)
            .filter(|exchange| !exchange.has_passive_request())
            .max_by_key(|exchange| exchange.start_time)
    }

    /// Returns the conversation ID that contains the given exchange ID, if any.
    /// Searches through all conversations for a given terminal view.
    pub fn conversation_id_for_exchange(
        &self,
        exchange_id: AIAgentExchangeId,
        terminal_view_id: EntityId,
    ) -> Option<AIConversationId> {
        self.live_conversation_ids_for_terminal_view
            .get(&terminal_view_id)?
            .iter()
            .find(|conversation_id| {
                self.conversations_by_id
                    .get(conversation_id)
                    .is_some_and(|conversation| {
                        conversation.exchange_with_id(exchange_id).is_some()
                    })
            })
            .copied()
    }

    /// Returns local conversation metadata,
    /// (excluding conversations from ambient agent runs).
    pub fn get_local_conversations_metadata(
        &self,
    ) -> impl Iterator<Item = &AIConversationMetadata> {
        self.all_conversations_metadata
            .values()
            .filter(|m| !m.is_ambient_agent_conversation())
    }

    /// Returns conversation metadata for a specific conversation ID.
    pub fn get_conversation_metadata(
        &self,
        conversation_id: &AIConversationId,
    ) -> Option<&AIConversationMetadata> {
        self.all_conversations_metadata.get(conversation_id)
    }

    /// Returns whether a conversation can be shared.
    ///
    /// A conversation can be shared if we have server metadata available
    /// (either from a loaded conversation or from conversation metadata).
    pub fn can_conversation_be_shared(&self, conversation_id: &AIConversationId) -> bool {
        self.get_server_conversation_metadata(conversation_id)
            .is_some()
    }

    /// Returns the server conversation metadata, used by the sharing dialog.
    ///
    /// This checks:
    /// 1. If the conversation is loaded in memory, returns from its server metadata
    /// 2. Otherwise, falls back to data stored in conversation metadata
    pub fn get_server_conversation_metadata(
        &self,
        conversation_id: &AIConversationId,
    ) -> Option<&ServerAIConversationMetadata> {
        // Check if conversation exists in memory and has server metadata
        if let Some(conversation) = self.conversation(conversation_id) {
            if let Some(m) = conversation.server_metadata() {
                return Some(m);
            }
        }

        // Fall back to conversation metadata
        if let Some(metadata) = self.get_conversation_metadata(conversation_id) {
            return metadata.server_conversation_metadata.as_ref();
        }

        None
    }

    /// Finds an AIConversationId by its server conversation token.
    ///
    /// O(1) lookup via `server_token_to_conversation_id`, which is maintained
    /// alongside every mutation of `conversations_by_id` /
    /// `all_conversations_metadata` that involves a token. Used to look up
    /// conversations for ambient agent tasks, which store the server token
    /// but not the AIConversationId.
    pub fn find_conversation_id_by_server_token(
        &self,
        server_token: &ServerConversationToken,
    ) -> Option<AIConversationId> {
        if let Some(id) = self.server_token_to_conversation_id.get(server_token) {
            return Some(*id);
        }

        // A token miss is the expected outcome whenever a task references a
        // conversation this client hasn't loaded (shared-session tasks from
        // other users, pre-sync state).
        log::debug!(
            "No conversation found for server token: {}",
            server_token.as_str()
        );
        None
    }

    /// Mark conversations as historical
    /// Historical conversations consist of non-live conversations that were read from the disk or server on startup,
    /// and conversations (recorded here) that were live this session but have now been cleared.
    pub fn mark_conversations_historical_for_terminal_view(&mut self, terminal_view_id: EntityId) {
        if self.is_terminal_view_conversation_transcript_viewer(terminal_view_id) {
            // We don't mark conversation transcript viewer conversations as historical,
            // as they are stored separately and should not be persisted/displayed as regular user conversations.
            return;
        }

        // There's a slight concern here that the conversations we're preserving might not have persisted succesfully
        // because of some unexpected error. Attempting to then restore these conversations would lead to unexpected behavior.
        // In the future it might be worthwhile to check that these conversations exist in the database before marking them as historical,
        // but for now this is an edge case that we don't need to worry about too much.
        let conversations_to_mark_historical: Vec<AIConversationMetadata> = self
            .all_live_conversations_for_terminal_view(terminal_view_id)
            .filter_map(|conversation| {
                let conversation_id = conversation.id();
                if !self.conversations_by_id.contains_key(&conversation_id)
                    || conversation.should_exclude_from_navigation()
                    || !blocklist_filter::conversation_would_render_in_blocklist(conversation)
                {
                    return None;
                }

                Some(conversation.into())
            })
            .collect();

        for metadata in conversations_to_mark_historical {
            if let Some(token) = &metadata.server_conversation_token {
                self.server_token_to_conversation_id
                    .insert(token.clone(), metadata.id);
            }
            self.all_conversations_metadata
                .insert(metadata.id, metadata);
        }
    }

    /// Inserts a conversation into memory by reconstructing exchanges from tasks.
    /// We use this when forking a conversation to ensure that the forked conversation
    /// is immediately available in memory before we try to restore it in a new tab.
    pub fn insert_forked_conversation_from_tasks(
        &mut self,
        conversation_id: AIConversationId,
        tasks: Vec<warp_multi_agent_api::Task>,
        conversation_data: AgentConversationData,
    ) -> anyhow::Result<AIConversation> {
        let mut conversation =
            AIConversation::new_restored(conversation_id, tasks, Some(conversation_data))?;

        // Assign fresh exchange IDs so persisted blocks do not collide.
        conversation.reassign_exchange_ids();

        if let Some(token) = conversation.server_conversation_token() {
            self.server_token_to_conversation_id
                .insert(token.clone(), conversation_id);
        }

        self.conversations_by_id
            .insert(conversation_id, conversation.clone());

        // This is harmless if we're opening the conversation immediately, but ensures it's in the conversation list right away if we fork in the background.
        let metadata = AIConversationMetadata::from(&conversation);
        if let Some(token) = &metadata.server_conversation_token {
            self.server_token_to_conversation_id
                .insert(token.clone(), conversation_id);
        }
        self.all_conversations_metadata
            .insert(conversation_id, metadata);

        Ok(conversation)
    }

    /// Clears all stored conversation-related data in memory.
    /// This is used when logging out to ensure no AI history persists across users.
    pub(crate) fn reset(&mut self) {
        self.live_conversation_ids_for_terminal_view.clear();
        self.cleared_conversation_ids_for_terminal_view.clear();
        self.conversations_by_id.clear();
        self.active_conversation_for_terminal_view.clear();
        self.ambient_agent_terminal_view_ids.clear();
        self.conversation_transcript_viewer_terminal_view_ids
            .clear();
        self.persisted_queries.clear();
        self.all_conversations_metadata.clear();
        self.agent_id_to_conversation_id.clear();
        self.server_token_to_conversation_id.clear();
    }
}

/// Returns the key to use in `agent_id_to_conversation_id` for the given
/// conversation.
fn agent_id_key(conversation: &AIConversation) -> Option<String> {
    conversation.orchestration_agent_id()
}

#[derive(Clone, Debug)]
pub enum BlocklistAIHistoryEvent {
    /// A new conversation was started.
    StartedNewConversation {
        new_conversation_id: AIConversationId,
        terminal_view_id: EntityId,
    },

    CreatedSubtask {
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        task_id: TaskId,
    },

    /// Emitted when the optimistically created task is "upgraded" to a server-backed task upon
    /// receiving a CreateTask client action.
    UpgradedTask {
        optimistic_id: TaskId,
        server_id: TaskId,
        terminal_view_id: EntityId,
    },

    AppendedExchange {
        exchange_id: AIAgentExchangeId,
        task_id: TaskId,
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        is_hidden: bool,

        // Populated if this exchange is appended as a result of an in-flight API request.
        response_stream_id: Option<ResponseStreamId>,
    },

    ReassignedExchange {
        exchange_id: AIAgentExchangeId,
        terminal_view_id: EntityId,
        new_task_id: TaskId,
        new_conversation_id: AIConversationId,
    },

    /// Includes the terminal view's [`EntityId`] so we can disambiguate the source of the event
    /// because this [`BlocklistAIHistoryModel`] is global.
    #[cfg_attr(not(feature = "local_fs"), allow(dead_code))]
    UpdatedStreamingExchange {
        exchange_id: AIAgentExchangeId,
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        is_hidden: bool,
    },

    UpdatedConversationStatus {
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
        is_restored: bool,
    },

    /// The active conversation was set to another conversation in the history.
    SetActiveConversation {
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
    },

    /// `conversation_id` is no longer marked as active for the given terminal view.
    ClearedActiveConversation {
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
    },

    ClearedConversationsInTerminalView {
        terminal_view_id: EntityId,
        active_conversation_id: Option<AIConversationId>,
    },

    UpdatedTodoList {
        terminal_view_id: EntityId,
    },

    UpdatedAutoexecuteOverride {
        terminal_view_id: EntityId,
    },

    /// Emitted when a conversation is split into two (on suggest starting new conversation)
    SplitConversation {
        terminal_view_id: EntityId,
        old_conversation_id: AIConversationId,
        new_conversation_id: AIConversationId,
    },

    /// This is emitted when an ephemeral/abandoned conversation is cleaned up
    /// (e.g. empty conversations the user never used, rejected passive code suggestions).
    RemoveConversation {
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
    },

    /// This is emitted when a user explicitly deletes an existing conversation.
    DeletedConversation {
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        conversation_title: Option<String>,
    },

    /// Emitted when conversations are restored in a terminal view.
    RestoredConversations {
        terminal_view_id: EntityId,
        conversation_ids: Vec<AIConversationId>,
    },

    /// Emitted when conversation metadata is updated.
    /// `terminal_view_id` is None when updating historical-only conversations.
    UpdatedConversationMetadata {
        terminal_view_id: Option<EntityId>,
        conversation_id: AIConversationId,
    },

    /// Emitted when conversation artifacts are updated (plans, PRs, etc.)
    UpdatedConversationArtifacts {
        terminal_view_id: EntityId,
        conversation_id: AIConversationId,
        artifact: Artifact,
    },

    /// Emitted when a conversation first receives its server-assigned conversation token
    /// (during StreamInit). Used by the StartAgentExecutor to resolve pending StartAgent
    /// actions for child agent conversations.
    ConversationServerTokenAssigned {
        conversation_id: AIConversationId,
        terminal_view_id: EntityId,
    },
}

impl BlocklistAIHistoryEvent {
    /// Returns the terminal view ID associated with this event, if any.
    /// Returns `None` for events that apply globally (e.g., historical conversation metadata updates).
    pub fn terminal_view_id(&self) -> Option<EntityId> {
        match self {
            BlocklistAIHistoryEvent::StartedNewConversation {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::AppendedExchange {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::UpdatedStreamingExchange {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::UpdatedConversationStatus {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::SetActiveConversation {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::ClearedActiveConversation {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::ClearedConversationsInTerminalView {
                terminal_view_id,
                ..
            }
            | BlocklistAIHistoryEvent::ReassignedExchange {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::UpdatedTodoList {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::UpdatedAutoexecuteOverride {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::SplitConversation {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::RemoveConversation {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::DeletedConversation {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::CreatedSubtask {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::RestoredConversations {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::UpgradedTask {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
                terminal_view_id, ..
            }
            | BlocklistAIHistoryEvent::ConversationServerTokenAssigned {
                terminal_view_id, ..
            } => Some(*terminal_view_id),
            // UpdatedConversationMetadata can have None when updating historical-only conversations
            BlocklistAIHistoryEvent::UpdatedConversationMetadata {
                terminal_view_id, ..
            } => *terminal_view_id,
        }
    }
}

impl Entity for BlocklistAIHistoryModel {
    type Event = BlocklistAIHistoryEvent;
}

impl SingletonEntity for BlocklistAIHistoryModel {}

/// Helper struct for showing AI history to the user. Guarantees that there is a user query and
/// contains less data than [`AIAgentExchange`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AIQueryHistory {
    /// The input originating from the user.
    pub query_text: String,

    /// The time the input was sent.
    pub start_time: DateTime<Local>,

    /// The status of the output streaming from the AI API.
    pub output_status: AIQueryHistoryOutputStatus,

    /// The working directory when the AI query was submitted.
    pub working_directory: Option<String>,

    /// The ordering category for this query in history.
    pub history_order: HistoryOrder,
}

impl AIQueryHistory {
    /// Creates a new [`AIQueryHistory`] for testing.
    #[cfg(test)]
    pub(crate) fn new_for_test(
        query_text: &str,
        start_time: DateTime<Local>,
        history_order: HistoryOrder,
    ) -> Self {
        Self {
            query_text: query_text.to_owned(),
            start_time,
            output_status: AIQueryHistoryOutputStatus::Pending,
            working_directory: None,
            history_order,
        }
    }
}

fn ai_exchange_to_query_history(
    value: &AIAgentExchange,
    history_order: HistoryOrder,
) -> Option<AIQueryHistory> {
    let query = value.input.iter().find_map(AIAgentInput::user_query)?;

    Some(AIQueryHistory {
        query_text: query,
        start_time: value.start_time,
        output_status: AIQueryHistoryOutputStatus::from(&value.output_status),
        working_directory: value.working_directory.clone(),
        history_order,
    })
}

fn persisted_ai_input_to_query_history(
    value: &PersistedAIInput,
    history_order: HistoryOrder,
) -> Option<AIQueryHistory> {
    // Extract the query text from the first Query input
    let query_text = value
        .inputs
        .iter()
        .map(|input| match input {
            PersistedAIInputType::Query { text, .. } => Some(text.clone()),
        })
        .next()
        .flatten()?;

    Some(AIQueryHistory {
        query_text,
        start_time: value.start_ts,
        output_status: value.output_status.clone(),
        working_directory: value.working_directory.clone(),
        history_order,
    })
}

/// Status of output streaming from the AI API.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum AIQueryHistoryOutputStatus {
    /// We are waiting to or are currently streaming output.
    Pending,
    /// The user manually cancelled output streaming.
    Cancelled,
    /// Output streaming failed.
    Failed,
    /// Output streaming completed successfully.
    Completed,
}

impl AIQueryHistoryOutputStatus {
    /// Returns a string representation of the output status.
    pub(crate) fn display_text(&self) -> &'static str {
        match self {
            AIQueryHistoryOutputStatus::Completed => "Completed successfully",
            AIQueryHistoryOutputStatus::Pending => "Pending",
            AIQueryHistoryOutputStatus::Cancelled => "Cancelled by user",
            AIQueryHistoryOutputStatus::Failed => "Failed",
        }
    }

    pub(crate) fn icon(&self) -> Icon {
        match self {
            AIQueryHistoryOutputStatus::Completed => Icon::Check,
            AIQueryHistoryOutputStatus::Pending => Icon::Loading,
            AIQueryHistoryOutputStatus::Cancelled => Icon::SlashCircle,
            AIQueryHistoryOutputStatus::Failed => Icon::AlertTriangle,
        }
    }
}

impl From<&AIAgentOutputStatus> for AIQueryHistoryOutputStatus {
    fn from(status: &AIAgentOutputStatus) -> Self {
        match status {
            AIAgentOutputStatus::Streaming { .. } => Self::Pending,
            AIAgentOutputStatus::Finished {
                finished_output, ..
            } => match finished_output {
                FinishedAIAgentOutput::Cancelled { .. } => Self::Cancelled,
                FinishedAIAgentOutput::Error { .. } => Self::Failed,
                FinishedAIAgentOutput::Success { .. } => Self::Completed,
            },
        }
    }
}

/// Updates the given tasks, which are presumed to be clones of tasks from a source conversation to be
/// used to back a fork or copy of the source conversation.
///
/// Reassigns new task IDs to each forked task to ensure task IDs remain globally unique and updates
/// description of the root task, prepending it with the given prefix.
fn update_forked_task_properties(
    tasks: Vec<warp_multi_agent_api::Task>,
    prefix: &str,
) -> Vec<warp_multi_agent_api::Task> {
    let mut old_to_new_task_ids = HashMap::new();
    fn get_new_task_id(new_ids: &mut HashMap<String, String>, old_task_id: &str) -> String {
        new_ids
            .entry(old_task_id.to_owned())
            .or_insert_with(|| Uuid::new_v4().to_string())
            .clone()
    }

    tasks
        .into_iter()
        .map(|mut t| {
            let new_id = get_new_task_id(&mut old_to_new_task_ids, &t.id);
            // Update task id to avoid duplicate tasks across conversations and ensure
            // all messages reference the new task id.
            t.id = new_id.clone();
            for message in &mut t.messages {
                message.task_id = new_id.clone();
                if let Some(subagent) = message.tool_call_mut().and_then(|tc| tc.subagent_mut()) {
                    subagent.task_id =
                        get_new_task_id(&mut old_to_new_task_ids, &subagent.task_id).clone();
                }
            }
            if let Some(deps) = t
                .dependencies
                .as_mut()
                .filter(|deps| !deps.parent_task_id.is_empty())
            {
                deps.parent_task_id =
                    get_new_task_id(&mut old_to_new_task_ids, &deps.parent_task_id).clone();
            } else {
                t.description = format!("{}{}", prefix, t.description);
            }
            t
        })
        .collect()
}

/// The default prefix used when forking a conversation.
pub const FORK_PREFIX: &str = "(Fork) ";

/// The prefix used when saving a conversation before a rewind operation.
pub const PRE_REWIND_PREFIX: &str = "(Pre-Rewind) ";

#[cfg(test)]
#[path = "history_model_test.rs"]
mod tests;