Files
galaxy/app/src/ai/agent/conversation.rs
T
Ryan WardandClaude Opus 4.6 f37a744692 Fix Warping indicator stuck after Bedrock LLM finishes responding
Bedrock was calling suggest_next_prompt tool which created a SuggestPrompt
action that waited forever on a oneshot channel for UI interaction that
never fires in the Bedrock path, keeping the conversation permanently
InProgress. Fixed by filtering the tool from the Bedrock tool list and
skipping it at the stream level when the LLM calls it from context history.

Also includes: Bedrock cache token tracking, cost estimation, LSP
improvements, conversation usage view updates, and external config support.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-05-18 12:13:05 -05:00

3856 lines
154 KiB
Rust

use crate::ai::agent::comment::CodeReview;
use crate::ai::agent::linearization::compute_task_depths;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::artifacts::Artifact;
use crate::ai::blocklist::{RequestInput, ResponseStreamId, SerializedBlockListItem};
use crate::ai::skills::SkillDescriptor;
use crate::code_review::CodeReviewTelemetryEvent;
use crate::notebooks::NotebookId;
use crate::persistence::model::{ConversationUsageMetadata, ModelTokenUsage, ToolUsageMetadata};
use crate::server::ids::ServerId;
use crate::terminal::general_settings::GeneralSettings;
use crate::terminal::model::block::{
AgentInteractionMetadata, AgentViewVisibility, BlockId, SerializedAIMetadata, SerializedBlock,
};
use crate::ai::agent::api::convert_conversation::{
compute_time_to_first_token_ms_from_messages, ConvertToExchanges,
};
use ai::document::AIDocumentId;
use chrono::{DateTime, Local, TimeZone};
use itertools::Itertools as _;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::{collections::HashMap, fmt::Display};
use super::task_store::TaskStore;
use galaxy_core::command::ExitCode;
use galaxy_core::execution_mode::AppExecutionMode;
use galaxy_core::features::FeatureFlag;
use galaxy_core::send_telemetry_from_ctx;
use galaxy_core::ui::appearance::Appearance;
use galaxy_core::ui::theme::color::internal_colors;
use galaxy_core::ui::theme::WarpTheme;
use galaxyui::color::ColorU;
use galaxyui::{EntityId, ModelContext, SingletonEntity};
use uuid::Uuid;
use vec1::{Size0Error, Vec1};
use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, response_event::stream_finished::TokenUsage};
use crate::ai::agent::{AIIdentifiers, CancellationReason};
use crate::{
ai::{
agent::{
icons::{
failed_icon, gray_stop_icon, in_progress_icon, succeeded_icon, yellow_stop_icon,
},
todos::AIAgentTodoList,
AIAgentOutputMessage, AIAgentOutputMessageType, MessageToAIAgentOutputMessageError,
},
blocklist::BlocklistAIHistoryEvent,
},
persistence::{
model::{AgentConversationData, PersistedAutoexecuteMode},
ModelEvent,
},
ui_components::icons::Icon,
BlocklistAIHistoryModel, GlobalResourceHandlesProvider,
};
use super::task::{ExtractMessagesError, UpdateTaskError, UpgradeOptimisticTaskError};
use super::{
api::ServerConversationToken,
task::{
derive_todo_lists_from_root_task,
helper::*,
transaction::{SavedTask, Transaction},
Task, TaskId,
},
AIAgentAction, AIAgentActionId, AIAgentContext, AIAgentExchange, AIAgentExchangeId,
AIAgentInput, AIAgentOutputStatus, AIAgentTodo, AIAgentTodoId, FinishedAIAgentOutput,
MessageId, RenderableAIError, RequestCost,
};
use super::{
AIAgentOutput, OutputModelInfo, ServerOutputId, Shared, SuggestedLoggingId, Suggestions,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TodoStatus {
Pending,
InProgress,
Completed,
Cancelled,
Stopped,
}
impl TodoStatus {
pub fn is_cancelled(&self) -> bool {
matches!(self, TodoStatus::Cancelled)
}
}
// basic info for creating a dummy command block based on an exchange's inputs
pub(crate) struct CommandBlockInfo {
pub(crate) command: String,
pub(crate) output: String,
pub(crate) exit_code: ExitCode,
pub(crate) ai_metadata: Option<String>,
/// The api message ID that this command block was extracted from.
/// Used to find the corresponding exchange for timestamp and PWD.
pub(crate) message_id: String,
}
#[derive(Debug, Clone)]
struct AddedExchange {
#[allow(dead_code)]
task_id: TaskId,
exchange_id: AIAgentExchangeId,
}
#[derive(thiserror::Error, Debug)]
pub enum RestoreConversationError {
#[error("Restored conversation has no root task")]
NoRootTask,
}
#[derive(thiserror::Error, Debug)]
#[error("Subagent task not found")]
pub struct SubagentTaskNotFound;
/// An Agent Mode conversation.
#[derive(Debug, Clone)]
pub struct AIConversation {
/// Unique ID for this conversation.
id: AIConversationId,
/// Whether this conversation is being shared from a different warp instance
/// (i.e. is not a local conversation).
is_viewing_shared_session: bool,
task_store: TaskStore,
optimistic_cli_subagent_subtask_id: Option<TaskId>,
/// TODO lists created during the conversation, ordered by creation time. The last list (if any) is the active list.
todo_lists: Vec<AIAgentTodoList>,
/// Current the code review in this conversation, `None` if the has never tried to address
/// comments in this conversation.
code_review: Option<CodeReview>,
status: ConversationStatus,
/// Optional detail for the current error status.
status_error_message: Option<String>,
/// Tracks whether the code review has been opened at least once for this conversation.
has_opened_code_review: bool,
/// Usage metadata for this conversation, including summarization status, context window usage,
/// credits spent, token usage, and tool usage.
conversation_usage_metadata: ConversationUsageMetadata,
/// The server-generated unique "token" for this conversation.
///
/// This must be roundtripped to the server when sending follow-ups within a given conversation.
server_conversation_token: Option<ServerConversationToken>,
/// The server-assigned task/run identifier (`ai_tasks.id`) for this
/// conversation, used for v2 orchestration.
///
/// For local conversations, parsed from `StreamInit.run_id` on the first
/// response. For remote child agents spawned via `POST /agent/run`, set
/// from `SpawnAgentResponse.task_id`.
///
/// Used for messaging API, events API, poller self-filtering, lifecycle
/// reports, parent↔child agent identity, and task status reporting.
/// The string form (for APIs that accept a run_id) is obtained via
/// `run_id()` which calls `.to_string()` on this field.
task_id: Option<AmbientAgentTaskId>,
/// The server conversation ID of the source conversation if this conversation was forked.
forked_from_server_conversation_token: Option<ServerConversationToken>,
/// Metadata from the server for this conversation (permissions, timestamps, etc.).
/// This is None for new conversations and gets populated after the first response completes.
/// TODO (roland): server_conversation_token, conversation_usage_metadata, and artifacts are duplicated in here.
/// Those are updated via stream events on init and finished respectively, while this is fetched via graphQL
/// Consider consolidating by having the stream events return this whole metadata
server_metadata: Option<ServerAIConversationMetadata>,
/// The active transaction for this conversation, if any.
transaction: Option<Transaction>,
/// The per-conversation override on the user's usual autonomy settings.
autoexecute_override: AIConversationAutoexecuteMode,
/// Map of new exchanges added keyed by ID of response stream corresponding to the MAA API
/// request.
added_exchanges_by_response: HashMap<ResponseStreamId, Vec1<AddedExchange>>,
/// A set of the hidden exchanges.
/// This is stored here instead of the AIAgentExchange because this is a view specific field.
/// We cache this here because we don't have access to the block everywhere we are updating the
/// persisted exchanges.
hidden_exchanges: HashSet<AIAgentExchangeId>,
/// A set of action IDs that have been reverted by the user.
reverted_action_ids: HashSet<AIAgentActionId>,
/// Accumulated suggestions received in the course of this conversation.
existing_suggestions: Option<Suggestions>,
/// A set of suggestion logging IDs that have been dismissed for this conversation.
dismissed_suggestion_ids: HashSet<SuggestedLoggingId>,
total_request_cost: RequestCost,
total_token_usage_by_model: HashMap<String, TokenUsage>,
last_block_token_usage_by_model: HashMap<String, TokenUsage>,
/// Fallback title used when no task description or initial query exists.
fallback_display_title: Option<String>,
/// Artifacts created during this conversation (plans, PRs, etc.).
artifacts: Vec<Artifact>,
/// Server-side identifier of the parent agent that spawned this child, if any.
/// In v1 this holds the parent's `server_conversation_token`; in v2 (OrchestrationV2)
/// it holds the parent's `run_id`. Persisted as `parent_agent_id` for serde compat.
parent_agent_id: Option<String>,
/// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator.
agent_name: Option<String>,
/// The local conversation ID of the parent that spawned this child, if any.
parent_conversation_id: Option<AIConversationId>,
/// True when this conversation is a placeholder for a child agent executing
/// on a remote worker. The parent's client does not drive execution for
/// these conversations — the remote worker's own client handles status
/// reporting. TaskStatusSyncModel skips status updates for these.
is_remote_child: bool,
/// The last event sequence number observed from the v2 orchestration
/// event log. Used on restore to resume event delivery without
/// re-delivering already-processed events.
last_event_sequence: Option<i64>,
/// Accumulated message history for direct Bedrock conversations.
/// Contains the full ordered sequence of user messages, assistant responses,
/// tool calls, and tool results sent to/received from Bedrock across all
/// request cycles. This is the source of truth for what Bedrock sees.
bedrock_message_history: Vec<crate::ai::bedrock::convert::ConversationMessage>,
}
pub(crate) fn artifact_from_fork_proto(
proto_artifact: &api::message::artifact_event::ConversationArtifact,
) -> Option<Artifact> {
use api::message::artifact_event::conversation_artifact::Artifact as ProtoArtifact;
match &proto_artifact.artifact {
Some(ProtoArtifact::PullRequest(pr)) => Some(Artifact::from(pr.clone())),
Some(ProtoArtifact::Screenshot(ss)) => Some(Artifact::from(ss.clone())),
Some(ProtoArtifact::Plan(plan)) => Some(Artifact::from(plan.clone())),
Some(ProtoArtifact::File(file)) => Some(Artifact::from(file.clone())),
None => None,
}
}
impl AIConversation {
pub fn new(is_viewing_shared_session: bool) -> Self {
let root_task = Task::new_optimistic_root();
Self {
id: AIConversationId::new(),
task_store: TaskStore::with_root_task(root_task),
optimistic_cli_subagent_subtask_id: None,
code_review: None,
is_viewing_shared_session,
todo_lists: vec![],
status: ConversationStatus::InProgress,
status_error_message: None,
has_opened_code_review: false,
conversation_usage_metadata: ConversationUsageMetadata::default(),
server_conversation_token: None,
task_id: None,
forked_from_server_conversation_token: None,
server_metadata: None,
transaction: None,
autoexecute_override: Default::default(),
added_exchanges_by_response: Default::default(),
hidden_exchanges: Default::default(),
reverted_action_ids: Default::default(),
existing_suggestions: None,
dismissed_suggestion_ids: Default::default(),
total_request_cost: RequestCost::new(0.),
total_token_usage_by_model: Default::default(),
last_block_token_usage_by_model: Default::default(),
fallback_display_title: None,
artifacts: Vec::new(),
parent_agent_id: None,
agent_name: None,
parent_conversation_id: None,
is_remote_child: false,
last_event_sequence: None,
bedrock_message_history: Vec::new(),
}
}
// TODO: derive todo list state from tasks instead of taking args.
//
// This would make it possible to fully restore a convo from tasks, instead of having to persist this additional data.
pub fn new_restored(
id: AIConversationId,
tasks: Vec<api::Task>,
conversation_data: Option<AgentConversationData>,
) -> Result<Self, RestoreConversationError> {
let api_tasks_by_id: HashMap<String, api::Task> =
tasks.into_iter().map(|t| (t.id.clone(), t)).collect();
// To process a task, we need to reference some of the data in its parent task. To
// avoid cloning, we process the task tree from deepest tasks to shallowest tasks. This
// ensures that children are always processed before their parents, avoiding any need to
// clone task data to ensure the parent is available when processing the child.
let depths = compute_task_depths(&api_tasks_by_id);
let mut task_ids: Vec<String> = api_tasks_by_id.keys().cloned().collect();
task_ids.sort_by(|a, b| {
depths
.get(b.as_str())
.unwrap_or(&0)
.cmp(depths.get(a.as_str()).unwrap_or(&0))
});
let mut api_tasks_and_exchanges_by_id: HashMap<_, _> = api_tasks_by_id
.into_iter()
.map(|(id, task)| {
let exchanges = task.into_exchanges();
(id, (task, exchanges))
})
.collect();
let mut tasks_by_id = HashMap::new();
let mut root_task = None;
for task_id in task_ids {
let Some((task, exchanges)) = api_tasks_and_exchanges_by_id.remove(&task_id) else {
continue;
};
if let Some(parent_id) = task.parent_id() {
if let Some((parent_task, _)) = api_tasks_and_exchanges_by_id.get(parent_id) {
tasks_by_id.insert(
TaskId::new(task.id.clone()),
Task::new_restored_subtask(task, parent_task, exchanges),
);
} else {
log::error!(
"Could not find parent task (id: {}) for task (id: {})",
parent_id,
task.id
);
}
} else if root_task.is_none() {
root_task = Some(Task::new_restored_root(task, exchanges.into_iter()));
}
}
let Some(root_task) = root_task else {
return Err(RestoreConversationError::NoRootTask);
};
// Derive todo lists from tasks by replaying UpdateTodos operations
let todo_lists = derive_todo_lists_from_root_task(&root_task);
let root_task_id = root_task.id().clone();
tasks_by_id.insert(root_task.id().clone(), root_task);
let (
server_conversation_token,
forked_from_server_conversation_token,
conversation_usage_metadata,
reverted_action_ids,
artifacts,
parent_agent_id,
agent_name,
parent_conversation_id,
run_id,
autoexecute_override,
last_event_sequence,
) = if let Some(data) = conversation_data {
let server_conversation_token = data
.server_conversation_token
.map(ServerConversationToken::new);
let conversation_usage_metadata = data.conversation_usage_metadata.unwrap_or_default();
let reverted_action_ids = data.reverted_action_ids.unwrap_or_default();
let forked_from_server_conversation_token = data
.forked_from_server_conversation_token
.map(ServerConversationToken::new);
let artifacts = data
.artifacts_json
.and_then(|json| {
serde_json::from_str(&json)
.map_err(|e| log::error!("Failed to deserialize artifacts: {e}"))
.ok()
})
.unwrap_or_default();
let parent_agent_id = data.parent_agent_id;
let agent_name = data.agent_name;
let parent_conversation_id = data
.parent_conversation_id
.and_then(|id| AIConversationId::try_from(id).ok());
let run_id = data.run_id;
let autoexecute_override = if FeatureFlag::RememberFastForwardState.is_enabled() {
data.autoexecute_override
.map(Into::into)
.unwrap_or_default()
} else {
AIConversationAutoexecuteMode::default()
};
let last_event_sequence = data.last_event_sequence;
(
server_conversation_token,
forked_from_server_conversation_token,
conversation_usage_metadata,
reverted_action_ids,
artifacts,
parent_agent_id,
agent_name,
parent_conversation_id,
run_id,
autoexecute_override,
last_event_sequence,
)
} else {
(
None,
None,
ConversationUsageMetadata::default(),
Default::default(),
Vec::new(),
None,
None,
None,
None,
AIConversationAutoexecuteMode::default(),
None,
)
};
// Convert these from the persistence type to the runtime one.
let reverted_action_ids = reverted_action_ids.into_iter().map_into().collect();
// Determine the correct status based on the exchanges before constructing
let status = Self::derive_status_from_root_task(&tasks_by_id.get(&root_task_id));
let task_store = TaskStore::from_tasks(tasks_by_id, root_task_id);
Ok(Self {
id,
is_viewing_shared_session: false,
task_store,
status,
status_error_message: None,
todo_lists,
// TODO(alokedesai): Support session restoration for code review comments.
code_review: None,
has_opened_code_review: false,
conversation_usage_metadata,
server_conversation_token,
task_id: run_id.as_deref().and_then(|id| id.parse().ok()),
forked_from_server_conversation_token,
server_metadata: None,
transaction: None,
autoexecute_override,
added_exchanges_by_response: Default::default(),
existing_suggestions: None,
hidden_exchanges: Default::default(),
reverted_action_ids,
dismissed_suggestion_ids: Default::default(),
total_request_cost: RequestCost::new(0.),
total_token_usage_by_model: Default::default(),
last_block_token_usage_by_model: Default::default(),
optimistic_cli_subagent_subtask_id: None,
fallback_display_title: None,
artifacts,
parent_agent_id,
agent_name,
parent_conversation_id,
is_remote_child: false,
last_event_sequence,
bedrock_message_history: Vec::new(),
})
}
pub fn id(&self) -> AIConversationId {
self.id
}
pub fn bedrock_message_history(&self) -> &[crate::ai::bedrock::convert::ConversationMessage] {
&self.bedrock_message_history
}
pub fn bedrock_message_history_mut(
&mut self,
) -> &mut Vec<crate::ai::bedrock::convert::ConversationMessage> {
&mut self.bedrock_message_history
}
pub fn append_to_bedrock_history(
&mut self,
messages: Vec<crate::ai::bedrock::convert::ConversationMessage>,
) {
self.bedrock_message_history.extend(messages);
}
/// Assigns fresh exchange IDs to all exchanges in this conversation.
/// Used when forking conversations to avoid ID collisions with persisted blocks.
pub fn reassign_exchange_ids(&mut self) {
let task_ids: Vec<TaskId> = self.task_store.tasks().map(|t| t.id().clone()).collect();
for task_id in task_ids {
self.task_store.modify_task(&task_id, |task| {
task.reassign_exchange_ids();
});
}
}
pub fn is_viewing_shared_session(&self) -> bool {
self.is_viewing_shared_session
}
pub fn set_is_viewing_shared_session(&mut self, is_viewing_shared_session: bool) {
self.is_viewing_shared_session = is_viewing_shared_session;
}
pub fn was_summarized(&self) -> bool {
self.conversation_usage_metadata.was_summarized
}
pub fn context_window_usage(&self) -> f32 {
self.conversation_usage_metadata.context_window_usage
}
pub fn credits_spent(&self) -> f32 {
(self.conversation_usage_metadata.credits_spent * 10.0).round() / 10.0
}
// Credits spent over the last block, where the block comprises
// all agent outputs since the most recent user input.
pub fn credits_spent_for_last_block(&self) -> Option<f32> {
self.conversation_usage_metadata
.credits_spent_for_last_block
.map(|credits| (credits * 10.0).round() / 10.0)
}
pub fn last_block_total_tokens(&self) -> u32 {
self.last_block_token_usage_by_model
.values()
.map(Self::total_tokens_for_usage)
.sum()
}
/// Time to first token for the last completed set of agent responses
/// since the most recent user query
pub fn time_to_first_token_for_last_user_query_ms(&self) -> i64 {
let exchanges = self.all_exchanges();
if exchanges.is_empty() {
return 0;
}
// Walk backwards from the end to find all exchanges in the last block
// (everything since the last user query).
for exchange in exchanges.iter().rev() {
if exchange.has_user_query() {
return exchange.time_to_first_token_ms.unwrap_or(0);
}
}
// If we never found a user query, return the time_to_first_token_ms from the first exchange
exchanges
.first()
.and_then(|ex| ex.time_to_first_token_ms)
.unwrap_or(0)
}
/// Helper to derive an exchange's finish time from its associated task messages.
fn finish_time_from_exchange_messages(
task: &Task,
exchange: &AIAgentExchange,
) -> Option<DateTime<Local>> {
task.messages()
.filter(|m| !m.id.is_empty())
.filter(|m| {
let id = MessageId::new(m.id.clone());
exchange.added_message_ids.contains(&id)
})
.filter_map(|m| {
m.timestamp.as_ref().and_then(|ts| {
let nanos = if ts.nanos < 0 { 0 } else { ts.nanos as u32 };
Local.timestamp_opt(ts.seconds, nanos).single()
})
})
.max()
}
/// Derive an exchange's start time from the latest input's context.
fn start_time_from_exchange_messages(exchange: &AIAgentExchange) -> Option<DateTime<Local>> {
exchange
.input
.last()
.and_then(|input| input.context())
.and_then(|contexts| {
contexts.iter().find_map(|context| match context {
AIAgentContext::CurrentTime { current_time } => Some(*current_time),
_ => None,
})
})
}
/// Derive the conversation status from the root task's exchanges.
/// Used when restoring conversations to determine if they were cancelled or completed successfully.
fn derive_status_from_root_task(root_task: &Option<&Task>) -> ConversationStatus {
let Some(root_task) = root_task else {
return ConversationStatus::Success;
};
// Check the last exchange's output status
if let Some(last_exchange) = root_task.last_exchange() {
match &last_exchange.output_status {
AIAgentOutputStatus::Finished {
finished_output: FinishedAIAgentOutput::Cancelled { .. },
} => return ConversationStatus::Cancelled,
AIAgentOutputStatus::Finished {
finished_output: FinishedAIAgentOutput::Error { .. },
} => return ConversationStatus::Error,
_ => {}
}
}
// If not cancelled or errored, it's successful
ConversationStatus::Success
}
/// Total agent response time for the last completed set of agent responses
/// since the most recent user query.
pub fn total_agent_response_time_since_last_user_query_ms(&self) -> i64 {
let exchanges = self.all_exchanges();
if exchanges.is_empty() {
return 0;
}
// Walk backwards, accumulating durations until we find a user query
let mut total_ms: i64 = 0;
for exchange in exchanges.iter().rev() {
total_ms += exchange
.duration()
.map(|duration| duration.num_milliseconds())
.unwrap_or(0);
if exchange.has_user_query() {
break;
}
}
total_ms
}
/// Wall-to-wall response time for the last completed set of agent responses.
pub fn wall_to_wall_response_time_since_last_query(&self) -> Option<i64> {
let exchanges = self.all_exchanges();
let last_exchange = exchanges.last().copied()?;
let finish_time = last_exchange.finish_time?;
// Walk backwards to find the most recent exchange with a user query
let start_time = exchanges.iter().rev().find_map(|exchange| {
if exchange.has_user_query() {
Some(exchange.start_time)
} else {
None
}
})?;
let duration = finish_time.signed_duration_since(start_time);
Some(duration.num_milliseconds())
}
pub fn token_usage(&self) -> &[ModelTokenUsage] {
&self.conversation_usage_metadata.token_usage
}
pub fn tool_usage_metadata(&self) -> &ToolUsageMetadata {
&self.conversation_usage_metadata.tool_usage_metadata
}
pub fn usage_metadata(&self) -> ConversationUsageMetadata {
self.conversation_usage_metadata.clone()
}
pub fn status(&self) -> &ConversationStatus {
&self.status
}
pub fn status_error_message(&self) -> Option<&str> {
self.status_error_message.as_deref()
}
pub fn update_status(
&mut self,
status: ConversationStatus,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
self.update_status_with_error_message(status, None, terminal_view_id, ctx);
}
pub fn update_status_with_error_message(
&mut self,
status: ConversationStatus,
error_message: Option<String>,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
self.status_error_message = if matches!(&status, ConversationStatus::Error) {
error_message.filter(|message| !message.trim().is_empty())
} else {
None
};
self.status = status;
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationStatus {
conversation_id: self.id,
terminal_view_id,
is_restored: false,
});
}
pub fn is_processing_response_stream(&self, stream_id: &ResponseStreamId) -> bool {
self.added_exchanges_by_response.contains_key(stream_id)
}
/// Removes the response stream tracking entry after the stream has fully completed.
pub fn cleanup_completed_response_stream(&mut self, stream_id: &ResponseStreamId) {
self.added_exchanges_by_response.remove(stream_id);
}
pub fn new_exchange_ids_for_response(
&self,
stream_id: &ResponseStreamId,
) -> impl Iterator<Item = AIAgentExchangeId> + '_ {
self.added_exchanges_by_response
.get(stream_id)
.into_iter()
.flat_map(|added_exchanges| {
added_exchanges
.iter()
.map(|new_exchange| new_exchange.exchange_id)
})
}
pub fn server_conversation_token(&self) -> Option<&ServerConversationToken> {
self.server_conversation_token.as_ref()
}
/// Returns the server-assigned run identifier as a string.
pub fn run_id(&self) -> Option<String> {
self.task_id.map(|id| id.to_string())
}
/// Sets the task ID by parsing a run_id string.
pub fn set_run_id(&mut self, id: String) {
self.task_id = id.parse().ok();
}
/// Returns the server-assigned task ID, if available.
pub fn task_id(&self) -> Option<AmbientAgentTaskId> {
self.task_id
}
/// Sets the task ID directly (used for child agents spawned via `SpawnAgentResponse`).
pub fn set_task_id(&mut self, id: AmbientAgentTaskId) {
self.task_id = Some(id);
}
/// Returns the server-side agent identifier appropriate for the active
/// orchestration version: `task_id` (as string) under v2,
/// `server_conversation_token` under v1.
pub fn orchestration_agent_id(&self) -> Option<String> {
if FeatureFlag::OrchestrationV2.is_enabled() {
self.run_id()
} else {
self.server_conversation_token
.as_ref()
.map(|t| t.as_str().to_string())
}
}
/// Updates the server conversation token for this conversation.
///
/// This is used internally for session sharing when a forked conversation receives
/// its new server-assigned token. The viewer needs to update the conversation's token
/// from the original (forked-from) token to the new token so subsequent messages can
/// be matched to the correct conversation.
///
/// This should only be called by session sharing viewer logic when linking forked conversations.
pub(crate) fn set_server_conversation_token(&mut self, token: String) {
self.server_conversation_token = Some(ServerConversationToken::new(token));
}
pub fn forked_from_server_conversation_token(&self) -> Option<&ServerConversationToken> {
self.forked_from_server_conversation_token.as_ref()
}
/// Clears the forked_from token after the first Init event has been sent to viewers.
/// This ensures we only send the forked_from token once during session sharing.
pub(crate) fn clear_forked_from_server_conversation_token(&mut self) {
self.forked_from_server_conversation_token = None;
}
pub fn server_id(&self) -> Option<ServerId> {
self.server_metadata
.as_ref()
.map(|metadata| metadata.metadata.uid)
}
pub fn server_metadata(&self) -> Option<&ServerAIConversationMetadata> {
self.server_metadata.as_ref()
}
pub fn set_server_metadata(&mut self, metadata: ServerAIConversationMetadata) {
self.server_metadata = Some(metadata);
}
pub fn parent_agent_id(&self) -> Option<&str> {
self.parent_agent_id.as_deref()
}
pub fn set_parent_agent_id(&mut self, id: String) {
self.parent_agent_id = Some(id);
}
pub fn agent_name(&self) -> Option<&str> {
self.agent_name.as_deref()
}
pub fn set_agent_name(&mut self, name: String) {
self.agent_name = Some(name);
}
pub fn parent_conversation_id(&self) -> Option<AIConversationId> {
self.parent_conversation_id
}
pub fn set_parent_conversation_id(&mut self, id: AIConversationId) {
self.parent_conversation_id = Some(id);
}
/// Returns the last observed v2 orchestration event sequence number, if any.
pub fn last_event_sequence(&self) -> Option<i64> {
self.last_event_sequence
}
/// Updates the last observed v2 orchestration event sequence number.
pub fn set_last_event_sequence(&mut self, sequence: i64) {
self.last_event_sequence = Some(sequence);
}
/// Returns true if this conversation was spawned by a parent orchestrator agent.
pub fn is_child_agent_conversation(&self) -> bool {
self.parent_conversation_id.is_some()
}
/// Returns true if this is a placeholder for a child agent executing on a
/// remote worker. The parent's client should not report task status for
/// these — the remote worker handles it.
pub fn is_remote_child(&self) -> bool {
self.is_remote_child
}
/// Marks this conversation as a remote child placeholder.
pub fn mark_as_remote_child(&mut self) {
self.is_remote_child = true;
}
/// Returns a flat list of linearized messages across all tasks, interpolating subtask messages
/// in between subagent tool calls and results, effectively corresponding to the order in which
/// the messages were created and added to the conversation.
pub fn all_linearized_messages(&self) -> Vec<&api::Message> {
self.task_store.all_linearized_messages()
}
/// Returns all the tasks in this conversation.
///
/// Note that until we've fully migrated to the multi-agent endpoint, in reality, each
/// conversation is comprised of a single task (the legacy endpoint `GenerateAIAgentOutput` does
/// not support multiple tasks within a conversation).
pub fn all_tasks(&self) -> impl Iterator<Item = &Task> {
self.task_store.tasks()
}
/// Returns the set of tasks that are still active (relevant to the agent).
///
/// This filters the full task list using DFS linearization to determine
/// which tasks have open subagent tool calls without corresponding results.
pub fn compute_active_tasks(&self) -> Vec<warp_multi_agent_api::Task> {
use std::collections::HashMap;
let root_task_id = self.get_root_task_id().to_string();
let all_tasks: HashMap<&str, &warp_multi_agent_api::Task> = self
.all_tasks()
.filter_map(|task| {
let source = task.source()?;
Some((source.id.as_str(), source))
})
.collect();
let active_task_ids =
crate::ai::agent::linearization::compute_active_task_ids(&root_task_id, &all_tasks);
all_tasks
.into_values()
.filter(|task| active_task_ids.contains(task.id.as_str()))
.cloned()
.collect()
}
/// Returns the titles from the CreateDocuments request corresponding to the given action ID (if any).
/// This is used by shared-session viewers to use the correct document titles from the original CreateDocuments action.
pub fn get_document_titles_for_action(
&self,
action_id: &AIAgentActionId,
) -> Option<Vec<String>> {
for exchange in self.all_exchanges() {
let Some(output) = exchange.output_status.output() else {
continue;
};
for message in &output.get().messages {
if let AIAgentOutputMessage {
message: AIAgentOutputMessageType::Action(action),
..
} = message
{
if &action.id == action_id {
if let super::AIAgentActionType::CreateDocuments(
super::CreateDocumentsRequest { documents },
) = &action.action
{
let titles = documents
.iter()
.map(|doc| doc.title.clone())
.collect::<Vec<_>>();
return Some(titles);
}
}
}
}
}
None
}
/// Returns the start timestamp of the earliest [`AIAgentExchange`] in the conversation, if
/// any.
pub fn start_ts(&self) -> Option<DateTime<Local>> {
self.root_task_exchanges()
.next()
.map(|exchange| exchange.start_time)
}
pub fn has_opened_code_review(&self) -> bool {
self.has_opened_code_review
}
pub fn mark_code_review_as_opened(&mut self) {
self.has_opened_code_review = true;
}
/// Returns the IDs of comments that have been addressed in this conversation.
pub fn addressed_comment_ids(&self) -> HashSet<crate::code_review::comments::CommentId> {
self.code_review
.as_ref()
.map(|cr| cr.addressed_comments.iter().map(|c| c.id).collect())
.unwrap_or_default()
}
pub fn is_entirely_passive_code_diff(&self) -> bool {
let mut has_passive_code_diff_exchange = false;
for exchange in self.root_task_exchanges() {
has_passive_code_diff_exchange |= exchange.has_passive_code_diff();
if exchange.has_user_query() {
return false;
}
}
has_passive_code_diff_exchange
}
pub fn is_entirely_passive(&self) -> bool {
let mut has_passive_exchange = false;
for exchange in self.root_task_exchanges() {
has_passive_exchange |= exchange.has_passive_request();
if exchange.has_user_query() {
return false;
}
}
has_passive_exchange
}
/// True if the conversation consists of just one exchange
/// and that exchange is a passive suggestion.
pub fn is_single_passive_exchange(&self) -> bool {
self.task_store.task_count() == 1
&& self.is_entirely_passive()
&& self
.get_root_task()
.is_some_and(|task| task.exchanges_len() == 1)
}
/// True if the conversation started with a CLI subagent and was never continued.
/// These conversations only have CLI subagent exchanges with no user queries,
/// meaning they never hit the primary agent.
pub fn is_orphaned_cli_subagent_conversation(&self) -> bool {
// Check if conversation has only 1 task (root task) and it's a CLI subagent
let started_with_cli_subagent = self.task_store.task_count() == 1
&& self
.get_root_task()
.is_some_and(|task| task.is_cli_subagent());
if !started_with_cli_subagent {
return false;
}
// Check if conversation was never continued (no user queries in any exchange)
let never_continued = self
.root_task_exchanges()
.all(|exchange| !exchange.has_user_query());
never_continued
}
/// Returns true if this conversation should be unconditionally excluded
/// from conversation navigation and history.
pub fn should_exclude_from_navigation(&self) -> bool {
// Passive-only suggestions without any follow-up requests shouldn't be presented as
// conversations.
self.is_entirely_passive()
// Orphaned CLI subagent conversations (invoked from within a terminal block) are
// internal and shouldn't appear in navigation.
|| self.is_orphaned_cli_subagent_conversation()
// Shared session viewer conversations are excluded because the shared session itself
// is visible/represented elsewhere.
|| self.is_viewing_shared_session()
// Child agent conversations spawned by an orchestrator are managed via the parent's
// status card and shouldn't clutter the navigation list.
|| self.is_child_agent_conversation()
}
pub fn existing_suggestions(&self) -> Option<&Suggestions> {
self.existing_suggestions.as_ref()
}
pub fn dismissed_suggestion_ids(&self) -> &HashSet<SuggestedLoggingId> {
&self.dismissed_suggestion_ids
}
pub fn dismiss_current_suggestions(&mut self) {
if let Some(suggestions) = &self.existing_suggestions {
self.dismissed_suggestion_ids
.extend(suggestions.rules.iter().map(|r| r.logging_id.clone()));
self.dismissed_suggestion_ids.extend(
suggestions
.agent_mode_workflows
.iter()
.map(|w| w.logging_id.clone()),
);
}
}
pub fn is_exchange_hidden(&self, exchange_id: AIAgentExchangeId) -> bool {
self.hidden_exchanges.contains(&exchange_id)
}
pub fn set_is_exchange_hidden(
&mut self,
exchange_id: AIAgentExchangeId,
is_hidden: bool,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
// If the status is not being modified, return.
if is_hidden == self.hidden_exchanges.contains(&exchange_id) {
return;
}
if is_hidden {
self.hidden_exchanges.insert(exchange_id);
} else {
self.hidden_exchanges.remove(&exchange_id);
}
// If the status is being toggled, set the persisted exchange hidden status.
// Find the exchange and the terminal view ID for the exchange and emit an event to update
// the exchange hidden state.
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden,
});
}
/// Returns an iterator over all exchanges in all tasks in this conversation.
pub fn all_exchanges(&self) -> Vec<&AIAgentExchange> {
self.task_store.all_exchanges().collect()
}
/// Returns a vector of vectors of exchanges, in linearized order as they appeared in the
/// conversation, grouped by task ID.
pub fn all_exchanges_by_task(&self) -> Vec<(TaskId, Vec<&AIAgentExchange>)> {
self.task_store.all_exchanges_by_task()
}
pub fn root_task_exchanges(&self) -> impl Iterator<Item = &AIAgentExchange> {
self.task_store
.root_task()
.into_iter()
.flat_map(|task| task.exchanges())
}
pub fn exchange_count(&self) -> usize {
self.task_store.exchange_count()
}
pub fn is_empty(&self) -> bool {
self.exchange_count() == 0
}
pub fn exchanges_reversed(&self) -> impl Iterator<Item = &AIAgentExchange> {
self.task_store
.root_task()
.into_iter()
.flat_map(|task| task.exchanges_reversed())
}
#[cfg_attr(target_family = "wasm", allow(unused))]
pub fn exchange_with_id(&self, exchange_id: AIAgentExchangeId) -> Option<&AIAgentExchange> {
for task in self.task_store.tasks() {
if let Some(exchange) = task.exchanges().find(|exchange| exchange.id == exchange_id) {
return Some(exchange);
}
}
None
}
/// Returns the exchange that preceded the exchange with the given id, if there is one.
pub fn previous_exchange(&self, exchange_id: &AIAgentExchangeId) -> Option<&AIAgentExchange> {
self.exchanges_reversed()
.skip_while(|e| e.id != *exchange_id)
.nth(1)
}
/// Returns the last exchange that didn't contain a passive request.
pub fn last_non_passive_exchange(&self) -> Option<&AIAgentExchange> {
self.exchanges_reversed()
.skip_while(|e| e.has_passive_request())
.nth(0)
}
pub fn first_exchange(&self) -> Option<&AIAgentExchange> {
self.task_store.first_exchange()
}
pub fn latest_exchange(&self) -> Option<&AIAgentExchange> {
self.task_store.latest_exchange()
}
pub fn latest_skills(&self) -> Option<Vec<SkillDescriptor>> {
self.task_store.latest_skills()
}
/// Get the auto-generated title of the given conversation
/// (falling back to the first query if the title is empty).
/// Get the title of the given conversation.
/// Priority: auto-generated task description > initial query > fallback_display_title.
pub fn title(&self) -> Option<String> {
self.task_store
.root_task()
.and_then(|task| {
if task.description().is_empty() {
self.initial_query()
} else {
Some(task.description().to_owned())
}
})
.or_else(|| self.fallback_display_title.clone())
}
/// Set a fallback title used when no task description or initial query exists.
pub fn set_fallback_display_title(&mut self, title: String) {
self.fallback_display_title = Some(title);
}
/// Returns the last time this conversation was modified (i.e., when the latest exchange was started).
pub fn last_modified_at(&self) -> Option<DateTime<Local>> {
self.latest_exchange()
.map(|e| e.finish_time.unwrap_or(e.start_time))
}
/// Returns artifacts created during this conversation.
pub fn artifacts(&self) -> &[Artifact] {
&self.artifacts
}
/// Adds an artifact to this conversation and persists the change.
pub fn add_artifact(
&mut self,
artifact: Artifact,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
self.artifacts.push(artifact.clone());
self.write_updated_conversation_state(ctx);
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
conversation_id: self.id,
artifact,
});
}
/// Updates the notebook_uid for a plan artifact when it's synced to Warp Drive.
pub fn update_plan_notebook_uid(
&mut self,
document_uid: AIDocumentId,
notebook_uid: NotebookId,
terminal_view_id: Option<EntityId>,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
let document_uid = document_uid.to_string();
for artifact in &mut self.artifacts {
if let Artifact::Plan {
document_uid: doc_uid,
notebook_uid: ref mut nb_uid,
..
} = artifact
{
if doc_uid == &document_uid {
*nb_uid = Some(notebook_uid);
let updated_artifact = artifact.clone();
self.write_updated_conversation_state(ctx);
if let Some(terminal_view_id) = terminal_view_id {
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
conversation_id: self.id,
artifact: updated_artifact,
});
}
return;
}
}
}
}
pub fn initial_query(&self) -> Option<String> {
self.root_task_exchanges()
.flat_map(|exchange| exchange.input.iter())
.find_map(|input| {
AIAgentInput::user_query(input)
.or_else(|| AIAgentInput::auto_code_diff_query(input).map(|s| s.to_string()))
.or_else(|| AIAgentInput::prompt_suggestion_result(input).cloned())
})
}
pub fn initial_user_query(&self) -> Option<String> {
self.root_task_exchanges()
.flat_map(|exchange| exchange.input.iter())
.find_map(AIAgentInput::user_query)
}
/// Export the conversation to markdown format.
/// This is used by both clipboard export and file export.
pub fn export_to_markdown(
&self,
action_model: Option<&crate::ai::blocklist::BlocklistAIActionModel>,
) -> String {
let mut result = Vec::new();
for exchange in self.all_exchanges() {
let formatted_exchange = exchange.format_for_copy(action_model);
if !formatted_exchange.is_empty() {
result.push(formatted_exchange);
}
}
result.join("\n\n")
}
pub fn has_auto_code_diff_query(&self) -> bool {
self.root_task_exchanges()
.flat_map(|exchange| exchange.input.iter())
.any(|input| input.is_auto_code_diff_query())
}
pub fn latest_user_query(&self) -> Option<String> {
self.exchanges_reversed().find_map(|exchange| {
exchange.input.iter().rev().find_map(|input| {
AIAgentInput::user_query(input)
.map(|query| query.trim().to_owned())
.filter(|query| !query.is_empty())
})
})
}
/// Returns an iterator over the IDs of all UseComputer actions across all exchanges
/// in this conversation.
pub fn use_computer_action_ids(&self) -> impl Iterator<Item = AIAgentActionId> + '_ {
self.all_exchanges().into_iter().flat_map(|exchange| {
exchange
.output_status
.output()
.into_iter()
.flat_map(|output| {
output
.get()
.actions()
.filter(|a| matches!(a.action, super::AIAgentActionType::UseComputer(_)))
.map(|a| a.id.clone())
.collect::<Vec<_>>()
})
})
}
pub fn contains_action(&self, action_id: &AIAgentActionId) -> bool {
self.task_store.tasks().any(|task| {
task.exchanges()
.any(|exchange| {
let Some(output) = exchange.output_status.output()
else {
return false;
};
output.get().messages.iter().any(|step| {
matches!(step, AIAgentOutputMessage{ message: AIAgentOutputMessageType::Action(AIAgentAction { id, .. }), .. } if id == action_id)
})
})
})
}
/// Returns the exchange ID that contains the given action ID, if any.
pub fn exchange_id_for_action(&self, action_id: &AIAgentActionId) -> Option<AIAgentExchangeId> {
for task in self.task_store.tasks() {
for exchange in task.exchanges() {
let Some(output) = exchange.output_status.output() else {
continue;
};
let contains_action = output.get().messages.iter().any(|step| {
matches!(step, AIAgentOutputMessage{ message: AIAgentOutputMessageType::Action(AIAgentAction { id, .. }), .. } if id == action_id)
});
if contains_action {
return Some(exchange.id);
}
}
}
None
}
/// Returns the `AIAgentContext` objects attached to the exchange with the given ID, if any.
pub fn context_for_exchange(
&self,
exchange_id: AIAgentExchangeId,
) -> impl Iterator<Item = &AIAgentContext> {
context_in_exchanges(self.exchange_with_id(exchange_id).into_iter())
}
pub fn update_for_new_request_input(
&mut self,
request_input: RequestInput,
stream_id: ResponseStreamId,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
if let Some(request_info) = self.added_exchanges_by_response.remove(&stream_id) {
log::error!(
"Existing response stream info for stream id {stream_id:?}: {request_info:?}"
);
}
let RequestInput {
input_messages,
working_directory,
model_id,
coding_model_id,
cli_agent_model_id,
computer_use_model_id,
shared_session_response_initiator,
request_start_ts,
..
} = request_input;
for (task_id, inputs) in input_messages.into_iter() {
let should_hide = inputs
.iter()
.any(|input| input.is_passive_suggestion_trigger());
let new_exchange = AIAgentExchange {
id: AIAgentExchangeId::new(),
input: inputs,
output_status: AIAgentOutputStatus::Streaming { output: None },
added_message_ids: HashSet::new(),
start_time: request_start_ts,
finish_time: None,
time_to_first_token_ms: None,
working_directory: working_directory.clone(),
// TODO(CORE-3546): fetch shell launch data from active session
model_id: model_id.clone(),
coding_model_id: coding_model_id.clone(),
cli_agent_model_id: cli_agent_model_id.clone(),
computer_use_model_id: computer_use_model_id.clone(),
request_cost: None,
// This will be None for non-shared sessions
response_initiator: shared_session_response_initiator.clone(),
};
let new_exchange_id = new_exchange.id;
self.append_exchange_to_task(&task_id, new_exchange)?;
self.added_exchanges_by_response.insert(
stream_id.clone(),
Vec1::new(AddedExchange {
task_id: task_id.clone(),
exchange_id: new_exchange_id,
}),
);
if should_hide {
self.hidden_exchanges.insert(new_exchange_id);
}
ctx.emit(BlocklistAIHistoryEvent::AppendedExchange {
exchange_id: new_exchange_id,
task_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: should_hide,
response_stream_id: Some(stream_id.clone()),
});
}
Ok(())
}
pub fn append_reassigned_exchange(
&mut self,
response_stream_id: &ResponseStreamId,
exchange: AIAgentExchange,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
let root_task_id = self.task_store.root_task_id().clone();
let exchange_id = exchange.id;
if exchange.output_status.is_streaming() {
if let Some(added_exchanges) =
self.added_exchanges_by_response.get_mut(response_stream_id)
{
added_exchanges.push(AddedExchange {
task_id: root_task_id.clone(),
exchange_id,
});
} else {
self.added_exchanges_by_response.insert(
response_stream_id.clone(),
Vec1::new(AddedExchange {
task_id: root_task_id.clone(),
exchange_id,
}),
);
}
}
self.append_exchange_to_task(&root_task_id, exchange)?;
ctx.emit(BlocklistAIHistoryEvent::ReassignedExchange {
exchange_id,
terminal_view_id,
new_task_id: root_task_id,
new_conversation_id: self.id,
});
Ok(())
}
fn append_exchange_to_task(
&mut self,
task_id: &TaskId,
exchange: AIAgentExchange,
) -> Result<(), UpdateConversationError> {
for input in exchange.input.iter() {
if let AIAgentInput::CodeReview {
review_comments, ..
} = input
{
let review_comments = review_comments
.comments
.clone()
.into_iter()
.map(|c| c.into())
.collect();
if let Some(code_review) = self.code_review.as_mut() {
code_review.pending_comments.extend(review_comments);
} else {
self.code_review = Some(CodeReview::new_with_pending_comments(review_comments));
}
}
}
if self.task_store.append_exchange(task_id, exchange) {
Ok(())
} else {
Err(UpdateConversationError::NoActiveTask)
}
}
pub fn remove_exchange(
&mut self,
exchange_id: AIAgentExchangeId,
) -> Result<AIAgentExchange, UpdateConversationError> {
let mut response_entries_to_remove = vec![];
for (stream_id, added_exchanges) in self.added_exchanges_by_response.iter_mut() {
if let Some(idx) = added_exchanges
.iter()
.position(|new_exchange| new_exchange.exchange_id == exchange_id)
{
if let Err(Size0Error) = added_exchanges.remove(idx) {
response_entries_to_remove.push(stream_id.clone());
}
}
}
for response_id in response_entries_to_remove.into_iter() {
self.added_exchanges_by_response.remove(&response_id);
}
// Find which task contains this exchange
let task_id = self.task_store.tasks().find_map(|task| {
task.exchanges()
.any(|e| e.id == exchange_id)
.then(|| task.id().clone())
});
if let Some(task_id) = task_id {
if let Some(exchange) = self.task_store.remove_task_exchange(&task_id, exchange_id) {
return Ok(exchange);
}
}
Err(UpdateConversationError::ExchangeNotFound)
}
pub fn initialize_output_for_response_stream(
&mut self,
stream_id: &ResponseStreamId,
init_event: warp_multi_agent_api::response_event::StreamInit,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
let Some(new_exchanges) = self.added_exchanges_by_response.get(stream_id).cloned() else {
return Err(UpdateConversationError::NoPendingRequest);
};
let request_id = init_event.request_id.clone();
for new_exchange_info in new_exchanges.iter() {
self.get_exchange_to_update(new_exchange_info.exchange_id)?
.init_output(ServerOutputId::new(request_id.clone()))?;
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id: new_exchange_info.exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self
.hidden_exchanges
.contains(&new_exchange_info.exchange_id),
});
}
self.server_conversation_token =
Some(ServerConversationToken::new(init_event.conversation_id));
let run_id = Some(init_event.run_id).filter(|s| !s.is_empty());
self.task_id = run_id.as_deref().and_then(|id| id.parse().ok());
Ok(())
}
pub fn update_cost_and_usage_for_request(
&mut self,
request_cost: Option<RequestCost>,
token_usage: Vec<TokenUsage>,
usage_metadata: Option<stream_finished::ConversationUsageMetadata>,
was_user_initiated_request: bool,
) -> Result<(), UpdateConversationError> {
if was_user_initiated_request {
self.last_block_token_usage_by_model.clear();
}
for usage in token_usage.into_iter() {
let entry = self
.total_token_usage_by_model
.entry(usage.model_id.clone())
.or_insert_with(|| TokenUsage {
model_id: usage.model_id.clone(),
total_input: 0,
output: 0,
input_cache_read: 0,
input_cache_write: 0,
cost_in_cents: 0.0,
});
entry.total_input += usage.total_input;
entry.output += usage.output;
entry.input_cache_read += usage.input_cache_read;
entry.input_cache_write += usage.input_cache_write;
entry.cost_in_cents += usage.cost_in_cents;
let last_block_entry = self
.last_block_token_usage_by_model
.entry(usage.model_id.clone())
.or_insert_with(|| TokenUsage {
model_id: usage.model_id.clone(),
total_input: 0,
output: 0,
input_cache_read: 0,
input_cache_write: 0,
cost_in_cents: 0.0,
});
last_block_entry.total_input += usage.total_input;
last_block_entry.output += usage.output;
last_block_entry.input_cache_read += usage.input_cache_read;
last_block_entry.input_cache_write += usage.input_cache_write;
last_block_entry.cost_in_cents += usage.cost_in_cents;
}
if let Some(request_cost) = request_cost {
let credits_spent_for_last_block = self
.conversation_usage_metadata
.credits_spent_for_last_block
.get_or_insert(0.0);
// If this exchange begins with a user input (implying it is initiating a new response),
// reset credits spent to only include credits for this new response.
if was_user_initiated_request {
*credits_spent_for_last_block = 0.;
}
// Accumulate response credit usage.
*credits_spent_for_last_block += request_cost.value() as f32;
self.total_request_cost += request_cost;
}
if let Some(usage_metadata) = usage_metadata {
self.conversation_usage_metadata.context_window_usage =
usage_metadata.context_window_usage;
self.conversation_usage_metadata.credits_spent = usage_metadata.credits_spent;
let mut token_usage: HashMap<_, ModelTokenUsage> = HashMap::new();
for (model_id, usage) in usage_metadata.warp_token_usage {
let entry = token_usage.entry(model_id.clone()).or_default();
entry.warp_tokens += usage.total_tokens;
for (category, tokens) in usage.token_usage_by_category {
*entry
.warp_token_usage_by_category
.entry(category)
.or_default() += tokens;
}
}
for (model_id, usage) in usage_metadata.byok_token_usage {
let entry = token_usage.entry(model_id.clone()).or_default();
entry.byok_tokens += usage.total_tokens;
for (category, tokens) in usage.token_usage_by_category {
*entry
.byok_token_usage_by_category
.entry(category)
.or_default() += tokens;
}
}
self.conversation_usage_metadata.token_usage = token_usage
.into_iter()
.map(|(name, mut usage)| {
usage.model_id = name;
usage
})
.collect();
self.conversation_usage_metadata.tool_usage_metadata = usage_metadata
.tool_usage_metadata
.as_ref()
.map(Into::into)
.unwrap_or_default();
// A conversation can never go from summarized to un-summarized,
// so we only update the summarized flag if it's going from false to true.
if usage_metadata.summarized && !self.conversation_usage_metadata.was_summarized {
self.conversation_usage_metadata.was_summarized = usage_metadata.summarized;
}
}
Ok(())
}
pub fn mark_request_completed(
&mut self,
stream_id: &ResponseStreamId,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
let Some(new_exchanges) = self.added_exchanges_by_response.get(stream_id).cloned() else {
log::error!("No pending request info for completed request.");
return Err(UpdateConversationError::NoPendingRequest);
};
let mut has_new_actions = false;
for AddedExchange {
exchange_id,
task_id,
} in new_exchanges.into_iter()
{
let completed_exchange = self.mark_exchange_completed(&task_id, exchange_id)?;
let output = completed_exchange
.output_status
.output()
.map(Shared::get_owned);
if let Some(output_shared) = output {
let output = output_shared.get();
has_new_actions |= output.actions().next().is_some();
if let Some(new_suggestions) = output.suggestions.clone() {
if let Some(existing_suggestions) = self.existing_suggestions.as_mut() {
existing_suggestions.rules.extend(new_suggestions.rules);
existing_suggestions
.agent_mode_workflows
.extend(new_suggestions.agent_mode_workflows);
} else {
self.existing_suggestions = Some(new_suggestions);
}
}
}
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self.is_exchange_hidden(exchange_id),
});
}
self.write_updated_conversation_state(ctx);
if !has_new_actions {
// Update conversation-level status to success if the output has no actions.
self.update_status(ConversationStatus::Success, terminal_view_id, ctx);
}
Ok(())
}
pub fn mark_completed_after_successful_split(
&mut self,
stream_id: &ResponseStreamId,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
// Remove the mapping between the response stream and this conversation, as the response stream is
// now associated with a different one.
if let Some(added_exchanges) = self.added_exchanges_by_response.remove(stream_id) {
for AddedExchange {
exchange_id,
task_id,
} in added_exchanges.into_iter()
{
let completed_exchange = self.mark_exchange_completed(&task_id, exchange_id)?;
let output = completed_exchange
.output_status
.output()
.map(Shared::get_owned);
if let Some(output_shared) = output {
let output = output_shared.get();
if let Some(new_suggestions) = output.suggestions.clone() {
if let Some(existing_suggestions) = self.existing_suggestions.as_mut() {
existing_suggestions.rules.extend(new_suggestions.rules);
existing_suggestions
.agent_mode_workflows
.extend(new_suggestions.agent_mode_workflows);
} else {
self.existing_suggestions = Some(new_suggestions);
}
}
}
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self.is_exchange_hidden(exchange_id),
});
}
}
self.write_updated_conversation_state(ctx);
// Update conversation-level status to success if the output has no actions.
self.update_status(ConversationStatus::Success, terminal_view_id, ctx);
Ok(())
}
pub fn mark_request_cancelled(
&mut self,
stream_id: &ResponseStreamId,
terminal_view_id: EntityId,
reason: CancellationReason,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
let Some(added_exchanges) = self.added_exchanges_by_response.get(stream_id).cloned() else {
log::error!("No pending request info for completed request.");
return Err(UpdateConversationError::NoPendingRequest);
};
if self.transaction.is_some() {
self.commit_transaction()
}
for AddedExchange {
exchange_id,
task_id,
} in added_exchanges.into_iter()
{
let is_viewing_shared_session = self.is_viewing_shared_session;
let task = self
.task_store
.get(&task_id)
.ok_or(UpdateConversationError::TaskNotFound)?
.clone();
let exchange = self.get_exchange_to_update(exchange_id)?;
let AIAgentOutputStatus::Streaming { output } = &exchange.output_status else {
// Skip exchanges that are already finished (e.g., a root task exchange
// that completed before a subagent exchange was cancelled).
continue;
};
exchange.output_status = AIAgentOutputStatus::Finished {
finished_output: FinishedAIAgentOutput::Cancelled {
output: output.as_ref().map(Shared::get_owned),
reason,
},
};
let finish_time = Self::finish_time_from_exchange_messages(&task, exchange)
.unwrap_or_else(Local::now);
// For shared-session viewers, derive start time and time to first token from server messages
// (in the same way we do when restoring/forking conversations).
if is_viewing_shared_session {
if let Some(start_time) = Self::start_time_from_exchange_messages(exchange) {
exchange.start_time = start_time;
}
exchange.time_to_first_token_ms = compute_time_to_first_token_ms_from_messages(
exchange.start_time,
task.messages().filter(|m| {
let id = MessageId::new(m.id.clone());
exchange.added_message_ids.contains(&id)
}),
);
}
exchange.finish_time = Some(finish_time);
let is_hidden = self.is_exchange_hidden(exchange_id);
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden,
});
}
self.write_updated_conversation_state(ctx);
// Don't mark the conversation as Cancelled if we're just cancelling to send a follow-up
// on the same conversation. The conversation will be immediately set back to InProgress.
if !reason.is_follow_up_for_same_conversation() {
self.update_status(ConversationStatus::Cancelled, terminal_view_id, ctx);
}
Ok(())
}
pub fn mark_request_cancelled_due_to_revert(
&mut self,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
if self.transaction.is_some() {
self.commit_transaction();
}
self.update_status(ConversationStatus::Success, terminal_view_id, ctx);
Ok(())
}
pub fn mark_request_completed_with_error(
&mut self,
stream_id: &ResponseStreamId,
error: RenderableAIError,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
let Some(added_exchanges) = self.added_exchanges_by_response.get(stream_id).cloned() else {
log::error!("No pending request info for completed request.");
return Err(UpdateConversationError::NoPendingRequest);
};
if self.transaction.is_some() {
self.commit_transaction()
}
let AddedExchange {
exchange_id: initial_exchange_id,
..
} = added_exchanges.first();
let identifiers = AIIdentifiers {
server_output_id: None,
server_conversation_id: self.server_conversation_token.clone().map(Into::into),
client_conversation_id: Some(self.id),
client_exchange_id: Some(*initial_exchange_id),
model_id: None,
};
let will_attempt_to_resume = matches!(
&error,
RenderableAIError::Other {
will_attempt_resume: true,
..
}
);
send_telemetry_from_ctx!(
crate::TelemetryEvent::AgentModeError {
identifiers,
error: error.to_string(),
is_user_visible: true,
will_attempt_to_resume,
},
ctx
);
for AddedExchange {
exchange_id,
task_id,
} in added_exchanges.into_iter()
{
let is_viewing_shared_session = self.is_viewing_shared_session;
let task = self
.task_store
.get(&task_id)
.ok_or(UpdateConversationError::TaskNotFound)?
.clone();
let exchange = self.get_exchange_to_update(exchange_id)?;
let AIAgentOutputStatus::Streaming { output } = &exchange.output_status else {
return Err(UpdateConversationError::OutputAlreadyFinished);
};
exchange.output_status = AIAgentOutputStatus::Finished {
finished_output: FinishedAIAgentOutput::Error {
output: output.as_ref().map(Shared::get_owned),
error: error.clone(),
},
};
let finish_time = Self::finish_time_from_exchange_messages(&task, exchange)
.unwrap_or_else(Local::now);
// For shared-session viewers, derive start time and time to first token from server messages
// (in the same way we do when restoring/forking conversations).
if is_viewing_shared_session {
if let Some(start_time) = Self::start_time_from_exchange_messages(exchange) {
exchange.start_time = start_time;
}
exchange.time_to_first_token_ms = compute_time_to_first_token_ms_from_messages(
exchange.start_time,
task.messages().filter(|m| {
let id = MessageId::new(m.id.clone());
exchange.added_message_ids.contains(&id)
}),
);
}
exchange.finish_time = Some(finish_time);
let is_hidden = self.is_exchange_hidden(exchange_id);
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden,
});
}
self.write_updated_conversation_state(ctx);
self.update_status_with_error_message(
ConversationStatus::Error,
Some(error.to_string()),
terminal_view_id,
ctx,
);
Ok(())
}
fn mark_exchange_completed(
&mut self,
task_id: &TaskId,
exchange_id: AIAgentExchangeId,
) -> Result<&AIAgentExchange, UpdateConversationError> {
let task = self
.task_store
.get(task_id)
.ok_or(UpdateConversationError::TaskNotFound)?
.clone();
let is_viewing_shared_session = self.is_viewing_shared_session;
let exchange = self.get_exchange_to_update(exchange_id)?;
let AIAgentOutputStatus::Streaming {
output: Some(output),
} = &exchange.output_status
else {
return Err(UpdateConversationError::OutputAlreadyFinished);
};
let output = output.get_owned();
exchange.output_status = AIAgentOutputStatus::Finished {
finished_output: FinishedAIAgentOutput::Success { output },
};
// Record finish time for this exchange based on the latest message timestamp associated
// with this exchange. Fallback to `Local::now()` if no timestamps are present so that
// duration calculations always have a sensible value.
let finish_time =
Self::finish_time_from_exchange_messages(&task, exchange).unwrap_or_else(Local::now);
// For shared-session viewers, derive start time and time to first token from server messages
// (in the same way we do when restoring/forking conversations).
if is_viewing_shared_session {
if let Some(start_time) = Self::start_time_from_exchange_messages(exchange) {
exchange.start_time = start_time;
}
exchange.time_to_first_token_ms = compute_time_to_first_token_ms_from_messages(
exchange.start_time,
task.messages().filter(|m| {
let id = MessageId::new(m.id.clone());
exchange.added_message_ids.contains(&id)
}),
);
}
exchange.finish_time = Some(finish_time);
let exchange = self
.exchange_with_id(exchange_id)
.ok_or(UpdateConversationError::ExchangeNotFound)?;
#[cfg(feature = "agent_mode_evals")]
{
// When running evals, log exchanges as they finish so there's a record if the container is killed due to timeout
// and there's no chance to gracefully export the whole conversation at the end.
let exchange_number = self.all_exchanges().len();
let token_usage = self.total_token_usage();
let token_usage_json: Vec<serde_json::Value> = token_usage
.iter()
.map(|usage| {
serde_json::json!({
"model_id": usage.model_id,
"total_input": usage.total_input,
"output": usage.output,
"input_cache_read": usage.input_cache_read,
"input_cache_write": usage.input_cache_write,
"cost_in_cents": usage.cost_in_cents
})
})
.collect();
println!(
"===== Exchange {exchange_number} - token_usage={}",
serde_json::to_string(&token_usage_json).unwrap_or_default()
);
for input in &exchange.input {
println!("\nInput:\n\n{input}\n");
}
println!("Output:\n{}\n", &exchange.output_status);
}
Ok(exchange)
}
pub fn apply_client_action(
&mut self,
response_stream_id: &ResponseStreamId,
terminal_view_id: EntityId,
action: warp_multi_agent_api::client_action::Action,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
use warp_multi_agent_api::client_action::*;
match action {
Action::BeginTransaction(_) => {
self.begin_transaction();
}
Action::CommitTransaction(_) => {
self.commit_transaction();
}
Action::RollbackTransaction(_) => {
log::debug!("Rollback transaction.");
self.rollback_transaction(response_stream_id);
}
Action::CreateTask(CreateTask { task: Some(task) }) => {
let task_id = TaskId::new(task.id.clone());
// Save an empty task to the transaction
self.checkpoint_task(&task_id);
if let Some(parent_id) = task.parent_id() {
// If we're expecting a server-created CLI subagent subtask, instead of creating
// a net-new subtask, we convert the optimistically-created CLI subtask into a
// server-backed one.
let optimistic_cli_subagent_subtask = self
.optimistic_cli_subagent_subtask_id
.as_ref()
.and_then(|id| self.task_store.remove(id));
let Some(parent_task) = self.task_store.get(&TaskId::new(parent_id.to_owned()))
else {
log::error!(
"Attempted to create task with parent id {parent_id} but no parent task found"
);
return Err(UpdateConversationError::TaskNotFound);
};
if let Some(optimistic_subtask) = optimistic_cli_subagent_subtask {
log::debug!(
"Upgrading optimistically created subtask with ID {:?} to server task with ID {:?}",
optimistic_subtask.id(),
task.id
);
self.optimistic_cli_subagent_subtask_id = None;
let optimistic_id = optimistic_subtask.id().clone();
let server_subtask = optimistic_subtask.into_server_created_task(
task,
parent_task.source(),
self.todo_lists.last(),
self.code_review.as_ref(),
)?;
ctx.emit(BlocklistAIHistoryEvent::UpgradedTask {
optimistic_id: optimistic_id.clone(),
server_id: server_subtask.id().clone(),
terminal_view_id,
});
for new_exchange in self
.added_exchanges_by_response
.get_mut(response_stream_id)
.into_iter()
.flat_map(|new_exchanges| new_exchanges.iter_mut())
{
if new_exchange.task_id == optimistic_id {
new_exchange.task_id = server_subtask.id().clone();
}
}
self.task_store.insert(server_subtask);
} else if let Some(existing_exchange) = self
.added_exchanges_by_response
.get(response_stream_id)
.map(|new_exchanges| new_exchanges.first())
.and_then(|new_exchange| {
self.task_store
.get(&new_exchange.task_id)
.and_then(|t| t.exchange(new_exchange.exchange_id))
})
{
let subtask = Task::new_subtask(
task,
parent_task
.source()
.ok_or(UpdateConversationError::TaskNotInitialized)?,
existing_exchange,
self.todo_lists.last(),
self.code_review.as_ref(),
// In shared-session viewers, we have to reconstruct what the original user input
// was using subsequent conversation messages (as the original input was not
// sent on this client). Once we reconstruct these inputs, we will insert them
// to mimic the normal conversation flow. (If this is not a shared session, the
// exchange inputs will already be populated).
self.is_viewing_shared_session,
);
// Subtasks can come pre-populated with messages (for example: an advice subagent
// or computer use subagent task created with an initial tool call already present
// in its task messages).
//
// In those cases, we need to ensure an AI block is created for the subtask's
// initial exchange; otherwise the first tool call/result can be "lost" from the
// block list because we only create AI blocks on AppendedExchange events.
//
// TODO(QUALITY-276): We should check if we can generally add exchanges from any
// subtask, or if that breaks things (e.g. in the CLI subagent).
let initial_exchange_ids: Vec<_> = if subtask.is_advice_subagent()
|| subtask.is_computer_use_subagent()
|| subtask.is_conversation_search_subagent()
{
subtask.exchanges().map(|e| e.id).collect()
} else {
Vec::new()
};
if self.is_viewing_shared_session {
// shared session viewers should move the current stream's new exchange from the root to the
// newly created subtask so there's exactly one "new" exchange and it
// belongs to the subtask (mirrors sharer semantics after optimistic upgrade).
let last_subtask_exchange_id = subtask
.exchanges()
.last()
.map(|e| e.id)
.ok_or(UpdateConversationError::ExchangeNotFound)?;
let new_exchanges = self
.added_exchanges_by_response
.get_mut(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?;
let first = new_exchanges.first_mut();
// we're updating first's id is because it should correspond with the newly generated subtask's new exchange
first.task_id = task_id.clone();
first.exchange_id = last_subtask_exchange_id;
} else {
let new_exchanges = self
.added_exchanges_by_response
.get_mut(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?;
new_exchanges.extend(subtask.exchanges().map(|exchange| {
AddedExchange {
task_id: task_id.clone(),
exchange_id: exchange.id,
}
}));
}
self.task_store.insert(subtask);
ctx.emit(BlocklistAIHistoryEvent::CreatedSubtask {
conversation_id: self.id,
terminal_view_id,
task_id: task_id.clone(),
});
for exchange_id in initial_exchange_ids {
let is_hidden = self.is_exchange_hidden(exchange_id);
ctx.emit(BlocklistAIHistoryEvent::AppendedExchange {
exchange_id,
task_id: task_id.clone(),
terminal_view_id,
conversation_id: self.id,
is_hidden,
response_stream_id: Some(response_stream_id.clone()),
});
}
}
} else {
let root_task_id = self.task_store.root_task_id().clone();
if let Some(mut root_task) = self.task_store.remove(&root_task_id) {
let old_id = root_task.id().clone();
root_task = root_task.into_server_created_task(
task,
None,
self.todo_lists.last(),
self.code_review.as_ref(),
)?;
ctx.emit(BlocklistAIHistoryEvent::UpgradedTask {
optimistic_id: old_id,
server_id: root_task.id().clone(),
terminal_view_id,
});
for AddedExchange {
ref mut task_id, ..
} in self
.added_exchanges_by_response
.get_mut(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?
.iter_mut()
{
if *task_id == root_task_id {
*task_id = root_task.id().clone();
}
}
self.task_store.set_root_task(root_task);
}
}
}
Action::UpdateTaskDescription(UpdateTaskDescription {
task_id,
description,
}) => {
let task_id = TaskId::new(task_id);
self.checkpoint_task(&task_id);
self.task_store
.modify_task(&task_id, |task| task.update_description(description))
.ok_or(UpdateConversationError::TaskNotFound)?;
}
Action::AddMessagesToTask(AddMessagesToTask { task_id, messages }) => {
for message in messages.iter() {
match message.message.as_ref() {
Some(api::message::Message::UpdateTodos(update)) => {
if let Some(todos_op) = update.operation.as_ref() {
update_todo_list_from_todo_op(
&mut self.todo_lists,
todos_op.clone(),
);
ctx.emit(BlocklistAIHistoryEvent::UpdatedTodoList {
terminal_view_id,
});
}
}
Some(api::message::Message::UpdateReviewComments(comments)) => {
if let Some(comments_op) = comments.operation.as_ref() {
if let Some(active_code_review) = self.code_review.as_mut() {
let resolved_count = update_comment_from_comment_operation(
active_code_review,
comments_op.clone(),
);
if resolved_count > 0 {
send_telemetry_from_ctx!(
CodeReviewTelemetryEvent::CommentResolved {
resolved_count
},
ctx
);
}
} else {
log::error!(
"Received an UpdateReviewComments message but there's no active code review state"
);
}
}
}
Some(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),
) => {
self.add_artifact(
Artifact::from(pr.clone()),
terminal_view_id,
ctx,
);
}
Some(
api::message::artifact_event::artifact_created::Artifact::Screenshot(screenshot),
) => {
self.add_artifact(
Artifact::from(screenshot.clone()),
terminal_view_id,
ctx,
);
}
Some(
api::message::artifact_event::artifact_created::Artifact::File(file),
) => {
self.add_artifact(
Artifact::from(file.clone()),
terminal_view_id,
ctx,
);
}
None => {}
}
}
Some(api::message::artifact_event::Event::ForkArtifacts(
fork_artifacts,
)) => {
for proto_artifact in &fork_artifacts.artifacts {
let Some(artifact) =
artifact_from_fork_proto(proto_artifact)
else {
continue;
};
self.add_artifact(artifact, terminal_view_id, ctx);
}
}
None => {}
}
}
Some(api::message::Message::ToolCallResult(tcr)) => {
// Clean up temp directories from conversation search subagents.
if let Some(api::message::tool_call_result::Result::Subagent(_)) =
&tcr.result
{
cleanup_conversation_search_temp_dir(
&tcr.tool_call_id,
&task_id,
&self.task_store,
);
}
}
Some(api::message::Message::ModelUsed(model_used)) => {
let exchange_id = self
.added_exchanges_by_response
.get(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?
.last()
.exchange_id;
let exchange = self.get_exchange_to_update(exchange_id)?;
if let Some(output) = exchange.output_status.output() {
let mut output = output.get_mut();
output.model_info = Some(OutputModelInfo {
model_id: model_used.model_id.clone().into(),
display_name: model_used.model_display_name.clone(),
is_fallback: model_used.is_fallback,
});
}
}
_ => {}
}
}
let task_id = TaskId::new(task_id);
self.checkpoint_task(&task_id);
let current_todo_list = self.todo_lists.last().cloned();
// Remove the task to relinquish mutable borrow on self, we add it back later.
let mut task = self
.task_store
.remove(&task_id)
.ok_or(UpdateConversationError::TaskNotFound)?;
let added_exchanges = self
.added_exchanges_by_response
.get(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?;
let exchange_id = if let Some(info) =
added_exchanges.iter().find(|info| info.task_id == task_id)
{
info.exchange_id
} else {
let existing_exchange = self
.get_task(&added_exchanges.last().task_id)
.ok_or(UpdateConversationError::TaskNotFound)?
.exchange(added_exchanges.last().exchange_id)
.ok_or(UpdateConversationError::ExchangeNotFound)?;
let new_exchange_id = task.append_new_exchange(existing_exchange);
if self.optimistic_cli_subagent_subtask_id.is_some() && task.is_root_task() {
// If we are lazily creating a new exchange at this point, this means we are updating
// a new task for the first time in this response stream.
//
// This is a bit of a hack, but if the optimistic CLI Subagent task is some and this is
// the root task, then this exchange corresponds to "setup" messages in the root task
// for bootstrapping the CLI subagent. In these cases, we don't care about
// surfacing the new root task messages in the UI (e.g. the blocklist) - there would basically
// be an empty AI Block corresponding to the CLI subagent tool call message added to the root
// task, with not user rendered output.
//
// The real fix here is to lazily create exchanges only when there are real messages to be
// rendered, or at the very least, lazily create AI blocks for an exchange only once the exchange
// actually has renderable content.
self.hidden_exchanges.insert(new_exchange_id);
}
new_exchange_id
};
let current_comment_state = self.code_review.as_ref().cloned();
task.add_messages(
messages,
exchange_id,
current_todo_list.as_ref(),
current_comment_state.as_ref(),
// In shared-session viewers, we have to reconstruct what the original user input
// was using subsequent conversation messages (as the original input was not
// sent on this client). Once we reconstruct these inputs, we will insert them
// to mimic the normal conversation flow. (If this is not a shared session, the
// exchange inputs will already be populated).
self.is_viewing_shared_session,
)?;
self.task_store.insert(task);
if !added_exchanges
.iter()
.any(|new_exchange_info| new_exchange_info.exchange_id == exchange_id)
{
self.added_exchanges_by_response
.get_mut(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?
.push(AddedExchange {
task_id: task_id.clone(),
exchange_id,
});
let is_hidden = self.hidden_exchanges.contains(&exchange_id);
ctx.emit(BlocklistAIHistoryEvent::AppendedExchange {
response_stream_id: Some(response_stream_id.clone()),
exchange_id,
task_id: task_id.clone(),
terminal_view_id,
conversation_id: self.id,
is_hidden,
});
}
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self.is_exchange_hidden(exchange_id),
});
}
Action::UpdateTaskServerData(UpdateTaskServerData {
task_id,
server_data,
}) => {
let task_id = TaskId::new(task_id);
self.task_store
.modify_task(&task_id, |task| task.update_task_server_data(server_data))
.ok_or(UpdateConversationError::TaskNotFound)?;
}
Action::UpdateTaskMessage(UpdateTaskMessage {
task_id,
message: Some(message),
mask: Some(mask),
}) => {
let task_id = TaskId::new(task_id);
let exchange_id = self
.added_exchanges_by_response
.get(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?
.iter()
.find_map(|new_exchange| {
(new_exchange.task_id == task_id).then_some(new_exchange.exchange_id)
})
.ok_or(UpdateConversationError::ExchangeNotFound)?;
let current_todo_list = self.todo_lists.last().cloned();
let current_comment_state = self.code_review.as_ref().cloned();
let is_viewing_shared_session = self.is_viewing_shared_session;
// In shared-session viewers, we have to reconstruct what the original user input
// was using subsequent conversation messages (as the original input was not
// sent on this client). Once we reconstruct these inputs, we will insert them
// to mimic the normal conversation flow. (If this is not a shared session, the
// exchange inputs will already be populated).
let todos_op = self
.task_store
.modify_task(&task_id, |task| {
task.upsert_message(
message,
exchange_id,
current_todo_list.as_ref(),
current_comment_state.as_ref(),
mask,
is_viewing_shared_session,
)
.map(|msg| msg.todos_op().cloned())
})
.ok_or(UpdateConversationError::TaskNotFound)??;
// Update todo list if needed
if let Some(todos_op) = todos_op {
update_todo_list_from_todo_op(&mut self.todo_lists, todos_op);
ctx.emit(BlocklistAIHistoryEvent::UpdatedTodoList { terminal_view_id });
}
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self.is_exchange_hidden(exchange_id),
});
}
Action::AppendToMessageContent(AppendToMessageContent {
task_id,
message: Some(message),
mask: Some(mask),
}) => {
let task_id = TaskId::new(task_id);
log::info!(
"[bedrock-debug] AppendToMessageContent: task_id={:?}, message_id={:?}",
task_id,
message.id
);
let exchange_id = match self.added_exchanges_by_response.get(response_stream_id) {
Some(exchanges) => {
log::info!(
"[bedrock-debug] AppendToMessageContent: found {} exchanges for stream",
exchanges.len()
);
for ex in exchanges.iter() {
log::info!(
"[bedrock-debug] exchange: task_id={:?}, exchange_id={:?}",
ex.task_id,
ex.exchange_id
);
}
match exchanges.iter().find_map(|new_exchange| {
(new_exchange.task_id == task_id).then_some(new_exchange.exchange_id)
}) {
Some(id) => id,
None => {
log::error!(
"[bedrock-debug] AppendToMessageContent: ExchangeNotFound - no exchange with matching task_id"
);
return Err(UpdateConversationError::ExchangeNotFound);
}
}
}
None => {
log::error!(
"[bedrock-debug] AppendToMessageContent: NoPendingRequest - no exchanges for this stream_id"
);
return Err(UpdateConversationError::NoPendingRequest);
}
};
log::info!(
"[bedrock-debug] AppendToMessageContent: found exchange_id={:?}",
exchange_id
);
let current_todo_list = self.todo_lists.last().cloned();
let current_comment_state = self.code_review.as_ref().cloned();
// Update the message and get the updated todos op, if any.
let todos_op = match self.task_store.modify_task(&task_id, |task| {
task.append_to_message_content(
message,
exchange_id,
current_todo_list.as_ref(),
current_comment_state.as_ref(),
mask,
)
.map(|msg| msg.todos_op().cloned())
}) {
Some(result) => match result {
Ok(todos_op) => {
log::info!("[bedrock-debug] AppendToMessageContent: append succeeded");
todos_op
}
Err(e) => {
log::error!(
"[bedrock-debug] AppendToMessageContent: append_to_message_content failed: {e:?}"
);
return Err(e.into());
}
},
None => {
log::error!(
"[bedrock-debug] AppendToMessageContent: TaskNotFound in task_store"
);
return Err(UpdateConversationError::TaskNotFound);
}
};
// Update todo list if needed
if let Some(todos_op) = todos_op {
update_todo_list_from_todo_op(&mut self.todo_lists, todos_op);
ctx.emit(BlocklistAIHistoryEvent::UpdatedTodoList { terminal_view_id });
}
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self.is_exchange_hidden(exchange_id),
});
}
Action::ShowSuggestions(suggestions) => {
let exchange_id = self
.added_exchanges_by_response
.get(response_stream_id)
.ok_or(UpdateConversationError::NoPendingRequest)?
.last()
.exchange_id;
let exchange_to_update = self.get_exchange_to_update(exchange_id)?;
exchange_to_update.update_suggestions(suggestions);
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_view_id,
conversation_id: self.id,
is_hidden: self.is_exchange_hidden(exchange_id),
});
}
Action::MoveMessagesToNewTask(MoveMessagesToNewTask {
source_task_id,
new_task: Some(mut new_task),
first_message_id,
last_message_id,
expected_message_count,
replacement_messages,
}) => {
let source_task_id = TaskId::new(source_task_id);
self.checkpoint_task(&source_task_id);
// Extract messages from the source task (this also inserts replacement messages).
let mut extracted_messages = self
.task_store
.modify_task(&source_task_id, |task| {
task.splice_messages(
&first_message_id,
&last_message_id,
expected_message_count,
replacement_messages,
)
})
.ok_or(UpdateConversationError::TaskNotFound)??;
// Update task_id on each extracted message to reference the new task.
for msg in &mut extracted_messages {
msg.task_id = new_task.id.clone();
}
// Append extracted messages to the new task.
new_task.messages.extend(extracted_messages);
// Get the source task's api::Task to look up subagent_params.
// At this point, the source task contains the replacement messages (including the
// subagent call referencing the new task), so new_summary_subtask can find them.
let source_api_task = self
.task_store
.get(&source_task_id)
.and_then(|t| t.source())
.cloned()
.ok_or(UpdateConversationError::TaskNotInitialized)?;
// Create the subtask and add it to the task store.
let subtask = Task::new_moved_messages_subtask(new_task, &source_api_task);
self.task_store.insert(subtask);
// Note: We do NOT emit any BlocklistAIHistoryEvent here because we
// intentionally keep the UI unchanged during a live session. The
// exchange's client representation (added_message_ids) remains
// unmodified, pointing to message IDs that now exist in a subtask.
}
Action::StartNewConversation(_) => {
// New conversations are handled at the BlocklistAIHistoryModel layer
}
_ => {
log::warn!("Received unsupported client action: {action:?}");
}
}
Ok(())
}
pub fn get_exchange_to_update(
&mut self,
exchange_id: AIAgentExchangeId,
) -> Result<&mut AIAgentExchange, UpdateConversationError> {
self.task_store
.exchange_mut(exchange_id)
.ok_or(UpdateConversationError::ExchangeNotFound)
}
pub fn get_root_task(&self) -> Option<&Task> {
self.task_store.root_task()
}
pub fn get_root_task_id(&self) -> &TaskId {
self.task_store.root_task_id()
}
pub fn get_task(&self, task_id: &TaskId) -> Option<&Task> {
self.task_store.get(task_id)
}
/// Optimistically creates a subtask for the CLISubagent task when a user query is sent while
/// the a command is running but no subagent has been spawned yet.
///
/// This is done in two scenarios:
///
/// 1) The user enters agent mode while a user-executed command is running, and sends a query.
/// 2) The agent has executed a long-running requested command, but before the response stream
/// finishes (in which the CLI subagent would be spawned), the user pre-empts with a query.
///
/// In both cases, we optimistically create a subtask for the query, and the next time we receive
/// a `CreateTask` client action for a subtask, we upgrade this optimistic subtask to a
/// server-backed task.
pub fn create_optimistic_cli_subagent_task(
&mut self,
block_id: &BlockId,
terminal_view_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> TaskId {
if self.optimistic_cli_subagent_subtask_id.take().is_some() {
log::error!(
"Tried to optimistically create new subtask for CLI agent when one exists already."
);
}
let new_task = Task::new_optimistic_cli_agent_subtask(block_id.clone());
let new_task_id = new_task.id().clone();
self.optimistic_cli_subagent_subtask_id = Some(new_task_id.clone());
self.task_store.insert(new_task);
ctx.emit(BlocklistAIHistoryEvent::CreatedSubtask {
conversation_id: self.id,
terminal_view_id,
task_id: new_task_id.clone(),
});
new_task_id
}
pub fn is_subagent_task_finished(
&self,
subagent_task_id: &TaskId,
) -> Result<bool, SubagentTaskNotFound> {
let subagent_task = self
.task_store
.get(subagent_task_id)
.ok_or(SubagentTaskNotFound)?;
let (Some(subagent_params), Some(parent_id)) =
(subagent_task.subagent_params(), subagent_task.parent_id())
else {
return Err(SubagentTaskNotFound);
};
let parent_task = self
.task_store
.get(&parent_id)
.ok_or(SubagentTaskNotFound)?;
Ok(parent_task
.source()
.into_iter()
.flat_map(|source| source.messages.iter())
.any(|message| {
message
.tool_call_result()
.is_some_and(|result| result.tool_call_id == subagent_params.tool_call_id)
}))
}
/// Returns true if any subagent task is currently active (not yet finished).
///
/// This covers both optimistic CLI subagent tasks (created before server
/// confirmation) and server-backed subagent tasks. Used to prevent
/// piggybacking orchestration events onto followup requests while a
/// subagent is active, since subagents cannot interpret those events and
/// inserting them breaks tool_use/tool_result ordering requirements.
pub fn has_active_subagent(&self) -> bool {
if self.optimistic_cli_subagent_subtask_id.is_some() {
return true;
}
self.all_tasks().any(|task| {
!task.is_root_task()
&& self
.is_subagent_task_finished(task.id())
.is_ok_and(|finished| !finished)
})
}
pub fn todo_lists(&self) -> &Vec<AIAgentTodoList> {
&self.todo_lists
}
pub fn active_todo_list(&self) -> Option<&AIAgentTodoList> {
self.todo_lists.last()
}
pub fn active_todo(&self) -> Option<&AIAgentTodo> {
self.active_todo_list()
.and_then(|todo_list| todo_list.in_progress_item())
}
pub fn todo_status(&self, todo_id: &AIAgentTodoId) -> Option<TodoStatus> {
for (i, list) in self.todo_lists.iter().rev().enumerate() {
let is_active_list = i == 0;
if let Some(pos) = list
.pending_items()
.iter()
.position(|item| &item.id == todo_id)
{
if is_active_list {
if pos == 0 {
return if self.status.is_in_progress() {
Some(TodoStatus::InProgress)
} else {
Some(TodoStatus::Stopped)
};
} else {
return Some(TodoStatus::Pending);
}
} else {
return Some(TodoStatus::Cancelled);
}
} else if list
.completed_items()
.iter()
.any(|item| &item.id == todo_id)
{
return Some(TodoStatus::Completed);
}
}
None
}
pub fn begin_transaction(&mut self) {
if self.transaction.is_some() {
log::error!("Transaction already in progress.");
return;
}
self.transaction = Some(Transaction::new());
}
fn commit_transaction(&mut self) {
// Clear the transaction if it exists.
if self.transaction.take().is_none() {
log::error!("No transaction in progress.");
}
}
pub(crate) fn write_updated_conversation_state(
&mut self,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
// We should not persist non-local conversations (e.g. shared sessions).
if self.is_viewing_shared_session {
return;
}
// Check if session restoration is enabled before writing any state.
if !*GeneralSettings::as_ref(ctx).restore_session
|| !AppExecutionMode::as_ref(ctx).can_save_session()
{
return;
}
let Some(sqlite_sender) = GlobalResourceHandlesProvider::as_ref(ctx)
.get()
.model_event_sender
.clone()
else {
return;
};
let reverted_action_ids = if self.reverted_action_ids.is_empty() {
None
} else {
Some(
self.reverted_action_ids
.clone()
.into_iter()
.map_into()
.collect(),
)
};
let artifacts_json = if self.artifacts.is_empty() {
None
} else {
match serde_json::to_string(&self.artifacts) {
Ok(json) => Some(json),
Err(e) => {
log::error!(
"Failed to serialize artifacts when persisting conversation data: {e}"
);
None
}
}
};
let event = ModelEvent::UpdateMultiAgentConversation {
conversation_id: self.id.to_string(),
updated_tasks: self
.all_tasks()
.filter_map(|task| task.source().cloned())
.collect(),
conversation_data: AgentConversationData {
server_conversation_token: self
.server_conversation_token
.clone()
.map(|token| token.into()),
conversation_usage_metadata: Some(self.conversation_usage_metadata.clone()),
reverted_action_ids,
forked_from_server_conversation_token: self
.forked_from_server_conversation_token
.clone()
.map(|token| token.into()),
artifacts_json,
parent_agent_id: self.parent_agent_id.clone(),
agent_name: self.agent_name.clone(),
parent_conversation_id: self.parent_conversation_id.map(|id| id.to_string()),
run_id: self.task_id.map(|id| id.to_string()),
autoexecute_override: Some(self.autoexecute_override.into()),
last_event_sequence: self.last_event_sequence,
},
};
ctx.spawn(
async move {
if let Err(e) = sqlite_sender.send(event) {
log::warn!("Failed to send updated AI tasks to sqlite writer thread: {e:?}");
}
},
|_, _, _| {},
);
}
pub fn rollback_transaction(&mut self, response_stream_id: &ResponseStreamId) {
let Some(transaction) = self.transaction.take() else {
log::error!("No transaction in progress.");
return;
};
let mut deleted_tasks = Vec::new();
let mut updated_tasks = Vec::new();
// For each saved task in the transaction:
for (_, saved_task) in transaction.saved_tasks() {
match saved_task {
SavedTask::New(id) => {
// The task was added during the transaction, so we need to delete it
deleted_tasks.push(id);
}
SavedTask::Existing(saved_task) => {
// The task was updated during the transaction, so we need to restore it
updated_tasks.push(*saved_task);
}
}
}
updated_tasks.into_iter().for_each(|task| {
log::debug!("Rolling back existing task: {:?}", task.id());
self.task_store.insert(task);
});
deleted_tasks.into_iter().for_each(|task_id| {
log::debug!("Rolling back new task: {task_id:?}");
self.task_store.remove(&task_id);
});
if let Some(added_exchanges) = self
.added_exchanges_by_response
.get(response_stream_id)
.cloned()
{
let mut updated_added_exchanges: Option<Vec1<AddedExchange>> = None;
for added_exchange in added_exchanges.into_iter() {
let does_exchange_exist = self
.task_store
.get(&added_exchange.task_id)
.and_then(|task| {
task.exchanges()
.find(|exchange| exchange.id == added_exchange.exchange_id)
})
.is_some();
if does_exchange_exist {
if let Some(updated_added_exchanges) = updated_added_exchanges.as_mut() {
updated_added_exchanges.push(added_exchange);
} else {
updated_added_exchanges = Some(Vec1::new(added_exchange));
}
}
}
if let Some(updated_added_exchanges) = updated_added_exchanges {
self.added_exchanges_by_response
.insert(response_stream_id.clone(), updated_added_exchanges);
}
}
}
pub fn checkpoint_task(&mut self, task_id: &TaskId) {
if let Some(transaction) = &mut self.transaction {
if let Some(task) = self.task_store.get(task_id) {
transaction.checkpoint_task(task);
} else {
transaction.checkpoint_new_task(task_id);
}
}
}
pub fn toggle_autoexecute_override(&mut self) {
self.autoexecute_override =
if self.autoexecute_override == AIConversationAutoexecuteMode::RespectUserSettings {
AIConversationAutoexecuteMode::RunToCompletion
} else {
AIConversationAutoexecuteMode::RespectUserSettings
};
}
pub fn autoexecute_override(&self) -> AIConversationAutoexecuteMode {
self.autoexecute_override
}
pub fn autoexecute_any_action(&self) -> bool {
self.autoexecute_override.is_autoexecute_any_action()
}
pub fn initial_working_directory(&self) -> Option<String> {
self.task_store
.root_task()
.and_then(Task::initial_working_directory)
}
/// Returns the current working directory from the most recent exchange that has one.
/// Scans exchanges in reverse order and returns the first populated working directory.
pub fn current_working_directory(&self) -> Option<String> {
self.task_store
.all_exchanges_rev()
.find_map(|exchange| exchange.working_directory.clone())
}
#[allow(dead_code)]
pub fn total_request_cost(&self) -> RequestCost {
self.total_request_cost
}
#[allow(dead_code)]
pub fn total_token_usage(&self) -> Vec<TokenUsage> {
self.total_token_usage_by_model.values().cloned().collect()
}
pub fn total_tokens(&self) -> u32 {
self.total_token_usage_by_model
.values()
.map(Self::total_tokens_for_usage)
.sum()
}
pub fn total_cost_cents(&self) -> f32 {
self.total_token_usage_by_model
.values()
.map(|u| u.cost_in_cents)
.sum()
}
fn total_tokens_for_usage(usage: &TokenUsage) -> u32 {
usage.total_input + usage.output + usage.input_cache_read + usage.input_cache_write
}
/// Normalize all newlines to CRLF so restored blocks render lines starting at column 0,
/// which is consistent with how we serialize real terminal blocks.
fn to_stylized_bytes(s: &str) -> Vec<u8> {
let s = s.replace("\r\n", "\n");
s.replace('\n', "\r\n").into_bytes()
}
/// Finds the RunShellCommand result for a given tool_call_id.
/// Returns both the result and the message ID of the result message.
pub(crate) fn find_run_shell_command_result(
&self,
tool_call_id: &str,
) -> Option<(api::RunShellCommandResult, String)> {
let root_task = self.get_root_task()?;
let api_task = root_task.source()?;
// Find the last tool call result with this tool call ID
api_task.messages.iter().rev().find_map(|msg| {
let result = msg.tool_call_result()?;
if result.tool_call_id == tool_call_id {
if let Some(api::message::tool_call_result::Result::RunShellCommand(cmd_result)) =
&result.result
{
return Some((cmd_result.clone(), msg.id.clone()));
}
}
None
})
}
/// Extracts all shell command blocks, in order, from the conversation's API task
/// messages.
///
/// This includes:
/// - RunShellCommand tool calls that completed
/// - Attachments from UserQuery/SystemQuery messages
/// - Context blocks from UserQuery/SystemQuery/ToolCallResult messages
///
/// Returns CommandBlockInfo with command, output, exit_code, and optional ai_metadata.
fn extract_command_blocks(&self) -> Vec<CommandBlockInfo> {
let mut command_blocks = Vec::new();
// Get the root task's API messages.
let Some(root_task) = self.get_root_task() else {
return command_blocks;
};
let Some(api_task) = root_task.source() else {
return command_blocks;
};
self.extract_command_blocks_from_messages(&api_task.messages, &mut command_blocks);
command_blocks
}
/// Extracts command blocks from a list of messages.
///
/// This recurses when it encounters a summarization subagent call, producing the list
/// of command blocks as it would have been had no summarization ever occurred.
fn extract_command_blocks_from_messages(
&self,
messages: &[api::Message],
command_blocks: &mut Vec<CommandBlockInfo>,
) {
// Build a map from tool_call_id to (RunShellCommandResult, result_message_id)
// for efficient lookup within this message set.
let tool_call_results: HashMap<&str, (&api::RunShellCommandResult, &str)> = messages
.iter()
.filter_map(|msg| {
let result = msg.tool_call_result()?;
if let Some(api::message::tool_call_result::Result::RunShellCommand(cmd_result)) =
&result.result
{
Some((result.tool_call_id.as_str(), (cmd_result, msg.id.as_str())))
} else {
None
}
})
.collect();
for message in messages {
let message_id = message.id.clone();
if let Some(tool_call) = message.tool_call() {
// Check if this is a moved-messages subtask (summarization subagent).
// If so, extract its command blocks here to maintain chronological order.
if let Some(subagent) = tool_call.subagent() {
if subagent.is_summarization() {
let subtask_id = TaskId::new(subagent.task_id.clone());
if let Some(subtask) = self.task_store.get(&subtask_id) {
if let Some(subtask_source) = subtask.source() {
// Recursively extract from subtask (in case of nested summarization).
self.extract_command_blocks_from_messages(
&subtask_source.messages,
command_blocks,
);
}
}
// Don't process this message further - it's just a subagent call.
continue;
}
}
// Extract from RunShellCommand tool calls.
if let Some(api::message::tool_call::Tool::RunShellCommand(run_cmd)) =
&tool_call.tool
{
let tool_call_id = &tool_call.tool_call_id;
let command = &run_cmd.command;
// Find the corresponding tool call result in this message set.
if let Some((cmd_result, result_message_id)) =
tool_call_results.get(tool_call_id.as_str())
{
if let Some(api::run_shell_command_result::Result::CommandFinished(
api::ShellCommandFinished {
output: command_output,
exit_code,
..
},
)) = &cmd_result.result
{
command_blocks.push(CommandBlockInfo {
command: command.clone(),
output: command_output.clone(),
exit_code: ExitCode::from(*exit_code),
ai_metadata: Some(
serde_json::to_string(&Some(
Into::<SerializedAIMetadata>::into(
AgentInteractionMetadata::new_hidden(
tool_call_id.clone().into(),
self.id(),
),
),
))
.unwrap_or_default(),
),
message_id: (*result_message_id).to_string(),
});
}
}
}
}
// Extract from UserQuery/SystemQuery attachments.
let attachments = match message.message.as_ref() {
Some(api::message::Message::UserQuery(user_query)) => user_query
.referenced_attachments
.values()
.collect::<Vec<_>>(),
Some(api::message::Message::SystemQuery(_)) => {
// SystemQuery doesn't have attachments currently.
vec![]
}
_ => vec![],
};
for attachment in attachments {
// Attachments have ExecutedShellCommand in their value oneof.
if let Some(api::attachment::Value::ExecutedShellCommand(cmd)) = &attachment.value {
command_blocks.push(CommandBlockInfo {
command: cmd.command.clone(),
output: cmd.output.clone(),
exit_code: ExitCode::from(cmd.exit_code),
ai_metadata: None,
message_id: message_id.clone(),
});
}
}
// Extract from UserQuery/SystemQuery context blocks.
let context_blocks = match message.message.as_ref() {
Some(api::message::Message::UserQuery(user_query)) => user_query.context.as_ref(),
Some(api::message::Message::SystemQuery(system_query)) => {
system_query.context.as_ref()
}
_ => None,
};
if let Some(context) = context_blocks {
#[allow(deprecated)]
for executed_shell_command in &context.executed_shell_commands {
if !executed_shell_command.command.is_empty() {
command_blocks.push(CommandBlockInfo {
command: executed_shell_command.command.clone(),
output: executed_shell_command.output.clone(),
exit_code: ExitCode::from(executed_shell_command.exit_code),
ai_metadata: None,
message_id: message_id.clone(),
});
}
}
}
}
}
/// Converts the conversation into a vector of serialized AI and command blocks.
/// When we open a new tab to restore a conversation in, we need to precompute this serialized list of blocks
/// to pass into the TerminalModel constructor since command blocks must be created
/// before the warp input block to not break bootstrapping.
/// Only the command blocks are actually created in the terminal model, but this sequencing is used later in the TerminalView
/// to know where to insert AI blocks relative to the command blocks.
pub fn to_serialized_blocklist_items(&self) -> Vec<SerializedBlockListItem> {
let mut serialized_blocks = Vec::new();
// Extract all command blocks from the task messages
let command_blocks = self.extract_command_blocks();
// Build a map from message ID to exchange for quick lookup
let mut message_id_to_exchange: HashMap<&str, &AIAgentExchange> = HashMap::new();
for exchange in self.root_task_exchanges() {
for message_id in &exchange.added_message_ids {
// MessageId derefs to str, so use &**message_id to get &str
message_id_to_exchange.insert(&**message_id, exchange);
}
}
// Get a fallback exchange for working directory and timestamp (used if message ID not found)
let first_exchange = self.root_task_exchanges().next();
let fallback_pwd = first_exchange.and_then(|e| e.working_directory.clone());
let fallback_time = first_exchange.map(|e| e.start_time).unwrap_or_default();
// Create serialized blocks from the extracted command blocks
for command_block in command_blocks {
// Find the exchange that contains this command block's message ID
let (pwd, timestamp) = message_id_to_exchange
.get(command_block.message_id.as_str())
.map(|exchange| (exchange.working_directory.clone(), exchange.start_time))
.unwrap_or((fallback_pwd.clone(), fallback_time));
let serialized_block = SerializedBlock {
id: BlockId::new(),
stylized_command: Self::to_stylized_bytes(&command_block.command),
stylized_output: Self::to_stylized_bytes(&command_block.output),
pwd,
git_head: None,
git_branch_name: None,
virtual_env: None,
conda_env: None,
node_version: None,
exit_code: command_block.exit_code,
did_execute: true,
start_ts: Some(timestamp),
completed_ts: Some(timestamp),
ps1: None,
rprompt: None,
honor_ps1: false,
session_id: None,
shell_host: None,
is_background: false,
prompt_snapshot: None,
ai_metadata: command_block.ai_metadata,
is_local: Some(true),
agent_view_visibility: Some(
AgentViewVisibility::new_from_conversation(self.id).into(),
),
};
serialized_blocks.push(SerializedBlockListItem::Command {
block: Box::new(serialized_block),
});
}
serialized_blocks
}
pub fn mark_action_as_reverted(
&mut self,
action_id: AIAgentActionId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) {
self.reverted_action_ids.insert(action_id);
self.write_updated_conversation_state(ctx);
}
pub fn is_action_reverted(&self, action_id: &AIAgentActionId) -> bool {
self.reverted_action_ids.contains(action_id)
}
pub fn reverted_action_ids(&self) -> &HashSet<AIAgentActionId> {
&self.reverted_action_ids
}
/// Truncates the conversation from the given exchange ID, removing all exchanges
/// from that exchange onwards (inclusive). This is a lossy operation - the removed
/// exchanges are permanently deleted from this conversation.
///
/// Returns the set of exchange IDs that were removed.
pub fn truncate_from_exchange(
&mut self,
from_exchange_id: AIAgentExchangeId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<HashSet<AIAgentExchangeId>, UpdateConversationError> {
let all_exchanges: Vec<AIAgentExchangeId> =
self.root_task_exchanges().map(|e| e.id).collect();
let truncate_from_idx = all_exchanges
.iter()
.position(|id| *id == from_exchange_id)
.ok_or(UpdateConversationError::ExchangeNotFound)?;
let exchanges_to_remove: HashSet<AIAgentExchangeId> =
all_exchanges[truncate_from_idx..].iter().copied().collect();
if exchanges_to_remove.is_empty() {
return Ok(exchanges_to_remove);
}
let message_ids_to_remove: HashSet<MessageId> = exchanges_to_remove
.iter()
.filter_map(|ex_id| self.exchange_with_id(*ex_id))
.flat_map(|ex| ex.added_message_ids.iter().cloned())
.collect();
if let Some(new_todo_lists) = self.task_store.modify_root_task(|root_task| {
root_task.truncate_exchanges_from(from_exchange_id);
root_task.remove_messages(&message_ids_to_remove);
// Return updated todo state
derive_todo_lists_from_root_task(root_task)
}) {
self.todo_lists = new_todo_lists;
}
// Make sure we don't have stale code review comment state
self.code_review = None;
self.added_exchanges_by_response
.retain(|_, added_exchanges| {
if added_exchanges
.iter()
.all(|added| exchanges_to_remove.contains(&added.exchange_id))
{
return false;
}
let _ = added_exchanges
.retain(|added| !exchanges_to_remove.contains(&added.exchange_id));
true
});
self.hidden_exchanges
.retain(|ex_id| !exchanges_to_remove.contains(ex_id));
// Stale ones are harmless, but might as well remove stale reverted action IDs
let mut new_reverted_action_ids = std::mem::take(&mut self.reverted_action_ids);
new_reverted_action_ids.retain(|id| self.contains_action(id));
self.reverted_action_ids = new_reverted_action_ids;
let root_task_is_empty = self
.task_store
.root_task()
.is_none_or(|task| task.exchanges_len() == 0);
// If all exchanges were removed, reset the root task to optimistic state.
// This allows the next message to go through the normal "first message" flow,
// where the server will create a new task and we'll upgrade the optimistic task.
if root_task_is_empty {
let root_task_id = self.task_store.root_task_id().clone();
self.task_store.remove(&root_task_id);
let new_root_task = Task::new_optimistic_root();
self.task_store.set_root_task(new_root_task);
self.server_conversation_token = None;
}
self.write_updated_conversation_state(ctx);
Ok(exchanges_to_remove)
}
}
pub(super) fn update_todo_list_from_todo_op(
todo_lists: &mut Vec<AIAgentTodoList>,
op: api::message::update_todos::Operation,
) {
use api::message::update_todos::Operation;
match op {
Operation::CreateTodoList(create_todo_list) => {
todo_lists.push(
AIAgentTodoList::default().with_pending_items(
create_todo_list
.initial_todos
.into_iter()
.map(Into::into)
.collect(),
),
);
}
Operation::UpdatePendingTodos(update_pending_todos) => {
let updated_todo_list = todo_lists.pop().unwrap_or_default().with_pending_items(
update_pending_todos
.updated_pending_todos
.into_iter()
.map(Into::into)
.collect(),
);
todo_lists.push(updated_todo_list);
}
Operation::MarkTodosCompleted(completed_items) => {
if let Some(todo_list) = todo_lists.last_mut() {
todo_list.mark_todos_complete(completed_items.todo_ids);
}
}
}
}
pub(super) fn update_comment_from_comment_operation(
current_comment_state: &mut CodeReview,
op: api::message::update_review_comments::Operation,
) -> usize {
use api::message::update_review_comments::Operation;
let mut resolved_count = 0usize;
match op {
Operation::AddressReviewComments(addressed_comments) => {
for comment_id in addressed_comments.comment_ids {
if let Some(item) = current_comment_state
.pending_comments
.iter()
.position(|item| item.id.to_string() == comment_id)
.map(|i| current_comment_state.pending_comments.remove(i))
{
current_comment_state.addressed_comments.push(item);
resolved_count += 1;
}
}
}
}
resolved_count
}
/// Cleans up temporary directories created by conversation search subagents.
///
/// When a SubagentResult comes back for a conversation_search subagent, the temp
/// directory containing materialized YAML files is no longer needed and should be removed.
fn cleanup_conversation_search_temp_dir(
tool_call_id: &str,
parent_task_id: &str,
task_store: &TaskStore,
) {
let parent_task_id = TaskId::new(parent_task_id.to_string());
let Some(parent_task) = task_store.get(&parent_task_id) else {
return;
};
// Find the Subagent tool call matching this tool_call_id.
let subtask_id = parent_task.messages().find_map(|m| {
let tc = m.tool_call()?;
if tc.tool_call_id != tool_call_id {
return None;
}
let sub = tc.subagent()?;
sub.is_conversation_search().then(|| sub.task_id.clone())
});
let Some(subtask_id) = subtask_id else {
return;
};
// Find the subtask and look for a FetchConversationResult with a directory_path.
let subtask_id = TaskId::new(subtask_id);
let Some(subtask) = task_store.get(&subtask_id) else {
return;
};
let base_dir = super::conversation_yaml::base_dir();
for msg in subtask.messages() {
if let Some(api::message::Message::ToolCallResult(tcr)) = &msg.message {
if let Some(api::message::tool_call_result::Result::FetchConversation(result)) =
&tcr.result
{
if let Some(api::fetch_conversation_result::Result::Success(success)) =
&result.result
{
let dir = std::path::Path::new(&success.directory_path);
if dir.starts_with(&base_dir) {
if let Err(e) = std::fs::remove_dir_all(dir) {
log::warn!(
"Failed to clean up conversation search temp dir {}: {e}",
dir.display(),
);
} else {
log::info!(
"Cleaned up conversation search temp dir: {}",
dir.display(),
);
}
}
}
}
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum UpdateConversationError {
#[error("Exchange not found.")]
ExchangeNotFound,
#[error("Could not update task: {0:?}")]
UpdateTask(#[from] UpdateTaskError),
#[error("Could not update upgrade optimistic task for server task: {0:?}")]
UpgradeOptimisticTask(#[from] UpgradeOptimisticTaskError),
#[error("Could not extract messages: {0:?}")]
ExtractMessages(#[from] ExtractMessagesError),
#[error("Task not found.")]
TaskNotFound,
#[error("Task never initialized with CreateTask client action.")]
TaskNotInitialized,
#[error("Message not found.")]
MessageNotFound,
#[error("Attempted to update already-finished output.")]
OutputAlreadyFinished,
#[error("Attempted to update output that was never initialized.")]
OutputNeverInitialized,
#[error("Failed to convert API message to client type: {0}")]
ConversionError(#[from] MessageToAIAgentOutputMessageError),
#[error("No active task")]
NoActiveTask,
#[error("No pending request.")]
NoPendingRequest,
}
/// A globally unique ID for a conversation with an AI agent.
#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct AIConversationId(Uuid);
impl Display for AIConversationId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl AIConversationId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
}
impl Default for AIConversationId {
fn default() -> Self {
Self::new()
}
}
impl TryFrom<String> for AIConversationId {
type Error = anyhow::Error;
fn try_from(value: String) -> Result<Self, Self::Error> {
Ok(Self(Uuid::try_parse(&value)?))
}
}
/// The harness that produced an agent conversation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AIAgentHarness {
Oz,
ClaudeCode,
Gemini,
Unknown,
}
/// Describes the format of the conversation transcript data.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AIAgentSerializedBlockFormat {
JsonV1,
}
/// Describes the format capabilities of a conversation.
#[derive(Debug, Clone)]
pub struct AIAgentConversationFormat {
/// Whether there is a Warp MAA task list available for this conversation.
pub has_task_list: bool,
/// The format of the TUI serialized block, if available.
pub block_snapshot: Option<AIAgentSerializedBlockFormat>,
}
/// Metadata for an AI conversation, containing all information from the GraphQL API
/// except the full task list data.
#[derive(Debug, Clone)]
pub struct ServerAIConversationMetadata {
/// The title of the conversation.
pub title: String,
/// The working directory where the conversation was started.
pub working_directory: Option<String>,
/// The harness that produced this conversation.
pub harness: AIAgentHarness,
/// Usage metadata including token counts, credits spent, etc.
pub usage: ConversationUsageMetadata,
/// Server metadata (revision, timestamps, creator info, etc.).
pub metadata: crate::cloud_object::ServerMetadata,
/// Permissions for this conversation (space, guests, link sharing).
pub permissions: crate::cloud_object::ServerPermissions,
/// The ID of the associated ambient agent task, if any.
pub ambient_agent_task_id: Option<crate::ai::ambient_agents::AmbientAgentTaskId>,
/// The server conversation token used to identify this conversation on the server.
pub server_conversation_token: ServerConversationToken,
/// Artifacts (plans, PRs) created during this conversation.
pub artifacts: Vec<Artifact>,
}
/// Returns an iterator over `AIAgentContext`s attached to inputs in the given `exchanges`, in the
/// same order in which they appeared.
pub(super) fn context_in_exchanges<'a>(
exchanges: impl Iterator<Item = &'a AIAgentExchange> + 'a,
) -> impl Iterator<Item = &'a AIAgentContext> + 'a {
exchanges.flat_map(|exchange| {
exchange
.input
.iter()
.filter_map(AIAgentInput::context)
.flatten()
})
}
impl AIAgentExchange {
/// Returns an error if the output was already initialized.
pub(super) fn init_output(
&mut self,
server_output_id: ServerOutputId,
) -> Result<(), UpdateTaskError> {
match &mut self.output_status {
AIAgentOutputStatus::Streaming { ref mut output } => {
if let Some(shared_output) = output {
// We expect to initialize output that has already been initialized if we retry
// after receiving a StreamInit event but before receiving any ClientActions.
shared_output.get_mut().server_output_id = Some(server_output_id);
} else {
*output = Some(Shared::new(AIAgentOutput {
messages: vec![],
citations: vec![],
server_output_id: Some(server_output_id),
api_metadata_bytes: None,
suggestions: None,
telemetry_events: vec![],
model_info: None,
request_cost: None,
}));
}
Ok(())
}
AIAgentOutputStatus::Finished { .. } => Err(UpdateTaskError::OutputAlreadyFinished),
}
}
fn update_suggestions(&self, suggestions: api::Suggestions) {
if let AIAgentOutputStatus::Streaming {
output: Some(output),
} = &self.output_status
{
let mut output = output.get_mut();
output.suggestions = Some(suggestions.into());
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum AIConversationAutoexecuteMode {
#[default]
RespectUserSettings,
RunToCompletion,
}
impl AIConversationAutoexecuteMode {
pub fn is_autoexecute_any_action(&self) -> bool {
matches!(self, AIConversationAutoexecuteMode::RunToCompletion)
}
}
impl From<PersistedAutoexecuteMode> for AIConversationAutoexecuteMode {
fn from(value: PersistedAutoexecuteMode) -> Self {
match value {
PersistedAutoexecuteMode::RespectUserSettings => Self::RespectUserSettings,
PersistedAutoexecuteMode::RunToCompletion => Self::RunToCompletion,
}
}
}
impl From<AIConversationAutoexecuteMode> for PersistedAutoexecuteMode {
fn from(value: AIConversationAutoexecuteMode) -> Self {
match value {
AIConversationAutoexecuteMode::RespectUserSettings => Self::RespectUserSettings,
AIConversationAutoexecuteMode::RunToCompletion => Self::RunToCompletion,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ConversationStatus {
/// Agent is running.
InProgress,
/// The last turn of the agent finished with success.
Success,
/// The last turn of the agent completed with error.
Error,
/// The last turn of the agent was cancelled by the user.
Cancelled,
/// The last turn of the agent resulted in an action whose execution is blocked by the user.
Blocked { blocked_action: String },
}
impl std::fmt::Display for ConversationStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ConversationStatus::InProgress => write!(f, "In progress"),
ConversationStatus::Success => write!(f, "Done"),
ConversationStatus::Error => write!(f, "Error"),
ConversationStatus::Cancelled => write!(f, "Cancelled"),
ConversationStatus::Blocked { .. } => write!(f, "Blocked"),
}
}
}
impl ConversationStatus {
pub fn render_icon(&self, appearance: &Appearance) -> galaxyui::elements::Icon {
match self {
ConversationStatus::InProgress => in_progress_icon(appearance),
ConversationStatus::Success => succeeded_icon(appearance),
ConversationStatus::Blocked { .. } => yellow_stop_icon(appearance),
ConversationStatus::Error => failed_icon(appearance),
ConversationStatus::Cancelled => gray_stop_icon(appearance),
}
}
pub fn status_icon_and_color(&self, theme: &WarpTheme) -> (Icon, ColorU) {
match self {
ConversationStatus::InProgress => (Icon::ClockLoader, theme.ansi_fg_magenta()),
ConversationStatus::Success => (Icon::Check, theme.ansi_fg_green()),
ConversationStatus::Error => (Icon::Triangle, theme.ansi_fg_red()),
ConversationStatus::Cancelled => (Icon::StopFilled, internal_colors::neutral_5(theme)),
ConversationStatus::Blocked { .. } => (Icon::StopFilled, theme.ansi_fg_yellow()),
}
}
pub fn is_in_progress(&self) -> bool {
matches!(self, ConversationStatus::InProgress)
}
pub fn is_blocked(&self) -> bool {
matches!(self, ConversationStatus::Blocked { .. })
}
pub fn is_cancelled(&self) -> bool {
matches!(self, ConversationStatus::Cancelled)
}
pub fn is_done(&self) -> bool {
matches!(
self,
ConversationStatus::Success | ConversationStatus::Error | ConversationStatus::Cancelled
)
}
pub fn is_error(&self) -> bool {
matches!(self, ConversationStatus::Error)
}
}
#[cfg(test)]
#[path = "conversation_tests.rs"]
mod tests;