Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

app/src/ai/agent/api.rs

@@ -0,0 +1,314 @@
+pub(crate) mod convert_conversation;
+mod convert_from;
+mod convert_to;
+mod r#impl;
+
+pub use ai::agent::convert::ConvertToAPITypeError;
+use ai::api_keys::ApiKeyManager;
+pub use convert_from::{
+    user_inputs_from_messages, ConversionParams, ConvertAPIMessageToClientOutputMessage,
+    MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError,
+};
+
+pub use r#impl::generate_multi_agent_output;
+
+use futures_lite::Stream;
+use serde::Serialize;
+use std::path::Path;
+use std::pin::Pin;
+use std::sync::Arc;
+use warp_core::channel::ChannelState;
+use warp_core::execution_mode::AppExecutionMode;
+use warp_core::features::FeatureFlag;
+
+use crate::ai::agent::conversation::AIConversationId;
+use crate::ai::ambient_agents::AmbientAgentTaskId;
+use crate::{
+    ai::{blocklist::SessionContext, llms::LLMId},
+    server::server_api::AIApiError,
+};
+
+use super::{AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions};
+use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput};
+use crate::ai::mcp::templatable_manager::TemplatableMCPServerInfo;
+use crate::ai::mcp::TemplatableMCPServerManager;
+use crate::settings::AISettings;
+use crate::terminal::safe_mode_settings::get_secret_obfuscation_mode;
+use crate::workspaces::user_workspaces::UserWorkspaces;
+use warp_core::user_preferences::GetUserPreferences;
+use warpui::{AppContext, EntityId, SingletonEntity as _};
+
+/// Unique, server-generated conversation-scoped token to be roundtripped to the API when sending
+/// requests that follow-up within a given conversation.
+#[derive(Serialize, Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ServerConversationToken(String);
+
+impl ServerConversationToken {
+    pub fn new(id: String) -> Self {
+        Self(id)
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+
+    pub fn debug_link(&self) -> String {
+        format!(
+            "{}/debug/maa/{}",
+            ChannelState::server_root_url(),
+            self.as_str()
+        )
+    }
+
+    pub fn conversation_link(&self) -> String {
+        format!(
+            "{}/conversation/{}",
+            ChannelState::server_root_url(),
+            self.as_str()
+        )
+    }
+}
+
+impl From<ServerConversationToken> for String {
+    fn from(value: ServerConversationToken) -> Self {
+        value.0
+    }
+}
+
+// Conversions between AI ServerConversationToken and protocol ServerConversationToken
+impl From<session_sharing_protocol::common::ServerConversationToken> for ServerConversationToken {
+    fn from(token: session_sharing_protocol::common::ServerConversationToken) -> Self {
+        Self(token.to_string())
+    }
+}
+
+impl TryFrom<ServerConversationToken>
+    for session_sharing_protocol::common::ServerConversationToken
+{
+    type Error = uuid::Error;
+
+    fn try_from(token: ServerConversationToken) -> Result<Self, Self::Error> {
+        token.as_str().parse()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RequestParams {
+    pub input: Vec<AIAgentInput>,
+    pub conversation_token: Option<ServerConversationToken>,
+    pub forked_from_conversation_token: Option<ServerConversationToken>,
+    pub ambient_agent_task_id: Option<AmbientAgentTaskId>,
+    pub tasks: Vec<warp_multi_agent_api::Task>,
+    pub existing_suggestions: Option<Suggestions>,
+    pub metadata: Option<RequestMetadata>,
+    pub session_context: SessionContext,
+    pub model: LLMId,
+    #[allow(unused)]
+    pub coding_model: LLMId,
+    pub cli_agent_model: LLMId,
+    pub computer_use_model: LLMId,
+    pub is_memory_enabled: bool,
+    pub warp_drive_context_enabled: bool,
+    pub mcp_context: Option<MCPContext>,
+    pub planning_enabled: bool,
+    should_redact_secrets: bool,
+
+    /// User-provided API keys for AI providers (BYO API Key).
+    pub api_keys: Option<warp_multi_agent_api::request::settings::ApiKeys>,
+    pub allow_use_of_warp_credits_with_byok: bool,
+    pub autonomy_level: warp_multi_agent_api::AutonomyLevel,
+    pub isolation_level: warp_multi_agent_api::IsolationLevel,
+    pub web_search_enabled: bool,
+    pub computer_use_enabled: bool,
+    pub ask_user_question_enabled: bool,
+    pub research_agent_enabled: bool,
+    pub orchestration_enabled: bool,
+    pub supported_tools_override: Option<Vec<warp_multi_agent_api::ToolType>>,
+    /// The conversation ID of the parent agent that spawned this child agent, if any.
+    pub parent_agent_id: Option<String>,
+    /// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator.
+    pub agent_name: Option<String>,
+}
+
+pub type Event = Result<warp_multi_agent_api::ResponseEvent, Arc<AIApiError>>;
+
+#[cfg(not(target_family = "wasm"))]
+pub type ResponseStream = Pin<Box<dyn Stream<Item = Event> + Send + 'static>>;
+
+// The WASM version of this type has no bound on `Send`, which is an unnecessary bound when
+// targeting wasm because the browser is single-threaded (and we don't leverage WebWorkers for async
+// execution in WoW).
+#[cfg(target_family = "wasm")]
+pub type ResponseStream = Pin<Box<dyn Stream<Item = Event>>>;
+
+#[derive(Debug, Clone)]
+pub struct ConversationData {
+    pub id: AIConversationId,
+    pub tasks: Vec<warp_multi_agent_api::Task>,
+    pub server_conversation_token: Option<ServerConversationToken>,
+    pub forked_from_conversation_token: Option<ServerConversationToken>,
+    pub ambient_agent_task_id: Option<AmbientAgentTaskId>,
+    pub existing_suggestions: Option<Suggestions>,
+}
+
+impl RequestParams {
+    pub fn new(
+        terminal_view_id: Option<EntityId>,
+        session_context: SessionContext,
+        request_input: &RequestInput,
+        conversation: ConversationData,
+        metadata: Option<RequestMetadata>,
+        app: &AppContext,
+    ) -> Self {
+        let ai_settings = AISettings::as_ref(app);
+        let is_memory_enabled = ai_settings.is_memory_enabled(app);
+        let warp_drive_context_enabled = ai_settings.is_warp_drive_context_enabled(app);
+
+        // Build MCP context - either grouped by server or flat lists based on feature flag
+        let mcp_context = if FeatureFlag::MCPGroupedServerContext.is_enabled() {
+            // Group MCP tools and resources by server
+            let templatable_manager = TemplatableMCPServerManager::as_ref(app);
+
+            let mut active_servers: Vec<&TemplatableMCPServerInfo> = templatable_manager
+                .get_active_templatable_servers()
+                .values()
+                .copied()
+                .collect();
+
+            // If file-based MCP servers are enabled, add active servers in scope of
+            // the user's current working directory
+            if let Some(cwd) = session_context.current_working_directory() {
+                active_servers.extend(
+                    templatable_manager
+                        .get_active_file_based_servers(Path::new(cwd), app)
+                        .values(),
+                );
+            }
+
+            // Include any ephemeral MCP servers started via the Oz CLI.
+            active_servers.extend(
+                templatable_manager
+                    .get_active_cli_spawned_servers()
+                    .values(),
+            );
+
+            let servers: Vec<MCPServer> = active_servers
+                .into_iter()
+                .map(|server| MCPServer {
+                    name: server.name().to_string(),
+                    description: server.description().unwrap_or_default().to_string(),
+                    id: server.installation_id().to_string(),
+                    resources: server.resources().to_vec(),
+                    tools: server.tools().to_vec(),
+                })
+                .collect();
+
+            if servers.is_empty() {
+                None
+            } else {
+                #[allow(deprecated)]
+                Some(MCPContext {
+                    resources: vec![],
+                    tools: vec![],
+                    servers,
+                })
+            }
+        } else {
+            // Flat lists of resources and tools
+            let templatable_mcp_manager = TemplatableMCPServerManager::as_ref(app);
+            let resources = templatable_mcp_manager
+                .resources()
+                .cloned()
+                .collect::<Vec<_>>();
+            let tools = templatable_mcp_manager.tools().cloned().collect::<Vec<_>>();
+
+            #[allow(deprecated)]
+            (!resources.is_empty() || !tools.is_empty()).then_some(MCPContext {
+                resources,
+                tools,
+                servers: vec![],
+            })
+        };
+
+        let should_redact_secrets = get_secret_obfuscation_mode(app).should_redact_secret();
+
+        let user_workspaces = UserWorkspaces::as_ref(app);
+        let api_keys = ApiKeyManager::as_ref(app).api_keys_for_request(
+            user_workspaces.is_byo_api_key_enabled(),
+            user_workspaces.is_aws_bedrock_credentials_enabled(app),
+        );
+        let allow_use_of_warp_credits_with_byok =
+            *AISettings::as_ref(app).can_use_warp_credits_with_byok;
+
+        let app_execution_mode = AppExecutionMode::as_ref(app);
+        let autonomy_level = if app_execution_mode.is_autonomous() {
+            warp_multi_agent_api::AutonomyLevel::Unsupervised
+        } else {
+            warp_multi_agent_api::AutonomyLevel::Supervised
+        };
+
+        let isolation_level = if app_execution_mode.is_sandboxed() {
+            warp_multi_agent_api::IsolationLevel::Sandbox
+        } else {
+            warp_multi_agent_api::IsolationLevel::None
+        };
+
+        let web_search_enabled =
+            BlocklistAIPermissions::as_ref(app).get_web_search_enabled(app, terminal_view_id);
+        let research_agent_enabled = app
+            .private_user_preferences()
+            .read_value("ResearchAgentEnabled")
+            .ok()
+            .flatten()
+            .and_then(|s| s.parse().ok())
+            .unwrap_or_default();
+        let is_ambient_agent = conversation.ambient_agent_task_id.is_some();
+        let computer_use_enabled = FeatureFlag::AgentModeComputerUse.is_enabled()
+            && BlocklistAIPermissions::as_ref(app)
+                .get_computer_use_setting(app, terminal_view_id)
+                .is_enabled()
+            && computer_use::is_supported_on_current_platform()
+            && (FeatureFlag::LocalComputerUse.is_enabled() || is_ambient_agent);
+        let ask_user_question_enabled = BlocklistAIPermissions::as_ref(app)
+            .get_ask_user_question_setting(app, terminal_view_id)
+            != crate::ai::execution_profiles::AskUserQuestionPermission::Never;
+
+        let orchestration_enabled = ai_settings.is_orchestration_enabled(app)
+            && session_context
+                .session_type()
+                .as_ref()
+                .is_none_or(|t| matches!(t, crate::terminal::model::session::SessionType::Local));
+
+        Self {
+            input: request_input.all_inputs().cloned().collect(),
+            conversation_token: conversation.server_conversation_token,
+            forked_from_conversation_token: conversation.forked_from_conversation_token,
+            ambient_agent_task_id: conversation.ambient_agent_task_id,
+            tasks: conversation.tasks,
+            existing_suggestions: conversation.existing_suggestions,
+            metadata,
+            session_context,
+            model: request_input.model_id.clone(),
+            coding_model: request_input.coding_model_id.clone(),
+            cli_agent_model: request_input.cli_agent_model_id.clone(),
+            computer_use_model: request_input.computer_use_model_id.clone(),
+            is_memory_enabled,
+            warp_drive_context_enabled,
+            mcp_context,
+            planning_enabled: true,
+            should_redact_secrets,
+            api_keys,
+            allow_use_of_warp_credits_with_byok,
+            autonomy_level,
+            isolation_level,
+            web_search_enabled,
+            computer_use_enabled,
+            ask_user_question_enabled,
+            research_agent_enabled,
+            orchestration_enabled,
+            supported_tools_override: request_input.supported_tools_override.clone(),
+            parent_agent_id: None,
+            agent_name: None,
+        }
+    }
+}

app/src/ai/agent/api/convert_from.rs

@@ -0,0 +1,944 @@
+//! Conversions from MAA API types to application types.
+use std::collections::HashMap;
+use std::time::Duration;
+
+use crate::ai::agent::api::convert_conversation::{
+    convert_input_context, convert_tool_call_result_to_input,
+};
+use crate::ai::agent::comment::CodeReview;
+use crate::ai::agent::task::TaskId;
+use crate::ai::agent::todos::AIAgentTodoList;
+use crate::ai::agent::{
+    util::parse_markdown_into_text_and_code_sections, AIAgentAction, AIAgentActionType,
+    AIAgentCitation, AIAgentInput, AIAgentOutputMessage, AIAgentText, AIAgentTodo,
+    ArtifactCreatedData, MessageId, StartAgentExecutionMode, SuggestedAgentModeWorkflow,
+    SuggestedRule, Suggestions, TodoOperation,
+};
+use crate::ai::agent::{
+    CloneRepositoryURL, SubagentCall, SubagentType, SummarizationType, WebFetchStatus,
+    WebSearchStatus,
+};
+use crate::ai::artifact_download::sanitized_basename;
+use crate::ai::document::ai_document_model::{AIDocumentId, AIDocumentVersion};
+use ai::agent::action::LifecycleEventType as StartAgentLifecycleEventType;
+use ai::agent::action_result::StartAgentVersion;
+use ai::agent::convert::ToolToAIAgentActionError;
+use ai::agent::UnknownCitationTypeError;
+use ai::skills::SkillReference;
+use api::ask_user_question::question::QuestionType;
+use warp_core::channel::ChannelState;
+use warp_multi_agent_api as api;
+
+use crate::ai::agent::{AIAgentAttachment, UserQueryMode};
+
+impl TryFrom<api::Attachment> for AIAgentAttachment {
+    type Error = anyhow::Error;
+
+    fn try_from(attachment: api::Attachment) -> Result<Self, Self::Error> {
+        match attachment.value {
+            Some(api::attachment::Value::FilePathReference(fpr)) => {
+                Ok(AIAgentAttachment::FilePathReference {
+                    file_id: String::new(),
+                    file_name: fpr
+                        .file_path
+                        .rsplit('/')
+                        .next()
+                        .unwrap_or(&fpr.file_path)
+                        .to_string(),
+                    file_path: fpr.file_path,
+                })
+            }
+            _ => anyhow::bail!("Unsupported attachment type for conversion"),
+        }
+    }
+}
+
+/// Converts proto UserQueryMode to the internal UserQueryMode type
+pub(crate) fn convert_user_query_mode(mode: Option<&api::UserQueryMode>) -> UserQueryMode {
+    let Some(mode) = mode else {
+        return UserQueryMode::default();
+    };
+
+    match &mode.r#type {
+        Some(api::user_query_mode::Type::Plan(_)) => UserQueryMode::Plan,
+        Some(api::user_query_mode::Type::Orchestrate(_)) => UserQueryMode::Orchestrate,
+        None => UserQueryMode::Normal,
+    }
+}
+
+fn convert_start_agent_lifecycle_event_type(
+    event_type: i32,
+) -> Option<StartAgentLifecycleEventType> {
+    let event_type = StartAgentLifecycleEventType::try_from(event_type).ok()?;
+    (event_type != StartAgentLifecycleEventType::Unspecified).then_some(event_type)
+}
+
+fn convert_start_agent_v2_harness_type(
+    harness: Option<api::start_agent_v2::execution_mode::Harness>,
+) -> Option<String> {
+    harness
+        .map(|harness| harness.r#type)
+        .filter(|harness_type| !harness_type.trim().is_empty())
+}
+
+fn convert_start_agent_execution_mode(
+    execution_mode: Option<api::start_agent::ExecutionMode>,
+) -> StartAgentExecutionMode {
+    match execution_mode.and_then(|execution_mode| execution_mode.mode) {
+        Some(api::start_agent::execution_mode::Mode::Remote(remote)) => {
+            StartAgentExecutionMode::remote_with_defaults(remote.environment_id)
+        }
+        Some(api::start_agent::execution_mode::Mode::Local(_)) | None => {
+            StartAgentExecutionMode::local_with_defaults()
+        }
+    }
+}
+
+fn convert_start_agent_v2_execution_mode(
+    execution_mode: Option<api::start_agent_v2::ExecutionMode>,
+) -> StartAgentExecutionMode {
+    match execution_mode.and_then(|execution_mode| execution_mode.mode) {
+        Some(api::start_agent_v2::execution_mode::Mode::Remote(remote)) => {
+            StartAgentExecutionMode::Remote {
+                environment_id: remote.environment_id,
+                skill_references: remote
+                    .skills
+                    .into_iter()
+                    .filter_map(convert_skill_reference)
+                    .collect(),
+                model_id: remote.model_id,
+                computer_use_enabled: remote.computer_use_enabled,
+                worker_host: remote.worker_host,
+                harness_type: convert_start_agent_v2_harness_type(remote.harness)
+                    .unwrap_or_default(),
+                title: remote.title,
+            }
+        }
+        Some(api::start_agent_v2::execution_mode::Mode::Local(local)) => {
+            convert_start_agent_v2_harness_type(local.harness)
+                .map(StartAgentExecutionMode::local_harness)
+                .unwrap_or_else(StartAgentExecutionMode::local_with_defaults)
+        }
+        None => StartAgentExecutionMode::local_with_defaults(),
+    }
+}
+
+fn convert_skill_reference(skill_ref: api::SkillRef) -> Option<SkillReference> {
+    match skill_ref.skill_reference {
+        Some(api::skill_ref::SkillReference::Path(path)) => Some(SkillReference::Path(path.into())),
+        Some(api::skill_ref::SkillReference::BundledSkillId(id)) => {
+            Some(SkillReference::BundledSkillId(id))
+        }
+        None => None,
+    }
+}
+
+/// Unexpected errors when trying to convert an [`api::Message`] to an [`AIAgentOutputMessage`].
+#[derive(Debug, thiserror::Error)]
+pub enum MessageToAIAgentOutputMessageError {
+    #[error("Missing expected message")]
+    MissingMessage,
+    #[error("Error converting tool to action: {0:?}")]
+    ToolError(#[from] ToolToAIAgentActionError),
+    #[error("Error converting citation: {0:?}")]
+    CitationError(#[from] UnknownCitationTypeError),
+}
+
+/// Successful result when trying to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
+#[allow(clippy::large_enum_variant)]
+pub enum MaybeAIAgentOutputMessage {
+    /// There is a mapping to a client output message.
+    Message(AIAgentOutputMessage),
+    /// We tried to parse a message that we don't care about.
+    NoClientRepresentation,
+}
+
+/// Successful result when trying to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
+#[allow(clippy::large_enum_variant)]
+enum MaybeAIAgentAction {
+    /// There is a mapping to a client action.
+    Action(AIAgentAction),
+    Subagent(SubagentCall),
+    /// We tried to parse a tool call that we don't care about.
+    NoClientRepresentation,
+}
+
+pub struct ConversionParams<'a> {
+    pub task_id: &'a TaskId,
+    pub current_todo_list: Option<&'a AIAgentTodoList>,
+    pub active_code_review: Option<&'a CodeReview>,
+}
+
+/// Trait for converting an [`api::Message`] to an [`AIAgentOutputMessage`].
+pub trait ConvertAPIMessageToClientOutputMessage {
+    fn to_client_output_message(
+        self,
+        params: ConversionParams,
+    ) -> Result<MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError>;
+}
+
+impl ConvertAPIMessageToClientOutputMessage for api::Message {
+    fn to_client_output_message(
+        self,
+        params: ConversionParams,
+    ) -> Result<MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError> {
+        let Some(message) = self.message else {
+            // In shared-session streams we can receive skeleton placeholder task messages without payloads.
+            // Treat them as having no client representation rather than erroring and aborting ingestion entirely.
+            return Ok(MaybeAIAgentOutputMessage::NoClientRepresentation);
+        };
+
+        let citations = self
+            .citations
+            .iter()
+            .map(|citation| (*citation).clone().try_into())
+            .collect::<Result<Vec<AIAgentCitation>, UnknownCitationTypeError>>()?;
+
+        match message {
+            api::message::Message::AgentOutput(output) => Ok(MaybeAIAgentOutputMessage::Message(
+                AIAgentOutputMessage::text(MessageId::new(self.id), output.into())
+                    .with_citations(citations),
+            )),
+            api::message::Message::AgentReasoning(reasoning) => {
+                let duration = reasoning
+                    .finished_duration
+                    .map(|d| Duration::from_secs(d.seconds as u64));
+                Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::reasoning(
+                        MessageId::new(self.id),
+                        reasoning.into(),
+                        duration,
+                    ),
+                ))
+            }
+            api::message::Message::ToolCall(tool_call) => match tool_call.to_action(params)? {
+                MaybeAIAgentAction::Action(action) => Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::action(MessageId::new(self.id), action)
+                        .with_citations(citations),
+                )),
+                MaybeAIAgentAction::Subagent(subagent) => Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::subagent(MessageId::new(self.id), subagent)
+                        .with_citations(citations),
+                )),
+                MaybeAIAgentAction::NoClientRepresentation => {
+                    Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                }
+            },
+            api::message::Message::WebSearch(web_search) => {
+                let status = match &web_search.status {
+                    Some(api::message::web_search::Status {
+                        r#type: Some(api::message::web_search::status::Type::Searching(searching)),
+                    }) => WebSearchStatus::Searching {
+                        query: if searching.query.is_empty() {
+                            None
+                        } else {
+                            Some(searching.query.clone())
+                        },
+                    },
+                    Some(api::message::web_search::Status {
+                        r#type: Some(api::message::web_search::status::Type::Success(success)),
+                    }) => WebSearchStatus::Success {
+                        query: success.query.clone(),
+                        pages: success
+                            .pages
+                            .iter()
+                            .map(|p| (p.url.clone(), p.title.clone()))
+                            .collect(),
+                    },
+                    Some(api::message::web_search::Status {
+                        r#type: Some(api::message::web_search::status::Type::Error(_)),
+                    }) => {
+                        // Error type doesn't have a query field currently, use empty string
+                        WebSearchStatus::Error {
+                            query: String::new(),
+                        }
+                    }
+                    _ => {
+                        // Unknown or missing status
+                        return Ok(MaybeAIAgentOutputMessage::NoClientRepresentation);
+                    }
+                };
+
+                Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::web_search(MessageId::new(self.id), status)
+                        .with_citations(citations),
+                ))
+            }
+            api::message::Message::WebFetch(web_fetch) => {
+                let status = match &web_fetch.status {
+                    Some(api::message::web_fetch::Status {
+                        r#type: Some(api::message::web_fetch::status::Type::Fetching(fetching)),
+                    }) => WebFetchStatus::Fetching {
+                        urls: fetching.urls.clone(),
+                    },
+                    Some(api::message::web_fetch::Status {
+                        r#type: Some(api::message::web_fetch::status::Type::Success(success)),
+                    }) => WebFetchStatus::Success {
+                        pages: success
+                            .pages
+                            .iter()
+                            .map(|p| (p.url.clone(), p.title.clone(), p.success))
+                            .collect(),
+                    },
+                    Some(api::message::web_fetch::Status {
+                        r#type: Some(api::message::web_fetch::status::Type::Error(_)),
+                    }) => WebFetchStatus::Error,
+                    _ => {
+                        // Unknown or missing status
+                        return Ok(MaybeAIAgentOutputMessage::NoClientRepresentation);
+                    }
+                };
+
+                Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::web_fetch(MessageId::new(self.id), status)
+                        .with_citations(citations),
+                ))
+            }
+            api::message::Message::ModelUsed(_) => {
+                Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+            }
+            api::message::Message::UpdateTodos(update_todos) => {
+                if let Some(operation) = update_todos.operation {
+                    match operation {
+                        api::message::update_todos::Operation::CreateTodoList(create_todo_list) => {
+                            Ok(MaybeAIAgentOutputMessage::Message(
+                                AIAgentOutputMessage::todo_operation(
+                                    MessageId::new(self.id),
+                                    TodoOperation::UpdateTodos {
+                                        todos: create_todo_list
+                                            .initial_todos
+                                            .into_iter()
+                                            .map(Into::into)
+                                            .collect(),
+                                    },
+                                )
+                                .with_citations(citations),
+                            ))
+                        }
+                        api::message::update_todos::Operation::UpdatePendingTodos(
+                            update_pending_todos,
+                        ) => Ok(MaybeAIAgentOutputMessage::Message(
+                            AIAgentOutputMessage::todo_operation(
+                                MessageId::new(self.id),
+                                TodoOperation::UpdateTodos {
+                                    todos: params
+                                        .current_todo_list
+                                        .iter()
+                                        .flat_map(|list| list.completed_items().iter().cloned())
+                                        .chain(
+                                            update_pending_todos
+                                                .updated_pending_todos
+                                                .into_iter()
+                                                .map(Into::into),
+                                        )
+                                        .collect(),
+                                },
+                            )
+                            .with_citations(citations),
+                        )),
+                        api::message::update_todos::Operation::MarkTodosCompleted(
+                            mark_todos_completed,
+                        ) => {
+                            if mark_todos_completed.todo_ids.is_empty() {
+                                Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                            } else {
+                                // This is a mark as completed operation
+                                Ok(MaybeAIAgentOutputMessage::Message(
+                                    AIAgentOutputMessage::todo_operation(
+                                        MessageId::new(self.id),
+                                        TodoOperation::MarkAsCompleted {
+                                            completed_todos: mark_todos_completed
+                                                .todo_ids
+                                                .into_iter()
+                                                .filter_map(|todo_id| {
+                                                    params.current_todo_list.and_then(|todo_list| {
+                                                        todo_list
+                                                            .completed_items()
+                                                            .iter()
+                                                            .find(|item| {
+                                                                item.id.as_ref() == todo_id.as_str()
+                                                            })
+                                                            .cloned()
+                                                    })
+                                                })
+                                                .collect(),
+                                        },
+                                    )
+                                    .with_citations(citations),
+                                ))
+                            }
+                        }
+                    }
+                } else {
+                    Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                }
+            }
+            api::message::Message::Summarization(summarization) => {
+                let duration = summarization
+                    .finished_duration
+                    .map(|d| Duration::from_secs(d.seconds as u64));
+                let (text, summarization_type, token_count) = match summarization.summary_type {
+                    Some(api::message::summarization::SummaryType::ConversationSummary(
+                        conv_summary,
+                    )) => {
+                        let token_count = if conv_summary.token_count > 0 {
+                            Some(conv_summary.token_count as u32)
+                        } else {
+                            None
+                        };
+                        let text = if !conv_summary.summary.is_empty() {
+                            AIAgentText {
+                                sections: parse_markdown_into_text_and_code_sections(
+                                    &conv_summary.summary,
+                                ),
+                            }
+                        } else {
+                            AIAgentText { sections: vec![] }
+                        };
+                        (text, SummarizationType::ConversationSummary, token_count)
+                    }
+                    Some(api::message::summarization::SummaryType::ToolCallResultSummary(_)) => (
+                        AIAgentText { sections: vec![] },
+                        SummarizationType::ToolCallResultSummary,
+                        None,
+                    ),
+                    None => {
+                        // Default to ConversationSummary if not specified
+                        (
+                            AIAgentText { sections: vec![] },
+                            SummarizationType::ConversationSummary,
+                            None,
+                        )
+                    }
+                };
+                Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::summarization(
+                        MessageId::new(self.id),
+                        text,
+                        duration,
+                        summarization_type,
+                        token_count,
+                    ),
+                ))
+            }
+            api::message::Message::UpdateReviewComments(update_comments) => {
+                if let Some(operation) = update_comments.operation {
+                    match operation {
+                        api::message::update_review_comments::Operation::AddressReviewComments(
+                            address_comments,
+                        ) => {
+                            if let Some(current_comments) = params.active_code_review {
+                                let addressed_comments = current_comments
+                                    .addressed_comments
+                                    .iter()
+                                    .filter(|comment| {
+                                        address_comments
+                                            .comment_ids
+                                            .iter()
+                                            .any(|id| id == &comment.id.to_string())
+                                    })
+                                    .cloned()
+                                    .collect();
+                                Ok(MaybeAIAgentOutputMessage::Message(
+                                    AIAgentOutputMessage::comments_addressed(
+                                        MessageId::new(self.id),
+                                        addressed_comments,
+                                    )
+                                    .with_citations(citations),
+                                ))
+                            } else {
+                                Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                            }
+                        }
+                    }
+                } else {
+                    Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                }
+            }
+            api::message::Message::DebugOutput(debug_output) => {
+                if ChannelState::enable_debug_features() {
+                    Ok(MaybeAIAgentOutputMessage::Message(
+                        AIAgentOutputMessage::debug_output(
+                            MessageId::new(self.id),
+                            debug_output.text,
+                        ),
+                    ))
+                } else {
+                    Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                }
+            }
+            api::message::Message::ArtifactEvent(artifact_event) => match artifact_event.event {
+                Some(api::message::artifact_event::Event::Created(artifact_created)) => {
+                    match artifact_created.artifact {
+                        Some(
+                            api::message::artifact_event::artifact_created::Artifact::PullRequest(
+                                pr,
+                            ),
+                        ) => Ok(MaybeAIAgentOutputMessage::Message(
+                            AIAgentOutputMessage::artifact_created(
+                                MessageId::new(self.id),
+                                ArtifactCreatedData::PullRequest {
+                                    url: pr.url,
+                                    branch: pr.branch,
+                                },
+                            )
+                            .with_citations(citations),
+                        )),
+                        Some(
+                            api::message::artifact_event::artifact_created::Artifact::Screenshot(
+                                screenshot,
+                            ),
+                        ) => Ok(MaybeAIAgentOutputMessage::Message(
+                            AIAgentOutputMessage::artifact_created(
+                                MessageId::new(self.id),
+                                ArtifactCreatedData::Screenshot {
+                                    artifact_uid: screenshot.artifact_uid,
+                                    mime_type: screenshot.mime_type,
+                                    description: if screenshot.description.is_empty() {
+                                        None
+                                    } else {
+                                        Some(screenshot.description)
+                                    },
+                                },
+                            )
+                            .with_citations(citations),
+                        )),
+                        Some(api::message::artifact_event::artifact_created::Artifact::File(
+                            file,
+                        )) => Ok(MaybeAIAgentOutputMessage::Message(
+                            AIAgentOutputMessage::artifact_created(
+                                MessageId::new(self.id),
+                                ArtifactCreatedData::File {
+                                    artifact_uid: file.artifact_uid,
+                                    filename: sanitized_basename(&file.filepath)
+                                        .unwrap_or_else(|| file.filepath.clone()),
+                                    filepath: file.filepath,
+                                    mime_type: file.mime_type,
+                                    description: if file.description.is_empty() {
+                                        None
+                                    } else {
+                                        Some(file.description)
+                                    },
+                                    size_bytes: file.size_bytes,
+                                },
+                            )
+                            .with_citations(citations),
+                        )),
+                        None => Ok(MaybeAIAgentOutputMessage::NoClientRepresentation),
+                    }
+                }
+                Some(api::message::artifact_event::Event::ForkArtifacts(_)) | None => {
+                    Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+                }
+            },
+            api::message::Message::MessagesReceivedFromAgents(messages_received_from_agents) => {
+                let messages = messages_received_from_agents
+                    .messages
+                    .into_iter()
+                    .map(|msg| crate::ai::agent::ReceivedMessageDisplay {
+                        message_id: msg.message_id,
+                        sender_agent_id: msg.sender_agent_id,
+                        addresses: msg.addresses,
+                        subject: msg.subject,
+                        message_body: msg.message_body,
+                    })
+                    .collect();
+                Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::messages_received_from_agents(
+                        MessageId::new(self.id),
+                        messages,
+                    )
+                    .with_citations(citations),
+                ))
+            }
+            api::message::Message::EventsFromAgents(events) => {
+                let event_ids = events
+                    .agent_events
+                    .iter()
+                    .map(|e| e.event_id.clone())
+                    .collect();
+                Ok(MaybeAIAgentOutputMessage::Message(
+                    AIAgentOutputMessage::events_from_agents(MessageId::new(self.id), event_ids)
+                        .with_citations(citations),
+                ))
+            }
+            // These messages don't indicate an error but they don't translate to a client-side output message.
+            api::message::Message::UserQuery(_)
+            | api::message::Message::SystemQuery(_)
+            | api::message::Message::ToolCallResult(_)
+            | api::message::Message::CodeReview(_)
+            | api::message::Message::ServerEvent(_)
+            | api::message::Message::InvokeSkill(_)
+            | api::message::Message::PassiveSuggestionResult(_) => {
+                Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
+            }
+        }
+    }
+}
+
+impl From<api::message::AgentOutput> for AIAgentText {
+    fn from(value: api::message::AgentOutput) -> Self {
+        AIAgentText {
+            sections: parse_markdown_into_text_and_code_sections(value.text.as_str()),
+        }
+    }
+}
+
+impl From<api::message::AgentReasoning> for AIAgentText {
+    fn from(value: api::message::AgentReasoning) -> Self {
+        AIAgentText {
+            sections: parse_markdown_into_text_and_code_sections(value.reasoning.as_str()),
+        }
+    }
+}
+
+/// Trait for converting an [`api::Message`] to an [`AIAgentOutputMessage`].
+trait ConvertAPIToolCallToAIAgentAction {
+    fn to_action(
+        self,
+        params: ConversionParams,
+    ) -> Result<MaybeAIAgentAction, ToolToAIAgentActionError>;
+}
+
+/// Trys to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
+///
+/// A [`Result::Error`] indicates an unexpected problem, while [`Ok(None)`]
+/// indicates a tool call that we aren't expected to parse.
+impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
+    fn to_action(
+        self,
+        params: ConversionParams,
+    ) -> Result<MaybeAIAgentAction, ToolToAIAgentActionError> {
+        let Some(tool) = self.tool else {
+            return Err(ToolToAIAgentActionError::MissingTool);
+        };
+
+        let create_standard_action = |action: AIAgentActionType| {
+            Ok(MaybeAIAgentAction::Action(AIAgentAction {
+                id: self.tool_call_id.clone().into(),
+                task_id: params.task_id.clone(),
+                action,
+                requires_result: true,
+            }))
+        };
+
+        match tool {
+            api::message::tool_call::Tool::RunShellCommand(run_shell_command) => {
+                create_standard_action(run_shell_command.into())
+            }
+            api::message::tool_call::Tool::WriteToLongRunningShellCommand(
+                write_to_long_running_shell_command,
+            ) => create_standard_action(write_to_long_running_shell_command.into()),
+            api::message::tool_call::Tool::ReadFiles(read_files) => {
+                create_standard_action(read_files.into())
+            }
+            api::message::tool_call::Tool::UploadFileArtifact(upload_file_artifact) => {
+                create_standard_action(upload_file_artifact.try_into()?)
+            }
+            api::message::tool_call::Tool::SearchCodebase(search_codebase) => {
+                create_standard_action(search_codebase.into())
+            }
+            api::message::tool_call::Tool::Grep(grep) => create_standard_action(grep.into()),
+            #[allow(deprecated)]
+            api::message::tool_call::Tool::FileGlob(glob) => create_standard_action(glob.into()),
+            api::message::tool_call::Tool::FileGlobV2(glob) => create_standard_action(glob.into()),
+            api::message::tool_call::Tool::ApplyFileDiffs(apply_file_diffs) => {
+                create_standard_action(apply_file_diffs.into())
+            }
+            api::message::tool_call::Tool::ReadMcpResource(read_mcp_resource) => {
+                create_standard_action(read_mcp_resource.into())
+            }
+            api::message::tool_call::Tool::CallMcpTool(call_mcp_tool) => {
+                match call_mcp_tool.try_into() {
+                    Ok(call_mcp_tool_action) => create_standard_action(call_mcp_tool_action),
+                    Err(error) => Err(error),
+                }
+            }
+            api::message::tool_call::Tool::SuggestNewConversation(suggest_new_conversation) => {
+                create_standard_action(suggest_new_conversation.into())
+            }
+            api::message::tool_call::Tool::SuggestPrompt(suggest_prompt) => {
+                match suggest_prompt.try_into() {
+                    Ok(suggest_prompt_action) => create_standard_action(suggest_prompt_action),
+                    Err(_) => Ok(MaybeAIAgentAction::NoClientRepresentation),
+                }
+            }
+            api::message::tool_call::Tool::OpenCodeReview(_) => {
+                create_standard_action(AIAgentActionType::OpenCodeReview)
+            }
+            api::message::tool_call::Tool::InitProject(_) => {
+                create_standard_action(AIAgentActionType::InitProject)
+            }
+            api::message::tool_call::Tool::ReadDocuments(read_documents) => {
+                create_standard_action(read_documents.into())
+            }
+            api::message::tool_call::Tool::EditDocuments(edit_documents) => {
+                create_standard_action(edit_documents.into())
+            }
+            api::message::tool_call::Tool::CreateDocuments(create_documents) => {
+                create_standard_action(create_documents.into())
+            }
+            api::message::tool_call::Tool::ReadShellCommandOutput(read_shell_command_output) => {
+                create_standard_action(read_shell_command_output.into())
+            }
+            api::message::tool_call::Tool::TransferShellCommandControlToUser(
+                transfer_shell_command_control_to_user,
+            ) => create_standard_action(transfer_shell_command_control_to_user.into()),
+            api::message::tool_call::Tool::UseComputer(use_computer) => {
+                create_standard_action(use_computer.try_into()?)
+            }
+            api::message::tool_call::Tool::RequestComputerUse(request_computer_use) => {
+                create_standard_action(request_computer_use.into())
+            }
+            api::message::tool_call::Tool::Subagent(subagent) => {
+                use api::message::tool_call::subagent::Metadata;
+                let subagent_type = match subagent.metadata {
+                    Some(Metadata::Cli(_)) => SubagentType::Cli,
+                    Some(Metadata::Research(_)) => SubagentType::Research,
+                    Some(Metadata::Advice(_)) => SubagentType::Advice,
+                    Some(Metadata::ComputerUse(_)) => SubagentType::ComputerUse,
+                    Some(Metadata::Summarization(_)) => SubagentType::Summarization,
+                    Some(Metadata::ConversationSearch(cs_meta)) => {
+                        let query = if cs_meta.query.is_empty() {
+                            None
+                        } else {
+                            Some(cs_meta.query)
+                        };
+                        let conversation_id = if cs_meta.conversation_id.is_empty() {
+                            None
+                        } else {
+                            Some(cs_meta.conversation_id)
+                        };
+                        SubagentType::ConversationSearch {
+                            query,
+                            conversation_id,
+                        }
+                    }
+                    Some(Metadata::WarpDocumentationSearch(_)) => {
+                        SubagentType::WarpDocumentationSearch
+                    }
+                    None => SubagentType::Unknown,
+                };
+                Ok(MaybeAIAgentAction::Subagent(SubagentCall {
+                    task_id: subagent.task_id,
+                    subagent_type,
+                }))
+            }
+            api::message::tool_call::Tool::StartAgent(start_agent) => {
+                create_standard_action(AIAgentActionType::StartAgent {
+                    version: StartAgentVersion::V1,
+                    name: start_agent.name,
+                    prompt: start_agent.prompt,
+                    execution_mode: convert_start_agent_execution_mode(start_agent.execution_mode),
+                    lifecycle_subscription: start_agent.lifecycle_subscription.map(
+                        |subscription| {
+                            subscription
+                                .event_types
+                                .into_iter()
+                                .filter_map(convert_start_agent_lifecycle_event_type)
+                                .collect()
+                        },
+                    ),
+                })
+            }
+            api::message::tool_call::Tool::StartAgentV2(start_agent) => {
+                create_standard_action(AIAgentActionType::StartAgent {
+                    version: StartAgentVersion::V2,
+                    name: start_agent.name,
+                    prompt: start_agent.prompt,
+                    execution_mode: convert_start_agent_v2_execution_mode(
+                        start_agent.execution_mode,
+                    ),
+                    lifecycle_subscription: start_agent.lifecycle_subscription.map(
+                        |subscription| {
+                            subscription
+                                .event_types
+                                .into_iter()
+                                .filter_map(convert_start_agent_lifecycle_event_type)
+                                .collect()
+                        },
+                    ),
+                })
+            }
+            api::message::tool_call::Tool::SendMessageToAgent(send_message) => {
+                create_standard_action(AIAgentActionType::SendMessageToAgent {
+                    addresses: send_message.addresses,
+                    subject: send_message.subject,
+                    message: send_message.message,
+                })
+            }
+            api::message::tool_call::Tool::InsertReviewComments(insert_review_comments) => {
+                create_standard_action(insert_review_comments.into())
+            }
+            api::message::tool_call::Tool::ReadSkill(read_skill) => {
+                create_standard_action(read_skill.try_into()?)
+            }
+            api::message::tool_call::Tool::FetchConversation(fetch_conversation) => {
+                create_standard_action(fetch_conversation.into())
+            }
+            api::message::tool_call::Tool::AskUserQuestion(ask) => {
+                let questions = ask
+                    .questions
+                    .into_iter()
+                    .filter_map(convert_api_question)
+                    .collect();
+                create_standard_action(AIAgentActionType::AskUserQuestion { questions })
+            }
+            // Clients do not need to know how to parse server tool-calls but receiving
+            // them is not an error.
+            api::message::tool_call::Tool::Server(_) => {
+                Ok(MaybeAIAgentAction::NoClientRepresentation)
+            }
+            _ => Err(ToolToAIAgentActionError::UnexpectedTool),
+        }
+    }
+}
+
+impl From<api::Suggestions> for Suggestions {
+    fn from(api_suggestions: api::Suggestions) -> Self {
+        Self {
+            rules: api_suggestions
+                .rules
+                .into_iter()
+                .map(|rule| SuggestedRule {
+                    name: rule.name,
+                    content: rule.content,
+                    logging_id: rule.logging_id.into(),
+                })
+                .collect(),
+            agent_mode_workflows: api_suggestions
+                .workflows
+                .into_iter()
+                .map(|workflow| SuggestedAgentModeWorkflow {
+                    name: workflow.name,
+                    prompt: workflow.prompt,
+                    logging_id: workflow.logging_id.into(),
+                })
+                .collect(),
+        }
+    }
+}
+
+impl From<api::TodoItem> for AIAgentTodo {
+    fn from(value: api::TodoItem) -> Self {
+        AIAgentTodo {
+            id: value.id.into(),
+            title: value.title,
+            description: value.description,
+        }
+    }
+}
+
+/// Reconstruct user inputs from the provided server messages
+/// (for use in shared agent exchanges where the input was not provided in this session)
+pub fn user_inputs_from_messages(messages: &[api::Message]) -> Vec<AIAgentInput> {
+    let mut inputs = Vec::new();
+    let mut document_versions: HashMap<AIDocumentId, AIDocumentVersion> = HashMap::new();
+    for m in messages {
+        let Some(inner) = &m.message else { continue };
+        match inner {
+            api::message::Message::UserQuery(uq) => {
+                let context = convert_input_context(uq.context.as_ref());
+                let referenced_attachments = uq
+                    .referenced_attachments
+                    .iter()
+                    .filter_map(|(key, attachment)| {
+                        AIAgentAttachment::try_from(attachment.clone())
+                            .ok()
+                            .map(|a| (key.clone(), a))
+                    })
+                    .collect();
+                inputs.push(AIAgentInput::UserQuery {
+                    query: uq.query.clone(),
+                    context,
+                    static_query_type: None,
+                    referenced_attachments,
+                    user_query_mode: convert_user_query_mode(uq.mode.as_ref()),
+                    running_command: None,
+                    intended_agent: Some(uq.intended_agent()),
+                });
+            }
+            api::message::Message::SystemQuery(sq) => {
+                let ctx = convert_input_context(sq.context.as_ref());
+                if let Some(t) = &sq.r#type {
+                    // These system queries appear as user inputs in ai blocks.
+                    match t {
+                        api::message::system_query::Type::CreateNewProject(p) => {
+                            inputs.push(AIAgentInput::CreateNewProject {
+                                query: p.query.clone(),
+                                context: ctx,
+                            });
+                        }
+                        api::message::system_query::Type::CloneRepository(p) => {
+                            inputs.push(AIAgentInput::CloneRepository {
+                                clone_repo_url: CloneRepositoryURL::new(p.url.clone()),
+                                context: ctx,
+                            });
+                        }
+                        api::message::system_query::Type::AutoCodeDiff(p) => {
+                            inputs.push(AIAgentInput::AutoCodeDiffQuery {
+                                query: p.query.clone(),
+                                context: ctx,
+                            });
+                        }
+                        api::message::system_query::Type::FetchReviewComments(fetch) => {
+                            inputs.push(AIAgentInput::FetchReviewComments {
+                                repo_path: fetch.repo_path.clone(),
+                                context: ctx,
+                            });
+                        }
+                        _ => {}
+                    }
+                }
+            }
+            api::message::Message::ToolCallResult(tcr) => {
+                let task_id = TaskId::new(m.task_id.clone());
+                if let Some(input) = convert_tool_call_result_to_input(
+                    &task_id,
+                    tcr,
+                    &HashMap::new(),
+                    &mut document_versions,
+                ) {
+                    inputs.push(input);
+                }
+            }
+            _ => {}
+        }
+    }
+    inputs
+}
+
+fn convert_api_question(
+    q: api::ask_user_question::Question,
+) -> Option<ai::agent::action::AskUserQuestionItem> {
+    let Some(QuestionType::MultipleChoice(mc)) = q.question_type else {
+        return None;
+    };
+
+    // Server sends -1 when there is no recommendation.
+    let recommended_idx = usize::try_from(mc.recommended_option_index)
+        .ok()
+        .filter(|idx| *idx < mc.options.len());
+    let options = mc
+        .options
+        .iter()
+        .enumerate()
+        .map(|(i, opt)| ai::agent::action::AskUserQuestionOption {
+            label: opt.label.clone(),
+            recommended: recommended_idx == Some(i),
+        })
+        .collect();
+    Some(ai::agent::action::AskUserQuestionItem {
+        question_id: q.question_id.clone(),
+        question: q.question,
+        question_type: ai::agent::action::AskUserQuestionType::MultipleChoice {
+            is_multiselect: mc.is_multiselect,
+            options,
+            supports_other: mc.supports_other,
+        },
+    })
+}
+
+#[cfg(test)]
+#[path = "convert_from_tests.rs"]
+mod tests;

app/src/ai/agent/api/convert_from_tests.rs

@@ -0,0 +1,644 @@
+use super::{
+    convert_api_question, ConversionParams, ConvertAPIMessageToClientOutputMessage,
+    MaybeAIAgentOutputMessage,
+};
+use crate::ai::agent::task::TaskId;
+use crate::ai::agent::{
+    AIAgentActionType, AIAgentOutputMessageType, LifecycleEventType, StartAgentExecutionMode,
+};
+use ai::agent::action::AskUserQuestionType;
+use ai::skills::SkillReference;
+use warp_multi_agent_api as api;
+
+fn start_agent_tool_call_message(
+    name: &str,
+    prompt: &str,
+    execution_mode: Option<api::start_agent::ExecutionMode>,
+    lifecycle_subscription_event_types: Option<Vec<i32>>,
+) -> api::Message {
+    api::Message {
+        id: "message-id".to_string(),
+        task_id: "task-id".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: "tool-call-id".to_string(),
+            tool: Some(api::message::tool_call::Tool::StartAgent(api::StartAgent {
+                name: name.to_string(),
+                prompt: prompt.to_string(),
+                execution_mode,
+                lifecycle_subscription: lifecycle_subscription_event_types
+                    .map(|event_types| api::start_agent::LifecycleSubscription { event_types }),
+            })),
+        })),
+        request_id: "request-id".to_string(),
+        timestamp: None,
+    }
+}
+
+fn local_start_agent_v2_execution_mode(harness_type: &str) -> api::start_agent_v2::ExecutionMode {
+    api::start_agent_v2::ExecutionMode {
+        mode: Some(api::start_agent_v2::execution_mode::Mode::Local(
+            api::start_agent_v2::execution_mode::Local {
+                harness: Some(api::start_agent_v2::execution_mode::Harness {
+                    r#type: harness_type.to_string(),
+                }),
+            },
+        )),
+    }
+}
+
+fn local_start_agent_v2_execution_mode_without_harness() -> api::start_agent_v2::ExecutionMode {
+    api::start_agent_v2::ExecutionMode {
+        mode: Some(api::start_agent_v2::execution_mode::Mode::Local(
+            api::start_agent_v2::execution_mode::Local { harness: None },
+        )),
+    }
+}
+
+fn start_agent_v2_tool_call_message(
+    name: &str,
+    prompt: &str,
+    execution_mode: Option<api::start_agent_v2::ExecutionMode>,
+    lifecycle_subscription_event_types: Option<Vec<i32>>,
+) -> api::Message {
+    api::Message {
+        id: "message-id".to_string(),
+        task_id: "task-id".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: "tool-call-id".to_string(),
+            tool: Some(api::message::tool_call::Tool::StartAgentV2(
+                api::StartAgentV2 {
+                    name: name.to_string(),
+                    prompt: prompt.to_string(),
+                    execution_mode,
+                    lifecycle_subscription: lifecycle_subscription_event_types.map(|event_types| {
+                        api::start_agent_v2::LifecycleSubscription { event_types }
+                    }),
+                },
+            )),
+        })),
+        request_id: "request-id".to_string(),
+        timestamp: None,
+    }
+}
+
+fn upload_artifact_tool_call_message(path: &str, description: &str) -> api::Message {
+    api::Message {
+        id: "message-id".to_string(),
+        task_id: "task-id".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: "tool-call-id".to_string(),
+            tool: Some(api::message::tool_call::Tool::UploadFileArtifact(
+                api::UploadFileArtifact {
+                    file: Some(api::FilePathReference {
+                        file_path: path.to_string(),
+                    }),
+                    description: description.to_string(),
+                },
+            )),
+        })),
+        request_id: "request-id".to_string(),
+        timestamp: None,
+    }
+}
+
+fn remote_start_agent_v2_execution_mode(
+    environment_id: &str,
+) -> api::start_agent_v2::ExecutionMode {
+    api::start_agent_v2::ExecutionMode {
+        mode: Some(api::start_agent_v2::execution_mode::Mode::Remote(
+            api::start_agent_v2::execution_mode::Remote {
+                environment_id: environment_id.to_string(),
+                skills: vec![
+                    api::SkillRef {
+                        skill_reference: Some(api::skill_ref::SkillReference::Path(
+                            "/tmp/SKILL.md".to_string(),
+                        )),
+                    },
+                    api::SkillRef {
+                        skill_reference: Some(api::skill_ref::SkillReference::BundledSkillId(
+                            "review-comments".to_string(),
+                        )),
+                    },
+                ],
+                model_id: "gpt-test".to_string(),
+                computer_use_enabled: true,
+                worker_host: "worker-host".to_string(),
+                harness: Some(api::start_agent_v2::execution_mode::Harness {
+                    r#type: "claude-code".to_string(),
+                }),
+                title: "Remote child".to_string(),
+            },
+        )),
+    }
+}
+
+fn file_artifact_created_message(filepath: &str, description: &str) -> api::Message {
+    api::Message {
+        id: "message-id".to_string(),
+        task_id: "task-id".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ArtifactEvent(
+            api::message::ArtifactEvent {
+                event: Some(api::message::artifact_event::Event::Created(
+                    api::message::artifact_event::ArtifactCreated {
+                        artifact: Some(
+                            api::message::artifact_event::artifact_created::Artifact::File(
+                                api::message::artifact_event::FileArtifact {
+                                    artifact_uid: "artifact-uid".to_string(),
+                                    filepath: filepath.to_string(),
+                                    mime_type: "text/plain".to_string(),
+                                    size_bytes: 42,
+                                    description: description.to_string(),
+                                },
+                            ),
+                        ),
+                    },
+                )),
+            },
+        )),
+        request_id: "request-id".to_string(),
+        timestamp: None,
+    }
+}
+
+fn remote_start_agent_execution_mode(environment_id: &str) -> api::start_agent::ExecutionMode {
+    api::start_agent::ExecutionMode {
+        mode: Some(api::start_agent::execution_mode::Mode::Remote(
+            api::start_agent::execution_mode::Remote {
+                environment_id: environment_id.to_string(),
+            },
+        )),
+    }
+}
+
+fn build_multiple_choice_question(
+    recommended_option_index: i32,
+) -> api::ask_user_question::Question {
+    api::ask_user_question::Question {
+        question_id: "q1".to_string(),
+        question: "Which option should we prefer?".to_string(),
+        question_type: Some(
+            api::ask_user_question::question::QuestionType::MultipleChoice(
+                api::ask_user_question::MultipleChoice {
+                    is_multiselect: false,
+                    options: vec![
+                        api::ask_user_question::Option {
+                            label: "First".to_string(),
+                        },
+                        api::ask_user_question::Option {
+                            label: "Second".to_string(),
+                        },
+                    ],
+                    recommended_option_index,
+                    supports_other: false,
+                },
+            ),
+        ),
+    }
+}
+
+#[test]
+fn convert_api_question_treats_negative_recommended_index_as_no_recommendation() {
+    let converted = convert_api_question(build_multiple_choice_question(-1))
+        .expect("multiple choice questions should convert");
+
+    let AskUserQuestionType::MultipleChoice { options, .. } = converted.question_type;
+    assert_eq!(options.len(), 2);
+    assert!(options.iter().all(|option| !option.recommended));
+}
+
+#[test]
+fn convert_api_question_uses_zero_based_recommended_index_when_present() {
+    let converted = convert_api_question(build_multiple_choice_question(0))
+        .expect("multiple choice questions should convert");
+
+    let AskUserQuestionType::MultipleChoice { options, .. } = converted.question_type;
+    assert_eq!(options.len(), 2);
+    assert!(options[0].recommended);
+    assert!(!options[1].recommended);
+}
+
+fn extract_start_agent_action(
+    output: MaybeAIAgentOutputMessage,
+) -> (
+    String,
+    String,
+    StartAgentExecutionMode,
+    Option<Vec<LifecycleEventType>>,
+) {
+    let MaybeAIAgentOutputMessage::Message(output_message) = output else {
+        panic!("expected output message");
+    };
+    let AIAgentOutputMessageType::Action(action) = output_message.message else {
+        panic!("expected action output message");
+    };
+    let AIAgentActionType::StartAgent {
+        version: _,
+        name,
+        prompt,
+        execution_mode,
+        lifecycle_subscription,
+    } = action.action
+    else {
+        panic!("expected StartAgent action");
+    };
+    (name, prompt, execution_mode, lifecycle_subscription)
+}
+
+fn extract_upload_artifact_action(output: MaybeAIAgentOutputMessage) -> (String, Option<String>) {
+    let MaybeAIAgentOutputMessage::Message(output_message) = output else {
+        panic!("expected output message");
+    };
+    let AIAgentOutputMessageType::Action(action) = output_message.message else {
+        panic!("expected action output message");
+    };
+    let AIAgentActionType::UploadArtifact(request) = action.action else {
+        panic!("expected UploadArtifact action");
+    };
+    (request.file_path, request.description)
+}
+
+fn extract_file_artifact_created(
+    output: MaybeAIAgentOutputMessage,
+) -> (String, String, Option<String>, i64) {
+    let MaybeAIAgentOutputMessage::Message(output_message) = output else {
+        panic!("expected output message");
+    };
+    let AIAgentOutputMessageType::ArtifactCreated(artifact) = output_message.message else {
+        panic!("expected artifact created output message");
+    };
+    let crate::ai::agent::ArtifactCreatedData::File {
+        filepath,
+        filename,
+        description,
+        size_bytes,
+        ..
+    } = artifact
+    else {
+        panic!("expected file artifact created output message");
+    };
+    (filepath, filename, description, size_bytes)
+}
+
+#[test]
+fn converts_start_agent_tool_call_to_action_with_prompt() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message =
+        start_agent_tool_call_message("Agent 1", "run tests and report failures", None, None);
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+
+    assert_eq!(name, "Agent 1");
+    assert_eq!(prompt, "run tests and report failures");
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::local_with_defaults()
+    );
+    assert_eq!(lifecycle_subscription, None);
+}
+
+#[test]
+fn converts_local_start_agent_v2_without_harness_type_to_defaults() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = start_agent_v2_tool_call_message(
+        "Agent 7",
+        "run in the default local harness",
+        Some(local_start_agent_v2_execution_mode_without_harness()),
+        None,
+    );
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+
+    assert_eq!(name, "Agent 7");
+    assert_eq!(prompt, "run in the default local harness");
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::local_with_defaults()
+    );
+    assert_eq!(lifecycle_subscription, None);
+}
+
+#[test]
+fn converts_upload_artifact_tool_call_to_action() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = upload_artifact_tool_call_message(
+        "/tmp/build/output.log",
+        "Build output for the latest run",
+    );
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (file_path, description) = extract_upload_artifact_action(output);
+
+    assert_eq!(file_path, "/tmp/build/output.log");
+    assert_eq!(
+        description.as_deref(),
+        Some("Build output for the latest run")
+    );
+}
+
+#[test]
+fn converts_file_artifact_created_message_with_filename() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message =
+        file_artifact_created_message("outputs/report.txt", "Build output for the latest run");
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (filepath, filename, description, size_bytes) = extract_file_artifact_created(output);
+
+    assert_eq!(filepath, "outputs/report.txt");
+    assert_eq!(filename, "report.txt");
+    assert_eq!(
+        description.as_deref(),
+        Some("Build output for the latest run")
+    );
+    assert_eq!(size_bytes, 42);
+}
+
+#[test]
+fn converts_start_agent_tool_calls_with_different_prompt_lengths() {
+    let task_id = TaskId::new("task-id".to_string());
+    let partial_message = start_agent_tool_call_message("Agent 2", "run tests", None, None);
+    let updated_message = start_agent_tool_call_message(
+        "Agent 2",
+        "run tests and then summarize failures",
+        None,
+        None,
+    );
+
+    let partial_output = partial_message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("partial conversion should succeed");
+    let updated_output = updated_message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("updated conversion should succeed");
+
+    let (_, partial_prompt, partial_execution_mode, _) = extract_start_agent_action(partial_output);
+    let (_, updated_prompt, updated_execution_mode, _) = extract_start_agent_action(updated_output);
+
+    assert_eq!(partial_prompt, "run tests");
+    assert_eq!(updated_prompt, "run tests and then summarize failures");
+    assert_eq!(
+        partial_execution_mode,
+        StartAgentExecutionMode::local_with_defaults()
+    );
+    assert_eq!(
+        updated_execution_mode,
+        StartAgentExecutionMode::local_with_defaults()
+    );
+}
+
+#[test]
+fn converts_start_agent_with_explicit_empty_lifecycle_subscription() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = start_agent_tool_call_message("Agent 3", "run tests", None, Some(vec![]));
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (_, _, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::local_with_defaults()
+    );
+    assert_eq!(lifecycle_subscription, Some(vec![]));
+}
+
+#[test]
+fn converts_start_agent_with_cancelled_and_blocked_lifecycle_subscription() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = start_agent_tool_call_message(
+        "Agent 4",
+        "wait for approval",
+        None,
+        Some(vec![
+            api::LifecycleEventType::Cancelled as i32,
+            api::LifecycleEventType::Blocked as i32,
+        ]),
+    );
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (_, _, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::local_with_defaults()
+    );
+    assert_eq!(
+        lifecycle_subscription,
+        Some(vec![
+            LifecycleEventType::Cancelled,
+            LifecycleEventType::Blocked
+        ])
+    );
+}
+
+#[test]
+fn converts_remote_start_agent_with_environment_id() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = start_agent_tool_call_message(
+        "Agent 5",
+        "run in the remote environment",
+        Some(remote_start_agent_execution_mode("env-123")),
+        None,
+    );
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+
+    assert_eq!(name, "Agent 5");
+    assert_eq!(prompt, "run in the remote environment");
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::Remote {
+            environment_id: "env-123".to_string(),
+            skill_references: vec![],
+            model_id: String::new(),
+            computer_use_enabled: false,
+            worker_host: String::new(),
+            harness_type: String::new(),
+            title: String::new(),
+        }
+    );
+    assert_eq!(lifecycle_subscription, None);
+}
+
+#[test]
+fn converts_remote_start_agent_v2_with_skill_references() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = start_agent_v2_tool_call_message(
+        "Agent 6",
+        "run in the remote environment",
+        Some(remote_start_agent_v2_execution_mode("env-123")),
+        None,
+    );
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+
+    assert_eq!(name, "Agent 6");
+    assert_eq!(prompt, "run in the remote environment");
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::Remote {
+            environment_id: "env-123".to_string(),
+            skill_references: vec![
+                SkillReference::Path("/tmp/SKILL.md".into()),
+                SkillReference::BundledSkillId("review-comments".to_string()),
+            ],
+            model_id: "gpt-test".to_string(),
+            computer_use_enabled: true,
+            worker_host: "worker-host".to_string(),
+            harness_type: "claude-code".to_string(),
+            title: "Remote child".to_string(),
+        }
+    );
+    assert_eq!(lifecycle_subscription, None);
+}
+
+#[test]
+fn converts_local_start_agent_v2_with_harness_type() {
+    let task_id = TaskId::new("task-id".to_string());
+    let message = start_agent_v2_tool_call_message(
+        "Agent 6",
+        "run in the local claude harness",
+        Some(local_start_agent_v2_execution_mode("claude-code")),
+        None,
+    );
+
+    let output = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("conversion should succeed");
+
+    let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
+
+    assert_eq!(name, "Agent 6");
+    assert_eq!(prompt, "run in the local claude harness");
+    assert_eq!(
+        execution_mode,
+        StartAgentExecutionMode::local_harness("claude-code".to_string())
+    );
+    assert_eq!(lifecycle_subscription, None);
+}
+
+#[test]
+fn transfer_control_tool_call_converts_to_action_message() {
+    let task_id = TaskId::new("task".to_string());
+    let reason = "Please finish the interactive flow".to_string();
+    let message = api::Message {
+        id: "message".to_string(),
+        task_id: "task".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: "tool_call".to_string(),
+            tool: Some(
+                api::message::tool_call::Tool::TransferShellCommandControlToUser(
+                    api::message::tool_call::TransferShellCommandControlToUser {
+                        reason: reason.clone(),
+                    },
+                ),
+            ),
+        })),
+        request_id: "req".to_string(),
+        timestamp: None,
+    };
+
+    let converted = message
+        .to_client_output_message(ConversionParams {
+            task_id: &task_id,
+            current_todo_list: None,
+            active_code_review: None,
+        })
+        .expect("transfer-control conversion should succeed");
+
+    match converted {
+        MaybeAIAgentOutputMessage::Message(output) => match output.message {
+            AIAgentOutputMessageType::Action(action) => {
+                assert_eq!(action.task_id, task_id);
+                assert_eq!(
+                    action.action,
+                    AIAgentActionType::TransferShellCommandControlToUser { reason }
+                );
+                assert!(action.requires_result);
+            }
+            other => panic!("Expected action message, got {other:?}"),
+        },
+        MaybeAIAgentOutputMessage::NoClientRepresentation => {
+            panic!("Expected transfer-control tool call to produce a client action")
+        }
+    }
+}

app/src/ai/agent/api/convert_to.rs

@@ -0,0 +1,977 @@
+//! Conversions from application types to MAA API types.
+
+use ai::agent::convert::ConvertToAPITypeError;
+use anyhow::anyhow;
+use chrono::{DateTime, Local, Timelike};
+use warp_multi_agent_api as api;
+
+use crate::ai::{
+    agent::{
+        AIAgentActionResult, AIAgentActionResultType, AIAgentAttachment, AIAgentContext,
+        AIAgentInput, DriveObjectPayload, MCPContext, PassiveSuggestionResultType,
+        PassiveSuggestionTrigger, RunningCommand, StaticQueryType, Suggestions, UserQueryMode,
+    },
+    block_context::BlockContext,
+};
+
+fn local_datetime_to_timestamp(timestamp: DateTime<Local>) -> prost_types::Timestamp {
+    prost_types::Timestamp {
+        seconds: timestamp.timestamp(),
+        nanos: timestamp.timestamp_subsec_nanos() as i32,
+    }
+}
+
+impl TryFrom<StaticQueryType> for api::request::input::query_with_canned_response::Type {
+    type Error = ConvertToAPITypeError;
+
+    fn try_from(value: StaticQueryType) -> Result<Self, Self::Error> {
+        match value {
+            StaticQueryType::Install => Ok(
+                api::request::input::query_with_canned_response::Type::Install(
+                    api::request::input::query_with_canned_response::Install {},
+                ),
+            ),
+            StaticQueryType::Code => {
+                Ok(api::request::input::query_with_canned_response::Type::Code(
+                    api::request::input::query_with_canned_response::Code {},
+                ))
+            }
+            StaticQueryType::Deploy => Ok(
+                api::request::input::query_with_canned_response::Type::Deploy(
+                    api::request::input::query_with_canned_response::Deploy {},
+                ),
+            ),
+            StaticQueryType::SomethingElse => Ok(
+                api::request::input::query_with_canned_response::Type::SomethingElse(
+                    api::request::input::query_with_canned_response::SomethingElse {},
+                ),
+            ),
+            StaticQueryType::CustomOnboardingRequest => Ok(
+                api::request::input::query_with_canned_response::Type::CustomOnboardingRequest(
+                    api::request::input::query_with_canned_response::CustomOnboardingRequest {},
+                ),
+            ),
+            StaticQueryType::EvaluationSuite => {
+                Err(anyhow::anyhow!("EvaluationSuite StaticQueryType not yet supported").into())
+            }
+        }
+    }
+}
+
+pub(super) fn convert_input(
+    mut inputs: Vec<AIAgentInput>,
+) -> Result<api::request::Input, ConvertToAPITypeError> {
+    if inputs.is_empty() {
+        return Err(anyhow!("Attempted to send multi-agent request with no input").into());
+    }
+    let api_context = inputs
+        .iter()
+        .rev()
+        .find_map(AIAgentInput::context)
+        .map(convert_context);
+
+    let mut api_inputs = vec![];
+    if inputs.len() == 1 {
+        match inputs.pop().expect("Input exists.") {
+            AIAgentInput::UserQuery {
+                query,
+                context,
+                static_query_type: Some(query_type),
+                ..
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::QueryWithCannedResponse(
+                        api::request::input::QueryWithCannedResponse {
+                            query,
+                            r#type: Some(query_type.try_into()?),
+                        },
+                    )),
+                });
+            }
+            AIAgentInput::AutoCodeDiffQuery { query, context } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::AutoCodeDiffQuery(
+                        api::request::input::AutoCodeDiffQuery { query },
+                    )),
+                });
+            }
+            AIAgentInput::ResumeConversation { context } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::ResumeConversation(
+                        api::request::input::ResumeConversation {},
+                    )),
+                });
+            }
+            AIAgentInput::InitProjectRules { context, .. } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::InitProjectRules(
+                        api::request::input::InitProjectRules {},
+                    )),
+                });
+            }
+            AIAgentInput::CreateEnvironment {
+                context,
+                repo_paths,
+                ..
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::CreateEnvironment(
+                        api::request::input::CreateEnvironment { repo_paths },
+                    )),
+                });
+            }
+            AIAgentInput::TriggerPassiveSuggestion {
+                context,
+                attachments,
+                trigger,
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::GeneratePassiveSuggestions(
+                        api::request::input::GeneratePassiveSuggestions {
+                            attachments: attachments
+                                .into_iter()
+                                .map(|attachment| attachment.into())
+                                .collect(),
+                            trigger: Some(trigger.into()),
+                        },
+                    )),
+                });
+            }
+            AIAgentInput::CreateNewProject { query, context } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::CreateNewProject(
+                        api::request::input::CreateNewProject { query },
+                    )),
+                });
+            }
+            AIAgentInput::CloneRepository {
+                clone_repo_url,
+                context,
+                ..
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::CloneRepository(
+                        api::request::input::CloneRepository {
+                            url: clone_repo_url.into_url(),
+                        },
+                    )),
+                });
+            }
+            AIAgentInput::CodeReview {
+                context,
+                review_comments,
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::CodeReview(
+                        api::request::input::CodeReview {
+                            operation: Some(
+                                api::request::input::code_review::Operation::InitialReviewComments(
+                                    api::request::input::code_review::InitialReviewComments {
+                                        review_comments: review_comments
+                                            .comments
+                                            .into_iter()
+                                            .map(Into::into)
+                                            .collect(),
+                                        diff_set: Some(api::DiffSet {
+                                            hunks: review_comments
+                                                .diff_set
+                                                .into_iter()
+                                                .flat_map(|(file_path, hunks)| {
+                                                    hunks.into_iter().map(move |hunk| {
+                                                        hunk.convert_to_api(file_path.clone())
+                                                    })
+                                                })
+                                                .collect(),
+                                            curr_ref: None,
+                                            base_ref: None,
+                                        }),
+                                    },
+                                ),
+                            ),
+                        },
+                    )),
+                });
+            }
+            AIAgentInput::FetchReviewComments { repo_path, context } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::FetchReviewComments(
+                        api::request::input::FetchReviewComments { repo_path },
+                    )),
+                });
+            }
+            AIAgentInput::SummarizeConversation { prompt } => {
+                return Ok(api::request::Input {
+                    context: None,
+                    r#type: Some(api::request::input::Type::SummarizeConversation(
+                        api::request::input::SummarizeConversation {
+                            prompt: prompt.unwrap_or_default(),
+                        },
+                    )),
+                });
+            }
+            AIAgentInput::InvokeSkill {
+                context,
+                skill,
+                user_query,
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::InvokeSkill(
+                        api::request::input::InvokeSkill {
+                            skill: Some(skill.into()),
+                            user_query: user_query.map(|user_query| {
+                                api::request::input::UserQuery {
+                                    query: user_query.query,
+                                    referenced_attachments: user_query
+                                        .referenced_attachments
+                                        .into_iter()
+                                        .map(|(k, attachment)| (k, attachment.into()))
+                                        .collect(),
+                                    mode: None,
+                                    intended_agent: Default::default(),
+                                }
+                            }),
+                        },
+                    )),
+                });
+            }
+            AIAgentInput::StartFromAmbientRunPrompt {
+                ambient_run_id,
+                context,
+                runtime_skill,
+                attachments_dir,
+            } => {
+                return Ok(api::request::Input {
+                    context: Some(convert_context(context.as_ref())),
+                    r#type: Some(api::request::input::Type::StartFromAmbientRunPrompt(
+                        api::request::input::StartFromAmbientRunPrompt {
+                            ambient_run_id,
+                            // Deprecated, we always resolve base_prompt from the stored task config.
+                            runtime_base_prompt: String::new(),
+
+                            runtime_skill: runtime_skill.map(|skill| skill.into()),
+                            attachments_dir: attachments_dir.unwrap_or_default(),
+                        },
+                    )),
+                });
+            }
+            other_input => match convert_input_to_user_input(other_input) {
+                Ok(api_input) => api_inputs.push(api_input),
+                Err(ConvertToAPITypeError::Ignore) => (),
+                Err(e) => return Err(e),
+            },
+        }
+    }
+
+    for input in inputs.into_iter() {
+        match convert_input_to_user_input(input) {
+            Ok(api_input) => api_inputs.push(api_input),
+            Err(ConvertToAPITypeError::Ignore) => continue,
+            Err(e) => return Err(e),
+        }
+    }
+
+    Ok(api::request::Input {
+        context: api_context,
+        r#type: Some(api::request::input::Type::UserInputs(
+            api::request::input::UserInputs {
+                inputs: api_inputs
+                    .into_iter()
+                    .map(|input| api::request::input::user_inputs::UserInput { input: Some(input) })
+                    .collect(),
+            },
+        )),
+    })
+}
+
+fn convert_input_to_user_input(
+    input: AIAgentInput,
+) -> Result<api::request::input::user_inputs::user_input::Input, ConvertToAPITypeError> {
+    match input {
+        AIAgentInput::UserQuery {
+            query,
+            static_query_type: None,
+            referenced_attachments,
+            user_query_mode,
+            running_command: None,
+            intended_agent,
+            ..
+        } => Ok(
+            api::request::input::user_inputs::user_input::Input::UserQuery(
+                api::request::input::UserQuery {
+                    query,
+                    referenced_attachments: referenced_attachments.into_iter().map(|(k, attachment)| (k, attachment.into())).collect(),
+                    mode: Some(user_query_mode.into()),
+                    intended_agent: intended_agent.map(|agent| agent.into()).unwrap_or_default(),
+                },
+            ),
+        ),
+        AIAgentInput::UserQuery {
+            query,
+            static_query_type: None,
+            referenced_attachments,
+            user_query_mode,
+            running_command: Some(RunningCommand{
+                command,
+                block_id,
+                grid_contents: output,
+                cursor,
+                requested_command_id,
+                is_alt_screen_active,
+            }),
+            ..
+        } => {
+            Ok(api::request::input::user_inputs::user_input::Input::CliAgentUserQuery(
+                api::request::input::CliAgentUserQuery {
+                    user_query: Some(api::request::input::UserQuery {
+                            query,
+                            referenced_attachments: referenced_attachments.into_iter().map(|(k, attachment)| (k, attachment.into())).collect(),
+                            mode: Some(user_query_mode.into()),
+                            intended_agent: api::AgentType::Cli.into(),
+                        }),
+                    running_command: Some(api::RunningShellCommand{
+                        command,
+                        snapshot: Some(api::LongRunningShellCommandSnapshot {
+                            output,
+                            cursor,
+                            command_id: block_id.as_str().to_owned(),
+                            is_alt_screen_active,
+                            is_preempted: false,
+                        }),
+                    }),
+                    run_shell_command_tool_call_id: requested_command_id.map(|id| id.to_string()).unwrap_or_default(),
+                }
+            ))
+        }
+        AIAgentInput::ActionResult { result, .. } => result.try_into(),
+        AIAgentInput::MessagesReceivedFromAgents { messages } => Ok(
+            api::request::input::user_inputs::user_input::Input::MessagesReceivedFromAgents(
+                api::request::input::user_inputs::MessagesReceivedFromAgents {
+                    messages: messages
+                        .into_iter()
+                        .map(
+                            |msg| api::request::input::user_inputs::messages_received_from_agents::ReceivedMessage {
+                                message_id: msg.message_id,
+                                sender_agent_id: msg.sender_agent_id,
+                                addresses: msg.addresses,
+                                subject: msg.subject,
+                                message_body: msg.message_body,
+                            },
+                        )
+                        .collect(),
+                },
+            ),
+        ),
+        AIAgentInput::EventsFromAgents { events } => Ok(
+            api::request::input::user_inputs::user_input::Input::EventsFromAgents(
+                api::request::input::user_inputs::EventsFromAgents {
+                    agent_events: events,
+                },
+            ),
+        ),
+        AIAgentInput::PassiveSuggestionResult {
+            trigger,
+            suggestion,
+            ..
+        } => {
+            let api_trigger = match trigger {
+                Some(PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger)) => Some(
+                    api::passive_suggestion_result_type::Trigger::ExecutedShellCommand(
+                        (*shell_trigger.executed_shell_command).into(),
+                    ),
+                ),
+                Some(PassiveSuggestionTrigger::AgentResponseCompleted { .. }) => Some(
+                    api::passive_suggestion_result_type::Trigger::AgentResponseCompleted(
+                        api::passive_suggestion_result_type::AgentResponseCompleted {},
+                    ),
+                ),
+                _ => None,
+            };
+            let api_suggestion = match suggestion {
+                PassiveSuggestionResultType::Prompt { prompt } => Some(
+                    api::passive_suggestion_result_type::Suggestion::Prompt(
+                        api::passive_suggestion_result_type::Prompt { prompt },
+                    ),
+                ),
+                PassiveSuggestionResultType::CodeDiff {
+                    diffs,
+                    summary,
+                    accepted,
+                } => Some(
+                    api::passive_suggestion_result_type::Suggestion::CodeDiff(
+                        api::passive_suggestion_result_type::CodeDiff {
+                            diffs: diffs
+                                .into_iter()
+                                .map(|d| api::passive_suggestion_result_type::code_diff::Diff {
+                                    file_path: d.file_path,
+                                    search: d.search,
+                                    replace: d.replace,
+                                })
+                                .collect(),
+                            summary,
+                            accepted,
+                        },
+                    ),
+                ),
+            };
+            Ok(
+                api::request::input::user_inputs::user_input::Input::PassiveSuggestionResult(
+                    api::request::input::user_inputs::PassiveSuggestionResultInput {
+                        result: Some(api::PassiveSuggestionResultType {
+                            trigger: api_trigger,
+                            suggestion: api_suggestion,
+                        }),
+                    },
+                ),
+            )
+        }
+        AIAgentInput::ResumeConversation { .. } => Err(ConvertToAPITypeError::Ignore),
+        AIAgentInput::InitProjectRules { .. } => Err(ConvertToAPITypeError::Ignore),
+        AIAgentInput::CodeReview { .. } => Err(ConvertToAPITypeError::Ignore),
+        AIAgentInput::FetchReviewComments { .. } => Err(ConvertToAPITypeError::Ignore),
+        AIAgentInput::CreateEnvironment { .. } => Err(ConvertToAPITypeError::Ignore),
+        AIAgentInput::InvokeSkill { .. } => Err(ConvertToAPITypeError::Ignore),
+        invalid_input => Err(anyhow!(
+            "Cannot convert non user query or action result input into API UserInput: {invalid_input:?}"
+        ).into()),
+    }
+}
+
+impl From<PassiveSuggestionTrigger> for api::request::input::generate_passive_suggestions::Trigger {
+    fn from(value: PassiveSuggestionTrigger) -> Self {
+        match value {
+            PassiveSuggestionTrigger::FilesChanged => {
+                api::request::input::generate_passive_suggestions::Trigger::FilesChanged(())
+            }
+            PassiveSuggestionTrigger::CommandRun => {
+                api::request::input::generate_passive_suggestions::Trigger::CommandRun(())
+            }
+            PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger) => {
+                api::request::input::generate_passive_suggestions::Trigger::ShellCommandCompleted(
+                    api::request::input::generate_passive_suggestions::ShellCommandCompleted {
+                        executed_shell_command: Some(
+                            (*shell_trigger.executed_shell_command).into(),
+                        ),
+                        relevant_files: shell_trigger
+                            .relevant_files
+                            .into_iter()
+                            .flat_map(|file| Vec::<api::AnyFileContent>::from(file).into_iter())
+                            .collect(),
+                    },
+                )
+            }
+            PassiveSuggestionTrigger::AgentResponseCompleted { .. } => {
+                api::request::input::generate_passive_suggestions::Trigger::AgentResponseCompleted(
+                    api::request::input::generate_passive_suggestions::AgentResponseCompleted {},
+                )
+            }
+        }
+    }
+}
+
+impl From<UserQueryMode> for warp_multi_agent_api::UserQueryMode {
+    fn from(value: UserQueryMode) -> Self {
+        match value {
+            UserQueryMode::Normal => warp_multi_agent_api::UserQueryMode { r#type: None },
+            UserQueryMode::Plan => warp_multi_agent_api::UserQueryMode {
+                r#type: Some(warp_multi_agent_api::user_query_mode::Type::Plan(())),
+            },
+            UserQueryMode::Orchestrate => warp_multi_agent_api::UserQueryMode {
+                r#type: Some(warp_multi_agent_api::user_query_mode::Type::Orchestrate(())),
+            },
+        }
+    }
+}
+
+impl From<AIAgentAttachment> for api::Attachment {
+    fn from(attachment: AIAgentAttachment) -> Self {
+        match attachment {
+            AIAgentAttachment::PlainText(text) => api::Attachment {
+                value: Some(api::attachment::Value::PlainText(text)),
+            },
+            AIAgentAttachment::Block(block) => api::Attachment {
+                value: Some(api::attachment::Value::ExecutedShellCommand(block.into())),
+            },
+            AIAgentAttachment::DriveObject { uid, payload } => api::Attachment {
+                value: Some(api::attachment::Value::DriveObject(api::DriveObject {
+                    uid,
+                    object_payload: payload.map(|p| match p {
+                        DriveObjectPayload::Workflow {
+                            name,
+                            description,
+                            command,
+                        } => api::drive_object::ObjectPayload::Workflow(api::Workflow {
+                            name,
+                            description,
+                            command,
+                        }),
+                        DriveObjectPayload::Notebook { title, content } => {
+                            api::drive_object::ObjectPayload::Notebook(api::Notebook {
+                                title,
+                                content,
+                            })
+                        }
+                        DriveObjectPayload::GenericStringObject {
+                            payload,
+                            object_type,
+                        } => api::drive_object::ObjectPayload::GenericStringObject(
+                            api::GenericStringObject {
+                                payload,
+                                object_type,
+                            },
+                        ),
+                    }),
+                })),
+            },
+            #[allow(deprecated)]
+            AIAgentAttachment::DiffHunk {
+                file_path,
+                line_range,
+                diff_content,
+                lines_added,
+                lines_removed,
+                current,
+                base,
+            } => api::Attachment {
+                value: Some(api::attachment::Value::DiffHunk(api::DiffHunk {
+                    file_path,
+                    line_range: Some(api::FileContentLineRange {
+                        start: line_range.start.as_usize() as u32,
+                        end: line_range.end.as_usize() as u32,
+                    }),
+                    diff_content,
+                    lines_added,
+                    lines_removed,
+                    current: current.map(Into::into),
+                    base: Some(base.into()),
+                })),
+            },
+            AIAgentAttachment::DocumentContent {
+                document_id,
+                content,
+                line_range,
+                // TODO: Add attachment source to API
+                ..
+            } => api::Attachment {
+                value: Some(api::attachment::Value::DocumentContent(
+                    api::DocumentContent {
+                        document_id,
+                        content,
+                        line_range: line_range.map(|range| api::FileContentLineRange {
+                            start: range.start.as_usize() as u32,
+                            end: range.end.as_usize() as u32,
+                        }),
+                    },
+                )),
+            },
+            AIAgentAttachment::DiffSet {
+                file_diffs,
+                current,
+                base,
+            } => api::Attachment {
+                value: Some(api::attachment::Value::DiffSet(api::DiffSet {
+                    hunks: file_diffs
+                        .into_iter()
+                        .flat_map(|(file_path, hunks)| {
+                            hunks
+                                .into_iter()
+                                .map(move |hunk| hunk.convert_to_api(file_path.clone()))
+                        })
+                        .collect(),
+                    curr_ref: current.map(Into::into),
+                    base_ref: Some(base.into()),
+                })),
+            },
+            AIAgentAttachment::FilePathReference { file_path, .. } => api::Attachment {
+                value: Some(api::attachment::Value::FilePathReference(
+                    api::FilePathReference { file_path },
+                )),
+            },
+        }
+    }
+}
+
+impl TryFrom<AIAgentActionResult> for api::request::input::user_inputs::user_input::Input {
+    type Error = ConvertToAPITypeError;
+
+    fn try_from(action_result: AIAgentActionResult) -> Result<Self, Self::Error> {
+        let result = match action_result.result {
+            AIAgentActionResultType::RequestCommandOutput(request_command_result) => {
+                Some(request_command_result.try_into()?)
+            }
+            AIAgentActionResultType::WriteToLongRunningShellCommand(result) => {
+                Some(result.try_into()?)
+            }
+            AIAgentActionResultType::ReadFiles(read_files_result) => {
+                Some(read_files_result.try_into()?)
+            }
+            AIAgentActionResultType::UploadArtifact(upload_artifact_result) => {
+                Some(upload_artifact_result.try_into()?)
+            }
+            AIAgentActionResultType::SearchCodebase(search_codebase_result) => {
+                Some(search_codebase_result.try_into()?)
+            }
+            AIAgentActionResultType::RequestFileEdits(request_file_edits_result) => {
+                Some(request_file_edits_result.try_into()?)
+            }
+            AIAgentActionResultType::Grep(grep_result) => Some(grep_result.try_into()?),
+            AIAgentActionResultType::FileGlob(file_glob_result) => {
+                Some(file_glob_result.try_into()?)
+            }
+            AIAgentActionResultType::FileGlobV2(file_glob_result) => {
+                Some(file_glob_result.try_into()?)
+            }
+            AIAgentActionResultType::ReadMCPResource(read_mcp_resource_result) => {
+                Some(read_mcp_resource_result.try_into()?)
+            }
+            AIAgentActionResultType::CallMCPTool(call_mcp_tool_result) => {
+                Some(call_mcp_tool_result.try_into()?)
+            }
+            AIAgentActionResultType::ReadSkill(read_skill_result) => {
+                Some(read_skill_result.try_into()?)
+            }
+            AIAgentActionResultType::SuggestNewConversation(suggest_new_conversation_result) => {
+                Some(suggest_new_conversation_result.try_into()?)
+            }
+            AIAgentActionResultType::SuggestPrompt(suggest_prompt_result) => {
+                Some(suggest_prompt_result.try_into()?)
+            }
+            AIAgentActionResultType::OpenCodeReview => Some(
+                warp_multi_agent_api::request::input::tool_call_result::Result::OpenCodeReview(
+                    warp_multi_agent_api::OpenCodeReviewResult {},
+                ),
+            ),
+            AIAgentActionResultType::InsertReviewComments(insert_review_comments_result) => {
+                Some(insert_review_comments_result.try_into()?)
+            }
+            AIAgentActionResultType::InitProject => Some(
+                warp_multi_agent_api::request::input::tool_call_result::Result::InitProject(
+                    warp_multi_agent_api::InitProjectResult {},
+                ),
+            ),
+            AIAgentActionResultType::ReadDocuments(read_documents_result) => {
+                Some(read_documents_result.try_into()?)
+            }
+            AIAgentActionResultType::EditDocuments(edit_documents_result) => {
+                Some(edit_documents_result.try_into()?)
+            }
+            AIAgentActionResultType::CreateDocuments(create_documents_result) => {
+                Some(create_documents_result.try_into()?)
+            }
+            AIAgentActionResultType::ReadShellCommandOutput(read_shell_command_output_result) => {
+                Some(read_shell_command_output_result.try_into()?)
+            }
+            AIAgentActionResultType::UseComputer(use_computer_result) => {
+                Some(use_computer_result.try_into()?)
+            }
+            AIAgentActionResultType::RequestComputerUse(request_computer_use_result) => {
+                Some(request_computer_use_result.try_into()?)
+            }
+            AIAgentActionResultType::FetchConversation(fetch_conversation_result) => {
+                Some(fetch_conversation_result.try_into()?)
+            }
+            AIAgentActionResultType::StartAgent(start_agent_result) => {
+                Some(start_agent_result.into())
+            }
+            AIAgentActionResultType::SendMessageToAgent(send_message_result) => {
+                Some(send_message_result.into())
+            }
+            AIAgentActionResultType::TransferShellCommandControlToUser(transfer_control_result) => {
+                Some(transfer_control_result.try_into()?)
+            }
+            AIAgentActionResultType::AskUserQuestion(ask_user_question_result) => {
+                Some(ask_user_question_result.into())
+            }
+        };
+        Ok(
+            api::request::input::user_inputs::user_input::Input::ToolCallResult(
+                api::request::input::ToolCallResult {
+                    tool_call_id: action_result.id.into(),
+                    result,
+                },
+            ),
+        )
+    }
+}
+
+fn convert_context(context: &[AIAgentContext]) -> api::InputContext {
+    let mut api_context = api::InputContext::default();
+    for context in context.iter().cloned() {
+        match context {
+            AIAgentContext::Block(block) => {
+                #[allow(deprecated)]
+                api_context.executed_shell_commands.push((*block).into());
+            }
+            AIAgentContext::Directory {
+                pwd,
+                home_dir,
+                are_file_symbols_indexed,
+            } => {
+                api_context.directory = Some(api::input_context::Directory {
+                    pwd: pwd.unwrap_or_default(),
+                    home: home_dir.unwrap_or_default(),
+                    pwd_file_symbols_indexed: are_file_symbols_indexed,
+                });
+            }
+            AIAgentContext::SelectedText(text) => {
+                api_context
+                    .selected_text
+                    .push(api::input_context::SelectedText { text });
+            }
+            AIAgentContext::ExecutionEnvironment(execution_ctx) => {
+                api_context.shell = Some(api::input_context::Shell {
+                    name: execution_ctx.shell_name,
+                    version: execution_ctx.shell_version.unwrap_or_default(),
+                });
+
+                if execution_ctx.os.category.is_none() && execution_ctx.os.distribution.is_none() {
+                    continue;
+                }
+                api_context.operating_system = Some(api::input_context::OperatingSystem {
+                    platform: execution_ctx.os.category.unwrap_or_default(),
+                    distribution: execution_ctx.os.distribution.unwrap_or_default(),
+                });
+            }
+            AIAgentContext::CurrentTime { current_time } => {
+                let utc_time = current_time.to_utc();
+                api_context.current_time = Some(prost_types::Timestamp {
+                    seconds: utc_time.timestamp(),
+                    nanos: utc_time.nanosecond() as i32,
+                });
+            }
+            AIAgentContext::Image(image_context) => {
+                api_context.images.push(api::input_context::Image {
+                    data: image_context.data.into(),
+                    mime_type: image_context.mime_type,
+                });
+            }
+            AIAgentContext::Codebase { path, name } => {
+                api_context
+                    .codebases
+                    .push(api::input_context::Codebase { path, name });
+            }
+            AIAgentContext::ProjectRules {
+                root_path,
+                active_rules,
+                additional_rule_paths,
+            } => {
+                api_context
+                    .project_rules
+                    .push(api::input_context::ProjectRules {
+                        root_path,
+                        active_rule_files: active_rules
+                            .into_iter()
+                            .flat_map(|rule| {
+                                let file_contents: Vec<api::FileContent> = rule.into();
+                                file_contents.into_iter()
+                            })
+                            .collect(),
+                        additional_rule_file_paths: additional_rule_paths,
+                    });
+            }
+            AIAgentContext::File(file_context) => {
+                let contents: Vec<api::FileContent> = file_context.into();
+
+                for content in contents {
+                    api_context.files.push(api::input_context::File {
+                        content: Some(content),
+                    });
+                }
+            }
+            AIAgentContext::Git { head, branch } => {
+                api_context.git = Some(api::input_context::Git {
+                    head,
+                    branch: branch.unwrap_or_default(),
+                });
+            }
+            AIAgentContext::Skills { skills } => {
+                api_context.updated_skills_context = Some(api::input_context::SkillsContext {
+                    available_skills: skills
+                        .into_iter()
+                        .map(|skill| api::SkillDescriptor {
+                            skill_reference: Some(skill.reference.into()),
+                            name: skill.name,
+                            description: skill.description,
+                            provider: Some(skill.provider.into()),
+                            scope: Some(skill.scope.into()),
+                        })
+                        .collect(),
+                });
+            }
+        }
+    }
+    api_context
+}
+
+impl From<Suggestions> for api::Suggestions {
+    fn from(value: Suggestions) -> Self {
+        Self {
+            rules: value
+                .rules
+                .into_iter()
+                .map(|rule| api::SuggestedRule {
+                    name: rule.name,
+                    content: rule.content,
+                    logging_id: rule.logging_id.to_string(),
+                })
+                .collect(),
+            workflows: value
+                .agent_mode_workflows
+                .into_iter()
+                .map(|workflow| api::SuggestedAgentModeWorkflow {
+                    name: workflow.name,
+                    prompt: workflow.prompt,
+                    logging_id: workflow.logging_id.to_string(),
+                })
+                .collect(),
+        }
+    }
+}
+
+// Convert rmcp resource to proto format.
+fn convert_mcp_resource(resource: rmcp::model::Resource) -> api::request::mcp_context::McpResource {
+    let rmcp::model::RawResource {
+        uri,
+        name,
+        description,
+        mime_type,
+        ..
+    } = resource.raw;
+    api::request::mcp_context::McpResource {
+        uri,
+        name,
+        description: description.unwrap_or_default(),
+        mime_type: mime_type.unwrap_or_default(),
+    }
+}
+
+// Convert rmcp tool to proto format, skipping tools with invalid schemas.
+fn convert_mcp_tool(tool: rmcp::model::Tool) -> Option<api::request::mcp_context::McpTool> {
+    let Ok(prost_types::Value {
+        kind: Some(prost_types::value::Kind::StructValue(input_schema)),
+    }) = serde_json_to_prost(tool.input_schema.as_ref().clone().into())
+    else {
+        return None;
+    };
+
+    Some(api::request::mcp_context::McpTool {
+        name: tool.name.to_string(),
+        description: tool.description.map(|d| d.to_string()).unwrap_or_default(),
+        input_schema: Some(input_schema),
+    })
+}
+
+impl From<MCPContext> for api::request::McpContext {
+    #[allow(deprecated)]
+    fn from(value: MCPContext) -> Self {
+        // Check if we're using the old flat structure (no servers)
+        // or the new grouped structure (servers populated)
+        if value.servers.is_empty() {
+            // Old behavior: use deprecated flat resources and tools lists
+            api::request::McpContext {
+                #[allow(deprecated)]
+                resources: value
+                    .resources
+                    .into_iter()
+                    .map(convert_mcp_resource)
+                    .collect(),
+                #[allow(deprecated)]
+                tools: value
+                    .tools
+                    .into_iter()
+                    .filter_map(convert_mcp_tool)
+                    .collect(),
+                servers: vec![], // Empty for old behavior
+            }
+        } else {
+            // New behavior: group by server
+            let servers: Vec<_> = value
+                .servers
+                .into_iter()
+                .map(|server| api::request::mcp_context::McpServer {
+                    id: server.id,
+                    name: server.name,
+                    description: server.description,
+                    resources: server
+                        .resources
+                        .into_iter()
+                        .map(convert_mcp_resource)
+                        .collect(),
+                    tools: server
+                        .tools
+                        .into_iter()
+                        .filter_map(convert_mcp_tool)
+                        .collect(),
+                })
+                .collect();
+
+            api::request::McpContext {
+                #[allow(deprecated)]
+                resources: vec![], // Empty - everything is grouped by server
+                #[allow(deprecated)]
+                tools: vec![], // Empty - everything is grouped by server
+                servers,
+            }
+        }
+    }
+}
+
+impl From<BlockContext> for api::ExecutedShellCommand {
+    fn from(block: BlockContext) -> Self {
+        api::ExecutedShellCommand {
+            command: block.command,
+            output: block.output,
+            exit_code: block.exit_code.value(),
+            command_id: block.id.into(),
+            is_auto_attached: block.is_auto_attached,
+            started_ts: block.started_ts.map(local_datetime_to_timestamp),
+            finished_ts: block.finished_ts.map(local_datetime_to_timestamp),
+        }
+    }
+}
+
+/// Trys to convert a [`serde_json::Value`] to a [`prost_types::Value`].
+#[cfg_attr(target_family = "wasm", allow(dead_code))]
+fn serde_json_to_prost(value: serde_json::Value) -> Result<prost_types::Value, String> {
+    use prost_types::value::Kind::*;
+    use serde_json::Value::*;
+    use std::collections::BTreeMap;
+
+    Ok(prost_types::Value {
+        kind: Some(match value {
+            Null => NullValue(0),
+            Bool(v) => BoolValue(v),
+            Number(n) => NumberValue(
+                n.as_f64()
+                    .ok_or_else(|| format!("float {n} is not valid JSON number"))?,
+            ),
+            String(s) => StringValue(s),
+            Array(a) => ListValue(prost_types::ListValue {
+                values: a
+                    .into_iter()
+                    .map(serde_json_to_prost)
+                    .collect::<Result<Vec<_>, std::string::String>>()?,
+            }),
+            Object(v) => StructValue(prost_types::Struct {
+                fields: v
+                    .into_iter()
+                    .map(|(k, v)| serde_json_to_prost(v).map(|v| (k, v)))
+                    .collect::<Result<BTreeMap<_, _>, std::string::String>>()?,
+            }),
+        }),
+    })
+}
+
+#[cfg(test)]
+#[path = "convert_to_tests.rs"]
+mod tests;

app/src/ai/agent/api/convert_to_tests.rs

@@ -0,0 +1,89 @@
+use crate::ai::agent::task::TaskId;
+use crate::ai::agent::{
+    AIAgentActionResult, AIAgentActionResultType, TransferShellCommandControlToUserResult,
+};
+use crate::terminal::model::block::BlockId;
+use warp_core::command::ExitCode;
+use warp_multi_agent_api as api;
+
+#[test]
+fn transfer_control_snapshot_result_converts_to_tool_call_result_input() {
+    let block_id = BlockId::default();
+    let input =
+        api::request::input::user_inputs::user_input::Input::try_from(AIAgentActionResult {
+            id: "tool_call".to_string().into(),
+            task_id: TaskId::new("task".to_string()),
+            result: AIAgentActionResultType::TransferShellCommandControlToUser(
+                TransferShellCommandControlToUserResult::Snapshot {
+                    block_id: block_id.clone(),
+                    grid_contents: "snapshot".to_string(),
+                    cursor: "<|cursor|>".to_string(),
+                    is_alt_screen_active: false,
+                    is_preempted: false,
+                },
+            ),
+        })
+        .unwrap();
+
+    match input {
+        api::request::input::user_inputs::user_input::Input::ToolCallResult(result) => {
+            assert_eq!(result.tool_call_id, "tool_call");
+            match result.result {
+                Some(api::request::input::tool_call_result::Result::TransferShellCommandControlToUser(
+                    api_result,
+                )) => match api_result.result {
+                    Some(
+                        api::transfer_shell_command_control_to_user_result::Result::LongRunningCommandSnapshot(snapshot),
+                    ) => {
+                        assert_eq!(snapshot.command_id, block_id.to_string());
+                        assert_eq!(snapshot.output, "snapshot");
+                        assert_eq!(snapshot.cursor, "<|cursor|>");
+                    }
+                    other => panic!("Expected snapshot result, got {other:?}"),
+                },
+                other => panic!("Expected transfer-control tool call result, got {other:?}"),
+            }
+        }
+        other => panic!("Expected tool-call-result input, got {other:?}"),
+    }
+}
+
+#[test]
+fn transfer_control_finished_result_converts_to_tool_call_result_input() {
+    let block_id = BlockId::default();
+    let input =
+        api::request::input::user_inputs::user_input::Input::try_from(AIAgentActionResult {
+            id: "tool_call".to_string().into(),
+            task_id: TaskId::new("task".to_string()),
+            result: AIAgentActionResultType::TransferShellCommandControlToUser(
+                TransferShellCommandControlToUserResult::CommandFinished {
+                    block_id: block_id.clone(),
+                    output: "done".to_string(),
+                    exit_code: ExitCode::from(17),
+                },
+            ),
+        })
+        .unwrap();
+
+    match input {
+        api::request::input::user_inputs::user_input::Input::ToolCallResult(result) => {
+            assert_eq!(result.tool_call_id, "tool_call");
+            match result.result {
+                Some(api::request::input::tool_call_result::Result::TransferShellCommandControlToUser(
+                    api_result,
+                )) => match api_result.result {
+                    Some(
+                        api::transfer_shell_command_control_to_user_result::Result::CommandFinished(finished),
+                    ) => {
+                        assert_eq!(finished.command_id, block_id.to_string());
+                        assert_eq!(finished.output, "done");
+                        assert_eq!(finished.exit_code, 17);
+                    }
+                    other => panic!("Expected command-finished result, got {other:?}"),
+                },
+                other => panic!("Expected transfer-control tool call result, got {other:?}"),
+            }
+        }
+        other => panic!("Expected tool-call-result input, got {other:?}"),
+    }
+}

app/src/ai/agent/api/impl.rs

@@ -0,0 +1,253 @@
+use std::{collections::HashMap, sync::Arc};
+
+use crate::{ai::agent::redaction, terminal::model::session::SessionType};
+use futures_util::StreamExt;
+use warp_core::features::FeatureFlag;
+use warp_multi_agent_api as api;
+
+use crate::server::server_api::ServerApi;
+
+use super::{convert_to::convert_input, ConvertToAPITypeError, RequestParams, ResponseStream};
+
+pub async fn generate_multi_agent_output(
+    server_api: Arc<ServerApi>,
+    mut params: RequestParams,
+    cancellation_rx: futures::channel::oneshot::Receiver<()>,
+) -> Result<ResponseStream, ConvertToAPITypeError> {
+    let supported_tools = params
+        .supported_tools_override
+        .take()
+        .unwrap_or_else(|| get_supported_tools(&params));
+    let supported_cli_agent_tools = get_supported_cli_agent_tools(&params);
+    let mut logging_metadata = HashMap::new();
+    if let Some(metadata) = params.metadata {
+        logging_metadata.insert(
+            "is_autodetected_user_query".to_owned(),
+            prost_types::Value {
+                kind: Some(prost_types::value::Kind::BoolValue(
+                    metadata.is_autodetected_user_query,
+                )),
+            },
+        );
+        logging_metadata.insert(
+            "entrypoint".to_owned(),
+            prost_types::Value {
+                kind: Some(prost_types::value::Kind::StringValue(
+                    metadata.entrypoint.entrypoint(),
+                )),
+            },
+        );
+        logging_metadata.insert(
+            "is_auto_resume_after_error".to_owned(),
+            prost_types::Value {
+                kind: Some(prost_types::value::Kind::BoolValue(
+                    metadata.is_auto_resume_after_error,
+                )),
+            },
+        );
+    }
+
+    if params.should_redact_secrets {
+        redaction::redact_inputs(&mut params.input);
+    }
+
+    let mut api_keys = params.api_keys;
+    if let Some(api_keys) = &mut api_keys {
+        api_keys.allow_use_of_warp_credits = params.allow_use_of_warp_credits_with_byok;
+    }
+
+    let request = api::Request {
+        task_context: Some(api::request::TaskContext {
+            tasks: params.tasks,
+        }),
+        input: Some(convert_input(params.input)?),
+        settings: Some(api::request::Settings {
+            model_config: Some(api::request::settings::ModelConfig {
+                base: params.model.into(),
+                cli_agent: params.cli_agent_model.into(),
+                computer_use_agent: params.computer_use_model.into(),
+                ..Default::default()
+            }),
+            rules_enabled: params.is_memory_enabled,
+            warp_drive_context_enabled: params.warp_drive_context_enabled,
+            web_context_retrieval_enabled: true,
+            supports_parallel_tool_calls: true,
+            use_anthropic_text_editor_tools: false,
+            planning_enabled: params.planning_enabled,
+            supports_create_files: true,
+            supported_tools: supported_tools.into_iter().map(Into::into).collect(),
+            supports_long_running_commands: true,
+            should_preserve_file_content_in_history: true,
+            supports_todos_ui: true,
+            supports_linked_code_blocks: FeatureFlag::LinkedCodeBlocks.is_enabled(),
+            supports_started_child_task_message: true,
+            supports_suggest_prompt: true,
+            supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(),
+            supports_reasoning_message: true,
+            api_keys,
+            autonomy_level: params.autonomy_level.into(),
+            isolation_level: params.isolation_level.into(),
+            web_search_enabled: params.web_search_enabled,
+            supported_cli_agent_tools: supported_cli_agent_tools
+                .into_iter()
+                .map(Into::into)
+                .collect(),
+            supports_v4a_file_diffs: FeatureFlag::V4AFileDiffs.is_enabled(),
+            supports_summarization_via_message_replacement:
+                FeatureFlag::SummarizationViaMessageReplacement.is_enabled(),
+            supports_bundled_skills: FeatureFlag::BundledSkills.is_enabled(),
+            supports_research_agent: params.research_agent_enabled,
+            supports_orchestration_v2: FeatureFlag::OrchestrationV2.is_enabled(),
+        }),
+        metadata: Some(api::request::Metadata {
+            logging: logging_metadata,
+            conversation_id: params
+                .conversation_token
+                .as_ref()
+                .map(|token| token.as_str().to_string())
+                .unwrap_or_default(),
+            ambient_agent_task_id: params
+                .ambient_agent_task_id
+                .map(|id| id.to_string())
+                .unwrap_or_default(),
+            forked_from_conversation_id: if params.conversation_token.is_none() {
+                // We only include this param on our initial request to the server
+                // (when the forked conversation has not been asigned a new id yet).
+                params
+                    .forked_from_conversation_token
+                    .map(|token| token.as_str().to_string())
+                    .unwrap_or_default()
+            } else {
+                String::new()
+            },
+            parent_agent_id: params.parent_agent_id.unwrap_or_default(),
+            agent_name: params.agent_name.unwrap_or_default(),
+        }),
+        existing_suggestions: params
+            .existing_suggestions
+            .map(|suggestions| suggestions.into()),
+        mcp_context: params.mcp_context.map(Into::into),
+    };
+
+    let response_stream = server_api.generate_multi_agent_output(&request).await;
+    match response_stream {
+        Ok(stream) => {
+            let output_stream = stream.take_until(cancellation_rx);
+            Ok(Box::pin(output_stream))
+        }
+        Err(e) => {
+            let (tx, rx) = async_channel::unbounded();
+            let _ = tx.send(Err(e)).await;
+            Ok(Box::pin(rx))
+        }
+    }
+}
+
+fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
+    let mut supported_tools = vec![
+        api::ToolType::Grep,
+        api::ToolType::FileGlob,
+        api::ToolType::FileGlobV2,
+        api::ToolType::ReadMcpResource,
+        api::ToolType::CallMcpTool,
+        api::ToolType::InitProject,
+        api::ToolType::OpenCodeReview,
+        api::ToolType::RunShellCommand,
+        api::ToolType::SuggestNewConversation,
+        api::ToolType::Subagent,
+        api::ToolType::WriteToLongRunningShellCommand,
+        api::ToolType::ReadShellCommandOutput,
+        api::ToolType::ReadDocuments,
+        api::ToolType::CreateDocuments,
+        api::ToolType::EditDocuments,
+        api::ToolType::SuggestPrompt,
+    ];
+
+    if FeatureFlag::ConversationsAsContext.is_enabled() {
+        supported_tools.push(api::ToolType::FetchConversation);
+    }
+
+    match params.session_context.session_type() {
+        None | Some(SessionType::Local) => {
+            supported_tools.extend(&[
+                api::ToolType::ReadFiles,
+                api::ToolType::ApplyFileDiffs,
+                api::ToolType::SearchCodebase,
+            ]);
+
+            if FeatureFlag::ArtifactCommand.is_enabled() {
+                supported_tools.push(api::ToolType::UploadFileArtifact);
+            }
+        }
+        Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
+            // Remote session with a known host — enable tools that route
+            // through RemoteServerClient. The host_id is only populated
+            // after a successful connection handshake, so its presence is a
+            // sufficient proxy for client availability.
+            // SearchCodebase remains disabled (follow-up work).
+            supported_tools.extend(&[api::ToolType::ReadFiles, api::ToolType::ApplyFileDiffs]);
+        }
+        Some(SessionType::WarpifiedRemote { host_id: None }) => {
+            // Feature flag off or not yet connected — no remote tools.
+        }
+    }
+
+    if FeatureFlag::AgentModeComputerUse.is_enabled() && params.computer_use_enabled {
+        supported_tools.extend(&[api::ToolType::UseComputer]);
+        supported_tools.extend(&[api::ToolType::RequestComputerUse])
+    }
+
+    if FeatureFlag::PRCommentsSlashCommand.is_enabled() {
+        supported_tools.push(api::ToolType::InsertReviewComments);
+    }
+
+    if FeatureFlag::ListSkills.is_enabled() {
+        supported_tools.push(api::ToolType::ReadSkill);
+    }
+
+    if params.orchestration_enabled {
+        supported_tools.push(if FeatureFlag::OrchestrationV2.is_enabled() {
+            api::ToolType::StartAgentV2
+        } else {
+            api::ToolType::StartAgent
+        });
+        supported_tools.push(api::ToolType::SendMessageToAgent);
+    }
+
+    if FeatureFlag::AskUserQuestion.is_enabled() && params.ask_user_question_enabled {
+        supported_tools.push(api::ToolType::AskUserQuestion);
+    }
+
+    supported_tools
+}
+
+fn get_supported_cli_agent_tools(params: &RequestParams) -> Vec<api::ToolType> {
+    let mut supported_cli_agent_tools = vec![
+        api::ToolType::WriteToLongRunningShellCommand,
+        api::ToolType::ReadShellCommandOutput,
+        api::ToolType::Grep,
+        api::ToolType::FileGlob,
+        api::ToolType::FileGlobV2,
+    ];
+
+    if FeatureFlag::TransferControlTool.is_enabled() {
+        supported_cli_agent_tools.push(api::ToolType::TransferShellCommandControlToUser);
+    }
+
+    match params.session_context.session_type() {
+        None | Some(SessionType::Local) => {
+            supported_cli_agent_tools
+                .extend(&[api::ToolType::ReadFiles, api::ToolType::SearchCodebase]);
+        }
+        Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
+            supported_cli_agent_tools.push(api::ToolType::ReadFiles);
+        }
+        Some(SessionType::WarpifiedRemote { host_id: None }) => {}
+    }
+
+    supported_cli_agent_tools
+}
+
+#[cfg(test)]
+#[path = "impl_tests.rs"]
+mod tests;

app/src/ai/agent/api/impl_tests.rs

@@ -0,0 +1,81 @@
+use crate::ai::agent::api::RequestParams;
+use crate::ai::blocklist::SessionContext;
+use crate::ai::llms::LLMId;
+use warp_core::features::FeatureFlag;
+use warp_multi_agent_api as api;
+
+use super::get_supported_tools;
+
+fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool) -> RequestParams {
+    let model = LLMId::from("test-model");
+
+    RequestParams {
+        input: vec![],
+        conversation_token: None,
+        forked_from_conversation_token: None,
+        ambient_agent_task_id: None,
+        tasks: vec![],
+        existing_suggestions: None,
+        metadata: None,
+        session_context: SessionContext::new_for_test(),
+        model: model.clone(),
+        coding_model: model.clone(),
+        cli_agent_model: model.clone(),
+        computer_use_model: model,
+        is_memory_enabled: false,
+        warp_drive_context_enabled: false,
+        mcp_context: None,
+        planning_enabled: true,
+        should_redact_secrets: false,
+        api_keys: None,
+        allow_use_of_warp_credits_with_byok: false,
+        autonomy_level: api::AutonomyLevel::Supervised,
+        isolation_level: api::IsolationLevel::None,
+        web_search_enabled: false,
+        computer_use_enabled: false,
+        ask_user_question_enabled,
+        research_agent_enabled: false,
+        orchestration_enabled: false,
+        supported_tools_override: None,
+        parent_agent_id: None,
+        agent_name: None,
+    }
+}
+
+#[test]
+fn supported_tools_omits_ask_user_question_when_disabled() {
+    let params = request_params_with_ask_user_question_enabled(false);
+    let supported_tools = get_supported_tools(&params);
+
+    assert!(!supported_tools.contains(&api::ToolType::AskUserQuestion));
+}
+
+#[test]
+fn supported_tools_includes_ask_user_question_when_enabled_and_feature_flag_is_enabled() {
+    if !FeatureFlag::AskUserQuestion.is_enabled() {
+        return;
+    }
+
+    let params = request_params_with_ask_user_question_enabled(true);
+    let supported_tools = get_supported_tools(&params);
+
+    assert!(supported_tools.contains(&api::ToolType::AskUserQuestion));
+}
+
+#[test]
+fn supported_tools_include_upload_artifact_when_feature_flag_is_enabled() {
+    let _flag = FeatureFlag::ArtifactCommand.override_enabled(true);
+    let params = request_params_with_ask_user_question_enabled(false);
+    let supported_tools = get_supported_tools(&params);
+
+    assert!(supported_tools.contains(&api::ToolType::UploadFileArtifact));
+}
+
+#[test]
+fn supported_tools_omit_upload_artifact_when_feature_flag_is_disabled() {
+    let _flag = FeatureFlag::ArtifactCommand.override_enabled(false);
+    let params = request_params_with_ask_user_question_enabled(false);
+    let supported_tools = get_supported_tools(&params);
+
+    assert!(!supported_tools.contains(&api::ToolType::UploadFileArtifact));
+}

app/src/ai/agent/comment.rs

@@ -0,0 +1,106 @@
+use crate::code_review::comments::CommentId;
+use std::path::PathBuf;
+
+/// The current state of a code review.
+#[derive(Debug, Clone, Default)]
+pub struct CodeReview {
+    /// Comments that are currently pending (have yet to be addressed).
+    pub pending_comments: Vec<ReviewComment>,
+    /// Comments that have been addressed.
+    pub addressed_comments: Vec<ReviewComment>,
+}
+
+impl CodeReview {
+    pub fn new_with_pending_comments(pending_comments: Vec<ReviewComment>) -> Self {
+        Self {
+            pending_comments,
+            ..Self::default()
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct ReviewComment {
+    pub id: CommentId,
+    pub content: String,
+    pub diff: ReviewDiff,
+    pub head_title: Option<String>,
+}
+
+impl ReviewComment {
+    pub fn title(&self) -> String {
+        match (&self.diff.file_path, self.diff.line_number) {
+            (Some(file_path), Some(line_number)) => {
+                let file_name = file_path
+                    .file_name()
+                    .and_then(|name| name.to_str())
+                    .unwrap_or("Invalid File Name");
+                let display_line = line_number + 1;
+                format!("{file_name}:{display_line}")
+            }
+            (Some(file_path), None) => {
+                let file_name = file_path
+                    .file_name()
+                    .and_then(|name| name.to_str())
+                    .unwrap_or("Invalid File Name");
+                file_name.to_string()
+            }
+            (None, _) => self
+                .head_title
+                .as_ref()
+                .cloned()
+                .unwrap_or_else(|| "Review Comment".to_string()),
+        }
+    }
+}
+
+impl From<crate::code_review::comments::AttachedReviewComment> for ReviewComment {
+    fn from(comment: crate::code_review::comments::AttachedReviewComment) -> Self {
+        let head_title = comment.head().map(|head| head.title());
+
+        ReviewComment {
+            id: comment.id,
+            content: comment.content,
+            diff: comment.target.into(),
+            head_title,
+        }
+    }
+}
+
+impl From<crate::code_review::comments::AttachedReviewCommentTarget> for ReviewDiff {
+    fn from(val: crate::code_review::comments::AttachedReviewCommentTarget) -> Self {
+        // Convert from the server format of a line number (which is zero indexed)
+        // to one that is one-indexed to display within the blocklist.
+        match val {
+            crate::code_review::comments::AttachedReviewCommentTarget::Line {
+                absolute_file_path,
+                line,
+                content: _,
+            } => {
+                let line_number = line
+                    .line_number()
+                    .map(|line_number| line_number.as_usize() + 1);
+                Self {
+                    file_path: Some(absolute_file_path),
+                    line_number,
+                }
+            }
+            crate::code_review::comments::AttachedReviewCommentTarget::File {
+                absolute_file_path,
+            } => Self {
+                file_path: Some(absolute_file_path),
+                line_number: None,
+            },
+            crate::code_review::comments::AttachedReviewCommentTarget::General => Self {
+                file_path: None,
+                line_number: None,
+            },
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct ReviewDiff {
+    pub file_path: Option<PathBuf>,
+    pub line_number: Option<usize>,
+}

app/src/ai/agent/conversation_tests.rs

@@ -0,0 +1,199 @@
+use std::collections::HashMap;
+
+use super::{
+    artifact_from_fork_proto, AIConversation, AIConversationAutoexecuteMode, AIConversationId,
+};
+use crate::ai::artifacts::Artifact;
+use crate::persistence::model::AgentConversationData;
+use warp_core::features::FeatureFlag;
+use warp_multi_agent_api as api;
+
+fn restored_conversation(conversation_data: Option<AgentConversationData>) -> AIConversation {
+    AIConversation::new_restored(
+        AIConversationId::new(),
+        vec![api::Task {
+            id: "root-task".to_string(),
+            messages: vec![],
+            dependencies: None,
+            description: String::new(),
+            summary: String::new(),
+            server_data: String::new(),
+        }],
+        conversation_data,
+    )
+    .unwrap()
+}
+
+fn user_query_message(id: &str, request_id: &str, query: &str) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: "root-task".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::UserQuery(api::message::UserQuery {
+            query: query.to_string(),
+            context: None,
+            referenced_attachments: HashMap::new(),
+            mode: None,
+            intended_agent: Default::default(),
+        })),
+        request_id: request_id.to_string(),
+        timestamp: None,
+    }
+}
+
+fn agent_output_message(id: &str, request_id: &str) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: "root-task".to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::AgentOutput(
+            api::message::AgentOutput {
+                text: "Done".to_string(),
+            },
+        )),
+        request_id: request_id.to_string(),
+        timestamp: None,
+    }
+}
+
+fn restored_conversation_with_queries(queries: &[&str]) -> AIConversation {
+    let messages = queries
+        .iter()
+        .enumerate()
+        .flat_map(|(index, query)| {
+            let request_id = format!("request-{index}");
+            [
+                user_query_message(&format!("user-{index}"), &request_id, query),
+                agent_output_message(&format!("agent-{index}"), &request_id),
+            ]
+        })
+        .collect();
+
+    AIConversation::new_restored(
+        AIConversationId::new(),
+        vec![api::Task {
+            id: "root-task".to_string(),
+            messages,
+            dependencies: None,
+            description: String::new(),
+            summary: String::new(),
+            server_data: String::new(),
+        }],
+        None,
+    )
+    .unwrap()
+}
+
+#[test]
+fn latest_user_query_returns_latest_non_empty_user_query() {
+    let conversation =
+        restored_conversation_with_queries(&["write unit tests", "fix the failing test"]);
+
+    assert_eq!(
+        conversation.latest_user_query(),
+        Some("fix the failing test".to_string())
+    );
+}
+
+#[test]
+fn latest_user_query_trims_and_skips_empty_queries() {
+    let conversation = restored_conversation_with_queries(&["  write unit tests  ", "  "]);
+
+    assert_eq!(
+        conversation.latest_user_query(),
+        Some("write unit tests".to_string())
+    );
+}
+
+#[test]
+fn restored_conversation_defaults_autoexecute_override_when_not_persisted() {
+    let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
+    let conversation_data: AgentConversationData =
+        serde_json::from_str(r#"{"server_conversation_token":null}"#).unwrap();
+
+    let conversation = restored_conversation(Some(conversation_data));
+
+    assert_eq!(
+        conversation.autoexecute_override(),
+        AIConversationAutoexecuteMode::RespectUserSettings
+    );
+}
+
+#[test]
+fn restored_conversation_defaults_unknown_persisted_autoexecute_override() {
+    let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
+    let conversation_data: AgentConversationData = serde_json::from_str(
+        r#"{"server_conversation_token":null,"autoexecute_override":"UnexpectedValue"}"#,
+    )
+    .unwrap();
+
+    let conversation = restored_conversation(Some(conversation_data));
+
+    assert_eq!(
+        conversation.autoexecute_override(),
+        AIConversationAutoexecuteMode::RespectUserSettings
+    );
+}
+
+#[test]
+fn restored_conversation_uses_persisted_autoexecute_override_when_enabled() {
+    let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
+    let conversation_data: AgentConversationData = serde_json::from_str(
+        r#"{"server_conversation_token":null,"autoexecute_override":"RunToCompletion"}"#,
+    )
+    .unwrap();
+
+    let conversation = restored_conversation(Some(conversation_data));
+
+    assert_eq!(
+        conversation.autoexecute_override(),
+        AIConversationAutoexecuteMode::RunToCompletion
+    );
+}
+
+#[test]
+fn restored_conversation_ignores_persisted_autoexecute_override_when_disabled() {
+    let _flag = FeatureFlag::RememberFastForwardState.override_enabled(false);
+    let conversation_data: AgentConversationData = serde_json::from_str(
+        r#"{"server_conversation_token":null,"autoexecute_override":"RunToCompletion"}"#,
+    )
+    .unwrap();
+
+    let conversation = restored_conversation(Some(conversation_data));
+
+    assert_eq!(
+        conversation.autoexecute_override(),
+        AIConversationAutoexecuteMode::RespectUserSettings
+    );
+}
+
+#[test]
+fn fork_artifacts_adds_file_artifacts_to_conversation() {
+    let proto_artifact = api::message::artifact_event::ConversationArtifact {
+        artifact: Some(
+            api::message::artifact_event::conversation_artifact::Artifact::File(
+                api::message::artifact_event::FileArtifact {
+                    artifact_uid: "artifact-file-1".to_string(),
+                    filepath: "outputs/report.txt".to_string(),
+                    mime_type: "text/plain".to_string(),
+                    size_bytes: 42,
+                    description: "Daily summary".to_string(),
+                },
+            ),
+        ),
+    };
+
+    assert_eq!(
+        artifact_from_fork_proto(&proto_artifact),
+        Some(Artifact::File {
+            artifact_uid: "artifact-file-1".to_string(),
+            filepath: "outputs/report.txt".to_string(),
+            filename: "report.txt".to_string(),
+            mime_type: "text/plain".to_string(),
+            description: Some("Daily summary".to_string()),
+            size_bytes: Some(42),
+        })
+    );
+}

app/src/ai/agent/conversation_yaml_tests.rs

@@ -0,0 +1,451 @@
+use std::fs;
+use std::path::Path;
+
+use warp_multi_agent_api as api;
+
+use crate::test_util::ai_agent_tasks::{
+    create_api_subtask, create_api_task, create_message, create_subagent_tool_call_message,
+};
+
+use super::{base_dir, materialize_tasks_to_yaml};
+
+/// Lists filenames (not full paths) in a directory, sorted.
+fn list_dir_sorted(dir: &Path) -> Vec<String> {
+    let mut entries: Vec<String> = fs::read_dir(dir)
+        .unwrap()
+        .filter_map(|e| e.ok())
+        .map(|e| e.file_name().to_string_lossy().into_owned())
+        .collect();
+    entries.sort();
+    entries
+}
+
+fn make_user_query_message(id: &str, task_id: &str, query: &str) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::UserQuery(api::message::UserQuery {
+            query: query.to_string(),
+            context: None,
+            mode: None,
+            referenced_attachments: Default::default(),
+            intended_agent: Default::default(),
+        })),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+fn make_tool_call_message(
+    id: &str,
+    task_id: &str,
+    tool_call_id: &str,
+    tool: api::message::tool_call::Tool,
+) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: tool_call_id.to_string(),
+            tool: Some(tool),
+        })),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+fn make_tool_call_result_message(
+    id: &str,
+    task_id: &str,
+    tool_call_id: &str,
+    result: api::message::tool_call_result::Result,
+) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCallResult(
+            api::message::ToolCallResult {
+                tool_call_id: tool_call_id.to_string(),
+                result: Some(result),
+                context: None,
+            },
+        )),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+fn cleanup_dir(path: &str) {
+    let _ = fs::remove_dir_all(path);
+}
+
+#[test]
+fn mixed_message_types_produce_sequentially_indexed_files() {
+    let task_id = "root";
+    let tasks = vec![create_api_task(
+        task_id,
+        vec![
+            make_user_query_message("m1", task_id, "hello"),
+            // AgentOutput via create_message helper
+            create_message("m2", task_id),
+            make_tool_call_message(
+                "m3",
+                task_id,
+                "tc1",
+                api::message::tool_call::Tool::Grep(api::message::tool_call::Grep {
+                    queries: vec!["foo".into()],
+                    path: "/src".into(),
+                }),
+            ),
+        ],
+    )];
+
+    let dir = materialize_tasks_to_yaml(&tasks).unwrap();
+    assert!(
+        Path::new(&dir).starts_with(base_dir()),
+        "returned path should be under temp_dir(), got: {dir}",
+    );
+    // Verify no mixed separators: on Windows the path should use only '\',
+    // on Unix only '/'. This catches the original bug where tempdir_in
+    // joined a forward-slash parent with a native backslash separator.
+    assert!(
+        !dir.contains('/') || !dir.contains('\\'),
+        "returned path has mixed separators: {dir}",
+    );
+    let files = list_dir_sorted(Path::new(&dir));
+
+    assert_eq!(files.len(), 3);
+    assert!(files[0].starts_with("000.m1.user_query"));
+    assert!(files[1].starts_with("001.m2.agent_output"));
+    assert!(files[2].starts_with("002.m3.tool_call.tc1.grep"));
+
+    // Verify user_query content is searchable.
+    let content = fs::read_to_string(Path::new(&dir).join(&files[0])).unwrap();
+    assert!(content.contains("type: user_query"));
+    assert!(content.contains("hello"));
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn subagent_file_and_subdirectory_share_same_index() {
+    let root_id = "root";
+    let subtask_id = "subtask1";
+
+    let root_task = create_api_task(
+        root_id,
+        vec![
+            make_user_query_message("m1", root_id, "search my conversation"),
+            create_subagent_tool_call_message(
+                "m2",
+                root_id,
+                subtask_id,
+                Some(
+                    api::message::tool_call::subagent::Metadata::ConversationSearch(
+                        Default::default(),
+                    ),
+                ),
+            ),
+        ],
+    );
+    let subtask = create_api_subtask(
+        subtask_id,
+        root_id,
+        vec![create_message("sub_m1", subtask_id)],
+    );
+
+    let dir = materialize_tasks_to_yaml(&[root_task, subtask]).unwrap();
+    let entries = list_dir_sorted(Path::new(&dir));
+
+    // Should have: 000.m1.user_query.yaml, 001.m2.subagent.*.yaml, 001.subtask1/ (directory)
+    assert_eq!(entries.len(), 3);
+
+    // The subagent YAML file and its subdirectory must share the same "001" prefix.
+    let subagent_file = entries
+        .iter()
+        .find(|e| e.contains("subagent") && e.ends_with(".yaml"))
+        .expect("should have subagent yaml file");
+    let subdir = entries
+        .iter()
+        .find(|e| e.contains(subtask_id) && !e.ends_with(".yaml"))
+        .expect("should have subtask directory");
+
+    let file_prefix: String = subagent_file.chars().take(3).collect();
+    let dir_prefix: String = subdir.chars().take(3).collect();
+    assert_eq!(
+        file_prefix, dir_prefix,
+        "subagent file ({subagent_file}) and directory ({subdir}) must share the same index prefix"
+    );
+    assert_eq!(file_prefix, "001");
+
+    // Verify subtask directory contains the subtask's messages.
+    let sub_entries = list_dir_sorted(&Path::new(&dir).join(subdir));
+    assert_eq!(sub_entries.len(), 1);
+    assert!(sub_entries[0].contains("sub_m1"));
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn missing_subtask_in_task_map_produces_file_but_no_directory() {
+    let root_id = "root";
+
+    // Subagent references subtask "missing_task" which is not in the task list.
+    let root_task = create_api_task(
+        root_id,
+        vec![create_subagent_tool_call_message(
+            "m1",
+            root_id,
+            "missing_task",
+            Some(api::message::tool_call::subagent::Metadata::Cli(
+                Default::default(),
+            )),
+        )],
+    );
+
+    let dir = materialize_tasks_to_yaml(&[root_task]).unwrap();
+    let entries = list_dir_sorted(Path::new(&dir));
+
+    // Should have just the YAML file, no subdirectory since the subtask is missing.
+    assert_eq!(entries.len(), 1);
+    assert!(entries[0].ends_with(".yaml"));
+    assert!(entries[0].contains("subagent"));
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn empty_task_list_returns_error() {
+    let result = materialize_tasks_to_yaml(&[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().contains("No root task found"));
+}
+
+#[test]
+fn tool_call_result_resolves_tool_name_from_matching_call() {
+    let task_id = "root";
+    let tasks = vec![create_api_task(
+        task_id,
+        vec![
+            make_tool_call_message(
+                "m1",
+                task_id,
+                "tc1",
+                api::message::tool_call::Tool::Grep(api::message::tool_call::Grep {
+                    queries: vec!["pattern".into()],
+                    path: "/src".into(),
+                }),
+            ),
+            make_tool_call_result_message(
+                "m2",
+                task_id,
+                "tc1",
+                api::message::tool_call_result::Result::Grep(api::GrepResult {
+                    result: Some(api::grep_result::Result::Success(
+                        api::grep_result::Success {
+                            matched_files: vec![api::grep_result::success::GrepFileMatch {
+                                file_path: "foo.rs".into(),
+                                matched_lines: vec![
+                                    api::grep_result::success::grep_file_match::GrepLineMatch {
+                                        line_number: 42,
+                                    },
+                                ],
+                            }],
+                        },
+                    )),
+                }),
+            ),
+        ],
+    )];
+
+    let dir = materialize_tasks_to_yaml(&tasks).unwrap();
+    let files = list_dir_sorted(Path::new(&dir));
+
+    assert_eq!(files.len(), 2);
+    // The result file should contain "grep" in its name, resolved from the tool call.
+    assert!(
+        files[1].contains("grep"),
+        "result filename should contain tool name 'grep', got: {}",
+        files[1]
+    );
+
+    // Verify line numbers are serialized.
+    let content = fs::read_to_string(Path::new(&dir).join(&files[1])).unwrap();
+    assert!(content.contains("foo.rs"), "should contain file path");
+    assert!(content.contains("42"), "should contain line number");
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn server_tool_calls_are_skipped() {
+    let task_id = "root";
+    let tasks = vec![create_api_task(
+        task_id,
+        vec![
+            make_user_query_message("m1", task_id, "hello"),
+            make_tool_call_message(
+                "m2",
+                task_id,
+                "tc_server",
+                api::message::tool_call::Tool::Server(api::message::tool_call::Server {
+                    payload: String::new(),
+                }),
+            ),
+            create_message("m3", task_id),
+        ],
+    )];
+
+    let dir = materialize_tasks_to_yaml(&tasks).unwrap();
+    let files = list_dir_sorted(Path::new(&dir));
+
+    // Server tool call should be skipped; only user_query and agent_output.
+    assert_eq!(files.len(), 2);
+    assert!(files[0].contains("user_query"));
+    assert!(files[1].contains("agent_output"));
+    // Index should still be sequential (000, 001) since server call was skipped.
+    assert!(files[0].starts_with("000"));
+    assert!(files[1].starts_with("001"));
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn start_agent_v2_tool_call_serializes_name_and_prompt() {
+    let task_id = "root";
+    let tasks = vec![create_api_task(
+        task_id,
+        vec![make_tool_call_message(
+            "m1",
+            task_id,
+            "tc_start_agent_v2",
+            api::message::tool_call::Tool::StartAgentV2(api::StartAgentV2 {
+                name: "Remote child".to_string(),
+                prompt: "Investigate the build failure".to_string(),
+                execution_mode: None,
+                lifecycle_subscription: None,
+            }),
+        )],
+    )];
+
+    let dir = materialize_tasks_to_yaml(&tasks).unwrap();
+    let files = list_dir_sorted(Path::new(&dir));
+    let content = fs::read_to_string(Path::new(&dir).join(&files[0])).unwrap();
+
+    assert!(content.contains("tool_name: start_agent"));
+    assert!(content.contains("name: \"Remote child\""));
+    assert!(content.contains("prompt: |"));
+    assert!(content.contains("Investigate the build failure"));
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn start_agent_v2_tool_call_result_serializes_agent_id_and_error() {
+    let task_id = "root";
+    let tasks = vec![create_api_task(
+        task_id,
+        vec![
+            make_tool_call_message(
+                "m1",
+                task_id,
+                "tc_start_agent_v2",
+                api::message::tool_call::Tool::StartAgentV2(api::StartAgentV2 {
+                    name: "Remote child".to_string(),
+                    prompt: "Investigate the build failure".to_string(),
+                    execution_mode: None,
+                    lifecycle_subscription: None,
+                }),
+            ),
+            make_tool_call_result_message(
+                "m2",
+                task_id,
+                "tc_start_agent_v2",
+                api::message::tool_call_result::Result::StartAgentV2(api::StartAgentV2Result {
+                    result: Some(api::start_agent_v2_result::Result::Success(
+                        api::start_agent_v2_result::Success {
+                            agent_id: "agent-123".to_string(),
+                        },
+                    )),
+                }),
+            ),
+            make_tool_call_result_message(
+                "m3",
+                task_id,
+                "tc_start_agent_v2",
+                api::message::tool_call_result::Result::StartAgentV2(api::StartAgentV2Result {
+                    result: Some(api::start_agent_v2_result::Result::Error(
+                        api::start_agent_v2_result::Error {
+                            error: "child failed".to_string(),
+                        },
+                    )),
+                }),
+            ),
+        ],
+    )];
+
+    let dir = materialize_tasks_to_yaml(&tasks).unwrap();
+    let files = list_dir_sorted(Path::new(&dir));
+    let success_content = fs::read_to_string(Path::new(&dir).join(&files[1])).unwrap();
+    let error_content = fs::read_to_string(Path::new(&dir).join(&files[2])).unwrap();
+
+    assert!(success_content.contains("agent_id: agent-123"));
+    assert!(error_content.contains("error: child failed"));
+
+    cleanup_dir(&dir);
+}
+
+#[test]
+fn upload_file_artifact_tool_call_result_serializes_only_supported_success_fields() {
+    let task_id = "root";
+    let tasks = vec![create_api_task(
+        task_id,
+        vec![
+            make_tool_call_message(
+                "m1",
+                task_id,
+                "tc_upload_file_artifact",
+                api::message::tool_call::Tool::UploadFileArtifact(api::UploadFileArtifact {
+                    file: Some(api::FilePathReference {
+                        file_path: "outputs/report.txt".to_string(),
+                    }),
+                    description: "Daily summary".to_string(),
+                }),
+            ),
+            make_tool_call_result_message(
+                "m2",
+                task_id,
+                "tc_upload_file_artifact",
+                api::message::tool_call_result::Result::UploadFileArtifact(
+                    api::UploadFileArtifactResult {
+                        result: Some(api::upload_file_artifact_result::Result::Success(
+                            api::upload_file_artifact_result::Success {
+                                artifact_uid: "artifact-123".to_string(),
+                                mime_type: "text/plain".to_string(),
+                                size_bytes: 42,
+                            },
+                        )),
+                    },
+                ),
+            ),
+        ],
+    )];
+
+    let dir = materialize_tasks_to_yaml(&tasks).unwrap();
+    let files = list_dir_sorted(Path::new(&dir));
+    let success_content = fs::read_to_string(Path::new(&dir).join(&files[1])).unwrap();
+
+    assert!(success_content.contains("artifact_uid: artifact-123"));
+    assert!(success_content.contains("mime_type: text/plain"));
+    assert!(success_content.contains("size_bytes: 42"));
+    assert!(!success_content.contains("filepath:"));
+    assert!(!success_content.contains("description:"));
+
+    cleanup_dir(&dir);
+}

app/src/ai/agent/icons.rs

@@ -0,0 +1,90 @@
+use warp_core::ui::{appearance::Appearance, theme::AnsiColorIdentifier};
+
+use crate::ui_components::{blended_colors, icons::Icon};
+
+pub fn todo_list_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::BulletedListBlock.into(),
+        blended_colors::neutral_7(appearance.theme()),
+    )
+}
+
+pub fn pending_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::Queued.into(),
+        blended_colors::neutral_5(appearance.theme()),
+    )
+}
+
+pub fn in_progress_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::Circle.into(),
+        AnsiColorIdentifier::Magenta.to_ansi_color(&appearance.theme().terminal_colors().normal),
+    )
+}
+
+pub fn succeeded_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::Check.into(),
+        AnsiColorIdentifier::Green.to_ansi_color(&appearance.theme().terminal_colors().normal),
+    )
+}
+
+pub fn addressed_comment_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::AddressedComment.into(),
+        AnsiColorIdentifier::Green.to_ansi_color(&appearance.theme().terminal_colors().normal),
+    )
+}
+
+pub fn failed_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::Triangle.into(),
+        AnsiColorIdentifier::Red.to_ansi_color(&appearance.theme().terminal_colors().normal),
+    )
+}
+
+/// Not running, does not need user's attention
+pub fn gray_stop_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::StopFilled.into(),
+        blended_colors::neutral_5(appearance.theme()),
+    )
+}
+
+/// Agent is waiting for user to follow-up with next prompt.
+pub fn gray_clock_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::ClockSnooze.into(),
+        blended_colors::neutral_5(appearance.theme()),
+    )
+}
+
+/// Loading but not actionable yet.
+pub fn gray_circle_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::Circle.into(),
+        blended_colors::neutral_5(appearance.theme()),
+    )
+}
+
+/// Not running, requires user's attention
+pub fn yellow_stop_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::StopFilled.into(),
+        AnsiColorIdentifier::Yellow.to_ansi_color(&appearance.theme().terminal_colors().normal),
+    )
+}
+
+/// To be used for actions (like running commands/reading files) that are long-running and executing.
+pub fn yellow_running_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(
+        Icon::Circle.into(),
+        AnsiColorIdentifier::Yellow.to_ansi_color(&appearance.theme().terminal_colors().normal),
+    )
+}
+
+/// Used for buttons that stop the current task
+pub fn red_stop_icon(appearance: &Appearance) -> warpui::elements::Icon {
+    warpui::elements::Icon::new(Icon::StopFilled.into(), appearance.theme().ansi_fg_red())
+}

app/src/ai/agent/linearization.rs

@@ -0,0 +1,108 @@
+//! Linearization utilities for task messages.
+//!
+//! This module provides pure functions for linearizing task messages in a conversation,
+//! following a DFS traversal that interleaves subtask messages at subagent tool calls.
+
+use std::collections::{HashMap, HashSet};
+
+use warp_multi_agent_api as api;
+
+use crate::ai::agent::task::helper::TaskExt as _;
+
+/// Computes the set of "active" task IDs in a task tree.
+///
+/// An active task is one that is still in progress. The algorithm:
+/// 1. Start with a queue containing the root task ID
+/// 2. For each task in the queue, walk through its messages:
+///    - When encountering a subagent ToolCall, add the subtask to the queue
+///    - When encountering a ToolCallResult matching a subagent call, remove from queue
+/// 3. After processing all messages, add the task to the active set
+/// 4. Repeat until the queue is empty
+pub fn compute_active_task_ids<'a>(
+    root_task_id: &str,
+    tasks: &HashMap<&str, &'a api::Task>,
+) -> HashSet<&'a str> {
+    let mut active_tasks = HashSet::new();
+    let mut visited = HashSet::new();
+    let mut queue = vec![root_task_id];
+
+    while let Some(task_id) = queue.pop() {
+        // Cycle protection: skip tasks we've already processed.
+        if !visited.insert(task_id) {
+            log::error!("Cycle detected in active task computation at task {task_id}");
+            continue;
+        }
+
+        let Some(task) = tasks.get(task_id) else {
+            // Task not found - skip it.
+            continue;
+        };
+
+        // Track subagent tool calls: tool_call_id -> subtask_id.
+        let mut pending_subagents: HashMap<&str, &str> = HashMap::new();
+
+        for message in &task.messages {
+            match &message.message {
+                Some(api::message::Message::ToolCall(tool_call)) => {
+                    // Check if this is a subagent call.
+                    if let Some(api::message::tool_call::Tool::Subagent(subagent)) = &tool_call.tool
+                    {
+                        if !subagent.task_id.is_empty() {
+                            // Add subtask to the queue.
+                            queue.push(subagent.task_id.as_str());
+                            // Track this subagent call so we can remove it when we see the result.
+                            pending_subagents
+                                .insert(tool_call.tool_call_id.as_str(), subagent.task_id.as_str());
+                        }
+                    }
+                }
+                Some(api::message::Message::ToolCallResult(result)) => {
+                    // If this result matches a pending subagent call, remove from queue.
+                    if let Some(subtask_id) = pending_subagents.remove(result.tool_call_id.as_str())
+                    {
+                        queue.retain(|id| *id != subtask_id);
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        // After processing all messages, add this task to the active set.
+        active_tasks.insert(task.id.as_str());
+    }
+
+    active_tasks
+}
+
+/// Computes the depth (distance from root) for each task in the map.
+///
+/// Tasks with no parent have depth 0. Tasks whose parent chain contains a cycle or leads to a
+/// missing task are assigned depth 0.
+pub fn compute_task_depths(tasks: &HashMap<String, api::Task>) -> HashMap<&str, usize> {
+    let mut depths = HashMap::new();
+    for (task_id, _) in tasks.iter() {
+        let mut depth = 0;
+        let mut current_id: &str = task_id;
+        let mut visited = HashSet::new();
+        while let Some(task) = tasks.get(current_id) {
+            if !visited.insert(current_id) {
+                // Cycle detected; treat as depth 0.
+                log::error!("Cycle detected in task parent chain starting from task {task_id}");
+                depth = 0;
+                break;
+            }
+            if let Some(parent_id) = task.parent_id() {
+                depth += 1;
+                current_id = parent_id;
+            } else {
+                break;
+            }
+        }
+        depths.insert(task_id.as_str(), depth);
+    }
+    depths
+}
+
+#[cfg(test)]
+#[path = "linearization_tests.rs"]
+mod tests;

app/src/ai/agent/linearization_tests.rs

@@ -0,0 +1,427 @@
+use std::collections::HashMap;
+
+use super::*;
+use warp_multi_agent_api as api;
+
+// Helper function to create a basic message
+fn create_message(id: &str, task_id: &str) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: "server_data".to_string(),
+        citations: vec![],
+        message: Some(api::message::Message::AgentOutput(
+            api::message::AgentOutput {
+                text: format!("Message content for {id}"),
+            },
+        )),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+fn create_subagent_tool_call_message(id: &str, task_id: &str, subtask_id: &str) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: "server_data".to_string(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: format!("{id}_tool_call"),
+            tool: Some(api::message::tool_call::Tool::Subagent(
+                api::message::tool_call::Subagent {
+                    task_id: subtask_id.to_string(),
+                    payload: String::new(),
+                    metadata: None,
+                },
+            )),
+        })),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+// Helper function to create a tool call result message.
+fn create_tool_call_result_message(id: &str, task_id: &str, tool_call_id: &str) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: "server_data".to_string(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCallResult(
+            api::message::ToolCallResult {
+                tool_call_id: tool_call_id.to_string(),
+                context: None,
+                result: None,
+            },
+        )),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+// Helper function to create a task with dependencies
+fn create_task(id: &str, messages: Vec<api::Message>, parent_task_id: Option<String>) -> api::Task {
+    let dependencies = parent_task_id.map(|parent_id| api::task::Dependencies {
+        parent_task_id: parent_id,
+    });
+
+    api::Task {
+        id: id.to_string(),
+        messages,
+        dependencies,
+        description: format!("Task {id}"),
+        summary: format!("Summary for task {id}"),
+        server_data: "".to_string(),
+    }
+}
+
+// Helper function to create a root task (no parent)
+fn create_root_task(id: &str, messages: Vec<api::Message>) -> api::Task {
+    create_task(id, messages, None)
+}
+
+// Helper function to create a child task
+fn create_child_task(id: &str, messages: Vec<api::Message>, parent_id: &str) -> api::Task {
+    create_task(id, messages, Some(parent_id.to_string()))
+}
+
+// Helper to build a task map from a slice of tasks.
+fn make_task_map(tasks: &[api::Task]) -> HashMap<&str, &api::Task> {
+    tasks.iter().map(|t| (t.id.as_str(), t)).collect()
+}
+
+#[test]
+fn test_compute_active_task_ids_single_root() {
+    let root = create_root_task("root", vec![create_message("m1", "root")]);
+    let tasks = vec![root];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    assert_eq!(active.len(), 1);
+    assert!(active.contains("root"));
+}
+
+#[test]
+fn test_compute_active_task_ids_subagent_in_progress() {
+    // Root calls a subagent but has not received the result yet.
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "child"),
+    ];
+    let child_messages = vec![create_message("child_m1", "child")];
+
+    let root = create_root_task("root", root_messages);
+    let child = create_child_task("child", child_messages, "root");
+    let tasks = vec![root, child];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    // Both root and child are active.
+    assert_eq!(active.len(), 2);
+    assert!(active.contains("root"));
+    assert!(active.contains("child"));
+}
+
+#[test]
+fn test_compute_active_task_ids_subagent_completed() {
+    // Root calls a subagent and has received the result.
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "child"),
+        create_tool_call_result_message("result1", "root", "call1_tool_call"),
+        create_message("m2", "root"),
+    ];
+    let child_messages = vec![create_message("child_m1", "child")];
+
+    let root = create_root_task("root", root_messages);
+    let child = create_child_task("child", child_messages, "root");
+    let tasks = vec![root, child];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    // Only root is active - child completed.
+    assert_eq!(active.len(), 1);
+    assert!(active.contains("root"));
+}
+
+#[test]
+fn test_compute_active_task_ids_nested_subagents() {
+    // root -> child (in progress) -> grandchild (in progress)
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "child"),
+    ];
+    let child_messages = vec![
+        create_message("child_m1", "child"),
+        create_subagent_tool_call_message("call2", "child", "grandchild"),
+    ];
+    let grandchild_messages = vec![create_message("gc_m1", "grandchild")];
+
+    let root = create_root_task("root", root_messages);
+    let child = create_child_task("child", child_messages, "root");
+    let grandchild = create_child_task("grandchild", grandchild_messages, "child");
+    let tasks = vec![root, child, grandchild];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    // All three are active.
+    assert_eq!(active.len(), 3);
+    assert!(active.contains("root"));
+    assert!(active.contains("child"));
+    assert!(active.contains("grandchild"));
+}
+
+#[test]
+fn test_compute_active_task_ids_nested_subagents_partial_completion() {
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "child"),
+    ];
+    let child_messages = vec![
+        create_message("child_m1", "child"),
+        create_subagent_tool_call_message("call2", "child", "grandchild"),
+        create_tool_call_result_message("result2", "child", "call2_tool_call"),
+    ];
+    let grandchild_messages = vec![create_message("gc_m1", "grandchild")];
+
+    let root = create_root_task("root", root_messages);
+    let child = create_child_task("child", child_messages, "root");
+    let grandchild = create_child_task("grandchild", grandchild_messages, "child");
+    let tasks = vec![root, child, grandchild];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    // Grandchild completed, but root and its child are still running.
+    assert_eq!(active.len(), 2);
+    assert!(active.contains("root"));
+    assert!(active.contains("child"));
+}
+
+#[test]
+fn test_compute_active_task_ids_multiple_parallel_subagents() {
+    // Root calls two subagents, one completed and one in progress.
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "child1"),
+        create_subagent_tool_call_message("call2", "root", "child2"),
+        create_tool_call_result_message("result1", "root", "call1_tool_call"),
+    ];
+    let child1_messages = vec![create_message("c1_m1", "child1")];
+    let child2_messages = vec![create_message("c2_m1", "child2")];
+
+    let root = create_root_task("root", root_messages);
+    let child1 = create_child_task("child1", child1_messages, "root");
+    let child2 = create_child_task("child2", child2_messages, "root");
+    let tasks = vec![root, child1, child2];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    // Root and child2 are active, child1 is completed.
+    assert_eq!(active.len(), 2);
+    assert!(active.contains("root"));
+    assert!(active.contains("child2"));
+}
+
+#[test]
+fn test_compute_active_task_ids_missing_subtask() {
+    // Root calls a subagent that doesn't exist in the task list.
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "nonexistent"),
+    ];
+
+    let root = create_root_task("root", root_messages);
+    let tasks = vec![root];
+    let tasks = make_task_map(&tasks);
+
+    let active = compute_active_task_ids("root", &tasks);
+
+    // Only root is active - the missing subtask is skipped.
+    assert_eq!(active.len(), 1);
+    assert!(active.contains("root"));
+}
+
+#[test]
+fn test_compute_active_task_ids_missing_root() {
+    // Root task ID doesn't exist in the map.
+    let tasks: HashMap<&str, &api::Task> = HashMap::new();
+
+    let active = compute_active_task_ids("nonexistent", &tasks);
+
+    assert!(active.is_empty());
+}
+
+#[test]
+fn test_compute_active_task_ids_cycle_protection() {
+    // Create a cycle: root -> child -> root (via subagent call).
+    // This should not cause an infinite loop.
+    let root_messages = vec![
+        create_message("m1", "root"),
+        create_subagent_tool_call_message("call1", "root", "child"),
+    ];
+    let child_messages = vec![
+        create_message("child_m1", "child"),
+        // Child calls back to root, creating a cycle.
+        create_subagent_tool_call_message("call2", "child", "root"),
+    ];
+
+    let root = create_root_task("root", root_messages);
+    let child = create_child_task("child", child_messages, "root");
+    let tasks = vec![root, child];
+    let tasks = make_task_map(&tasks);
+
+    // This should complete without infinite looping.
+    let active = compute_active_task_ids("root", &tasks);
+
+    // Both root and child are active (cycle is broken by visited check).
+    assert_eq!(active.len(), 2);
+    assert!(active.contains("root"));
+    assert!(active.contains("child"));
+}
+
+// ============================================================================
+// compute_task_depths tests
+// ============================================================================
+
+/// Creates a task with the given ID and optional parent, without any messages.
+fn create_task_for_depth(id: &str, parent_id: Option<&str>) -> api::Task {
+    create_task(id, vec![], parent_id.map(str::to_string))
+}
+
+#[test]
+fn test_compute_task_depths_empty() {
+    let tasks = HashMap::new();
+    let depths = compute_task_depths(&tasks);
+    assert!(depths.is_empty());
+}
+
+#[test]
+fn test_compute_task_depths_single_root() {
+    let tasks: HashMap<String, _> =
+        [("root".to_string(), create_task_for_depth("root", None))].into();
+
+    let depths = compute_task_depths(&tasks);
+
+    assert_eq!(depths.get("root"), Some(&0));
+}
+
+#[test]
+fn test_compute_task_depths_linear_chain() {
+    // root -> child -> grandchild
+    let tasks: HashMap<String, _> = [
+        ("root".to_string(), create_task_for_depth("root", None)),
+        (
+            "child".to_string(),
+            create_task_for_depth("child", Some("root")),
+        ),
+        (
+            "grandchild".to_string(),
+            create_task_for_depth("grandchild", Some("child")),
+        ),
+    ]
+    .into();
+
+    let depths = compute_task_depths(&tasks);
+
+    assert_eq!(depths.get("root"), Some(&0));
+    assert_eq!(depths.get("child"), Some(&1));
+    assert_eq!(depths.get("grandchild"), Some(&2));
+}
+
+#[test]
+fn test_compute_task_depths_tree_structure() {
+    // root -> child1 -> grandchild1
+    //      -> child2
+    let tasks: HashMap<String, _> = [
+        ("root".to_string(), create_task_for_depth("root", None)),
+        (
+            "child1".to_string(),
+            create_task_for_depth("child1", Some("root")),
+        ),
+        (
+            "child2".to_string(),
+            create_task_for_depth("child2", Some("root")),
+        ),
+        (
+            "grandchild1".to_string(),
+            create_task_for_depth("grandchild1", Some("child1")),
+        ),
+    ]
+    .into();
+
+    let depths = compute_task_depths(&tasks);
+
+    assert_eq!(depths.get("root"), Some(&0));
+    assert_eq!(depths.get("child1"), Some(&1));
+    assert_eq!(depths.get("child2"), Some(&1));
+    assert_eq!(depths.get("grandchild1"), Some(&2));
+}
+
+#[test]
+fn test_compute_task_depths_orphan() {
+    // Task with parent that doesn't exist in the map.
+    let tasks: HashMap<String, _> = [(
+        "orphan".to_string(),
+        create_task_for_depth("orphan", Some("missing_parent")),
+    )]
+    .into();
+
+    let depths = compute_task_depths(&tasks);
+
+    // Orphan's parent is missing, so the chain breaks immediately after computing depth 1.
+    assert_eq!(depths.get("orphan"), Some(&1));
+}
+
+#[test]
+fn test_compute_task_depths_cycle_two_tasks() {
+    // a -> b -> a (cycle)
+    let tasks: HashMap<String, _> = [
+        ("a".to_string(), create_task_for_depth("a", Some("b"))),
+        ("b".to_string(), create_task_for_depth("b", Some("a"))),
+    ]
+    .into();
+
+    let depths = compute_task_depths(&tasks);
+
+    // Both tasks are in a cycle, so they should get depth 0.
+    assert_eq!(depths.get("a"), Some(&0));
+    assert_eq!(depths.get("b"), Some(&0));
+}
+
+#[test]
+fn test_compute_task_depths_self_referential_cycle() {
+    // a -> a (self-referential)
+    let tasks: HashMap<String, _> =
+        [("a".to_string(), create_task_for_depth("a", Some("a")))].into();
+
+    let depths = compute_task_depths(&tasks);
+
+    // Self-referential task should get depth 0.
+    assert_eq!(depths.get("a"), Some(&0));
+}
+
+#[test]
+fn test_compute_task_depths_cycle_with_tail() {
+    // Task c points to a cycle: c -> a -> b -> a
+    let tasks: HashMap<String, _> = [
+        ("a".to_string(), create_task_for_depth("a", Some("b"))),
+        ("b".to_string(), create_task_for_depth("b", Some("a"))),
+        ("c".to_string(), create_task_for_depth("c", Some("a"))),
+    ]
+    .into();
+
+    let depths = compute_task_depths(&tasks);
+
+    // a and b are in a cycle, so depth 0.
+    assert_eq!(depths.get("a"), Some(&0));
+    assert_eq!(depths.get("b"), Some(&0));
+    // c's parent chain leads into a cycle, so it also gets depth 0.
+    assert_eq!(depths.get("c"), Some(&0));
+}

app/src/ai/agent/mod_test.rs

@@ -0,0 +1,191 @@
+use std::ops::Range;
+use std::sync::Arc;
+
+use warp_multi_agent_api::{FileContent, FileContentLineRange};
+
+use crate::ai::agent::{
+    AIAgentOutput, AIAgentOutputMessage, AIAgentOutputMessageType, AIAgentText, AIAgentTextSection,
+    AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram, AnyFileContent,
+    FileContext, FormattedTextWrapper, MessageId, ProgrammingLanguage,
+};
+use crate::terminal::shell::ShellType;
+use markdown_parser::{FormattedText, FormattedTextFragment, FormattedTextLine};
+
+fn to_range(range: Range<u32>) -> Option<FileContentLineRange> {
+    Some(FileContentLineRange {
+        start: range.start,
+        end: range.end,
+    })
+}
+
+#[test]
+fn formatted_text_wrapper_shares_arc_across_calls() {
+    let text = FormattedText::new([FormattedTextLine::Line(vec![
+        FormattedTextFragment::plain_text("hello world"),
+    ])]);
+    let wrapper = FormattedTextWrapper::from(text);
+    let arc1 = wrapper.formatted_text_arc();
+    let arc2 = wrapper.formatted_text_arc();
+    // Both calls must return the same allocation — not independent deep copies.
+    assert!(Arc::ptr_eq(&arc1, &arc2));
+}
+
+#[test]
+fn formatted_text_wrapper_preserves_content() {
+    let text = FormattedText::new([
+        FormattedTextLine::Line(vec![FormattedTextFragment::plain_text("line one")]),
+        FormattedTextLine::Line(vec![FormattedTextFragment::plain_text("line two")]),
+    ]);
+    let wrapper = FormattedTextWrapper::from(text);
+    // lines() metadata matches the cached Arc
+    assert_eq!(wrapper.lines().len(), 2);
+    assert_eq!(wrapper.lines()[0].raw_text(), "line one\n");
+    assert_eq!(wrapper.lines()[1].raw_text(), "line two\n");
+    // Arc contains the same lines
+    let ft = wrapper.formatted_text_arc();
+    assert_eq!(ft.lines.len(), 2);
+}
+
+#[test]
+fn test_convert_files() {
+    let a = FileContext::new(
+        "a.txt".to_string(),
+        AnyFileContent::StringContent("hey\nyou".to_string()),
+        None,
+        None,
+    );
+
+    assert_eq!(
+        Into::<Vec<FileContent>>::into(a),
+        vec![FileContent {
+            file_path: "a.txt".to_string(),
+            content: "hey\nyou".to_string(),
+            line_range: None,
+        }]
+    );
+}
+
+#[test]
+fn test_convert_files_range() {
+    // Content is pre-sliced to match the line range.
+    let a = FileContext::new(
+        "a.txt".to_string(),
+        AnyFileContent::StringContent("hey\nyou".to_string()),
+        Some(1..2),
+        None,
+    );
+
+    assert_eq!(
+        Into::<Vec<FileContent>>::into(a),
+        vec![FileContent {
+            file_path: "a.txt".to_string(),
+            content: "hey\nyou".to_string(),
+            line_range: to_range(1..2),
+        }]
+    );
+}
+
+#[test]
+fn test_convert_files_range_out_of_bounds() {
+    // Even with an out-of-bounds range, content is passed through as-is.
+    let a = FileContext::new(
+        "a.txt".to_string(),
+        AnyFileContent::StringContent(String::new()),
+        Some(10..20),
+        None,
+    );
+
+    assert_eq!(
+        Into::<Vec<FileContent>>::into(a),
+        vec![FileContent {
+            file_path: "a.txt".to_string(),
+            content: String::new(),
+            line_range: to_range(10..20),
+        }]
+    );
+}
+
+#[test]
+fn test_programming_language_from_string() {
+    // Shell language specifiers should produce Shell variants
+    assert_eq!(
+        ProgrammingLanguage::from("bash".to_string()),
+        ProgrammingLanguage::Shell(ShellType::Bash)
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("shell".to_string()),
+        ProgrammingLanguage::Shell(ShellType::Bash)
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("sh".to_string()),
+        ProgrammingLanguage::Shell(ShellType::Bash)
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("zsh".to_string()),
+        ProgrammingLanguage::Shell(ShellType::Zsh)
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("fish".to_string()),
+        ProgrammingLanguage::Shell(ShellType::Fish)
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("powershell".to_string()),
+        ProgrammingLanguage::Shell(ShellType::PowerShell)
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("pwsh".to_string()),
+        ProgrammingLanguage::Shell(ShellType::PowerShell)
+    );
+
+    // Non-shell languages should produce Other variants
+    assert_eq!(
+        ProgrammingLanguage::from("python".to_string()),
+        ProgrammingLanguage::Other("python".to_string())
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("rust".to_string()),
+        ProgrammingLanguage::Other("rust".to_string())
+    );
+    assert_eq!(
+        ProgrammingLanguage::from("javascript".to_string()),
+        ProgrammingLanguage::Other("javascript".to_string())
+    );
+}
+
+#[test]
+fn format_for_copy_preserves_visual_markdown_sections() {
+    let output = AIAgentOutput {
+        messages: vec![AIAgentOutputMessage {
+            id: MessageId::new("message-1".to_string()),
+            message: AIAgentOutputMessageType::Text(AIAgentText {
+                sections: vec![
+                    AIAgentTextSection::PlainText {
+                        text: "Intro".to_string().into(),
+                    },
+                    AIAgentTextSection::Image {
+                        image: AgentOutputImage {
+                            alt_text: "Diagram".to_string(),
+                            source: "./diagram.png".to_string(),
+                            title: None,
+                            markdown_source: "![Diagram](./diagram.png)".to_string(),
+                            layout: AgentOutputImageLayout::Block,
+                        },
+                    },
+                    AIAgentTextSection::MermaidDiagram {
+                        diagram: AgentOutputMermaidDiagram {
+                            source: "graph TD\nA --> B".to_string(),
+                            markdown_source: "```mermaid\ngraph TD\nA --> B\n```".to_string(),
+                        },
+                    },
+                ],
+            }),
+            citations: Vec::new(),
+        }],
+        ..Default::default()
+    };
+
+    assert_eq!(
+        output.format_for_copy(None),
+        "Intro\n![Diagram](./diagram.png)\n```mermaid\ngraph TD\nA --> B\n```"
+    );
+}

app/src/ai/agent/redaction.rs

@@ -0,0 +1,409 @@
+use std::sync::Arc;
+
+use crate::ai::agent::{
+    AIAgentActionResultType, AIAgentAttachment, AIAgentContext, AIAgentInput, AnyFileContent,
+    AskUserQuestionAnswerItem, AskUserQuestionResult, BlockContext, PassiveSuggestionResultType,
+    PassiveSuggestionTrigger, RequestCommandOutputResult, TransferShellCommandControlToUserResult,
+};
+
+use super::super::blocklist::block::secret_redaction::{
+    find_secrets_in_text, SECRET_REDACTION_REPLACEMENT_CHARACTER,
+};
+
+/// Redact all detected secrets in-place within the given string.
+pub(crate) fn redact_secrets(input: &mut String) {
+    let mut secrets: Vec<_> = find_secrets_in_text(input)
+        .into_iter()
+        .map(|r| r.byte_range)
+        .collect();
+    // Replace from the end to preserve indices
+    secrets.sort_by_key(|range| range.start);
+    for range in secrets.into_iter().rev() {
+        let replacement =
+            SECRET_REDACTION_REPLACEMENT_CHARACTER.repeat(range.end.saturating_sub(range.start));
+        input.replace_range(range.start..range.end, &replacement);
+    }
+}
+
+/// Redact secrets in-place for all user-provided text fields inside the inputs that will be
+/// sent to the server.
+pub(crate) fn redact_inputs(inputs: &mut [AIAgentInput]) {
+    for input in inputs.iter_mut() {
+        match input {
+            AIAgentInput::UserQuery {
+                query,
+                context,
+                referenced_attachments,
+                ..
+            } => {
+                redact_secrets(query);
+                redact_context(Arc::make_mut(context));
+                referenced_attachments
+                    .values_mut()
+                    .for_each(redact_attachment);
+            }
+            AIAgentInput::AutoCodeDiffQuery { query, context, .. } => {
+                redact_secrets(query);
+                redact_context(Arc::make_mut(context));
+            }
+            AIAgentInput::CreateNewProject { context, .. }
+            | AIAgentInput::CloneRepository { context, .. }
+            | AIAgentInput::ResumeConversation { context }
+            | AIAgentInput::InitProjectRules { context, .. }
+            | AIAgentInput::StartFromAmbientRunPrompt { context, .. } => {
+                redact_context(Arc::make_mut(context));
+            }
+            AIAgentInput::SummarizeConversation { prompt } => {
+                if let Some(p) = prompt {
+                    redact_secrets(p);
+                }
+            }
+            AIAgentInput::CreateEnvironment { context, .. } => {
+                redact_context(Arc::make_mut(context));
+            }
+            AIAgentInput::TriggerPassiveSuggestion {
+                context,
+                attachments,
+                trigger,
+            } => {
+                redact_context(Arc::make_mut(context));
+                attachments.iter_mut().for_each(redact_attachment);
+                if let PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger) = trigger {
+                    redact_secrets(&mut shell_trigger.executed_shell_command.command);
+                    redact_secrets(&mut shell_trigger.executed_shell_command.output);
+                    for file in shell_trigger.relevant_files.iter_mut() {
+                        if let AnyFileContent::StringContent(content) = &mut file.content {
+                            redact_secrets(content);
+                        }
+                    }
+                }
+            }
+            AIAgentInput::CodeReview {
+                context,
+                review_comments,
+            } => {
+                redact_context(Arc::make_mut(context));
+                for comment in review_comments.comments.iter_mut() {
+                    redact_secrets(&mut comment.content);
+                    match &mut comment.target {
+                        crate::code_review::comments::AttachedReviewCommentTarget::Line {
+                            content,
+                            ..
+                        } => {
+                            redact_secrets(&mut content.content);
+                        }
+                        crate::code_review::comments::AttachedReviewCommentTarget::File {
+                            ..
+                        }
+                        | crate::code_review::comments::AttachedReviewCommentTarget::General => {}
+                    }
+                }
+
+                for diff in review_comments.diff_set.values_mut().flatten() {
+                    redact_secrets(&mut diff.diff_content);
+                }
+            }
+            // No user-provided text to redact in inter-agent relay inputs.
+            AIAgentInput::MessagesReceivedFromAgents { .. }
+            | AIAgentInput::EventsFromAgents { .. } => {}
+            AIAgentInput::ActionResult { result, context } => {
+                redact_context(Arc::make_mut(context));
+                match &mut result.result {
+                    AIAgentActionResultType::RequestCommandOutput(output) => {
+                        if let RequestCommandOutputResult::Completed { output, .. } = output {
+                            redact_secrets(output);
+                        }
+                    }
+                    AIAgentActionResultType::WriteToLongRunningShellCommand(result) => {
+                        use crate::ai::agent::WriteToLongRunningShellCommandResult::*;
+                        match result {
+                            Snapshot { grid_contents, .. } => redact_secrets(grid_contents),
+                            CommandFinished { output, .. } => redact_secrets(output),
+                            Error(_) | Cancelled => {}
+                        }
+                    }
+                    AIAgentActionResultType::ReadShellCommandOutput(result) => {
+                        use crate::ai::agent::ReadShellCommandOutputResult::*;
+                        match result {
+                            CommandFinished { output, .. } => redact_secrets(output),
+                            LongRunningCommandSnapshot { grid_contents, .. } => {
+                                redact_secrets(grid_contents)
+                            }
+                            Error(_) | Cancelled => {}
+                        }
+                    }
+                    AIAgentActionResultType::ReadFiles(read_files_result) => {
+                        if let crate::ai::agent::ReadFilesResult::Success { files } =
+                            read_files_result
+                        {
+                            for file in files {
+                                if let AnyFileContent::StringContent(content) = &mut file.content {
+                                    redact_secrets(content);
+                                }
+                            }
+                        }
+                    }
+                    AIAgentActionResultType::UploadArtifact(upload_result) => {
+                        use crate::ai::agent::UploadArtifactResult;
+                        match upload_result {
+                            UploadArtifactResult::Success {
+                                filepath,
+                                description,
+                                ..
+                            } => {
+                                if let Some(filepath) = filepath {
+                                    redact_secrets(filepath);
+                                }
+                                if let Some(description) = description {
+                                    redact_secrets(description);
+                                }
+                            }
+                            UploadArtifactResult::Error(error) => redact_secrets(error),
+                            UploadArtifactResult::Cancelled => {}
+                        }
+                    }
+                    AIAgentActionResultType::SearchCodebase(search_codebase_result) => {
+                        if let crate::ai::agent::SearchCodebaseResult::Success { files } =
+                            search_codebase_result
+                        {
+                            for file in files {
+                                if let AnyFileContent::StringContent(content) = &mut file.content {
+                                    redact_secrets(content);
+                                }
+                            }
+                        }
+                    }
+                    AIAgentActionResultType::RequestFileEdits(request_file_edits_result) => {
+                        if let crate::ai::agent::RequestFileEditsResult::Success {
+                            diff,
+                            updated_files,
+                            deleted_files,
+                            ..
+                        } = request_file_edits_result
+                        {
+                            redact_secrets(diff);
+                            for file in updated_files {
+                                if let AnyFileContent::StringContent(content) =
+                                    &mut file.file_context.content
+                                {
+                                    redact_secrets(content);
+                                }
+                            }
+                            for file_path in deleted_files {
+                                redact_secrets(file_path);
+                            }
+                        }
+                    }
+                    AIAgentActionResultType::InsertReviewComments(result) => {
+                        use crate::ai::agent::InsertReviewCommentsResult::*;
+                        match result {
+                            Success { repo_path } => redact_secrets(repo_path),
+                            Error { repo_path, message } => {
+                                redact_secrets(repo_path);
+                                redact_secrets(message);
+                            }
+                            Cancelled => {}
+                        }
+                    }
+
+                    // These are effectively flow control and don't contain secrets
+                    AIAgentActionResultType::SuggestNewConversation { .. }
+                    | AIAgentActionResultType::OpenCodeReview
+                    | AIAgentActionResultType::InitProject => {}
+
+                    // Contains only file path/line number information
+                    AIAgentActionResultType::Grep(_)
+                    | AIAgentActionResultType::FileGlob(_)
+                    | AIAgentActionResultType::FileGlobV2(_) => {}
+
+                    // TODO: Redact MCP-related results
+                    AIAgentActionResultType::CallMCPTool { .. }
+                    | AIAgentActionResultType::ReadSkill { .. }
+                    | AIAgentActionResultType::ReadMCPResource { .. }
+                    | AIAgentActionResultType::SuggestPrompt { .. }
+                    | AIAgentActionResultType::ReadDocuments(_)
+                    | AIAgentActionResultType::EditDocuments(_)
+                    | AIAgentActionResultType::CreateDocuments(_) => {}
+
+                    // TODO(AGENT-2282): figure out whether there's any reasonable way to
+                    // do redaction here (probably not).
+                    AIAgentActionResultType::UseComputer(_) => {}
+
+                    // Request computer use just contains screen dimensions, no secrets
+                    AIAgentActionResultType::RequestComputerUse(_) => {}
+
+                    // FetchConversation results contain tasks returned from the server,
+                    // which were already redacted before being sent as client inputs.
+                    // (client inputs -> redaction -> server request -> task messages)
+                    AIAgentActionResultType::FetchConversation(_) => {}
+
+                    // StartAgent results contain only an agent ID string, no secrets
+                    AIAgentActionResultType::StartAgent(_) => {}
+
+                    // SendMessageToAgent results contain only a message ID or error string, no secrets
+                    AIAgentActionResultType::SendMessageToAgent(_) => {}
+                    // TransferShellCommandControlToUser result - similar to WriteToLongRunningShellCommand
+                    AIAgentActionResultType::TransferShellCommandControlToUser(result) => {
+                        match result {
+                            TransferShellCommandControlToUserResult::Snapshot {
+                                grid_contents,
+                                ..
+                            } => redact_secrets(grid_contents),
+                            TransferShellCommandControlToUserResult::CommandFinished {
+                                output,
+                                ..
+                            } => redact_secrets(output),
+                            TransferShellCommandControlToUserResult::Error(_)
+                            | TransferShellCommandControlToUserResult::Cancelled => {}
+                        }
+                    }
+                    AIAgentActionResultType::AskUserQuestion(result) => {
+                        redact_ask_user_question_result(result);
+                    }
+                }
+            }
+            AIAgentInput::FetchReviewComments { repo_path, context } => {
+                redact_secrets(repo_path);
+                redact_context(Arc::make_mut(context));
+            }
+            AIAgentInput::InvokeSkill {
+                context,
+                skill,
+                user_query,
+            } => {
+                redact_context(Arc::make_mut(context));
+                redact_secrets(&mut skill.content);
+                if let Some(user_query) = user_query {
+                    redact_secrets(&mut user_query.query);
+                    for attachment in user_query.referenced_attachments.values_mut() {
+                        redact_attachment(attachment);
+                    }
+                }
+            }
+            AIAgentInput::PassiveSuggestionResult {
+                trigger,
+                suggestion,
+                context,
+            } => {
+                redact_context(Arc::make_mut(context));
+                match suggestion {
+                    PassiveSuggestionResultType::Prompt { prompt } => redact_secrets(prompt),
+                    PassiveSuggestionResultType::CodeDiff { diffs, .. } => {
+                        for diff in diffs {
+                            redact_secrets(&mut diff.file_path);
+                            redact_secrets(&mut diff.search);
+                            redact_secrets(&mut diff.replace);
+                        }
+                    }
+                }
+                if let Some(PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger)) =
+                    trigger
+                {
+                    redact_secrets(&mut shell_trigger.executed_shell_command.command);
+                    redact_secrets(&mut shell_trigger.executed_shell_command.output);
+                    for file in shell_trigger.relevant_files.iter_mut() {
+                        if let AnyFileContent::StringContent(content) = &mut file.content {
+                            redact_secrets(content);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+fn redact_ask_user_question_result(result: &mut AskUserQuestionResult) {
+    match result {
+        AskUserQuestionResult::Success { answers } => {
+            for answer in answers {
+                if let AskUserQuestionAnswerItem::Answered { other_text, .. } = answer {
+                    redact_secrets(other_text);
+                }
+            }
+        }
+        AskUserQuestionResult::SkippedByAutoApprove { .. } => {}
+        AskUserQuestionResult::Error(message) => redact_secrets(message),
+        AskUserQuestionResult::Cancelled => {}
+    }
+}
+fn redact_context(context: &mut [AIAgentContext]) {
+    for context_item in context {
+        match context_item {
+            AIAgentContext::Block(context) => {
+                redact_secrets(&mut context.command);
+                redact_secrets(&mut context.output);
+            }
+            AIAgentContext::SelectedText(text) => {
+                redact_secrets(text);
+            }
+            // Other context types don't contain user-provided text that needs redaction
+            AIAgentContext::Directory { .. }
+            | AIAgentContext::ExecutionEnvironment(_)
+            | AIAgentContext::CurrentTime { .. }
+            | AIAgentContext::Image(_)
+            | AIAgentContext::Codebase { .. }
+            | AIAgentContext::ProjectRules { .. }
+            | AIAgentContext::Git { .. }
+            | AIAgentContext::File(_)
+            | AIAgentContext::Skills { .. } => {}
+        }
+    }
+}
+
+fn redact_attachment(attachment: &mut AIAgentAttachment) {
+    match attachment {
+        AIAgentAttachment::PlainText(text) => {
+            redact_secrets(text);
+        }
+        AIAgentAttachment::Block(BlockContext {
+            command, output, ..
+        }) => {
+            redact_secrets(command);
+            redact_secrets(output);
+        }
+        AIAgentAttachment::DriveObject { payload, .. } => {
+            if let Some(drive_payload) = payload {
+                match drive_payload {
+                    crate::ai::agent::DriveObjectPayload::Workflow {
+                        name,
+                        description,
+                        command,
+                    } => {
+                        redact_secrets(name);
+                        redact_secrets(description);
+                        redact_secrets(command);
+                    }
+                    crate::ai::agent::DriveObjectPayload::Notebook { title, content } => {
+                        redact_secrets(title);
+                        redact_secrets(content);
+                    }
+                    crate::ai::agent::DriveObjectPayload::GenericStringObject {
+                        payload, ..
+                    } => {
+                        redact_secrets(payload);
+                    }
+                }
+            }
+        }
+        AIAgentAttachment::DiffHunk {
+            file_path,
+            diff_content,
+            ..
+        } => {
+            redact_secrets(file_path);
+            redact_secrets(diff_content);
+        }
+        AIAgentAttachment::DiffSet { file_diffs, .. } => {
+            for hunks in file_diffs.values_mut() {
+                for hunk in hunks {
+                    redact_secrets(&mut hunk.diff_content);
+                }
+            }
+        }
+        AIAgentAttachment::DocumentContent { content, .. } => {
+            redact_secrets(content);
+        }
+        // FilePathReference only contains a file ID and filename, no user secrets.
+        AIAgentAttachment::FilePathReference { .. } => {}
+    }
+}

app/src/ai/agent/suggestion_test.rs

@@ -0,0 +1,159 @@
+use crate::ai::agent::{
+    SuggestedAgentModeWorkflow, SuggestedLoggingId, SuggestedRule, Suggestions,
+};
+
+#[test]
+fn test_extend_suggestions() {
+    // Create base suggestions
+    let mut base_suggestions = Suggestions {
+        rules: vec![
+            SuggestedRule {
+                name: "rule1".into(),
+                content: "content1".into(),
+                logging_id: SuggestedLoggingId::from("id1".to_string()),
+            },
+            SuggestedRule {
+                name: "rule2".into(),
+                content: "content2".into(),
+                logging_id: SuggestedLoggingId::from("id2".to_string()),
+            },
+        ],
+        agent_mode_workflows: vec![
+            SuggestedAgentModeWorkflow {
+                name: "workflow1".into(),
+                prompt: "prompt1".into(),
+                logging_id: SuggestedLoggingId::from("wid1".to_string()),
+            },
+            SuggestedAgentModeWorkflow {
+                name: "workflow2".into(),
+                prompt: "prompt2".into(),
+                logging_id: SuggestedLoggingId::from("wid2".to_string()),
+            },
+        ],
+    };
+
+    // Create additional suggestions with both unique and duplicate logging_ids
+    let additional_suggestions = Suggestions {
+        rules: vec![
+            // Duplicate logging_id but different name/content
+            SuggestedRule {
+                name: "rule1_modified".into(),
+                content: "content1_modified".into(),
+                logging_id: SuggestedLoggingId::from("id1".to_string()),
+            },
+            // New unique rule
+            SuggestedRule {
+                name: "rule3".into(),
+                content: "content3".into(),
+                logging_id: SuggestedLoggingId::from("id3".to_string()),
+            },
+            // Another new unique rule
+            SuggestedRule {
+                name: "rule4".into(),
+                content: "content4".into(),
+                logging_id: SuggestedLoggingId::from("id4".to_string()),
+            },
+        ],
+        agent_mode_workflows: vec![
+            // Duplicate workflow logging_id but different name/prompt
+            SuggestedAgentModeWorkflow {
+                name: "workflow1_modified".into(),
+                prompt: "prompt1_modified".into(),
+                logging_id: SuggestedLoggingId::from("wid1".to_string()),
+            },
+            // New unique workflow
+            SuggestedAgentModeWorkflow {
+                name: "workflow3".into(),
+                prompt: "prompt3".into(),
+                logging_id: SuggestedLoggingId::from("wid3".to_string()),
+            },
+            // Another new unique workflow
+            SuggestedAgentModeWorkflow {
+                name: "workflow4".into(),
+                prompt: "prompt4".into(),
+                logging_id: SuggestedLoggingId::from("wid4".to_string()),
+            },
+        ],
+    };
+
+    // Extend base suggestions with additional ones
+    base_suggestions.extend(&additional_suggestions);
+
+    // Verify rules
+
+    // Verify the length (should be 4 because one was a duplicate)
+    assert_eq!(base_suggestions.rules.len(), 4);
+
+    // Verify that original rules with id1 and id2 are still present and unchanged
+    assert!(base_suggestions
+        .rules
+        .iter()
+        .any(|r| r.logging_id.to_string() == "id1"
+            && r.name == "rule1"
+            && r.content == "content1"));
+    assert!(base_suggestions
+        .rules
+        .iter()
+        .any(|r| r.logging_id.to_string() == "id2"
+            && r.name == "rule2"
+            && r.content == "content2"));
+
+    // Verify that new unique rules (id3 and id4) were added
+    assert!(base_suggestions
+        .rules
+        .iter()
+        .any(|r| r.logging_id.to_string() == "id3"
+            && r.name == "rule3"
+            && r.content == "content3"));
+    assert!(base_suggestions
+        .rules
+        .iter()
+        .any(|r| r.logging_id.to_string() == "id4"
+            && r.name == "rule4"
+            && r.content == "content4"));
+
+    // Verify that the modified version of id1 was not added (deduplication worked)
+    assert!(!base_suggestions
+        .rules
+        .iter()
+        .any(|r| r.logging_id.to_string() == "id1" && r.name == "rule1_modified"));
+
+    // Verify workflows
+
+    // Verify the length (should be 4 because one was a duplicate)
+    assert_eq!(base_suggestions.agent_mode_workflows.len(), 4);
+
+    // Verify that original workflows with wid1 and wid2 are still present and unchanged
+    assert!(base_suggestions
+        .agent_mode_workflows
+        .iter()
+        .any(|w| w.logging_id.to_string() == "wid1"
+            && w.name == "workflow1"
+            && w.prompt == "prompt1"));
+    assert!(base_suggestions
+        .agent_mode_workflows
+        .iter()
+        .any(|w| w.logging_id.to_string() == "wid2"
+            && w.name == "workflow2"
+            && w.prompt == "prompt2"));
+
+    // Verify that new unique workflows (wid3 and wid4) were added
+    assert!(base_suggestions
+        .agent_mode_workflows
+        .iter()
+        .any(|w| w.logging_id.to_string() == "wid3"
+            && w.name == "workflow3"
+            && w.prompt == "prompt3"));
+    assert!(base_suggestions
+        .agent_mode_workflows
+        .iter()
+        .any(|w| w.logging_id.to_string() == "wid4"
+            && w.name == "workflow4"
+            && w.prompt == "prompt4"));
+
+    // Verify that the modified version of wid1 was not added (deduplication worked)
+    assert!(!base_suggestions
+        .agent_mode_workflows
+        .iter()
+        .any(|w| w.logging_id.to_string() == "wid1" && w.name == "workflow1_modified"));
+}

app/src/ai/agent/task/helper.rs

@@ -0,0 +1,214 @@
+//! This module contains traits and trait implementations for exposing helper methods for accessing
+//! proto fields.
+use warp_multi_agent_api as api;
+
+pub trait TaskExt {
+    fn parent_id(&self) -> Option<&str>;
+}
+
+impl TaskExt for api::Task {
+    fn parent_id(&self) -> Option<&str> {
+        self.dependencies
+            .as_ref()
+            .map(|deps| deps.parent_task_id.as_str())
+            .filter(|id| !id.is_empty())
+    }
+}
+
+pub trait MessageExt {
+    fn todos_op(&self) -> Option<&api::message::update_todos::Operation>;
+    fn tool_call(&self) -> Option<&api::message::ToolCall>;
+    fn tool_call_mut(&mut self) -> Option<&mut api::message::ToolCall>;
+    fn tool_call_result(&self) -> Option<&api::message::ToolCallResult>;
+}
+
+pub trait ToolCallExt {
+    fn subagent(&self) -> Option<&api::message::tool_call::Subagent>;
+    fn subagent_mut(&mut self) -> Option<&mut api::message::tool_call::Subagent>;
+}
+
+pub trait ToolExt {
+    fn name(&self) -> &'static str;
+}
+
+pub trait SubagentExt {
+    fn is_cli(&self) -> bool;
+    fn is_advice(&self) -> bool;
+    fn is_computer_use(&self) -> bool;
+    fn is_summarization(&self) -> bool;
+    fn is_conversation_search(&self) -> bool;
+    fn is_warp_documentation_search(&self) -> bool;
+    fn type_name(&self) -> &'static str;
+}
+
+impl MessageExt for api::Message {
+    fn todos_op(&self) -> Option<&api::message::update_todos::Operation> {
+        self.message.as_ref().and_then(|message| {
+            if let api::message::Message::UpdateTodos(update) = message {
+                update.operation.as_ref()
+            } else {
+                None
+            }
+        })
+    }
+
+    fn tool_call(&self) -> Option<&api::message::ToolCall> {
+        self.message.as_ref().and_then(|message| {
+            if let api::message::Message::ToolCall(tool_call) = message {
+                Some(tool_call)
+            } else {
+                None
+            }
+        })
+    }
+
+    fn tool_call_mut(&mut self) -> Option<&mut api::message::ToolCall> {
+        self.message.as_mut().and_then(|message| {
+            if let api::message::Message::ToolCall(tool_call) = message {
+                Some(tool_call)
+            } else {
+                None
+            }
+        })
+    }
+
+    fn tool_call_result(&self) -> Option<&api::message::ToolCallResult> {
+        self.message.as_ref().and_then(|message| {
+            if let api::message::Message::ToolCallResult(result) = message {
+                Some(result)
+            } else {
+                None
+            }
+        })
+    }
+}
+
+impl ToolCallExt for api::message::ToolCall {
+    fn subagent(&self) -> Option<&api::message::tool_call::Subagent> {
+        match self.tool.as_ref() {
+            Some(api::message::tool_call::Tool::Subagent(subagent)) => Some(subagent),
+            _ => None,
+        }
+    }
+
+    fn subagent_mut(&mut self) -> Option<&mut api::message::tool_call::Subagent> {
+        match self.tool.as_mut() {
+            Some(api::message::tool_call::Tool::Subagent(subagent)) => Some(subagent),
+            _ => None,
+        }
+    }
+}
+
+impl ToolExt for api::message::tool_call::Tool {
+    fn name(&self) -> &'static str {
+        use api::message::tool_call::Tool;
+        match self {
+            Tool::RunShellCommand(_) => "run_shell_command",
+            Tool::SearchCodebase(_) => "search_codebase",
+            Tool::ReadFiles(_) => "read_files",
+            Tool::UploadFileArtifact(_) => "upload_artifact",
+            Tool::ApplyFileDiffs(_) => "apply_file_diffs",
+            Tool::Grep(_) => "grep",
+            #[allow(deprecated)]
+            Tool::FileGlob(_) => "file_glob",
+            Tool::FileGlobV2(_) => "file_glob_v2",
+            Tool::ReadMcpResource(_) => "read_mcp_resource",
+            Tool::CallMcpTool(_) => "call_mcp_tool",
+            Tool::WriteToLongRunningShellCommand(_) => "write_to_lrc",
+            Tool::ReadDocuments(_) => "read_documents",
+            Tool::EditDocuments(_) => "edit_documents",
+            Tool::CreateDocuments(_) => "create_documents",
+            Tool::ReadShellCommandOutput(_) => "read_shell_command_output",
+            Tool::UseComputer(_) => "use_computer",
+            Tool::RequestComputerUse(_) => "request_computer_use",
+            Tool::FetchConversation(_) => "fetch_conversation",
+            Tool::InsertReviewComments(_) => "insert_review_comments",
+            Tool::ReadSkill(_) => "read_skill",
+            Tool::SuggestPlan(_) => "suggest_plan",
+            Tool::SuggestCreatePlan(_) => "suggest_create_plan",
+            Tool::SuggestNewConversation(_) => "suggest_new_conversation",
+            Tool::SuggestPrompt(_) => "suggest_prompt",
+            Tool::OpenCodeReview(_) => "open_code_review",
+            Tool::InitProject(_) => "init_project",
+            Tool::StartAgent(_) => "start_agent",
+            // Keep the logical tool name stable across the v1/v2 schema split so analytics,
+            // history, and UI handling continue to treat both as the same tool.
+            Tool::StartAgentV2(_) => "start_agent",
+            Tool::Server(_) => "server",
+            Tool::Subagent(_) => "subagent",
+            Tool::AskUserQuestion(_) => "ask_user_question",
+            Tool::SendMessageToAgent(_) => "send_message_to_agent",
+            Tool::TransferShellCommandControlToUser(_) => "transfer_shell_command_control",
+        }
+    }
+}
+
+impl SubagentExt for api::message::tool_call::Subagent {
+    fn is_cli(&self) -> bool {
+        self.metadata.as_ref().is_some_and(|metadata| {
+            matches!(
+                metadata,
+                api::message::tool_call::subagent::Metadata::Cli(_)
+            )
+        })
+    }
+
+    fn is_advice(&self) -> bool {
+        self.metadata.as_ref().is_some_and(|metadata| {
+            matches!(
+                metadata,
+                api::message::tool_call::subagent::Metadata::Advice(_)
+            )
+        })
+    }
+
+    fn is_computer_use(&self) -> bool {
+        self.metadata.as_ref().is_some_and(|metadata| {
+            matches!(
+                metadata,
+                api::message::tool_call::subagent::Metadata::ComputerUse(_)
+            )
+        })
+    }
+
+    fn is_summarization(&self) -> bool {
+        self.metadata.as_ref().is_some_and(|metadata| {
+            matches!(
+                metadata,
+                api::message::tool_call::subagent::Metadata::Summarization(_)
+            )
+        })
+    }
+
+    fn is_conversation_search(&self) -> bool {
+        self.metadata.as_ref().is_some_and(|metadata| {
+            matches!(
+                metadata,
+                api::message::tool_call::subagent::Metadata::ConversationSearch(_)
+            )
+        })
+    }
+
+    fn is_warp_documentation_search(&self) -> bool {
+        self.metadata.as_ref().is_some_and(|metadata| {
+            matches!(
+                metadata,
+                api::message::tool_call::subagent::Metadata::WarpDocumentationSearch(_)
+            )
+        })
+    }
+
+    fn type_name(&self) -> &'static str {
+        use api::message::tool_call::subagent::Metadata;
+        match &self.metadata {
+            Some(Metadata::Cli(_)) => "cli",
+            Some(Metadata::Research(_)) => "research",
+            Some(Metadata::Advice(_)) => "advice",
+            Some(Metadata::ComputerUse(_)) => "computer_use",
+            Some(Metadata::Summarization(_)) => "summarization",
+            Some(Metadata::ConversationSearch(_)) => "conversation_search",
+            Some(Metadata::WarpDocumentationSearch(_)) => "warp_documentation_search",
+            None => "unknown",
+        }
+    }
+}

app/src/ai/agent/task/transaction.rs

@@ -0,0 +1,54 @@
+use std::collections::HashMap;
+
+use crate::ai::agent::task::TaskId;
+
+use super::Task;
+
+/// Keeps track of the state of tasks before they are modified.
+/// Messages are assumed to be only updated during the same transaction
+/// in which they were added, so we can clean up message by simply
+/// deleting them.
+#[derive(Debug, Clone)]
+pub struct Transaction {
+    saved_tasks: HashMap<TaskId, SavedTask>,
+}
+
+/// Saves state for either a newly added task or a pre-existing task
+/// modified during a transaction.
+#[derive(Debug, Clone)]
+pub enum SavedTask {
+    New(TaskId),
+    Existing(Box<Task>),
+}
+
+impl Transaction {
+    pub fn new() -> Self {
+        Self {
+            saved_tasks: HashMap::new(),
+        }
+    }
+
+    /// A map of the tasks modified in this transaction.
+    pub fn saved_tasks(self) -> HashMap<TaskId, SavedTask> {
+        self.saved_tasks
+    }
+
+    /// Saves a SavedTask::New to the transaction, representing a newly added task.
+    pub fn checkpoint_new_task(&mut self, task_id: &TaskId) {
+        if !self.saved_tasks.contains_key(task_id) {
+            let task = SavedTask::New(task_id.clone());
+            self.saved_tasks.insert(task_id.clone(), task);
+        }
+    }
+
+    /// Saves a SavedTask::Existing to the transaction, representing an existing
+    /// task which is being modified.
+    pub fn checkpoint_task(&mut self, task: &Task) {
+        if !self.saved_tasks.contains_key(task.id()) {
+            self.saved_tasks.insert(
+                task.id().clone(),
+                SavedTask::Existing(Box::new(task.clone())),
+            );
+        }
+    }
+}

app/src/ai/agent/task_store.rs

@@ -0,0 +1,346 @@
+use std::collections::HashMap;
+
+use warp_multi_agent_api as api;
+
+use crate::ai::{
+    agent::{AIAgentContext, AIAgentInput},
+    skills::SkillDescriptor,
+};
+
+use super::{
+    task::{
+        helper::{MessageExt, ToolCallExt},
+        Task, TaskId,
+    },
+    AIAgentExchange, AIAgentExchangeId, AIAgentOutputMessageType,
+};
+
+#[derive(Debug, Clone)]
+struct ExchangeRef {
+    task_id: TaskId,
+    exchange_index: usize,
+}
+
+/// Task storage with a linearized exchange index for O(1) first/last access.
+#[derive(Debug, Clone)]
+pub struct TaskStore {
+    root_task_id: TaskId,
+    tasks: HashMap<TaskId, Task>,
+    linearized_refs: Vec<ExchangeRef>,
+}
+
+impl TaskStore {
+    pub fn with_root_task(root_task: Task) -> Self {
+        let root_task_id = root_task.id().clone();
+        let mut store = Self {
+            tasks: HashMap::new(),
+            linearized_refs: Vec::new(),
+            root_task_id: root_task_id.clone(),
+        };
+        store.tasks.insert(root_task_id, root_task);
+        store.rebuild_linearized_refs_index();
+        store
+    }
+
+    /// Creates a TaskStore from an existing HashMap of tasks.
+    /// Rebuilds the linearized index after construction.
+    pub fn from_tasks(tasks: HashMap<TaskId, Task>, root_task_id: TaskId) -> Self {
+        let mut store = Self {
+            tasks,
+            linearized_refs: Vec::new(),
+            root_task_id,
+        };
+        store.rebuild_linearized_refs_index();
+        store
+    }
+
+    pub fn root_task_id(&self) -> &TaskId {
+        &self.root_task_id
+    }
+
+    pub fn get(&self, task_id: &TaskId) -> Option<&Task> {
+        self.tasks.get(task_id)
+    }
+
+    pub fn tasks(&self) -> impl Iterator<Item = &Task> {
+        self.tasks.values()
+    }
+
+    pub fn task_count(&self) -> usize {
+        self.tasks.len()
+    }
+
+    /// Appends an exchange to a task and rebuilds the index.
+    /// Returns true if the task was found and the exchange was appended.
+    pub fn append_exchange(&mut self, task_id: &TaskId, exchange: AIAgentExchange) -> bool {
+        let Some(task) = self.tasks.get_mut(task_id) else {
+            return false;
+        };
+        task.append_exchange(exchange);
+        self.rebuild_linearized_refs_index();
+        true
+    }
+
+    /// Removes an exchange from a task and rebuilds the index.
+    /// Returns the removed exchange if found.
+    pub fn remove_task_exchange(
+        &mut self,
+        task_id: &TaskId,
+        exchange_id: AIAgentExchangeId,
+    ) -> Option<AIAgentExchange> {
+        let task = self.tasks.get_mut(task_id)?;
+        let exchange = task.remove_exchange(exchange_id)?;
+        self.rebuild_linearized_refs_index();
+        Some(exchange)
+    }
+
+    /// Returns a mutable reference to an exchange by its ID, searching all tasks.
+    pub fn exchange_mut(&mut self, exchange_id: AIAgentExchangeId) -> Option<&mut AIAgentExchange> {
+        for task in self.tasks.values_mut() {
+            if let Some(exchange) = task.exchange_mut(exchange_id) {
+                return Some(exchange);
+            }
+        }
+        None
+    }
+
+    /// Modifies a task via the provided closure and rebuilds the exchange index
+    /// if exchanges changed.
+    pub fn modify_task<R>(
+        &mut self,
+        task_id: &TaskId,
+        f: impl FnOnce(&mut Task) -> R,
+    ) -> Option<R> {
+        let exchange_count_before = self.tasks.get(task_id)?.exchanges_len();
+        let task = self.tasks.get_mut(task_id)?;
+        let result = f(task);
+        let exchange_count_after = self
+            .tasks
+            .get(task_id)
+            .map(|t| t.exchanges_len())
+            .unwrap_or(0);
+        if exchange_count_before != exchange_count_after {
+            self.rebuild_linearized_refs_index();
+        }
+        Some(result)
+    }
+
+    /// Modifies the root task via the provided closure and rebuilds the exchange index if exchanges changed.
+    pub fn modify_root_task<R>(&mut self, f: impl FnOnce(&mut Task) -> R) -> Option<R> {
+        let root_task_id = self.root_task_id.clone();
+        self.modify_task(&root_task_id, f)
+    }
+
+    pub fn root_task(&self) -> Option<&Task> {
+        self.tasks.get(&self.root_task_id)
+    }
+
+    /// Sets or replaces the root task, removing any previous root if it exists.
+    pub fn set_root_task(&mut self, root_task: Task) {
+        // Remove the old root task and its exchange refs
+        let old_root_id = self.root_task_id.clone();
+        self.remove(&old_root_id);
+
+        let new_root_id = root_task.id().clone();
+        self.root_task_id = new_root_id;
+        self.insert(root_task);
+    }
+
+    pub fn first_exchange(&self) -> Option<&AIAgentExchange> {
+        self.linearized_refs
+            .first()
+            .and_then(|r| self.lookup_exchange(r))
+    }
+
+    pub fn latest_exchange(&self) -> Option<&AIAgentExchange> {
+        self.linearized_refs
+            .last()
+            .and_then(|r| self.lookup_exchange(r))
+    }
+
+    pub fn exchange_count(&self) -> usize {
+        self.linearized_refs.len()
+    }
+
+    pub fn all_exchanges(&self) -> impl Iterator<Item = &AIAgentExchange> {
+        self.linearized_refs
+            .iter()
+            .filter_map(|r| self.lookup_exchange(r))
+    }
+
+    pub fn all_exchanges_rev(&self) -> impl Iterator<Item = &AIAgentExchange> {
+        self.linearized_refs
+            .iter()
+            .rev()
+            .filter_map(|r| self.lookup_exchange(r))
+    }
+
+    pub fn all_exchanges_by_task(&self) -> Vec<(TaskId, Vec<&AIAgentExchange>)> {
+        let mut result: Vec<(TaskId, Vec<&AIAgentExchange>)> = Vec::new();
+
+        for exchange_ref in &self.linearized_refs {
+            let Some(exchange) = self.lookup_exchange(exchange_ref) else {
+                continue;
+            };
+
+            // Check if we should append to the last group or start a new one
+            if let Some((last_task_id, exchanges)) = result.last_mut() {
+                if last_task_id == &exchange_ref.task_id {
+                    exchanges.push(exchange);
+                    continue;
+                }
+            }
+
+            // Start a new group
+            result.push((exchange_ref.task_id.clone(), vec![exchange]));
+        }
+
+        result
+    }
+
+    pub fn latest_skills(&self) -> Option<Vec<SkillDescriptor>> {
+        self.linearized_refs.iter().rev().find_map(|exchange_ref| {
+            let exchange = self.lookup_exchange(exchange_ref);
+
+            if let Some(exchange) = exchange {
+                let skills = exchange.input.iter().find_map(|input| {
+                    let context = match input {
+                        AIAgentInput::UserQuery { context, .. } => Some(context),
+                        AIAgentInput::ResumeConversation { context, .. } => Some(context),
+                        AIAgentInput::ActionResult { context, .. } => Some(context),
+                        AIAgentInput::TriggerPassiveSuggestion { context, .. } => Some(context),
+                        _ => None,
+                    };
+
+                    context.and_then(|ctx| {
+                        ctx.iter().find_map(|context| {
+                            if let AIAgentContext::Skills { skills } = context {
+                                Some(skills)
+                            } else {
+                                None
+                            }
+                        })
+                    })
+                });
+
+                skills.cloned()
+            } else {
+                None
+            }
+        })
+    }
+
+    /// Returns all messages in linearized DFS order, interleaving subtask messages
+    /// immediately after their parent subagent call messages.
+    pub fn all_linearized_messages(&self) -> Vec<&api::Message> {
+        fn collect_messages_dfs<'a>(
+            me: &'a TaskStore,
+            messages: &mut Vec<&'a api::Message>,
+            task: &'a Task,
+        ) {
+            for message in task.messages() {
+                messages.push(message);
+                // If this message is a subagent call, recursively add subtask messages
+                if let Some(subagent_call) = message
+                    .tool_call()
+                    .and_then(|tc: &api::message::ToolCall| tc.subagent())
+                {
+                    if let Some(subtask) = me.get(&TaskId::new(subagent_call.task_id.clone())) {
+                        collect_messages_dfs(me, messages, subtask);
+                    }
+                }
+            }
+        }
+
+        let mut messages = Vec::new();
+        if let Some(root_task) = self.root_task() {
+            collect_messages_dfs(self, &mut messages, root_task);
+        }
+        messages
+    }
+
+    pub fn insert(&mut self, task: Task) {
+        self.tasks.insert(task.id().clone(), task);
+        self.rebuild_linearized_refs_index();
+    }
+
+    pub fn remove(&mut self, task_id: &TaskId) -> Option<Task> {
+        let task = self.tasks.remove(task_id)?;
+        self.linearized_refs.retain(|r| &r.task_id != task_id);
+        Some(task)
+    }
+
+    fn lookup_exchange(&self, r: &ExchangeRef) -> Option<&AIAgentExchange> {
+        self.tasks
+            .get(&r.task_id)?
+            .exchanges()
+            .nth(r.exchange_index)
+    }
+
+    /// Rebuilds the linearized index from scratch using DFS traversal.
+    fn rebuild_linearized_refs_index(&mut self) {
+        self.linearized_refs = Self::build_linearized_refs(&self.tasks, &self.root_task_id);
+    }
+
+    /// Builds linearized exchange refs via DFS traversal without mutating self.
+    /// This allows us to borrow `tasks` immutably throughout the traversal.
+    fn build_linearized_refs(
+        tasks: &HashMap<TaskId, Task>,
+        root_task_id: &TaskId,
+    ) -> Vec<ExchangeRef> {
+        let mut refs = Vec::new();
+
+        fn append_refs_for_task(
+            tasks: &HashMap<TaskId, Task>,
+            refs: &mut Vec<ExchangeRef>,
+            task: &Task,
+        ) {
+            let task_id = task.id().clone();
+
+            for (exchange_index, exchange) in task.exchanges().enumerate() {
+                refs.push(ExchangeRef {
+                    task_id: task_id.clone(),
+                    exchange_index,
+                });
+
+                // Check for subagent calls in the exchange output.
+                if let Some(output) = exchange.output_status.output() {
+                    for output_message in output.get().messages.iter() {
+                        if let AIAgentOutputMessageType::Subagent(subagent_call) =
+                            &output_message.message
+                        {
+                            if let Some(subtask) =
+                                tasks.get(&TaskId::new(subagent_call.task_id.clone()))
+                            {
+                                append_refs_for_task(tasks, refs, subtask);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        if let Some(root_task) = tasks.get(root_task_id) {
+            append_refs_for_task(tasks, &mut refs, root_task);
+        }
+
+        refs
+    }
+}
+
+#[cfg(test)]
+mod testing {
+    use crate::ai::agent::task::TaskId;
+
+    use super::TaskStore;
+
+    impl TaskStore {
+        pub fn contains(&self, task_id: &TaskId) -> bool {
+            self.tasks.contains_key(task_id)
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "task_store_tests.rs"]
+mod tests;

app/src/ai/agent/task_store_tests.rs

@@ -0,0 +1,625 @@
+use std::collections::HashSet;
+
+use chrono::Local;
+use uuid::Uuid;
+
+use crate::ai::{
+    agent::{
+        task::{Task, TaskId},
+        AIAgentExchange, AIAgentExchangeId, AIAgentOutput, AIAgentOutputMessage,
+        AIAgentOutputMessageType, AIAgentOutputStatus, FinishedAIAgentOutput, MessageId, Shared,
+        SubagentCall,
+    },
+    llms::LLMId,
+};
+
+use super::TaskStore;
+
+fn create_test_exchange() -> AIAgentExchange {
+    AIAgentExchange {
+        id: AIAgentExchangeId::new(),
+        input: vec![],
+        output_status: AIAgentOutputStatus::Streaming { output: None },
+        added_message_ids: HashSet::new(),
+        start_time: Local::now(),
+        finish_time: None,
+        time_to_first_token_ms: None,
+        working_directory: None,
+        model_id: LLMId::from(""),
+        request_cost: None,
+        coding_model_id: LLMId::from(""),
+        cli_agent_model_id: LLMId::from(""),
+        computer_use_model_id: LLMId::from(""),
+        response_initiator: None,
+    }
+}
+
+fn create_test_task_with_exchanges(exchange_count: usize) -> Task {
+    let mut task = Task::new_optimistic_root();
+    for _ in 0..exchange_count {
+        task.append_exchange(create_test_exchange());
+    }
+    task
+}
+
+fn create_test_subtask_with_exchanges(exchange_count: usize) -> Task {
+    use crate::terminal::model::block::BlockId;
+    let mut task = Task::new_optimistic_cli_agent_subtask(BlockId::new());
+    for _ in 0..exchange_count {
+        task.append_exchange(create_test_exchange());
+    }
+    task
+}
+
+/// Creates an exchange with a finished output containing a subagent call to the given task_id.
+fn create_exchange_with_subagent_call(subtask_id: &TaskId) -> AIAgentExchange {
+    let output = AIAgentOutput {
+        messages: vec![AIAgentOutputMessage {
+            id: MessageId::new(Uuid::new_v4().to_string()),
+            message: AIAgentOutputMessageType::Subagent(SubagentCall {
+                task_id: subtask_id.to_string(),
+                subagent_type: crate::ai::agent::SubagentType::Unknown,
+            }),
+            citations: vec![],
+        }],
+        citations: vec![],
+        server_output_id: None,
+        api_metadata_bytes: None,
+        suggestions: None,
+        telemetry_events: vec![],
+        model_info: None,
+        request_cost: None,
+    };
+
+    AIAgentExchange {
+        id: AIAgentExchangeId::new(),
+        input: vec![],
+        output_status: AIAgentOutputStatus::Finished {
+            finished_output: FinishedAIAgentOutput::Success {
+                output: Shared::new(output),
+            },
+        },
+        added_message_ids: HashSet::new(),
+        start_time: Local::now(),
+        finish_time: None,
+        time_to_first_token_ms: None,
+        working_directory: None,
+        model_id: LLMId::from(""),
+        request_cost: None,
+        coding_model_id: LLMId::from(""),
+        cli_agent_model_id: LLMId::from(""),
+        computer_use_model_id: LLMId::from(""),
+        response_initiator: None,
+    }
+}
+
+#[test]
+fn test_with_root_task() {
+    let task = create_test_task_with_exchanges(3);
+    let task_id = task.id().clone();
+    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+
+    let store = TaskStore::with_root_task(task);
+
+    assert_eq!(store.task_count(), 1);
+    assert_eq!(store.exchange_count(), 3);
+    assert_eq!(store.root_task_id(), &task_id);
+    assert_eq!(store.root_task().expect("task exists").id(), &task_id);
+    assert_eq!(store.first_exchange().map(|e| e.id), Some(exchange_ids[0]));
+    assert_eq!(store.latest_exchange().map(|e| e.id), Some(exchange_ids[2]));
+}
+
+#[test]
+fn test_first_and_latest_exchange_o1() {
+    let task = create_test_task_with_exchanges(5);
+    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+    let store = TaskStore::with_root_task(task);
+
+    // These should be O(1) operations
+    let first = store.first_exchange().expect("has exchanges");
+    let latest = store.latest_exchange().expect("has exchanges");
+
+    assert_eq!(first.id, exchange_ids[0]);
+    assert_eq!(latest.id, exchange_ids[4]);
+}
+
+#[test]
+fn test_insert_subtask() {
+    // Create root task with 1 exchange, then we'll add a subagent call exchange
+    let root_task = create_test_task_with_exchanges(1);
+    let root_task_id = root_task.id().clone();
+    let mut store = TaskStore::with_root_task(root_task);
+
+    // Create subtask with 1 exchange
+    let subtask = create_test_subtask_with_exchanges(1);
+    let subtask_id = subtask.id().clone();
+
+    // Add exchange with subagent call to root task BEFORE inserting subtask
+    let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
+    store.append_exchange(&root_task_id, subagent_exchange);
+
+    // Now insert the subtask - its exchanges should be included via the subagent call
+    store.insert(subtask);
+
+    assert_eq!(store.task_count(), 2);
+    // 1 root exchange + 1 subagent call exchange + 1 subtask exchange = 3
+    assert_eq!(store.exchange_count(), 3);
+    assert!(store.get(&root_task_id).is_some());
+    assert!(store.get(&subtask_id).is_some());
+    assert!(store.contains(&subtask_id));
+}
+
+#[test]
+fn test_remove_task() {
+    let task = create_test_task_with_exchanges(3);
+    let task_id = task.id().clone();
+    let mut store = TaskStore::with_root_task(task);
+
+    assert_eq!(store.task_count(), 1);
+    assert_eq!(store.exchange_count(), 3);
+
+    let removed = store.remove(&task_id);
+    assert!(removed.is_some());
+    assert_eq!(store.task_count(), 0);
+    assert_eq!(store.exchange_count(), 0);
+    assert!(store.get(&task_id).is_none());
+}
+
+#[test]
+fn test_remove_nonexistent_task() {
+    let task = create_test_task_with_exchanges(2);
+    let mut store = TaskStore::with_root_task(task);
+
+    let nonexistent_id = TaskId::new(Uuid::new_v4().to_string());
+    let removed = store.remove(&nonexistent_id);
+    assert!(removed.is_none());
+    assert_eq!(store.task_count(), 1);
+}
+
+#[test]
+fn test_all_exchanges_iteration() {
+    let task = create_test_task_with_exchanges(4);
+    let expected_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+    let store = TaskStore::with_root_task(task);
+
+    let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+
+    assert_eq!(actual_ids, expected_ids);
+}
+
+#[test]
+fn test_all_exchanges_by_task() {
+    let task = create_test_task_with_exchanges(3);
+    let task_id = task.id().clone();
+    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+    let store = TaskStore::with_root_task(task);
+
+    let by_task = store.all_exchanges_by_task();
+    assert_eq!(by_task.len(), 1);
+    assert_eq!(by_task[0].0, task_id);
+    assert_eq!(by_task[0].1.len(), 3);
+
+    let actual_ids: Vec<_> = by_task[0].1.iter().map(|e| e.id).collect();
+    assert_eq!(actual_ids, exchange_ids);
+}
+
+#[test]
+fn test_set_root_task_replaces_old() {
+    let task1 = create_test_task_with_exchanges(2);
+    let task1_id = task1.id().clone();
+    let mut store = TaskStore::with_root_task(task1);
+
+    let task2 = create_test_task_with_exchanges(3);
+    let task2_id = task2.id().clone();
+    let task2_exchange_ids: Vec<_> = task2.exchanges().map(|e| e.id).collect();
+    store.set_root_task(task2);
+
+    assert_eq!(store.task_count(), 1);
+    assert_eq!(store.exchange_count(), 3);
+    assert!(store.get(&task1_id).is_none());
+    assert!(store.get(&task2_id).is_some());
+    assert_eq!(store.root_task_id(), &task2_id);
+
+    let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(actual_ids, task2_exchange_ids);
+}
+
+#[test]
+fn test_append_exchange() {
+    let task = create_test_task_with_exchanges(2);
+    let task_id = task.id().clone();
+    let mut store = TaskStore::with_root_task(task);
+
+    let new_exchange = create_test_exchange();
+    let new_exchange_id = new_exchange.id;
+
+    let result = store.append_exchange(&task_id, new_exchange);
+    assert!(result);
+    assert_eq!(store.exchange_count(), 3);
+
+    // Verify the new exchange is accessible
+    assert_eq!(store.latest_exchange().map(|e| e.id), Some(new_exchange_id));
+}
+
+#[test]
+fn test_remove_task_exchange() {
+    let task = create_test_task_with_exchanges(3);
+    let task_id = task.id().clone();
+    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+    let mut store = TaskStore::with_root_task(task);
+
+    // First verify initial state
+    assert_eq!(store.exchange_count(), 3);
+
+    // Remove the middle exchange
+    let removed = store.remove_task_exchange(&task_id, exchange_ids[1]);
+    assert!(removed.is_some());
+    assert_eq!(removed.unwrap().id, exchange_ids[1]);
+
+    // After removal, we should have 2 exchanges
+    assert_eq!(store.exchange_count(), 2);
+
+    // Verify the remaining exchanges are correct
+    let remaining_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(remaining_ids, vec![exchange_ids[0], exchange_ids[2]]);
+}
+
+#[test]
+fn test_append_exchange_to_nonexistent_task() {
+    let task = create_test_task_with_exchanges(1);
+    let mut store = TaskStore::with_root_task(task);
+
+    let nonexistent_id = TaskId::new(Uuid::new_v4().to_string());
+    let result = store.append_exchange(&nonexistent_id, create_test_exchange());
+    assert!(!result);
+    assert_eq!(store.exchange_count(), 1);
+}
+
+#[test]
+fn test_remove_exchange_from_nonexistent_task() {
+    let task = create_test_task_with_exchanges(1);
+    let exchange_id = task.exchanges().next().unwrap().id;
+    let mut store = TaskStore::with_root_task(task);
+
+    let nonexistent_task_id = TaskId::new(Uuid::new_v4().to_string());
+    let result = store.remove_task_exchange(&nonexistent_task_id, exchange_id);
+    assert!(result.is_none());
+    assert_eq!(store.exchange_count(), 1);
+}
+
+#[test]
+fn test_tasks_iteration() {
+    let task1 = create_test_task_with_exchanges(2);
+    let task1_id = task1.id().clone();
+    let mut store = TaskStore::with_root_task(task1);
+
+    let task2 = create_test_subtask_with_exchanges(1);
+    let task2_id = task2.id().clone();
+    store.insert(task2);
+
+    let task_ids: HashSet<_> = store.tasks().map(|t| t.id().clone()).collect();
+    assert_eq!(task_ids.len(), 2);
+    assert!(task_ids.contains(&task1_id));
+    assert!(task_ids.contains(&task2_id));
+}
+
+#[test]
+fn test_multiple_tasks_exchange_order() {
+    // Create root task with 1 exchange, then add subagent call exchange
+    let root_task = create_test_task_with_exchanges(1);
+    let root_task_id = root_task.id().clone();
+    let first_root_exchange_id = root_task.exchanges().next().unwrap().id;
+    let mut store = TaskStore::with_root_task(root_task);
+
+    // Create subtask with 2 exchanges
+    let subtask = create_test_subtask_with_exchanges(2);
+    let subtask_id = subtask.id().clone();
+    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();
+
+    // Add subagent call exchange to root
+    let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
+    let subagent_exchange_id = subagent_exchange.id;
+    store.append_exchange(&root_task_id, subagent_exchange);
+
+    // Insert subtask - now linked via subagent call
+    store.insert(subtask);
+
+    // 1 root + 1 subagent call + 2 subtask = 4 exchanges
+    assert_eq!(store.exchange_count(), 4);
+
+    // First/last should still work
+    assert_eq!(
+        store.first_exchange().map(|e| e.id),
+        Some(first_root_exchange_id)
+    );
+    assert_eq!(
+        store.latest_exchange().map(|e| e.id),
+        Some(subtask_exchange_ids[1])
+    );
+
+    // Order: root[0], root[subagent_call], subtask[0], subtask[1]
+    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(all_ids.len(), 4);
+    assert_eq!(all_ids[0], first_root_exchange_id);
+    assert_eq!(all_ids[1], subagent_exchange_id);
+    assert_eq!(all_ids[2], subtask_exchange_ids[0]);
+    assert_eq!(all_ids[3], subtask_exchange_ids[1]);
+}
+
+#[test]
+fn test_empty_task_handling() {
+    let task = create_test_task_with_exchanges(0);
+    let task_id = task.id().clone();
+    let store = TaskStore::with_root_task(task);
+
+    assert_eq!(store.task_count(), 1);
+    assert_eq!(store.exchange_count(), 0);
+    assert!(store.first_exchange().is_none());
+    assert!(store.latest_exchange().is_none());
+    assert!(store.get(&task_id).is_some());
+}
+
+#[test]
+fn test_from_tasks() {
+    use std::collections::HashMap;
+
+    let task = create_test_task_with_exchanges(3);
+    let task_id = task.id().clone();
+    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+
+    let mut tasks = HashMap::new();
+    tasks.insert(task_id.clone(), task);
+
+    let store = TaskStore::from_tasks(tasks, task_id.clone());
+
+    assert_eq!(store.task_count(), 1);
+    assert_eq!(store.exchange_count(), 3);
+    assert_eq!(store.root_task_id(), &task_id);
+
+    let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(actual_ids, exchange_ids);
+}
+
+#[test]
+fn test_exchange_mut() {
+    let task = create_test_task_with_exchanges(2);
+    let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+    let mut store = TaskStore::with_root_task(task);
+
+    // Can find existing exchange
+    let exchange = store.exchange_mut(exchange_ids[0]);
+    assert!(exchange.is_some());
+    assert_eq!(exchange.unwrap().id, exchange_ids[0]);
+
+    // Returns None for non-existent exchange
+    let fake_id = AIAgentExchangeId::new();
+    assert!(store.exchange_mut(fake_id).is_none());
+}
+
+#[test]
+fn test_modify_task_conditional_rebuild() {
+    let task = create_test_task_with_exchanges(2);
+    let task_id = task.id().clone();
+    let original_exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
+    let mut store = TaskStore::with_root_task(task);
+
+    // Verify initial state
+    assert_eq!(store.exchange_count(), 2);
+
+    // modify_task with no exchange change should still work
+    let result = store.modify_task(&task_id, |task| {
+        assert_eq!(task.exchanges_len(), 2);
+        "no change"
+    });
+    assert_eq!(result, Some("no change"));
+    assert_eq!(store.exchange_count(), 2);
+
+    // modify_task that adds an exchange should update the index
+    let new_exchange = create_test_exchange();
+    let new_exchange_id = new_exchange.id;
+    store.modify_task(&task_id, |task| {
+        task.append_exchange(new_exchange);
+    });
+    assert_eq!(store.exchange_count(), 3);
+    assert_eq!(store.latest_exchange().map(|e| e.id), Some(new_exchange_id));
+
+    // Verify all exchanges are in the index
+    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(all_ids.len(), 3);
+    assert_eq!(all_ids[0], original_exchange_ids[0]);
+    assert_eq!(all_ids[1], original_exchange_ids[1]);
+    assert_eq!(all_ids[2], new_exchange_id);
+}
+
+// =============================================================================
+// Subtask Linearization Tests
+// =============================================================================
+
+#[test]
+fn test_linearization_parent_with_one_subtask() {
+    // Create a subtask first so we have its ID
+    let subtask = create_test_subtask_with_exchanges(2);
+    let subtask_id = subtask.id().clone();
+    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();
+
+    // Create root task: exchange1, exchange_with_subagent_call, exchange3
+    let mut root_task = Task::new_optimistic_root();
+    let exchange1 = create_test_exchange();
+    let exchange1_id = exchange1.id;
+    root_task.append_exchange(exchange1);
+
+    let exchange_with_call = create_exchange_with_subagent_call(&subtask_id);
+    let exchange_with_call_id = exchange_with_call.id;
+    root_task.append_exchange(exchange_with_call);
+
+    let exchange3 = create_test_exchange();
+    let exchange3_id = exchange3.id;
+    root_task.append_exchange(exchange3);
+
+    // Build the store
+    let mut store = TaskStore::with_root_task(root_task);
+    store.insert(subtask);
+
+    // Total exchanges: 3 root + 2 subtask = 5
+    assert_eq!(store.exchange_count(), 5);
+
+    // Expected order: root[0], root[1] (with subagent call), subtask[0], subtask[1], root[2]
+    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(all_ids.len(), 5);
+    assert_eq!(all_ids[0], exchange1_id);
+    assert_eq!(all_ids[1], exchange_with_call_id);
+    assert_eq!(all_ids[2], subtask_exchange_ids[0]);
+    assert_eq!(all_ids[3], subtask_exchange_ids[1]);
+    assert_eq!(all_ids[4], exchange3_id);
+
+    // Verify first and last
+    assert_eq!(store.first_exchange().map(|e| e.id), Some(exchange1_id));
+    assert_eq!(store.latest_exchange().map(|e| e.id), Some(exchange3_id));
+}
+
+#[test]
+fn test_linearization_nested_subtasks() {
+    // Create nested subtask (grandchild) first
+    let grandchild_subtask = create_test_subtask_with_exchanges(1);
+    let grandchild_id = grandchild_subtask.id().clone();
+    let grandchild_exchange_id = grandchild_subtask.exchanges().next().unwrap().id;
+
+    // Create child subtask with a call to grandchild
+    use crate::terminal::model::block::BlockId;
+    let mut child_subtask = Task::new_optimistic_cli_agent_subtask(BlockId::new());
+    let child_id = child_subtask.id().clone();
+
+    let child_exchange1 = create_test_exchange();
+    let child_exchange1_id = child_exchange1.id;
+    child_subtask.append_exchange(child_exchange1);
+
+    let child_call_to_grandchild = create_exchange_with_subagent_call(&grandchild_id);
+    let child_call_exchange_id = child_call_to_grandchild.id;
+    child_subtask.append_exchange(child_call_to_grandchild);
+
+    // Create root task with a call to child
+    let mut root_task = Task::new_optimistic_root();
+
+    let root_exchange1 = create_test_exchange();
+    let root_exchange1_id = root_exchange1.id;
+    root_task.append_exchange(root_exchange1);
+
+    let root_call_to_child = create_exchange_with_subagent_call(&child_id);
+    let root_call_exchange_id = root_call_to_child.id;
+    root_task.append_exchange(root_call_to_child);
+
+    // Build the store
+    let mut store = TaskStore::with_root_task(root_task);
+    store.insert(child_subtask);
+    store.insert(grandchild_subtask);
+
+    // Total: 2 root + 2 child + 1 grandchild = 5
+    assert_eq!(store.exchange_count(), 5);
+
+    // Expected DFS order:
+    // root[0], root[1] (calls child) -> child[0], child[1] (calls grandchild) -> grandchild[0]
+    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(all_ids.len(), 5);
+    assert_eq!(all_ids[0], root_exchange1_id);
+    assert_eq!(all_ids[1], root_call_exchange_id);
+    assert_eq!(all_ids[2], child_exchange1_id);
+    assert_eq!(all_ids[3], child_call_exchange_id);
+    assert_eq!(all_ids[4], grandchild_exchange_id);
+}
+
+#[test]
+fn test_linearization_multiple_subtasks_same_parent() {
+    // Create two subtasks
+    let subtask1 = create_test_subtask_with_exchanges(1);
+    let subtask1_id = subtask1.id().clone();
+    let subtask1_exchange_id = subtask1.exchanges().next().unwrap().id;
+
+    let subtask2 = create_test_subtask_with_exchanges(2);
+    let subtask2_id = subtask2.id().clone();
+    let subtask2_exchange_ids: Vec<_> = subtask2.exchanges().map(|e| e.id).collect();
+
+    // Create root task with calls to both subtasks in separate exchanges
+    let mut root_task = Task::new_optimistic_root();
+
+    let call_to_subtask1 = create_exchange_with_subagent_call(&subtask1_id);
+    let call_to_subtask1_id = call_to_subtask1.id;
+    root_task.append_exchange(call_to_subtask1);
+
+    let middle_exchange = create_test_exchange();
+    let middle_exchange_id = middle_exchange.id;
+    root_task.append_exchange(middle_exchange);
+
+    let call_to_subtask2 = create_exchange_with_subagent_call(&subtask2_id);
+    let call_to_subtask2_id = call_to_subtask2.id;
+    root_task.append_exchange(call_to_subtask2);
+
+    // Build the store
+    let mut store = TaskStore::with_root_task(root_task);
+    store.insert(subtask1);
+    store.insert(subtask2);
+
+    // Total: 3 root + 1 subtask1 + 2 subtask2 = 6
+    assert_eq!(store.exchange_count(), 6);
+
+    // Expected order:
+    // root[0] (calls subtask1) -> subtask1[0], root[1], root[2] (calls subtask2) -> subtask2[0], subtask2[1]
+    let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
+    assert_eq!(all_ids.len(), 6);
+    assert_eq!(all_ids[0], call_to_subtask1_id);
+    assert_eq!(all_ids[1], subtask1_exchange_id);
+    assert_eq!(all_ids[2], middle_exchange_id);
+    assert_eq!(all_ids[3], call_to_subtask2_id);
+    assert_eq!(all_ids[4], subtask2_exchange_ids[0]);
+    assert_eq!(all_ids[5], subtask2_exchange_ids[1]);
+}
+
+#[test]
+fn test_all_exchanges_by_task_with_subtasks() {
+    // Create a subtask
+    let subtask = create_test_subtask_with_exchanges(2);
+    let subtask_id = subtask.id().clone();
+    let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();
+
+    // Create root task with a call to subtask
+    let mut root_task = Task::new_optimistic_root();
+    let root_id = root_task.id().clone();
+
+    let root_exchange1 = create_test_exchange();
+    let root_exchange1_id = root_exchange1.id;
+    root_task.append_exchange(root_exchange1);
+
+    let call_to_subtask = create_exchange_with_subagent_call(&subtask_id);
+    let call_exchange_id = call_to_subtask.id;
+    root_task.append_exchange(call_to_subtask);
+
+    let root_exchange3 = create_test_exchange();
+    let root_exchange3_id = root_exchange3.id;
+    root_task.append_exchange(root_exchange3);
+
+    // Build the store
+    let mut store = TaskStore::with_root_task(root_task);
+    store.insert(subtask);
+
+    // Check all_exchanges_by_task grouping
+    let by_task = store.all_exchanges_by_task();
+
+    // Should have 3 groups: root[0-1], subtask[0-1], root[2]
+    assert_eq!(by_task.len(), 3);
+
+    // First group: root task's first two exchanges
+    assert_eq!(by_task[0].0, root_id);
+    assert_eq!(by_task[0].1.len(), 2);
+    assert_eq!(by_task[0].1[0].id, root_exchange1_id);
+    assert_eq!(by_task[0].1[1].id, call_exchange_id);
+
+    // Second group: subtask's exchanges
+    assert_eq!(by_task[1].0, subtask_id);
+    assert_eq!(by_task[1].1.len(), 2);
+    assert_eq!(by_task[1].1[0].id, subtask_exchange_ids[0]);
+    assert_eq!(by_task[1].1[1].id, subtask_exchange_ids[1]);
+
+    // Third group: root task's last exchange
+    assert_eq!(by_task[2].0, root_id);
+    assert_eq!(by_task[2].1.len(), 1);
+    assert_eq!(by_task[2].1[0].id, root_exchange3_id);
+}

app/src/ai/agent/task_tests.rs

@@ -0,0 +1,586 @@
+use std::collections::HashSet;
+
+use crate::ai::agent::{
+    AIAgentActionType, AIAgentExchange, AIAgentOutput, AIAgentOutputMessageType,
+    AIAgentOutputStatus, MessageId, Shared,
+};
+use crate::ai::llms::LLMId;
+use crate::test_util::ai_agent_tasks::{
+    create_api_subtask, create_api_task, create_message, create_subagent_tool_call_message,
+};
+use chrono::Local;
+use prost_types::FieldMask;
+use warp_multi_agent_api as api;
+
+use super::{ExtractMessagesError, Task};
+
+/// Creates a Task backed by server data from the given api::Task.
+fn create_server_task(api_task: api::Task) -> Task {
+    Task::new_restored_root(api_task, std::iter::empty())
+}
+
+fn create_streaming_exchange_with_output() -> AIAgentExchange {
+    AIAgentExchange {
+        id: Default::default(),
+        input: vec![],
+        output_status: AIAgentOutputStatus::Streaming {
+            output: Some(Shared::new(AIAgentOutput::default())),
+        },
+        added_message_ids: HashSet::new(),
+        start_time: Local::now(),
+        finish_time: None,
+        time_to_first_token_ms: None,
+        working_directory: None,
+        model_id: LLMId::from(""),
+        request_cost: None,
+        coding_model_id: LLMId::from(""),
+        cli_agent_model_id: LLMId::from(""),
+        computer_use_model_id: LLMId::from(""),
+        response_initiator: None,
+    }
+}
+
+fn create_start_agent_tool_call_message(
+    id: &str,
+    task_id: &str,
+    name: &str,
+    prompt: &str,
+) -> api::Message {
+    api::Message {
+        id: id.to_string(),
+        task_id: task_id.to_string(),
+        server_message_data: String::new(),
+        citations: vec![],
+        message: Some(api::message::Message::ToolCall(api::message::ToolCall {
+            tool_call_id: format!("{id}_tool_call"),
+            tool: Some(api::message::tool_call::Tool::StartAgent(api::StartAgent {
+                name: name.to_string(),
+                prompt: prompt.to_string(),
+                execution_mode: None,
+                lifecycle_subscription: None,
+            })),
+        })),
+        request_id: String::new(),
+        timestamp: None,
+    }
+}
+
+fn assert_start_agent_prompt(
+    task: &Task,
+    exchange_id: crate::ai::agent::AIAgentExchangeId,
+    prompt: &str,
+) {
+    let exchange = task.exchange(exchange_id).expect("exchange should exist");
+    let output = exchange
+        .output_status
+        .output()
+        .expect("output should be initialized");
+    let output = output.get();
+    let output_message = output
+        .messages
+        .iter()
+        .find(|message| message.id == MessageId::new("start_agent_message".to_string()))
+        .expect("start agent output message should exist");
+
+    let AIAgentOutputMessageType::Action(action) = &output_message.message else {
+        panic!("expected action output message");
+    };
+    let AIAgentActionType::StartAgent {
+        prompt: current_prompt,
+        ..
+    } = &action.action
+    else {
+        panic!("expected StartAgent action");
+    };
+
+    assert_eq!(current_prompt, prompt);
+}
+
+#[test]
+fn test_upsert_message_adds_start_agent_prompt_to_output() {
+    let task_id = "task1";
+    let mut task = create_server_task(create_api_task(task_id, vec![]));
+
+    let exchange = create_streaming_exchange_with_output();
+    let exchange_id = exchange.id;
+    task.append_exchange(exchange);
+
+    task.upsert_message(
+        create_start_agent_tool_call_message(
+            "start_agent_message",
+            task_id,
+            "Agent 1",
+            "run tests",
+        ),
+        exchange_id,
+        None,
+        None,
+        FieldMask {
+            paths: vec!["message.tool_call".to_string()],
+        },
+        false,
+    )
+    .expect("initial upsert should succeed");
+    assert_start_agent_prompt(&task, exchange_id, "run tests");
+}
+
+// =============================================================================
+// Tests for Task::splice_messages()
+// =============================================================================
+
+#[test]
+fn test_splice_messages_happy_path() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+            create_message("m4", task_id),
+            create_message("m5", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract m2, m3, m4 (middle 3 messages).
+    let replacement = vec![create_message("replacement", task_id)];
+    let result = task.splice_messages("m2", "m4", 3, replacement);
+
+    assert!(result.is_ok());
+    let extracted = result.unwrap();
+    assert_eq!(extracted.len(), 3);
+    assert_eq!(extracted[0].id, "m2");
+    assert_eq!(extracted[1].id, "m3");
+    assert_eq!(extracted[2].id, "m4");
+
+    // Verify the task now has: m1, replacement, m5.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["m1", "replacement", "m5"]);
+}
+
+#[test]
+fn test_splice_messages_single_message() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract just m2.
+    let replacement = vec![create_message("replacement", task_id)];
+    let result = task.splice_messages("m2", "m2", 1, replacement);
+
+    assert!(result.is_ok());
+    let extracted = result.unwrap();
+    assert_eq!(extracted.len(), 1);
+    assert_eq!(extracted[0].id, "m2");
+
+    // Verify the task now has: m1, replacement, m3.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["m1", "replacement", "m3"]);
+}
+
+#[test]
+fn test_splice_messages_all_messages() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract all messages.
+    let replacement = vec![create_message("replacement", task_id)];
+    let result = task.splice_messages("m1", "m3", 3, replacement);
+
+    assert!(result.is_ok());
+    let extracted = result.unwrap();
+    assert_eq!(extracted.len(), 3);
+
+    // Verify the task now only has the replacement.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["replacement"]);
+}
+
+#[test]
+fn test_splice_messages_empty_replacement() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract m2 with no replacement (pure deletion).
+    let result = task.splice_messages("m2", "m2", 1, vec![]);
+
+    assert!(result.is_ok());
+    let extracted = result.unwrap();
+    assert_eq!(extracted.len(), 1);
+    assert_eq!(extracted[0].id, "m2");
+
+    // Verify the task now has: m1, m3.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["m1", "m3"]);
+}
+
+#[test]
+fn test_splice_messages_multiple_replacements() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract m2 and replace with two messages.
+    let replacement = vec![create_message("r1", task_id), create_message("r2", task_id)];
+    let result = task.splice_messages("m2", "m2", 1, replacement);
+
+    assert!(result.is_ok());
+
+    // Verify the task now has: m1, r1, r2, m3.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["m1", "r1", "r2", "m3"]);
+}
+
+#[test]
+fn test_splice_messages_first_message_not_found() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![create_message("m1", task_id), create_message("m2", task_id)],
+    );
+    let mut task = create_server_task(api_task);
+
+    let result = task.splice_messages("nonexistent", "m2", 1, vec![]);
+
+    assert!(matches!(
+        result,
+        Err(ExtractMessagesError::FirstMessageNotFound(id)) if id == "nonexistent"
+    ));
+}
+
+#[test]
+fn test_splice_messages_last_message_not_found() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![create_message("m1", task_id), create_message("m2", task_id)],
+    );
+    let mut task = create_server_task(api_task);
+
+    let result = task.splice_messages("m1", "nonexistent", 1, vec![]);
+
+    assert!(matches!(
+        result,
+        Err(ExtractMessagesError::LastMessageNotFound(id)) if id == "nonexistent"
+    ));
+}
+
+#[test]
+fn test_splice_messages_invalid_range() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // first_message_id appears after last_message_id.
+    let result = task.splice_messages("m3", "m1", 3, vec![]);
+
+    assert!(matches!(result, Err(ExtractMessagesError::InvalidRange)));
+}
+
+#[test]
+fn test_splice_messages_checksum_mismatch_too_few() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Claim there are 5 messages when there are only 3 in the range.
+    let result = task.splice_messages("m1", "m3", 5, vec![]);
+
+    assert!(matches!(
+        result,
+        Err(ExtractMessagesError::ChecksumMismatch {
+            expected: 5,
+            actual: 3
+        })
+    ));
+}
+
+#[test]
+fn test_splice_messages_checksum_mismatch_too_many() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Claim there is 1 message when there are 3 in the range.
+    let result = task.splice_messages("m1", "m3", 1, vec![]);
+
+    assert!(matches!(
+        result,
+        Err(ExtractMessagesError::ChecksumMismatch {
+            expected: 1,
+            actual: 3
+        })
+    ));
+}
+
+#[test]
+fn test_splice_messages_optimistic_task_not_initialized() {
+    let mut task = Task::new_optimistic_root();
+
+    let result = task.splice_messages("m1", "m2", 2, vec![]);
+
+    assert!(matches!(
+        result,
+        Err(ExtractMessagesError::TaskNotInitialized)
+    ));
+}
+
+#[test]
+fn test_splice_messages_from_beginning() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+            create_message("m4", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract from the beginning.
+    let replacement = vec![create_message("replacement", task_id)];
+    let result = task.splice_messages("m1", "m2", 2, replacement);
+
+    assert!(result.is_ok());
+    let extracted = result.unwrap();
+    assert_eq!(extracted.len(), 2);
+
+    // Verify the task now has: replacement, m3, m4.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["replacement", "m3", "m4"]);
+}
+
+#[test]
+fn test_splice_messages_from_end() {
+    let task_id = "task1";
+    let api_task = create_api_task(
+        task_id,
+        vec![
+            create_message("m1", task_id),
+            create_message("m2", task_id),
+            create_message("m3", task_id),
+            create_message("m4", task_id),
+        ],
+    );
+    let mut task = create_server_task(api_task);
+
+    // Extract from the end.
+    let replacement = vec![create_message("replacement", task_id)];
+    let result = task.splice_messages("m3", "m4", 2, replacement);
+
+    assert!(result.is_ok());
+    let extracted = result.unwrap();
+    assert_eq!(extracted.len(), 2);
+
+    // Verify the task now has: m1, m2, replacement.
+    let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(remaining_ids, vec!["m1", "m2", "replacement"]);
+}
+
+// =============================================================================
+// Tests for Task::new_moved_messages_subtask()
+// =============================================================================
+
+#[test]
+fn test_new_moved_messages_subtask_basic() {
+    let parent_id = "parent";
+    let subtask_id = "subtask";
+
+    // Create parent task with a subagent call referencing the subtask.
+    let parent_api_task = create_api_task(
+        parent_id,
+        vec![
+            create_message("m1", parent_id),
+            create_subagent_tool_call_message("subagent_call", parent_id, subtask_id, None),
+            create_message("m2", parent_id),
+        ],
+    );
+
+    // Create the subtask api::Task with some messages.
+    let subtask_api_task = create_api_task(
+        subtask_id,
+        vec![
+            create_message("s1", subtask_id),
+            create_message("s2", subtask_id),
+        ],
+    );
+
+    let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
+
+    assert_eq!(subtask.id().to_string(), subtask_id);
+    assert!(subtask.exchanges().next().is_none()); // No exchanges.
+    assert_eq!(subtask.messages().count(), 2);
+
+    // Should have subagent_params extracted from parent.
+    let subagent_params = subtask.subagent_params();
+    assert!(subagent_params.is_some());
+    assert_eq!(
+        subagent_params.unwrap().tool_call_id,
+        "subagent_call_tool_call"
+    );
+}
+
+#[test]
+fn test_new_moved_messages_subtask_with_summarization_metadata() {
+    let parent_id = "parent";
+    let subtask_id = "subtask";
+
+    // Create parent task with a summarization subagent call.
+    let parent_api_task = create_api_task(
+        parent_id,
+        vec![create_subagent_tool_call_message(
+            "summary_call",
+            parent_id,
+            subtask_id,
+            Some(api::message::tool_call::subagent::Metadata::Summarization(
+                (),
+            )),
+        )],
+    );
+
+    let subtask_api_task = create_api_task(subtask_id, vec![create_message("s1", subtask_id)]);
+
+    let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
+
+    // Check that subagent_params has the summarization metadata.
+    let subagent_params = subtask.subagent_params();
+    assert!(subagent_params.is_some());
+
+    let call = &subagent_params.unwrap().call;
+    assert!(matches!(
+        call.metadata,
+        Some(api::message::tool_call::subagent::Metadata::Summarization(
+            _
+        ))
+    ));
+}
+
+#[test]
+fn test_new_moved_messages_subtask_no_matching_subagent_call() {
+    let parent_id = "parent";
+    let subtask_id = "subtask";
+
+    // Parent task has no subagent call to this subtask.
+    let parent_api_task = create_api_task(
+        parent_id,
+        vec![
+            create_message("m1", parent_id),
+            // Subagent call references a different task.
+            create_subagent_tool_call_message("other_call", parent_id, "other_task", None),
+        ],
+    );
+
+    let subtask_api_task = create_api_task(subtask_id, vec![create_message("s1", subtask_id)]);
+
+    let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
+
+    // No subagent_params since no matching call was found.
+    assert!(subtask.subagent_params().is_none());
+}
+
+#[test]
+fn test_new_moved_messages_subtask_preserves_messages() {
+    let parent_id = "parent";
+    let subtask_id = "subtask";
+
+    let parent_api_task = create_api_task(
+        parent_id,
+        vec![create_subagent_tool_call_message(
+            "call", parent_id, subtask_id, None,
+        )],
+    );
+
+    // Subtask with multiple messages.
+    let subtask_api_task = create_api_task(
+        subtask_id,
+        vec![
+            create_message("s1", subtask_id),
+            create_message("s2", subtask_id),
+            create_message("s3", subtask_id),
+        ],
+    );
+
+    let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
+
+    // All messages should be preserved.
+    let message_ids: Vec<_> = subtask.messages().map(|m| m.id.as_str()).collect();
+    assert_eq!(message_ids, vec!["s1", "s2", "s3"]);
+}
+
+// =============================================================================
+// Tests for Warp docs subagent classification
+// =============================================================================
+
+#[test]
+fn test_is_warp_documentation_search_subagent() {
+    let parent_id = "parent";
+    let subtask_id = "subtask";
+    let parent_api_task = create_api_task(
+        parent_id,
+        vec![create_subagent_tool_call_message(
+            "docs_call",
+            parent_id,
+            subtask_id,
+            Some(api::message::tool_call::subagent::Metadata::WarpDocumentationSearch(())),
+        )],
+    );
+    let subtask_api_task = create_api_subtask(subtask_id, parent_id, vec![]);
+    let subtask = Task::new_restored_subtask(subtask_api_task, &parent_api_task, vec![]);
+
+    assert!(subtask.is_warp_documentation_search_subagent());
+    assert!(!subtask.is_conversation_search_subagent());
+}

app/src/ai/agent/telemetry.rs

@@ -0,0 +1,123 @@
+use serde::Serialize;
+use warpui::{AppContext, SingletonEntity};
+
+use crate::ai::llms::LLMId;
+use crate::CloudModel;
+use crate::{
+    server::telemetry::AgentModeCitation as CitationForTelemetry,
+    terminal::view::block_onboarding::onboarding_agentic_suggestions_block::OnboardingChipType,
+};
+
+use super::conversation::AIConversationId;
+use super::{
+    AIAgentCitation, AIAgentExchangeId, EntrypointType, PassiveSuggestionTriggerType,
+    ServerOutputId,
+};
+
+pub trait ForTelemetry {
+    type Output;
+
+    fn for_telemetry(&self, ctx: &AppContext) -> Option<Self::Output>;
+}
+
+impl ForTelemetry for AIAgentCitation {
+    type Output = CitationForTelemetry;
+
+    fn for_telemetry(&self, ctx: &AppContext) -> Option<Self::Output> {
+        match self {
+            Self::WarpDriveObject { uid } => {
+                CloudModel::as_ref(ctx).get_by_uid(uid).map(|object| {
+                    CitationForTelemetry::WarpDriveObject {
+                        object_type: object.object_type(),
+                        uid: object.uid(),
+                    }
+                })
+            }
+            Self::WarpDocumentation { path } => {
+                Some(CitationForTelemetry::WarpDocs { page: path.clone() })
+            }
+            Self::WebPage { url } => Some(CitationForTelemetry::WebPage { url: url.clone() }),
+        }
+    }
+}
+
+impl EntrypointType {
+    pub fn entrypoint(&self) -> String {
+        match self {
+            Self::Onboarding { chip_type } => {
+                format!(
+                    "ONBOARDING.{}",
+                    match chip_type {
+                        OnboardingChipType::FixAnIssue => "FIX_AN_ISSUE",
+                        OnboardingChipType::PullCloudLogs => "PULL_CLOUD_LOGS",
+                        OnboardingChipType::StartAFeature => "START_A_FEATURE",
+                        OnboardingChipType::PythonSnakeGame => "PYTHON_SNAKE_GAME",
+                        OnboardingChipType::ExploreGitHistory => "EXPLORE_GIT_HISTORY",
+                        OnboardingChipType::MatrixThemePicker => "MATRIX_THEME_PICKER",
+                        OnboardingChipType::Other => "OTHER",
+                    }
+                )
+            }
+            Self::PromptSuggestion {
+                is_static,
+                is_coding,
+            } => match (is_static, is_coding) {
+                (true, true) => "PROMPT_SUGGESTION.CODING_STATIC".to_string(),
+                (true, false) => "PROMPT_SUGGESTION.STATIC".to_string(),
+                (false, true) => "PROMPT_SUGGESTION.CODING".to_string(),
+                (false, false) => "PROMPT_SUGGESTION.SIMPLE".to_string(),
+            },
+            Self::ZeroStateAgentModePromptSuggestion => {
+                "ZERO_STATE_AGENT_MODE_PROMPT_SUGGESTION".to_string()
+            }
+            Self::InitProjectRules => "INIT_PROJECT_RULES".to_string(),
+            Self::UserInitiated => "USER_INITIATED".to_string(),
+            Self::AgentInitiated => "AGENT_INITIATED".to_string(),
+            Self::TriggerPassiveSuggestion { trigger } => {
+                let trigger_name = match trigger {
+                    Some(PassiveSuggestionTriggerType::FilesChanged) => "FILES_CHANGED",
+                    Some(PassiveSuggestionTriggerType::CommandRun) => "COMMAND_RUN",
+                    Some(PassiveSuggestionTriggerType::ShellCommandCompleted) => {
+                        "SHELL_COMMAND_COMPLETED"
+                    }
+                    Some(PassiveSuggestionTriggerType::AgentResponseCompleted) => {
+                        "AGENT_RESPONSE_COMPLETED"
+                    }
+                    None => "NONE",
+                };
+                format!("TRIGGER_SUGGEST_PROMPT.{trigger_name}")
+            }
+            Self::CloneRepository => "CLONE_REPOSITORY".to_string(),
+            Self::SharedSession => "SHARED_SESSION".to_string(),
+            Self::ResumeConversation => "RESUME_CONVERSATION".to_string(),
+        }
+    }
+}
+
+#[derive(Clone, Default, Debug, Serialize)]
+pub struct AIIdentifiers {
+    /// Useful for joining to client-side telemetry.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub client_conversation_id: Option<AIConversationId>,
+    /// A stable ID to relate failures for the same underlying request.
+    #[serde(rename = "exchange_id", skip_serializing_if = "Option::is_none")]
+    pub client_exchange_id: Option<AIAgentExchangeId>,
+    /// Unique ID for this output coming from the AI API. Generated by the server. Only passed in
+    /// the initial response chunk.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub server_output_id: Option<ServerOutputId>,
+    /// The conversation ID included in the response chunk.
+    ///
+    /// Once this is set, it is never updated. That shouldn't be an issue because the conversation
+    /// ID is only expected to be passed in the initial response chunk.
+    ///
+    /// Note that this conversation ID is server-scoped; it is _not_ related to the
+    /// `AIConversationId`, which is entirely a client-side abstraction.
+    ///
+    /// This is mainly used for server-side logging and analytics.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub server_conversation_id: Option<String>,
+    /// The ID of the model actually used to generate the output. This may differ from the requested model.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model_id: Option<LLMId>,
+}

app/src/ai/agent/todos/mod.rs

@@ -0,0 +1,91 @@
+use crate::ai::agent::AIAgentTodo;
+
+use super::AIAgentTodoId;
+pub(crate) mod popup;
+
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct AIAgentTodoList {
+    completed_items: Vec<AIAgentTodo>,
+    pending_items: Vec<AIAgentTodo>,
+}
+
+impl AIAgentTodoList {
+    pub fn with_pending_items(mut self, pending_items: Vec<AIAgentTodo>) -> Self {
+        self.pending_items = pending_items;
+        self
+    }
+
+    pub fn with_completed_items(mut self, completed_items: Vec<AIAgentTodo>) -> Self {
+        self.completed_items = completed_items;
+        self
+    }
+
+    pub fn update_pending_items(&mut self, pending_items: Vec<AIAgentTodo>) {
+        self.pending_items = pending_items;
+    }
+
+    pub fn clear_pending_items(&mut self) {
+        self.pending_items.clear();
+    }
+
+    pub fn len(&self) -> usize {
+        self.pending_items.len() + self.completed_items.len()
+    }
+
+    pub fn is_finished(&self) -> bool {
+        self.pending_items.is_empty() && !self.completed_items.is_empty()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.pending_items.is_empty() && self.completed_items.is_empty()
+    }
+
+    pub fn in_progress_item(&self) -> Option<&AIAgentTodo> {
+        self.pending_items.first()
+    }
+
+    pub fn pending_items(&self) -> &[AIAgentTodo] {
+        &self.pending_items
+    }
+
+    pub fn completed_items(&self) -> &[AIAgentTodo] {
+        &self.completed_items
+    }
+
+    pub fn is_pending(&self, todo_id: &AIAgentTodoId) -> bool {
+        self.pending_items.iter().any(|item| &item.id == todo_id)
+    }
+
+    pub fn is_completed(&self, todo_id: &AIAgentTodoId) -> bool {
+        self.completed_items.iter().any(|item| &item.id == todo_id)
+    }
+
+    pub fn get_item(&self, todo_id: &AIAgentTodoId) -> Option<&AIAgentTodo> {
+        self.items().find(|item| &item.id == todo_id)
+    }
+
+    pub fn get_item_index(&self, todo_id: &AIAgentTodoId) -> Option<usize> {
+        self.items().position(|item| &item.id == todo_id)
+    }
+
+    fn items(&self) -> impl Iterator<Item = &AIAgentTodo> {
+        self.completed_items.iter().chain(self.pending_items.iter())
+    }
+
+    pub fn update_pending_todos(&mut self, todos: Vec<AIAgentTodo>) {
+        self.pending_items = todos;
+    }
+
+    pub fn mark_todos_complete(&mut self, completed_todo_ids: Vec<String>) {
+        for completed_todo_id in completed_todo_ids.into_iter() {
+            if let Some(item) = self
+                .pending_items
+                .iter()
+                .position(|item| item.id == completed_todo_id.clone().into())
+                .map(|i| self.pending_items.remove(i))
+            {
+                self.completed_items.push(item);
+            }
+        }
+    }
+}

app/src/ai/agent/todos/popup.rs

@@ -0,0 +1,319 @@
+use crate::ai::blocklist::{BlocklistAIContextEvent, BlocklistAIContextModel};
+use pathfinder_color::ColorU;
+use warp_core::ui::appearance::Appearance;
+use warp_core::ui::theme::Fill;
+use warpui::elements::{
+    ClippedScrollStateHandle, ClippedScrollable, Dismiss, Empty, Expanded, ParentElement,
+    SavePosition, ScrollTarget, ScrollToPositionMode, ScrollbarWidth, Shrinkable,
+};
+use warpui::fonts::FamilyId;
+use warpui::ModelHandle;
+use warpui::SingletonEntity;
+use warpui::{
+    elements::{
+        Border, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, DropShadow, Flex,
+        MainAxisSize, Radius, Text,
+    },
+    fonts::{Properties, Weight},
+    keymap::FixedBinding,
+    AppContext, Element, Entity, EntityId, TypedActionView, View, ViewContext,
+};
+
+use crate::ai::agent::icons::{in_progress_icon, pending_icon, succeeded_icon};
+use crate::ai::agent::todos::AIAgentTodoList;
+use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel};
+use crate::ui_components::blended_colors;
+
+pub struct AgentTodosPopupView {
+    terminal_view_id: EntityId,
+    ai_context_model: ModelHandle<BlocklistAIContextModel>,
+    scroll_state: ClippedScrollStateHandle,
+}
+
+const IN_PROGRESS_POSITION_ID: &str = "AgentTodosPopup-in-progress";
+
+#[derive(Debug, Clone, Copy)]
+pub enum AgentTodosPopupAction {
+    ClosePopup,
+}
+
+pub enum AgentTodosPopupEvent {
+    Close,
+}
+
+struct Styles {
+    ui_font_family: FamilyId,
+    background: Fill,
+    main_text_color: ColorU,
+    sub_text_color: ColorU,
+    detail_font_size: f32,
+}
+
+pub fn init(app: &mut AppContext) {
+    use warpui::keymap::macros::*;
+
+    app.register_fixed_bindings([FixedBinding::new(
+        "escape",
+        AgentTodosPopupAction::ClosePopup,
+        id!(AgentTodosPopupView::ui_name()),
+    )]);
+}
+
+impl AgentTodosPopupView {
+    pub fn new(
+        terminal_view_id: EntityId,
+        ai_context_model: ModelHandle<BlocklistAIContextModel>,
+        ctx: &mut ViewContext<Self>,
+    ) -> Self {
+        let blocklist_history_model = BlocklistAIHistoryModel::handle(ctx);
+        ctx.subscribe_to_model(&blocklist_history_model, move |me, _, event, ctx| {
+            me.handle_blocklist_history_event(event, ctx);
+        });
+        ctx.subscribe_to_model(&ai_context_model, move |_, _, event, ctx| {
+            if let BlocklistAIContextEvent::PendingQueryStateUpdated = event {
+                ctx.notify();
+            }
+        });
+        Self {
+            terminal_view_id,
+            ai_context_model,
+            scroll_state: Default::default(),
+        }
+    }
+
+    fn handle_blocklist_history_event(
+        &mut self,
+        event: &BlocklistAIHistoryEvent,
+        ctx: &mut ViewContext<Self>,
+    ) {
+        if let BlocklistAIHistoryEvent::UpdatedTodoList { terminal_view_id } = event {
+            if *terminal_view_id == self.terminal_view_id {
+                ctx.notify();
+            }
+        }
+    }
+
+    /// Scroll to the in-progress item, if not currently visible.
+    pub fn scroll_to_in_progress_item(&self) {
+        self.scroll_state.scroll_to_position(ScrollTarget {
+            position_id: IN_PROGRESS_POSITION_ID.to_string(),
+            mode: ScrollToPositionMode::FullyIntoView,
+        });
+    }
+
+    fn close(&mut self, ctx: &mut ViewContext<Self>) {
+        ctx.emit(AgentTodosPopupEvent::Close);
+    }
+
+    fn styles(&self, appearance: &Appearance) -> Styles {
+        let theme = appearance.theme();
+        let background = theme.surface_1();
+        let main_text_color = blended_colors::text_main(theme, background);
+        let sub_text_color = blended_colors::text_sub(theme, background);
+        let detail_font_size = appearance.ui_font_size();
+        let ui_font_family = appearance.ui_font_family();
+
+        Styles {
+            ui_font_family,
+            background,
+            main_text_color,
+            sub_text_color,
+            detail_font_size,
+        }
+    }
+
+    fn render_header(
+        &self,
+        app: &warpui::AppContext,
+        todo_list: &AIAgentTodoList,
+    ) -> Box<dyn warpui::Element> {
+        let appearance = Appearance::as_ref(app);
+        let styles = self.styles(appearance);
+        let theme = appearance.theme();
+
+        let completed_count = todo_list.completed_items().len();
+        let total_count = todo_list.pending_items().len() + completed_count;
+
+        let mut header_row = Flex::row().with_cross_axis_alignment(CrossAxisAlignment::Center);
+        let mut header = Text::new(
+            "Tasks".to_string(),
+            appearance.header_font_family(),
+            styles.detail_font_size + 2.,
+        )
+        .with_color(styles.main_text_color)
+        .with_style(Properties::default().weight(Weight::Semibold));
+
+        header.add_text_with_highlights(
+            format!(" {completed_count}/{total_count}"),
+            theme.sub_text_color(theme.surface_1()).into(),
+            Properties::default().weight(Weight::Semibold),
+        );
+
+        header_row.add_child(header.finish());
+        header_row.finish()
+    }
+}
+
+impl View for AgentTodosPopupView {
+    fn ui_name() -> &'static str {
+        "AgentTodosPopup"
+    }
+
+    fn render(&self, app: &warpui::AppContext) -> Box<dyn warpui::Element> {
+        let Some(todo_list) = self
+            .ai_context_model
+            .as_ref(app)
+            .selected_conversation_todolist(app)
+        else {
+            // We don't have an empty state.
+            // Assume the popup will only be shown if there are todos.
+            return Empty::new().finish();
+        };
+
+        let appearance = Appearance::as_ref(app);
+        let styles = self.styles(appearance);
+        let theme = appearance.theme();
+
+        let background = styles.background;
+        let main_text_color = styles.main_text_color;
+        let sub_text_color = styles.sub_text_color;
+        let detail_font_size = styles.detail_font_size;
+        let ui_font_family = styles.ui_font_family;
+
+        let mut list_col = Flex::column()
+            .with_cross_axis_alignment(CrossAxisAlignment::Stretch)
+            .with_spacing(12.);
+
+        let items_with_icons = todo_list
+            .completed_items()
+            .iter()
+            .map(|item| (item, succeeded_icon(appearance)))
+            .chain(
+                todo_list
+                    .pending_items()
+                    .iter()
+                    .enumerate()
+                    .map(|(i, item)| {
+                        (
+                            item,
+                            if i == 0 {
+                                in_progress_icon(appearance)
+                            } else {
+                                pending_icon(appearance)
+                            },
+                        )
+                    }),
+            );
+
+        for (item, status_icon) in items_with_icons {
+            let mut row = Flex::row()
+                .with_cross_axis_alignment(CrossAxisAlignment::Center)
+                .with_main_axis_size(MainAxisSize::Max);
+
+            // Status icon
+            row.add_child(
+                Container::new(
+                    ConstrainedBox::new(status_icon.finish())
+                        .with_width(16.)
+                        .with_height(16.)
+                        .finish(),
+                )
+                .with_margin_right(8.)
+                .finish(),
+            );
+
+            let is_in_progress = todo_list
+                .in_progress_item()
+                .map(|t| t.id.clone())
+                .as_ref()
+                .map(|id| &item.id == id)
+                .unwrap_or(false);
+
+            // Title and status
+            let mut text_col =
+                Flex::column().with_cross_axis_alignment(CrossAxisAlignment::Stretch);
+
+            let text_color = if is_in_progress {
+                main_text_color
+            } else {
+                sub_text_color
+            };
+
+            text_col.add_child(
+                Text::new(item.title.clone(), ui_font_family, detail_font_size)
+                    .with_color(text_color)
+                    .finish(),
+            );
+
+            row.add_child(Expanded::new(1.0, text_col.finish()).finish());
+
+            let row = if is_in_progress {
+                SavePosition::new(row.finish(), IN_PROGRESS_POSITION_ID).finish()
+            } else {
+                row.finish()
+            };
+
+            list_col.add_child(row);
+        }
+
+        let header = Container::new(self.render_header(app, todo_list))
+            .with_padding_top(16.)
+            .with_horizontal_padding(16.)
+            .with_padding_bottom(8.)
+            .finish();
+
+        let scrollable_body = ClippedScrollable::vertical(
+            self.scroll_state.clone(),
+            Container::new(list_col.finish())
+                .with_horizontal_padding(16.)
+                .with_padding_bottom(16.)
+                .finish(),
+            ScrollbarWidth::Auto,
+            theme.nonactive_ui_detail().into(),
+            theme.active_ui_detail().into(),
+            warpui::elements::Fill::None,
+        )
+        .with_overlayed_scrollbar()
+        .finish();
+
+        let panel_col = Flex::column()
+            .with_cross_axis_alignment(CrossAxisAlignment::Stretch)
+            .with_child(header)
+            .with_child(Shrinkable::new(1.0, scrollable_body).finish());
+
+        Dismiss::new(
+            ConstrainedBox::new(
+                Container::new(panel_col.finish())
+                    .with_background(background)
+                    .with_corner_radius(CornerRadius::with_all(Radius::Pixels(8.)))
+                    .with_border(Border::all(1.).with_border_fill(theme.outline()))
+                    .with_drop_shadow(DropShadow::default())
+                    .finish(),
+            )
+            .with_width(300.)
+            .with_max_height(420.)
+            .finish(),
+        )
+        .prevent_interaction_with_other_elements()
+        .on_dismiss(|ctx, _app| {
+            ctx.dispatch_typed_action(AgentTodosPopupAction::ClosePopup);
+        })
+        .finish()
+    }
+}
+
+impl Entity for AgentTodosPopupView {
+    type Event = AgentTodosPopupEvent;
+}
+
+impl TypedActionView for AgentTodosPopupView {
+    type Action = AgentTodosPopupAction;
+
+    fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
+        match action {
+            AgentTodosPopupAction::ClosePopup => {
+                self.close(ctx);
+            }
+        }
+    }
+}

app/src/ai/agent/util.rs

@@ -0,0 +1,301 @@
+use super::{
+    AIAgentTextSection, AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram,
+    AgentOutputTable, ProgrammingLanguage,
+};
+use crate::code::editor_management::CodeSource;
+use crate::features::FeatureFlag;
+use ai::gfm_table::{format_gfm_table, maybe_collect_gfm_table_lines};
+use itertools::Itertools;
+use lazy_static::lazy_static;
+use markdown_parser::{
+    parse_image_run_line, parse_markdown_with_gfm_tables, FormattedImage, FormattedTextLine,
+};
+use mermaid_to_svg::is_mermaid_diagram;
+use regex::Regex;
+use std::{collections::HashMap, path::PathBuf};
+use warp_util::path::LineAndColumnArg;
+
+lazy_static! {
+    /// Markdown prefix for code blocks. Matches on triple backticks followed by a language.
+    /// Importantly, parameters for linked code blocks are captured into their own group.
+    static ref CODE_START_REGEX: Regex = Regex::new(r"^\s*```([\w+-]*)(.*)$").expect("Regex is valid");
+
+    /// Markdown suffix for code blocks.
+    static ref CODE_END_REGEX: Regex = Regex::new(r"^\s*```\s*$").expect("Regex is valid");
+
+    /// Extracts key-value parameters from text in the format: key=value, used for code block metadata.
+    /// Expects to match on text with format path=/path/to/file start=<line_number>
+    static ref CODE_PARAMS_REGEX: Regex = Regex::new(r"(\w+)=([^\s]+)").expect("Regex is valid");
+}
+
+/// Converts the given `markdown_text` into corresponding `Text` and `Code` `AIAgentOutputStep`s.
+pub(super) fn parse_markdown_into_text_and_code_sections(
+    markdown_text: &str,
+) -> Vec<AIAgentTextSection> {
+    let mut sections = vec![];
+    let mut current_section = CurrentSection::PlainText(String::new());
+
+    let mut lines = markdown_text.lines().peekable();
+    while let Some(line) = lines.next() {
+        match &mut current_section {
+            CurrentSection::PlainText(text) => {
+                // Detect tables and render them as formatted table sections.
+                if let Some(table_lines) = maybe_collect_gfm_table_lines(line, &mut lines, |l| {
+                    CODE_START_REGEX.is_match(l)
+                }) {
+                    let markdown_source = table_lines.join("\n");
+                    let table_section = if FeatureFlag::BlocklistMarkdownTableRendering.is_enabled()
+                    {
+                        parse_agent_output_table(&markdown_source)
+                    } else {
+                        Some(AgentOutputTable::legacy(format_gfm_table(&table_lines)))
+                    };
+                    if let Some(table_section) = table_section {
+                        if !text.is_empty() {
+                            flush_plain_text_sections(text, &mut sections);
+                            text.clear();
+                        }
+                        sections.push(AIAgentTextSection::Table {
+                            table: table_section,
+                        });
+                        continue;
+                    }
+
+                    if !text.is_empty() {
+                        text.push('\n');
+                    }
+                    text.push_str(&markdown_source);
+                    continue;
+                }
+
+                if let Some((_, [language, param_str])) = CODE_START_REGEX
+                    .captures(line)
+                    .map(|capture_group| capture_group.extract())
+                {
+                    if !text.is_empty() {
+                        flush_plain_text_sections(text, &mut sections);
+                    }
+
+                    let source = {
+                        let mut params = HashMap::new();
+                        for (_, [key, value]) in CODE_PARAMS_REGEX
+                            .captures_iter(param_str)
+                            .map(|c| c.extract())
+                        {
+                            params.insert(key, value);
+                        }
+                        match (params.get("path"), params.get("start")) {
+                            (Some(path), Some(start)) => {
+                                start
+                                    .parse::<usize>()
+                                    .ok()
+                                    .map(|line_num| CodeSource::Link {
+                                        path: PathBuf::from(path),
+                                        range_start: Some(LineAndColumnArg {
+                                            line_num,
+                                            column_num: None,
+                                        }),
+                                        range_end: None,
+                                    })
+                            }
+                            _ => None,
+                        }
+                    };
+                    current_section = CurrentSection::Code {
+                        code: String::new(),
+                        language_token: Some(language.to_owned()).filter(|l| !l.is_empty()),
+                        language: Some(language)
+                            .filter(|l| !l.is_empty())
+                            .map(|l| l.to_owned().into()),
+                        source,
+                    };
+                } else {
+                    if !text.is_empty() {
+                        text.push('\n');
+                    }
+                    text.push_str(line);
+                }
+            }
+            CurrentSection::Code {
+                code,
+                language,
+                language_token,
+                source,
+            } => {
+                if CODE_END_REGEX.is_match(line) {
+                    if !code.is_empty() {
+                        if let Some(CodeSource::Link {
+                            range_start: Some(start),
+                            range_end,
+                            ..
+                        }) = source.as_mut()
+                        {
+                            *range_end = Some(LineAndColumnArg {
+                                line_num: start.line_num + code.lines().count() - 1,
+                                column_num: None,
+                            });
+                        }
+                        push_code_or_mermaid_section(
+                            std::mem::take(code),
+                            language.clone(),
+                            language_token.as_deref(),
+                            source.take(),
+                            &mut sections,
+                        );
+                    }
+                    current_section = CurrentSection::PlainText(String::new());
+                } else {
+                    if !code.is_empty() {
+                        code.push('\n');
+                    }
+                    code.push_str(line);
+                }
+            }
+        }
+    }
+
+    match current_section {
+        CurrentSection::PlainText(text) => {
+            flush_plain_text_sections(&text, &mut sections);
+        }
+        CurrentSection::Code {
+            code,
+            language,
+            language_token,
+            source,
+        } => {
+            push_code_or_mermaid_section(
+                code,
+                language,
+                language_token.as_deref(),
+                source,
+                &mut sections,
+            );
+        }
+    }
+
+    if sections.is_empty() {
+        sections.push(AIAgentTextSection::PlainText {
+            text: String::new().into(),
+        });
+    }
+    sections
+}
+
+fn parse_agent_output_table(markdown_source: &str) -> Option<AgentOutputTable> {
+    let formatted_text = parse_markdown_with_gfm_tables(markdown_source).ok()?;
+    let table = formatted_text
+        .lines
+        .into_iter()
+        .exactly_one()
+        .ok()
+        .and_then(|line| match line {
+            FormattedTextLine::Table(table) => Some(table),
+            _ => None,
+        })?;
+    Some(AgentOutputTable::structured(
+        markdown_source.to_owned(),
+        table,
+    ))
+}
+
+enum CurrentSection {
+    PlainText(String),
+    Code {
+        code: String,
+        language_token: Option<String>,
+        language: Option<ProgrammingLanguage>,
+        source: Option<CodeSource>,
+    },
+}
+
+fn flush_plain_text_sections(markdown_text: &str, sections: &mut Vec<AIAgentTextSection>) {
+    if markdown_text.is_empty() {
+        return;
+    }
+
+    let mut plain_text = String::new();
+    for line in markdown_text.split_inclusive('\n') {
+        if let Some(images) = parse_image_run_line(line) {
+            if !plain_text.is_empty() {
+                sections.push(AIAgentTextSection::PlainText {
+                    text: std::mem::take(&mut plain_text).into(),
+                });
+            }
+
+            if images.len() == 1 {
+                if let Some(image) = images.into_iter().next() {
+                    sections.push(image_section(image, AgentOutputImageLayout::Block));
+                }
+            } else {
+                sections.extend(
+                    images
+                        .into_iter()
+                        .map(|image| image_section(image, AgentOutputImageLayout::Inline)),
+                );
+            }
+        } else {
+            plain_text.push_str(line);
+        }
+    }
+
+    if !plain_text.is_empty() {
+        sections.push(AIAgentTextSection::PlainText {
+            text: plain_text.into(),
+        });
+    }
+}
+fn image_section(image: FormattedImage, layout: AgentOutputImageLayout) -> AIAgentTextSection {
+    AIAgentTextSection::Image {
+        image: AgentOutputImage {
+            markdown_source: markdown_source_for_image(&image),
+            alt_text: image.alt_text,
+            source: image.source,
+            title: image.title,
+            layout,
+        },
+    }
+}
+
+fn markdown_source_for_image(image: &FormattedImage) -> String {
+    warp_editor::content::text::format_image_markdown(
+        &image.alt_text,
+        &image.source,
+        image.title.as_deref(),
+    )
+}
+
+fn markdown_source_for_mermaid(source: &str) -> String {
+    format!("```mermaid\n{source}\n```")
+}
+
+fn push_code_or_mermaid_section(
+    code: String,
+    language: Option<ProgrammingLanguage>,
+    language_token: Option<&str>,
+    source: Option<CodeSource>,
+    sections: &mut Vec<AIAgentTextSection>,
+) {
+    if code.is_empty() {
+        return;
+    }
+
+    if language_token.is_some_and(is_mermaid_diagram) {
+        sections.push(AIAgentTextSection::MermaidDiagram {
+            diagram: AgentOutputMermaidDiagram {
+                markdown_source: markdown_source_for_mermaid(&code),
+                source: code,
+            },
+        });
+    } else {
+        sections.push(AIAgentTextSection::Code {
+            code,
+            language,
+            source,
+        });
+    }
+}
+
+#[cfg(test)]
+#[path = "util_tests.rs"]
+mod tests;

app/src/ai/agent/util_tests.rs

@@ -0,0 +1,280 @@
+use super::parse_markdown_into_text_and_code_sections;
+use crate::ai::agent::{AIAgentTextSection, AgentOutputImageLayout, AgentOutputTableRendering};
+use crate::features::FeatureFlag;
+
+#[test]
+fn extracts_gfm_pipe_table_into_table_section() {
+    let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(true);
+    let input = "Intro\n\n| A | B |\n| --- | --- |\n| 1 | 2 |\n\nOutro";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 3);
+
+    match &sections[0] {
+        AIAgentTextSection::PlainText { text } => {
+            assert!(text.text().contains("Intro"));
+        }
+        _ => panic!("expected first section to be PlainText"),
+    }
+
+    match &sections[1] {
+        AIAgentTextSection::Table { table } => {
+            assert_eq!(table.markdown_source, "| A | B |\n| --- | --- |\n| 1 | 2 |");
+            match &table.rendering {
+                AgentOutputTableRendering::Legacy { .. } => {
+                    panic!("expected structured table rendering")
+                }
+                AgentOutputTableRendering::Structured { table } => {
+                    assert_eq!(table.headers.len(), 2);
+                    assert_eq!(table.rows.len(), 1);
+                }
+            }
+            assert_eq!(
+                table.rendered_lines(),
+                vec!["A\tB".to_string(), "1\t2".to_string()]
+            );
+        }
+        _ => panic!("expected second section to be Table"),
+    }
+
+    match &sections[2] {
+        AIAgentTextSection::PlainText { text } => {
+            assert!(text.text().contains("Outro"));
+        }
+        _ => panic!("expected third section to be PlainText"),
+    }
+}
+
+#[test]
+fn does_not_extract_pipe_text_without_separator_row() {
+    let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(true);
+    let input = "a | b\nc | d";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 1);
+    assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
+}
+
+#[test]
+fn table_can_be_followed_immediately_by_text() {
+    let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(true);
+    let input = "| A | B |\n|---|---|\n| 1 | 2 |\nAfter";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 2);
+    assert!(matches!(sections[0], AIAgentTextSection::Table { .. }));
+    match &sections[1] {
+        AIAgentTextSection::PlainText { text } => {
+            assert!(text.text().contains("After"));
+        }
+        _ => panic!("expected second section to be PlainText"),
+    }
+}
+
+#[test]
+fn extracts_gfm_pipe_table_into_legacy_table_section_when_flag_disabled() {
+    let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(false);
+    let input = "Intro\n\n| A | B |\n|---|---|\n| 1 | 2 |\n\nOutro";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 3);
+
+    match &sections[1] {
+        AIAgentTextSection::Table { table } => {
+            assert_eq!(
+                table.markdown_source,
+                "| A   | B   |\n| --- | --- |\n| 1   | 2   |"
+            );
+            assert_eq!(
+                table.rendered_lines(),
+                vec![
+                    "| A   | B   |".to_string(),
+                    "| --- | --- |".to_string(),
+                    "| 1   | 2   |".to_string()
+                ]
+            );
+            assert!(matches!(
+                &table.rendering,
+                AgentOutputTableRendering::Legacy { .. }
+            ));
+        }
+        _ => panic!("expected second section to be Table"),
+    }
+}
+
+#[test]
+fn extracts_markdown_image_into_image_section() {
+    let input = "Intro\n\n![Diagram](./diagram.png)\n\nOutro";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 3);
+    assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
+    match &sections[1] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.alt_text, "Diagram");
+            assert_eq!(image.source, "./diagram.png");
+            assert_eq!(image.markdown_source, "![Diagram](./diagram.png)");
+            assert_eq!(image.layout, AgentOutputImageLayout::Block);
+        }
+        _ => panic!("expected second section to be Image"),
+    }
+    assert!(matches!(sections[2], AIAgentTextSection::PlainText { .. }));
+}
+
+#[test]
+fn extracts_multiple_markdown_images_in_order() {
+    let input = "![One](one.png)\n![Two](two.png)\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 2);
+    match &sections[0] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.markdown_source, "![One](one.png)");
+            assert_eq!(image.layout, AgentOutputImageLayout::Block);
+        }
+        _ => panic!("expected first section to be Image"),
+    }
+    match &sections[1] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.markdown_source, "![Two](two.png)");
+            assert_eq!(image.layout, AgentOutputImageLayout::Block);
+        }
+        _ => panic!("expected second section to be Image"),
+    }
+}
+
+#[test]
+fn extracts_same_line_markdown_images_into_inline_image_sections() {
+    let input = "![One](one.png) ![Two](two.png)\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 2);
+    for (section, expected_markdown) in sections.iter().zip(["![One](one.png)", "![Two](two.png)"])
+    {
+        match section {
+            AIAgentTextSection::Image { image } => {
+                assert_eq!(image.markdown_source, expected_markdown);
+                assert_eq!(image.layout, AgentOutputImageLayout::Inline);
+            }
+            _ => panic!("expected inline image section"),
+        }
+    }
+}
+
+#[test]
+fn does_not_extract_inline_image_run_from_mixed_text_line() {
+    let input = "Intro ![One](one.png) ![Two](two.png)\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 1);
+    assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
+}
+
+#[test]
+fn extracts_block_image_with_commonmark_title() {
+    let input = "![Rex](./rex.png \"My dog Rex\")\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 1);
+    match &sections[0] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.alt_text, "Rex");
+            assert_eq!(image.source, "./rex.png");
+            assert_eq!(image.title.as_deref(), Some("My dog Rex"));
+            // Right-click copy uses `markdown_source`, so it must round-trip
+            // the authored title (product invariant 9).
+            assert_eq!(image.markdown_source, "![Rex](./rex.png \"My dog Rex\")");
+        }
+        _ => panic!("expected image section"),
+    }
+}
+
+#[test]
+fn extracts_inline_image_run_with_partial_title() {
+    // The inline run contains two images where only the second carries a title.
+    let input = "![One](one.png) ![Two](two.png \"caption\")\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 2);
+    match &sections[0] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.title, None);
+            assert_eq!(image.markdown_source, "![One](one.png)");
+            assert_eq!(image.layout, AgentOutputImageLayout::Inline);
+        }
+        _ => panic!("expected first inline image"),
+    }
+    match &sections[1] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.title.as_deref(), Some("caption"));
+            assert_eq!(image.markdown_source, "![Two](two.png \"caption\")");
+            assert_eq!(image.layout, AgentOutputImageLayout::Inline);
+        }
+        _ => panic!("expected second inline image"),
+    }
+}
+
+#[test]
+fn block_image_with_empty_title_normalizes_to_none() {
+    let input = "![Alt](image.png \"\")\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 1);
+    match &sections[0] {
+        AIAgentTextSection::Image { image } => {
+            assert_eq!(image.title, None);
+            // Empty titles normalize away, so `markdown_source` is the
+            // canonical untitled form, not the original source text.
+            assert_eq!(image.markdown_source, "![Alt](image.png)");
+        }
+        _ => panic!("expected image section"),
+    }
+}
+
+#[test]
+fn block_image_with_unclosed_title_falls_back_to_plain_text() {
+    let input = "![Alt](image.png \"unterminated)\n";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 1);
+    // Unclosed titles cause the whole image to render as plain text.
+    assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
+}
+
+#[test]
+fn extracts_mermaid_code_block_into_mermaid_section() {
+    let input = "```mermaid\ngraph TD\nA[Start] --> B[Finish]\n```";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 1);
+    match &sections[0] {
+        AIAgentTextSection::MermaidDiagram { diagram } => {
+            assert_eq!(diagram.source, "graph TD\nA[Start] --> B[Finish]");
+            assert_eq!(
+                diagram.markdown_source,
+                "```mermaid\ngraph TD\nA[Start] --> B[Finish]\n```"
+            );
+        }
+        _ => panic!("expected mermaid diagram section"),
+    }
+}
+
+#[test]
+fn extracts_multiple_mermaid_code_blocks_in_order() {
+    let input = "```mermaid\ngraph TD\nA --> B\n```\n\n```mermaid\ngraph TD\nB --> C\n```";
+    let sections = parse_markdown_into_text_and_code_sections(input);
+
+    assert_eq!(sections.len(), 2);
+    match &sections[0] {
+        AIAgentTextSection::MermaidDiagram { diagram } => {
+            assert_eq!(diagram.source, "graph TD\nA --> B");
+        }
+        _ => panic!("expected first section to be MermaidDiagram"),
+    }
+    match &sections[1] {
+        AIAgentTextSection::MermaidDiagram { diagram } => {
+            assert_eq!(diagram.source, "graph TD\nB --> C");
+        }
+        _ => panic!("expected second section to be MermaidDiagram"),
+    }
+}