Initial public release of Warp.
Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
This commit is contained in:
David Stern
@@ -0,0 +1,509 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::ai::blocklist::agent_view::{AgentViewController, AgentViewControllerEvent};
|
||||
use crate::ai::blocklist::BlocklistAIHistoryModel;
|
||||
use crate::terminal::model::session::active_session::ActiveSession;
|
||||
use warpui::{
|
||||
AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity, WeakModelHandle,
|
||||
WindowId,
|
||||
};
|
||||
|
||||
/// Contains the handles needed to track an active agent view.
|
||||
struct ActiveAgentViewHandles {
|
||||
controller: WeakModelHandle<AgentViewController>,
|
||||
active_session: WeakModelHandle<ActiveSession>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum ActiveAgentViewsEvent {
|
||||
/// A conversation was closed (exited from the agent view or its pane was removed).
|
||||
ConversationClosed { conversation_id: AIConversationId },
|
||||
/// A conversation was entered within a terminal view.
|
||||
TerminalViewFocused,
|
||||
/// An ambient agent session was opened in a tab.
|
||||
AmbientSessionOpened {
|
||||
#[allow(dead_code)]
|
||||
task_id: AmbientAgentTaskId,
|
||||
},
|
||||
/// An ambient agent session tab was closed.
|
||||
AmbientSessionClosed {
|
||||
#[allow(dead_code)]
|
||||
task_id: AmbientAgentTaskId,
|
||||
},
|
||||
/// A window was closed and its focused state was removed.
|
||||
WindowClosed,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum ConversationOrTaskId {
|
||||
ConversationId(AIConversationId),
|
||||
TaskId(AmbientAgentTaskId),
|
||||
}
|
||||
|
||||
impl ConversationOrTaskId {
|
||||
pub fn conversation_id(&self) -> Option<AIConversationId> {
|
||||
match self {
|
||||
ConversationOrTaskId::ConversationId(conversation_id) => Some(*conversation_id),
|
||||
ConversationOrTaskId::TaskId(..) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// State of the focused terminal view and the active conversation in that terminal view.
|
||||
#[derive(Clone)]
|
||||
struct FocusedTerminalState {
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
focused_terminal_id: EntityId,
|
||||
active_conversation_id: Option<ConversationOrTaskId>,
|
||||
}
|
||||
|
||||
/// ActiveAgentViewsModel tracks which agent conversations are currently "active" - meaning either:
|
||||
/// - An interactive conversation whose agent view is expanded in a pane
|
||||
/// - An ambient conversation that is open in a tab
|
||||
/// This model also tracks which conversation is focused (i.e. active in the currently focused pane).
|
||||
pub struct ActiveAgentViewsModel {
|
||||
/// Per-window focused terminal state, keyed by WindowId.
|
||||
focused_terminal_states: HashMap<WindowId, FocusedTerminalState>,
|
||||
last_focused_terminal_state: Option<FocusedTerminalState>,
|
||||
/// Map from terminal_view_id to agent view handles (for interactive conversations).
|
||||
agent_view_handles: HashMap<EntityId, ActiveAgentViewHandles>,
|
||||
/// Map from terminal_view_id to ambient task ID (for open ambient sessions).
|
||||
ambient_sessions: HashMap<EntityId, AmbientAgentTaskId>,
|
||||
/// Tracks when each conversation was last opened/focused for sorting purposes.
|
||||
last_opened_times: HashMap<ConversationOrTaskId, DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl Entity for ActiveAgentViewsModel {
|
||||
type Event = ActiveAgentViewsEvent;
|
||||
}
|
||||
|
||||
impl SingletonEntity for ActiveAgentViewsModel {}
|
||||
|
||||
impl ActiveAgentViewsModel {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
focused_terminal_states: HashMap::new(),
|
||||
last_focused_terminal_state: None,
|
||||
agent_view_handles: HashMap::new(),
|
||||
ambient_sessions: HashMap::new(),
|
||||
last_opened_times: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register an agent view controller to track when the agent view is entered/exited.
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
pub fn register_agent_view_controller(
|
||||
&mut self,
|
||||
controller: &ModelHandle<AgentViewController>,
|
||||
active_session: &ModelHandle<ActiveSession>,
|
||||
terminal_view_id: EntityId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
// Skip registering this controller if it is already registered.
|
||||
if let Some(existing) = self.agent_view_handles.get(&terminal_view_id) {
|
||||
if existing
|
||||
.controller
|
||||
.upgrade(ctx)
|
||||
.is_some_and(|c| c.id() == controller.id())
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
self.agent_view_handles.insert(
|
||||
terminal_view_id,
|
||||
ActiveAgentViewHandles {
|
||||
controller: controller.downgrade(),
|
||||
active_session: active_session.downgrade(),
|
||||
},
|
||||
);
|
||||
|
||||
ctx.subscribe_to_model(controller, move |model, event, ctx| match event {
|
||||
AgentViewControllerEvent::EnteredAgentView {
|
||||
conversation_id, ..
|
||||
} => {
|
||||
let conv_id = ConversationOrTaskId::ConversationId(*conversation_id);
|
||||
model.last_opened_times.insert(conv_id, Utc::now());
|
||||
|
||||
// Update the focused conversation in whichever window owns this terminal view.
|
||||
// We ignore agent view changes if we are focused on an ambient conversation,
|
||||
// as ambient conversation navigation operates at the task level instead of the conversation level.
|
||||
for focused_terminal_state in model.focused_terminal_states.values_mut() {
|
||||
if focused_terminal_state.focused_terminal_id == terminal_view_id
|
||||
&& !matches!(
|
||||
focused_terminal_state.active_conversation_id,
|
||||
Some(ConversationOrTaskId::TaskId(_))
|
||||
)
|
||||
{
|
||||
focused_terminal_state.active_conversation_id = Some(conv_id);
|
||||
}
|
||||
}
|
||||
// Emit so subscribers can move this conversation to the Active section.
|
||||
ctx.emit(ActiveAgentViewsEvent::TerminalViewFocused);
|
||||
}
|
||||
AgentViewControllerEvent::ExitedAgentView {
|
||||
conversation_id, ..
|
||||
} => {
|
||||
model
|
||||
.last_opened_times
|
||||
.remove(&ConversationOrTaskId::ConversationId(*conversation_id));
|
||||
|
||||
// Clear the focused conversation in whichever window owns this terminal view.
|
||||
for state in model.focused_terminal_states.values_mut() {
|
||||
if state.focused_terminal_id == terminal_view_id
|
||||
&& !matches!(
|
||||
state.active_conversation_id,
|
||||
Some(ConversationOrTaskId::TaskId(_))
|
||||
)
|
||||
{
|
||||
state.active_conversation_id = None;
|
||||
}
|
||||
}
|
||||
// Emit so subscribers can move this conversation to the Past section.
|
||||
ctx.emit(ActiveAgentViewsEvent::ConversationClosed {
|
||||
conversation_id: *conversation_id,
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
});
|
||||
}
|
||||
|
||||
/// Unregister an agent view controller
|
||||
/// (called when the controller's terminal pane is hidden or closed).
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
pub fn unregister_agent_view_controller(
|
||||
&mut self,
|
||||
terminal_pane_id: EntityId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if let Some(handles) = self.agent_view_handles.remove(&terminal_pane_id) {
|
||||
let closed_conversation_id = handles
|
||||
.controller
|
||||
.upgrade(ctx)
|
||||
.and_then(|c| c.as_ref(ctx).agent_view_state().active_conversation_id());
|
||||
|
||||
// If the focused terminal is the one being unregistered, clear the focused state.
|
||||
self.focused_terminal_states
|
||||
.retain(|_, state| state.focused_terminal_id != terminal_pane_id);
|
||||
if self
|
||||
.last_focused_terminal_state
|
||||
.as_ref()
|
||||
.is_some_and(|state| state.focused_terminal_id == terminal_pane_id)
|
||||
{
|
||||
self.last_focused_terminal_state = None;
|
||||
}
|
||||
|
||||
if let Some(conversation_id) = closed_conversation_id {
|
||||
ctx.emit(ActiveAgentViewsEvent::ConversationClosed { conversation_id });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn handle_pane_focus_change(
|
||||
&mut self,
|
||||
window_id: WindowId,
|
||||
focused_terminal_view_id: Option<EntityId>,
|
||||
focused_task_id: Option<AmbientAgentTaskId>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let old_focused = self.get_focused_conversation(window_id);
|
||||
|
||||
if let Some(terminal_view_id) = focused_terminal_view_id {
|
||||
// Task ID takes precedence if viewing a shared ambient agent session.
|
||||
let active_conversation_id = if let Some(task_id) = focused_task_id {
|
||||
Some(ConversationOrTaskId::TaskId(task_id))
|
||||
} else {
|
||||
self.agent_view_handles
|
||||
.get(&terminal_view_id)
|
||||
.and_then(|handles| handles.controller.upgrade(ctx))
|
||||
.and_then(|controller| {
|
||||
controller
|
||||
.as_ref(ctx)
|
||||
.agent_view_state()
|
||||
.active_conversation_id()
|
||||
})
|
||||
.map(ConversationOrTaskId::ConversationId)
|
||||
};
|
||||
|
||||
let new_state = FocusedTerminalState {
|
||||
focused_terminal_id: terminal_view_id,
|
||||
active_conversation_id,
|
||||
};
|
||||
self.last_focused_terminal_state = Some(new_state.clone());
|
||||
self.focused_terminal_states.insert(window_id, new_state);
|
||||
} else {
|
||||
self.focused_terminal_states.remove(&window_id);
|
||||
}
|
||||
|
||||
if old_focused != self.get_focused_conversation(window_id) {
|
||||
ctx.emit(ActiveAgentViewsEvent::TerminalViewFocused);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the focused conversation for a specific window.
|
||||
/// Returns None if the window doesn't have an active agent view or ambient conversation.
|
||||
pub fn get_focused_conversation(&self, window_id: WindowId) -> Option<ConversationOrTaskId> {
|
||||
self.focused_terminal_states
|
||||
.get(&window_id)
|
||||
.and_then(|state| state.active_conversation_id)
|
||||
}
|
||||
|
||||
/// Get the last focused terminal view id (persisted across non-terminal focus changes).
|
||||
pub fn get_last_focused_terminal_id(&self) -> Option<EntityId> {
|
||||
self.last_focused_terminal_state
|
||||
.as_ref()
|
||||
.map(|state| state.focused_terminal_id)
|
||||
}
|
||||
|
||||
/// Returns the focused conversation ID if it's a new/empty conversation view.
|
||||
/// Only returns Some if the focused agent view was just created to start a new
|
||||
/// conversation (i.e. has no exchanges yet).
|
||||
pub fn maybe_get_focused_new_conversation(
|
||||
&self,
|
||||
window_id: WindowId,
|
||||
ctx: &AppContext,
|
||||
) -> Option<AIConversationId> {
|
||||
let state = self.focused_terminal_states.get(&window_id)?;
|
||||
let terminal_id = state.focused_terminal_id;
|
||||
|
||||
let is_new = self
|
||||
.agent_view_handles
|
||||
.get(&terminal_id)
|
||||
.and_then(|handles| handles.controller.upgrade(ctx))
|
||||
.map(|c| c.as_ref(ctx).agent_view_state().is_new())
|
||||
.unwrap_or(false);
|
||||
|
||||
if is_new {
|
||||
match state.active_conversation_id {
|
||||
Some(ConversationOrTaskId::ConversationId(id)) => Some(id),
|
||||
_ => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove the focused state for a window
|
||||
/// (called when said window is closed and cleaned up from the undo stack).
|
||||
pub fn remove_focused_state_for_window(
|
||||
&mut self,
|
||||
window_id: WindowId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if self.focused_terminal_states.remove(&window_id).is_some() {
|
||||
ctx.emit(ActiveAgentViewsEvent::WindowClosed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Register an ambient session (open in a tab).
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
pub fn register_ambient_session(
|
||||
&mut self,
|
||||
terminal_view_id: EntityId,
|
||||
task_id: AmbientAgentTaskId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let existing = self.ambient_sessions.insert(terminal_view_id, task_id);
|
||||
if existing != Some(task_id) {
|
||||
self.last_opened_times
|
||||
.insert(ConversationOrTaskId::TaskId(task_id), Utc::now());
|
||||
ctx.emit(ActiveAgentViewsEvent::AmbientSessionOpened { task_id });
|
||||
}
|
||||
}
|
||||
|
||||
/// Unregister an ambient session when the tab is closed.
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
pub fn unregister_ambient_session(
|
||||
&mut self,
|
||||
terminal_view_id: EntityId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if let Some(task_id) = self.ambient_sessions.remove(&terminal_view_id) {
|
||||
self.last_opened_times
|
||||
.remove(&ConversationOrTaskId::TaskId(task_id));
|
||||
ctx.emit(ActiveAgentViewsEvent::AmbientSessionClosed { task_id });
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the terminal view ID for a conversation if it's currently active
|
||||
/// (i.e., has an expanded agent view in some pane).
|
||||
pub fn terminal_view_id_for_conversation(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &AppContext,
|
||||
) -> Option<EntityId> {
|
||||
self.agent_view_handles
|
||||
.iter()
|
||||
.find_map(|(terminal_view_id, handles)| {
|
||||
let controller = handles.controller.upgrade(ctx)?;
|
||||
controller
|
||||
.as_ref(ctx)
|
||||
.agent_view_state()
|
||||
.active_conversation_id()
|
||||
.is_some_and(|id| id == conversation_id)
|
||||
.then_some(*terminal_view_id)
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns true if the conversation is currently open
|
||||
/// (i.e., has an expanded agent view in some pane).
|
||||
pub fn is_conversation_open(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &AppContext,
|
||||
) -> bool {
|
||||
self.terminal_view_id_for_conversation(conversation_id, ctx)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
/// Returns the active session for a conversation if it's currently active
|
||||
/// (i.e., has an expanded agent view).
|
||||
pub fn get_active_session_for_conversation(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &AppContext,
|
||||
) -> Option<ModelHandle<ActiveSession>> {
|
||||
for handles in self.agent_view_handles.values() {
|
||||
let Some(controller) = handles.controller.upgrade(ctx) else {
|
||||
continue;
|
||||
};
|
||||
let is_active = controller
|
||||
.as_ref(ctx)
|
||||
.agent_view_state()
|
||||
.active_conversation_id()
|
||||
.is_some_and(|id| id == conversation_id);
|
||||
if is_active {
|
||||
return handles.active_session.upgrade(ctx);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Returns the controller for a conversation if it's currently active
|
||||
/// (i.e., has an expanded agent view).
|
||||
pub fn get_controller_for_conversation(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &AppContext,
|
||||
) -> Option<ModelHandle<AgentViewController>> {
|
||||
for handles in self.agent_view_handles.values() {
|
||||
if let Some(controller) = handles.controller.upgrade(ctx) {
|
||||
let is_active = controller
|
||||
.as_ref(ctx)
|
||||
.agent_view_state()
|
||||
.active_conversation_id()
|
||||
.is_some_and(|id| id == conversation_id);
|
||||
if is_active {
|
||||
return Some(controller);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Returns the last opened time for a conversation, used for sorting active conversations.
|
||||
pub fn get_last_opened_time(&self, id: &ConversationOrTaskId) -> Option<DateTime<Utc>> {
|
||||
self.last_opened_times.get(id).copied()
|
||||
}
|
||||
|
||||
/// Returns the terminal view ID that has an active ambient session with the given task ID.
|
||||
pub fn get_terminal_view_id_for_ambient_task(
|
||||
&self,
|
||||
task_id: AmbientAgentTaskId,
|
||||
) -> Option<EntityId> {
|
||||
self.ambient_sessions
|
||||
.iter()
|
||||
.find_map(|(view_id, id)| (*id == task_id).then_some(*view_id))
|
||||
}
|
||||
|
||||
/// Returns the terminal view ID that has an active conversation with the given ID.
|
||||
pub fn get_terminal_view_id_for_conversation(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &AppContext,
|
||||
) -> Option<EntityId> {
|
||||
for (terminal_view_id, handles) in &self.agent_view_handles {
|
||||
let Some(controller) = handles.controller.upgrade(ctx) else {
|
||||
continue;
|
||||
};
|
||||
let is_active = controller
|
||||
.as_ref(ctx)
|
||||
.agent_view_state()
|
||||
.active_conversation_id()
|
||||
.is_some_and(|id| id == conversation_id);
|
||||
if is_active {
|
||||
return Some(*terminal_view_id);
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Get all currently active conversation IDs.
|
||||
/// A conversation is active if it is open and a query has been sent since it was last opened.
|
||||
/// New (empty) conversations and ambient sessions are always considered active when open.
|
||||
pub fn get_all_active_conversation_ids(
|
||||
&self,
|
||||
ctx: &AppContext,
|
||||
) -> HashSet<ConversationOrTaskId> {
|
||||
let history_model = BlocklistAIHistoryModel::as_ref(ctx);
|
||||
let mut ids = HashSet::new();
|
||||
|
||||
for handles in self.agent_view_handles.values() {
|
||||
if let Some(controller) = handles.controller.upgrade(ctx) {
|
||||
let state = controller.as_ref(ctx).agent_view_state();
|
||||
if let Some(conversation_id) = state.active_conversation_id() {
|
||||
let Some(conversation) = history_model.conversation(&conversation_id) else {
|
||||
continue;
|
||||
};
|
||||
if !conversation.is_entirely_passive()
|
||||
&& state.was_conversation_modified_since_opening(history_model)
|
||||
{
|
||||
ids.insert(ConversationOrTaskId::ConversationId(conversation_id));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Ambient sessions are always considered active when open.
|
||||
for task_id in self.ambient_sessions.values() {
|
||||
ids.insert(ConversationOrTaskId::TaskId(*task_id));
|
||||
}
|
||||
|
||||
ids
|
||||
}
|
||||
|
||||
/// Get all currently open conversation IDs.
|
||||
/// A conversation is considered open if it is in an expanded agent view.
|
||||
pub fn get_all_open_conversation_ids(&self, ctx: &AppContext) -> HashSet<ConversationOrTaskId> {
|
||||
let mut ids = HashSet::new();
|
||||
|
||||
// Collect from interactive agent views (expanded).
|
||||
for handles in self.agent_view_handles.values() {
|
||||
if let Some(controller) = handles.controller.upgrade(ctx) {
|
||||
if let Some(conversation_id) = controller
|
||||
.as_ref(ctx)
|
||||
.agent_view_state()
|
||||
.active_conversation_id()
|
||||
{
|
||||
ids.insert(ConversationOrTaskId::ConversationId(conversation_id));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Collect from ambient sessions (open in tabs)
|
||||
for task_id in self.ambient_sessions.values() {
|
||||
ids.insert(ConversationOrTaskId::TaskId(*task_id));
|
||||
}
|
||||
|
||||
ids
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "active_agent_views_model_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,213 @@
|
||||
use std::str::FromStr;
|
||||
|
||||
use warpui::{App, EntityId, WindowId};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn setup_model(app: &mut App) -> ModelHandle<ActiveAgentViewsModel> {
|
||||
app.add_singleton_model(|_| ActiveAgentViewsModel::new())
|
||||
}
|
||||
|
||||
fn new_task_id() -> AmbientAgentTaskId {
|
||||
AmbientAgentTaskId::from_str(&uuid::Uuid::new_v4().to_string()).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn per_window_focused_state_is_independent() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window_a = WindowId::new();
|
||||
let window_b = WindowId::new();
|
||||
let terminal_a = EntityId::new();
|
||||
let terminal_b = EntityId::new();
|
||||
let task_a = new_task_id();
|
||||
let task_b = new_task_id();
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_a, Some(terminal_a), Some(task_a), ctx);
|
||||
model.handle_pane_focus_change(window_b, Some(terminal_b), Some(task_b), ctx);
|
||||
});
|
||||
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(
|
||||
model.get_focused_conversation(window_a),
|
||||
Some(ConversationOrTaskId::TaskId(task_a))
|
||||
);
|
||||
assert_eq!(
|
||||
model.get_focused_conversation(window_b),
|
||||
Some(ConversationOrTaskId::TaskId(task_b))
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clearing_one_window_does_not_affect_other() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window_a = WindowId::new();
|
||||
let window_b = WindowId::new();
|
||||
let terminal_a = EntityId::new();
|
||||
let terminal_b = EntityId::new();
|
||||
let task_a = new_task_id();
|
||||
let task_b = new_task_id();
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_a, Some(terminal_a), Some(task_a), ctx);
|
||||
model.handle_pane_focus_change(window_b, Some(terminal_b), Some(task_b), ctx);
|
||||
});
|
||||
|
||||
// Clear window A's focus by passing None for terminal_view_id.
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_a, None, None, ctx);
|
||||
});
|
||||
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_focused_conversation(window_a), None);
|
||||
assert_eq!(
|
||||
model.get_focused_conversation(window_b),
|
||||
Some(ConversationOrTaskId::TaskId(task_b))
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn last_focused_terminal_tracks_most_recent_globally() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window_a = WindowId::new();
|
||||
let window_b = WindowId::new();
|
||||
let terminal_a = EntityId::new();
|
||||
let terminal_b = EntityId::new();
|
||||
let task_a = new_task_id();
|
||||
let task_b = new_task_id();
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_a, Some(terminal_a), Some(task_a), ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal_a));
|
||||
});
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_b, Some(terminal_b), Some(task_b), ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal_b));
|
||||
});
|
||||
|
||||
// Clearing window B's focus should NOT clear last_focused (it persists).
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_b, None, None, ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal_b));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_window_returns_none() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window_a = WindowId::new();
|
||||
let unknown_window = WindowId::new();
|
||||
let terminal = EntityId::new();
|
||||
let task = new_task_id();
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_a, Some(terminal), Some(task), ctx);
|
||||
});
|
||||
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_focused_conversation(unknown_window), None);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn focus_change_without_task_id_has_no_conversation() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window = WindowId::new();
|
||||
let terminal = EntityId::new();
|
||||
|
||||
// No agent view handles registered, no task_id → active_conversation_id should be None.
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window, Some(terminal), None, ctx);
|
||||
});
|
||||
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_focused_conversation(window), None);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_focused_state_for_window_cleans_up() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window_a = WindowId::new();
|
||||
let window_b = WindowId::new();
|
||||
let terminal_a = EntityId::new();
|
||||
let terminal_b = EntityId::new();
|
||||
let task_a = new_task_id();
|
||||
let task_b = new_task_id();
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window_a, Some(terminal_a), Some(task_a), ctx);
|
||||
model.handle_pane_focus_change(window_b, Some(terminal_b), Some(task_b), ctx);
|
||||
});
|
||||
|
||||
// Remove window A's state (simulating undo-close expiry).
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.remove_focused_state_for_window(window_a, ctx);
|
||||
});
|
||||
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.get_focused_conversation(window_a), None);
|
||||
assert_eq!(
|
||||
model.get_focused_conversation(window_b),
|
||||
Some(ConversationOrTaskId::TaskId(task_b))
|
||||
);
|
||||
});
|
||||
|
||||
// Removing again is a no-op.
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.remove_focused_state_for_window(window_a, ctx);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overwriting_same_window_updates_state() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = setup_model(&mut app);
|
||||
let window = WindowId::new();
|
||||
let terminal_1 = EntityId::new();
|
||||
let terminal_2 = EntityId::new();
|
||||
let task_1 = new_task_id();
|
||||
let task_2 = new_task_id();
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window, Some(terminal_1), Some(task_1), ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(
|
||||
model.get_focused_conversation(window),
|
||||
Some(ConversationOrTaskId::TaskId(task_1))
|
||||
);
|
||||
});
|
||||
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.handle_pane_focus_change(window, Some(terminal_2), Some(task_2), ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(
|
||||
model.get_focused_conversation(window),
|
||||
Some(ConversationOrTaskId::TaskId(task_2))
|
||||
);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
pub(crate) mod convert_conversation;
|
||||
mod convert_from;
|
||||
mod convert_to;
|
||||
mod r#impl;
|
||||
|
||||
pub use ai::agent::convert::ConvertToAPITypeError;
|
||||
use ai::api_keys::ApiKeyManager;
|
||||
pub use convert_from::{
|
||||
user_inputs_from_messages, ConversionParams, ConvertAPIMessageToClientOutputMessage,
|
||||
MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError,
|
||||
};
|
||||
|
||||
pub use r#impl::generate_multi_agent_output;
|
||||
|
||||
use futures_lite::Stream;
|
||||
use serde::Serialize;
|
||||
use std::path::Path;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use warp_core::channel::ChannelState;
|
||||
use warp_core::execution_mode::AppExecutionMode;
|
||||
use warp_core::features::FeatureFlag;
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::{
|
||||
ai::{blocklist::SessionContext, llms::LLMId},
|
||||
server::server_api::AIApiError,
|
||||
};
|
||||
|
||||
use super::{AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions};
|
||||
use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput};
|
||||
use crate::ai::mcp::templatable_manager::TemplatableMCPServerInfo;
|
||||
use crate::ai::mcp::TemplatableMCPServerManager;
|
||||
use crate::settings::AISettings;
|
||||
use crate::terminal::safe_mode_settings::get_secret_obfuscation_mode;
|
||||
use crate::workspaces::user_workspaces::UserWorkspaces;
|
||||
use warp_core::user_preferences::GetUserPreferences;
|
||||
use warpui::{AppContext, EntityId, SingletonEntity as _};
|
||||
|
||||
/// Unique, server-generated conversation-scoped token to be roundtripped to the API when sending
|
||||
/// requests that follow-up within a given conversation.
|
||||
#[derive(Serialize, Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct ServerConversationToken(String);
|
||||
|
||||
impl ServerConversationToken {
|
||||
pub fn new(id: String) -> Self {
|
||||
Self(id)
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub fn debug_link(&self) -> String {
|
||||
format!(
|
||||
"{}/debug/maa/{}",
|
||||
ChannelState::server_root_url(),
|
||||
self.as_str()
|
||||
)
|
||||
}
|
||||
|
||||
pub fn conversation_link(&self) -> String {
|
||||
format!(
|
||||
"{}/conversation/{}",
|
||||
ChannelState::server_root_url(),
|
||||
self.as_str()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ServerConversationToken> for String {
|
||||
fn from(value: ServerConversationToken) -> Self {
|
||||
value.0
|
||||
}
|
||||
}
|
||||
|
||||
// Conversions between AI ServerConversationToken and protocol ServerConversationToken
|
||||
impl From<session_sharing_protocol::common::ServerConversationToken> for ServerConversationToken {
|
||||
fn from(token: session_sharing_protocol::common::ServerConversationToken) -> Self {
|
||||
Self(token.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<ServerConversationToken>
|
||||
for session_sharing_protocol::common::ServerConversationToken
|
||||
{
|
||||
type Error = uuid::Error;
|
||||
|
||||
fn try_from(token: ServerConversationToken) -> Result<Self, Self::Error> {
|
||||
token.as_str().parse()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RequestParams {
|
||||
pub input: Vec<AIAgentInput>,
|
||||
pub conversation_token: Option<ServerConversationToken>,
|
||||
pub forked_from_conversation_token: Option<ServerConversationToken>,
|
||||
pub ambient_agent_task_id: Option<AmbientAgentTaskId>,
|
||||
pub tasks: Vec<warp_multi_agent_api::Task>,
|
||||
pub existing_suggestions: Option<Suggestions>,
|
||||
pub metadata: Option<RequestMetadata>,
|
||||
pub session_context: SessionContext,
|
||||
pub model: LLMId,
|
||||
#[allow(unused)]
|
||||
pub coding_model: LLMId,
|
||||
pub cli_agent_model: LLMId,
|
||||
pub computer_use_model: LLMId,
|
||||
pub is_memory_enabled: bool,
|
||||
pub warp_drive_context_enabled: bool,
|
||||
pub mcp_context: Option<MCPContext>,
|
||||
pub planning_enabled: bool,
|
||||
should_redact_secrets: bool,
|
||||
|
||||
/// User-provided API keys for AI providers (BYO API Key).
|
||||
pub api_keys: Option<warp_multi_agent_api::request::settings::ApiKeys>,
|
||||
pub allow_use_of_warp_credits_with_byok: bool,
|
||||
pub autonomy_level: warp_multi_agent_api::AutonomyLevel,
|
||||
pub isolation_level: warp_multi_agent_api::IsolationLevel,
|
||||
pub web_search_enabled: bool,
|
||||
pub computer_use_enabled: bool,
|
||||
pub ask_user_question_enabled: bool,
|
||||
pub research_agent_enabled: bool,
|
||||
pub orchestration_enabled: bool,
|
||||
pub supported_tools_override: Option<Vec<warp_multi_agent_api::ToolType>>,
|
||||
/// The conversation ID of the parent agent that spawned this child agent, if any.
|
||||
pub parent_agent_id: Option<String>,
|
||||
/// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator.
|
||||
pub agent_name: Option<String>,
|
||||
}
|
||||
|
||||
pub type Event = Result<warp_multi_agent_api::ResponseEvent, Arc<AIApiError>>;
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub type ResponseStream = Pin<Box<dyn Stream<Item = Event> + Send + 'static>>;
|
||||
|
||||
// The WASM version of this type has no bound on `Send`, which is an unnecessary bound when
|
||||
// targeting wasm because the browser is single-threaded (and we don't leverage WebWorkers for async
|
||||
// execution in WoW).
|
||||
#[cfg(target_family = "wasm")]
|
||||
pub type ResponseStream = Pin<Box<dyn Stream<Item = Event>>>;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConversationData {
|
||||
pub id: AIConversationId,
|
||||
pub tasks: Vec<warp_multi_agent_api::Task>,
|
||||
pub server_conversation_token: Option<ServerConversationToken>,
|
||||
pub forked_from_conversation_token: Option<ServerConversationToken>,
|
||||
pub ambient_agent_task_id: Option<AmbientAgentTaskId>,
|
||||
pub existing_suggestions: Option<Suggestions>,
|
||||
}
|
||||
|
||||
impl RequestParams {
|
||||
pub fn new(
|
||||
terminal_view_id: Option<EntityId>,
|
||||
session_context: SessionContext,
|
||||
request_input: &RequestInput,
|
||||
conversation: ConversationData,
|
||||
metadata: Option<RequestMetadata>,
|
||||
app: &AppContext,
|
||||
) -> Self {
|
||||
let ai_settings = AISettings::as_ref(app);
|
||||
let is_memory_enabled = ai_settings.is_memory_enabled(app);
|
||||
let warp_drive_context_enabled = ai_settings.is_warp_drive_context_enabled(app);
|
||||
|
||||
// Build MCP context - either grouped by server or flat lists based on feature flag
|
||||
let mcp_context = if FeatureFlag::MCPGroupedServerContext.is_enabled() {
|
||||
// Group MCP tools and resources by server
|
||||
let templatable_manager = TemplatableMCPServerManager::as_ref(app);
|
||||
|
||||
let mut active_servers: Vec<&TemplatableMCPServerInfo> = templatable_manager
|
||||
.get_active_templatable_servers()
|
||||
.values()
|
||||
.copied()
|
||||
.collect();
|
||||
|
||||
// If file-based MCP servers are enabled, add active servers in scope of
|
||||
// the user's current working directory
|
||||
if let Some(cwd) = session_context.current_working_directory() {
|
||||
active_servers.extend(
|
||||
templatable_manager
|
||||
.get_active_file_based_servers(Path::new(cwd), app)
|
||||
.values(),
|
||||
);
|
||||
}
|
||||
|
||||
// Include any ephemeral MCP servers started via the Oz CLI.
|
||||
active_servers.extend(
|
||||
templatable_manager
|
||||
.get_active_cli_spawned_servers()
|
||||
.values(),
|
||||
);
|
||||
|
||||
let servers: Vec<MCPServer> = active_servers
|
||||
.into_iter()
|
||||
.map(|server| MCPServer {
|
||||
name: server.name().to_string(),
|
||||
description: server.description().unwrap_or_default().to_string(),
|
||||
id: server.installation_id().to_string(),
|
||||
resources: server.resources().to_vec(),
|
||||
tools: server.tools().to_vec(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
if servers.is_empty() {
|
||||
None
|
||||
} else {
|
||||
#[allow(deprecated)]
|
||||
Some(MCPContext {
|
||||
resources: vec![],
|
||||
tools: vec![],
|
||||
servers,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
// Flat lists of resources and tools
|
||||
let templatable_mcp_manager = TemplatableMCPServerManager::as_ref(app);
|
||||
let resources = templatable_mcp_manager
|
||||
.resources()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let tools = templatable_mcp_manager.tools().cloned().collect::<Vec<_>>();
|
||||
|
||||
#[allow(deprecated)]
|
||||
(!resources.is_empty() || !tools.is_empty()).then_some(MCPContext {
|
||||
resources,
|
||||
tools,
|
||||
servers: vec![],
|
||||
})
|
||||
};
|
||||
|
||||
let should_redact_secrets = get_secret_obfuscation_mode(app).should_redact_secret();
|
||||
|
||||
let user_workspaces = UserWorkspaces::as_ref(app);
|
||||
let api_keys = ApiKeyManager::as_ref(app).api_keys_for_request(
|
||||
user_workspaces.is_byo_api_key_enabled(),
|
||||
user_workspaces.is_aws_bedrock_credentials_enabled(app),
|
||||
);
|
||||
let allow_use_of_warp_credits_with_byok =
|
||||
*AISettings::as_ref(app).can_use_warp_credits_with_byok;
|
||||
|
||||
let app_execution_mode = AppExecutionMode::as_ref(app);
|
||||
let autonomy_level = if app_execution_mode.is_autonomous() {
|
||||
warp_multi_agent_api::AutonomyLevel::Unsupervised
|
||||
} else {
|
||||
warp_multi_agent_api::AutonomyLevel::Supervised
|
||||
};
|
||||
|
||||
let isolation_level = if app_execution_mode.is_sandboxed() {
|
||||
warp_multi_agent_api::IsolationLevel::Sandbox
|
||||
} else {
|
||||
warp_multi_agent_api::IsolationLevel::None
|
||||
};
|
||||
|
||||
let web_search_enabled =
|
||||
BlocklistAIPermissions::as_ref(app).get_web_search_enabled(app, terminal_view_id);
|
||||
let research_agent_enabled = app
|
||||
.private_user_preferences()
|
||||
.read_value("ResearchAgentEnabled")
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or_default();
|
||||
let is_ambient_agent = conversation.ambient_agent_task_id.is_some();
|
||||
let computer_use_enabled = FeatureFlag::AgentModeComputerUse.is_enabled()
|
||||
&& BlocklistAIPermissions::as_ref(app)
|
||||
.get_computer_use_setting(app, terminal_view_id)
|
||||
.is_enabled()
|
||||
&& computer_use::is_supported_on_current_platform()
|
||||
&& (FeatureFlag::LocalComputerUse.is_enabled() || is_ambient_agent);
|
||||
let ask_user_question_enabled = BlocklistAIPermissions::as_ref(app)
|
||||
.get_ask_user_question_setting(app, terminal_view_id)
|
||||
!= crate::ai::execution_profiles::AskUserQuestionPermission::Never;
|
||||
|
||||
let orchestration_enabled = ai_settings.is_orchestration_enabled(app)
|
||||
&& session_context
|
||||
.session_type()
|
||||
.as_ref()
|
||||
.is_none_or(|t| matches!(t, crate::terminal::model::session::SessionType::Local));
|
||||
|
||||
Self {
|
||||
input: request_input.all_inputs().cloned().collect(),
|
||||
conversation_token: conversation.server_conversation_token,
|
||||
forked_from_conversation_token: conversation.forked_from_conversation_token,
|
||||
ambient_agent_task_id: conversation.ambient_agent_task_id,
|
||||
tasks: conversation.tasks,
|
||||
existing_suggestions: conversation.existing_suggestions,
|
||||
metadata,
|
||||
session_context,
|
||||
model: request_input.model_id.clone(),
|
||||
coding_model: request_input.coding_model_id.clone(),
|
||||
cli_agent_model: request_input.cli_agent_model_id.clone(),
|
||||
computer_use_model: request_input.computer_use_model_id.clone(),
|
||||
is_memory_enabled,
|
||||
warp_drive_context_enabled,
|
||||
mcp_context,
|
||||
planning_enabled: true,
|
||||
should_redact_secrets,
|
||||
api_keys,
|
||||
allow_use_of_warp_credits_with_byok,
|
||||
autonomy_level,
|
||||
isolation_level,
|
||||
web_search_enabled,
|
||||
computer_use_enabled,
|
||||
ask_user_question_enabled,
|
||||
research_agent_enabled,
|
||||
orchestration_enabled,
|
||||
supported_tools_override: request_input.supported_tools_override.clone(),
|
||||
parent_agent_id: None,
|
||||
agent_name: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,944 @@
|
||||
//! Conversions from MAA API types to application types.
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::ai::agent::api::convert_conversation::{
|
||||
convert_input_context, convert_tool_call_result_to_input,
|
||||
};
|
||||
use crate::ai::agent::comment::CodeReview;
|
||||
use crate::ai::agent::task::TaskId;
|
||||
use crate::ai::agent::todos::AIAgentTodoList;
|
||||
use crate::ai::agent::{
|
||||
util::parse_markdown_into_text_and_code_sections, AIAgentAction, AIAgentActionType,
|
||||
AIAgentCitation, AIAgentInput, AIAgentOutputMessage, AIAgentText, AIAgentTodo,
|
||||
ArtifactCreatedData, MessageId, StartAgentExecutionMode, SuggestedAgentModeWorkflow,
|
||||
SuggestedRule, Suggestions, TodoOperation,
|
||||
};
|
||||
use crate::ai::agent::{
|
||||
CloneRepositoryURL, SubagentCall, SubagentType, SummarizationType, WebFetchStatus,
|
||||
WebSearchStatus,
|
||||
};
|
||||
use crate::ai::artifact_download::sanitized_basename;
|
||||
use crate::ai::document::ai_document_model::{AIDocumentId, AIDocumentVersion};
|
||||
use ai::agent::action::LifecycleEventType as StartAgentLifecycleEventType;
|
||||
use ai::agent::action_result::StartAgentVersion;
|
||||
use ai::agent::convert::ToolToAIAgentActionError;
|
||||
use ai::agent::UnknownCitationTypeError;
|
||||
use ai::skills::SkillReference;
|
||||
use api::ask_user_question::question::QuestionType;
|
||||
use warp_core::channel::ChannelState;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::ai::agent::{AIAgentAttachment, UserQueryMode};
|
||||
|
||||
impl TryFrom<api::Attachment> for AIAgentAttachment {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(attachment: api::Attachment) -> Result<Self, Self::Error> {
|
||||
match attachment.value {
|
||||
Some(api::attachment::Value::FilePathReference(fpr)) => {
|
||||
Ok(AIAgentAttachment::FilePathReference {
|
||||
file_id: String::new(),
|
||||
file_name: fpr
|
||||
.file_path
|
||||
.rsplit('/')
|
||||
.next()
|
||||
.unwrap_or(&fpr.file_path)
|
||||
.to_string(),
|
||||
file_path: fpr.file_path,
|
||||
})
|
||||
}
|
||||
_ => anyhow::bail!("Unsupported attachment type for conversion"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts proto UserQueryMode to the internal UserQueryMode type
|
||||
pub(crate) fn convert_user_query_mode(mode: Option<&api::UserQueryMode>) -> UserQueryMode {
|
||||
let Some(mode) = mode else {
|
||||
return UserQueryMode::default();
|
||||
};
|
||||
|
||||
match &mode.r#type {
|
||||
Some(api::user_query_mode::Type::Plan(_)) => UserQueryMode::Plan,
|
||||
Some(api::user_query_mode::Type::Orchestrate(_)) => UserQueryMode::Orchestrate,
|
||||
None => UserQueryMode::Normal,
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_start_agent_lifecycle_event_type(
|
||||
event_type: i32,
|
||||
) -> Option<StartAgentLifecycleEventType> {
|
||||
let event_type = StartAgentLifecycleEventType::try_from(event_type).ok()?;
|
||||
(event_type != StartAgentLifecycleEventType::Unspecified).then_some(event_type)
|
||||
}
|
||||
|
||||
fn convert_start_agent_v2_harness_type(
|
||||
harness: Option<api::start_agent_v2::execution_mode::Harness>,
|
||||
) -> Option<String> {
|
||||
harness
|
||||
.map(|harness| harness.r#type)
|
||||
.filter(|harness_type| !harness_type.trim().is_empty())
|
||||
}
|
||||
|
||||
fn convert_start_agent_execution_mode(
|
||||
execution_mode: Option<api::start_agent::ExecutionMode>,
|
||||
) -> StartAgentExecutionMode {
|
||||
match execution_mode.and_then(|execution_mode| execution_mode.mode) {
|
||||
Some(api::start_agent::execution_mode::Mode::Remote(remote)) => {
|
||||
StartAgentExecutionMode::remote_with_defaults(remote.environment_id)
|
||||
}
|
||||
Some(api::start_agent::execution_mode::Mode::Local(_)) | None => {
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_start_agent_v2_execution_mode(
|
||||
execution_mode: Option<api::start_agent_v2::ExecutionMode>,
|
||||
) -> StartAgentExecutionMode {
|
||||
match execution_mode.and_then(|execution_mode| execution_mode.mode) {
|
||||
Some(api::start_agent_v2::execution_mode::Mode::Remote(remote)) => {
|
||||
StartAgentExecutionMode::Remote {
|
||||
environment_id: remote.environment_id,
|
||||
skill_references: remote
|
||||
.skills
|
||||
.into_iter()
|
||||
.filter_map(convert_skill_reference)
|
||||
.collect(),
|
||||
model_id: remote.model_id,
|
||||
computer_use_enabled: remote.computer_use_enabled,
|
||||
worker_host: remote.worker_host,
|
||||
harness_type: convert_start_agent_v2_harness_type(remote.harness)
|
||||
.unwrap_or_default(),
|
||||
title: remote.title,
|
||||
}
|
||||
}
|
||||
Some(api::start_agent_v2::execution_mode::Mode::Local(local)) => {
|
||||
convert_start_agent_v2_harness_type(local.harness)
|
||||
.map(StartAgentExecutionMode::local_harness)
|
||||
.unwrap_or_else(StartAgentExecutionMode::local_with_defaults)
|
||||
}
|
||||
None => StartAgentExecutionMode::local_with_defaults(),
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_skill_reference(skill_ref: api::SkillRef) -> Option<SkillReference> {
|
||||
match skill_ref.skill_reference {
|
||||
Some(api::skill_ref::SkillReference::Path(path)) => Some(SkillReference::Path(path.into())),
|
||||
Some(api::skill_ref::SkillReference::BundledSkillId(id)) => {
|
||||
Some(SkillReference::BundledSkillId(id))
|
||||
}
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Unexpected errors when trying to convert an [`api::Message`] to an [`AIAgentOutputMessage`].
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum MessageToAIAgentOutputMessageError {
|
||||
#[error("Missing expected message")]
|
||||
MissingMessage,
|
||||
#[error("Error converting tool to action: {0:?}")]
|
||||
ToolError(#[from] ToolToAIAgentActionError),
|
||||
#[error("Error converting citation: {0:?}")]
|
||||
CitationError(#[from] UnknownCitationTypeError),
|
||||
}
|
||||
|
||||
/// Successful result when trying to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
pub enum MaybeAIAgentOutputMessage {
|
||||
/// There is a mapping to a client output message.
|
||||
Message(AIAgentOutputMessage),
|
||||
/// We tried to parse a message that we don't care about.
|
||||
NoClientRepresentation,
|
||||
}
|
||||
|
||||
/// Successful result when trying to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
enum MaybeAIAgentAction {
|
||||
/// There is a mapping to a client action.
|
||||
Action(AIAgentAction),
|
||||
Subagent(SubagentCall),
|
||||
/// We tried to parse a tool call that we don't care about.
|
||||
NoClientRepresentation,
|
||||
}
|
||||
|
||||
pub struct ConversionParams<'a> {
|
||||
pub task_id: &'a TaskId,
|
||||
pub current_todo_list: Option<&'a AIAgentTodoList>,
|
||||
pub active_code_review: Option<&'a CodeReview>,
|
||||
}
|
||||
|
||||
/// Trait for converting an [`api::Message`] to an [`AIAgentOutputMessage`].
|
||||
pub trait ConvertAPIMessageToClientOutputMessage {
|
||||
fn to_client_output_message(
|
||||
self,
|
||||
params: ConversionParams,
|
||||
) -> Result<MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError>;
|
||||
}
|
||||
|
||||
impl ConvertAPIMessageToClientOutputMessage for api::Message {
|
||||
fn to_client_output_message(
|
||||
self,
|
||||
params: ConversionParams,
|
||||
) -> Result<MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError> {
|
||||
let Some(message) = self.message else {
|
||||
// In shared-session streams we can receive skeleton placeholder task messages without payloads.
|
||||
// Treat them as having no client representation rather than erroring and aborting ingestion entirely.
|
||||
return Ok(MaybeAIAgentOutputMessage::NoClientRepresentation);
|
||||
};
|
||||
|
||||
let citations = self
|
||||
.citations
|
||||
.iter()
|
||||
.map(|citation| (*citation).clone().try_into())
|
||||
.collect::<Result<Vec<AIAgentCitation>, UnknownCitationTypeError>>()?;
|
||||
|
||||
match message {
|
||||
api::message::Message::AgentOutput(output) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::text(MessageId::new(self.id), output.into())
|
||||
.with_citations(citations),
|
||||
)),
|
||||
api::message::Message::AgentReasoning(reasoning) => {
|
||||
let duration = reasoning
|
||||
.finished_duration
|
||||
.map(|d| Duration::from_secs(d.seconds as u64));
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::reasoning(
|
||||
MessageId::new(self.id),
|
||||
reasoning.into(),
|
||||
duration,
|
||||
),
|
||||
))
|
||||
}
|
||||
api::message::Message::ToolCall(tool_call) => match tool_call.to_action(params)? {
|
||||
MaybeAIAgentAction::Action(action) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::action(MessageId::new(self.id), action)
|
||||
.with_citations(citations),
|
||||
)),
|
||||
MaybeAIAgentAction::Subagent(subagent) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::subagent(MessageId::new(self.id), subagent)
|
||||
.with_citations(citations),
|
||||
)),
|
||||
MaybeAIAgentAction::NoClientRepresentation => {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
},
|
||||
api::message::Message::WebSearch(web_search) => {
|
||||
let status = match &web_search.status {
|
||||
Some(api::message::web_search::Status {
|
||||
r#type: Some(api::message::web_search::status::Type::Searching(searching)),
|
||||
}) => WebSearchStatus::Searching {
|
||||
query: if searching.query.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(searching.query.clone())
|
||||
},
|
||||
},
|
||||
Some(api::message::web_search::Status {
|
||||
r#type: Some(api::message::web_search::status::Type::Success(success)),
|
||||
}) => WebSearchStatus::Success {
|
||||
query: success.query.clone(),
|
||||
pages: success
|
||||
.pages
|
||||
.iter()
|
||||
.map(|p| (p.url.clone(), p.title.clone()))
|
||||
.collect(),
|
||||
},
|
||||
Some(api::message::web_search::Status {
|
||||
r#type: Some(api::message::web_search::status::Type::Error(_)),
|
||||
}) => {
|
||||
// Error type doesn't have a query field currently, use empty string
|
||||
WebSearchStatus::Error {
|
||||
query: String::new(),
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// Unknown or missing status
|
||||
return Ok(MaybeAIAgentOutputMessage::NoClientRepresentation);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::web_search(MessageId::new(self.id), status)
|
||||
.with_citations(citations),
|
||||
))
|
||||
}
|
||||
api::message::Message::WebFetch(web_fetch) => {
|
||||
let status = match &web_fetch.status {
|
||||
Some(api::message::web_fetch::Status {
|
||||
r#type: Some(api::message::web_fetch::status::Type::Fetching(fetching)),
|
||||
}) => WebFetchStatus::Fetching {
|
||||
urls: fetching.urls.clone(),
|
||||
},
|
||||
Some(api::message::web_fetch::Status {
|
||||
r#type: Some(api::message::web_fetch::status::Type::Success(success)),
|
||||
}) => WebFetchStatus::Success {
|
||||
pages: success
|
||||
.pages
|
||||
.iter()
|
||||
.map(|p| (p.url.clone(), p.title.clone(), p.success))
|
||||
.collect(),
|
||||
},
|
||||
Some(api::message::web_fetch::Status {
|
||||
r#type: Some(api::message::web_fetch::status::Type::Error(_)),
|
||||
}) => WebFetchStatus::Error,
|
||||
_ => {
|
||||
// Unknown or missing status
|
||||
return Ok(MaybeAIAgentOutputMessage::NoClientRepresentation);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::web_fetch(MessageId::new(self.id), status)
|
||||
.with_citations(citations),
|
||||
))
|
||||
}
|
||||
api::message::Message::ModelUsed(_) => {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
api::message::Message::UpdateTodos(update_todos) => {
|
||||
if let Some(operation) = update_todos.operation {
|
||||
match operation {
|
||||
api::message::update_todos::Operation::CreateTodoList(create_todo_list) => {
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::todo_operation(
|
||||
MessageId::new(self.id),
|
||||
TodoOperation::UpdateTodos {
|
||||
todos: create_todo_list
|
||||
.initial_todos
|
||||
.into_iter()
|
||||
.map(Into::into)
|
||||
.collect(),
|
||||
},
|
||||
)
|
||||
.with_citations(citations),
|
||||
))
|
||||
}
|
||||
api::message::update_todos::Operation::UpdatePendingTodos(
|
||||
update_pending_todos,
|
||||
) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::todo_operation(
|
||||
MessageId::new(self.id),
|
||||
TodoOperation::UpdateTodos {
|
||||
todos: params
|
||||
.current_todo_list
|
||||
.iter()
|
||||
.flat_map(|list| list.completed_items().iter().cloned())
|
||||
.chain(
|
||||
update_pending_todos
|
||||
.updated_pending_todos
|
||||
.into_iter()
|
||||
.map(Into::into),
|
||||
)
|
||||
.collect(),
|
||||
},
|
||||
)
|
||||
.with_citations(citations),
|
||||
)),
|
||||
api::message::update_todos::Operation::MarkTodosCompleted(
|
||||
mark_todos_completed,
|
||||
) => {
|
||||
if mark_todos_completed.todo_ids.is_empty() {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
} else {
|
||||
// This is a mark as completed operation
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::todo_operation(
|
||||
MessageId::new(self.id),
|
||||
TodoOperation::MarkAsCompleted {
|
||||
completed_todos: mark_todos_completed
|
||||
.todo_ids
|
||||
.into_iter()
|
||||
.filter_map(|todo_id| {
|
||||
params.current_todo_list.and_then(|todo_list| {
|
||||
todo_list
|
||||
.completed_items()
|
||||
.iter()
|
||||
.find(|item| {
|
||||
item.id.as_ref() == todo_id.as_str()
|
||||
})
|
||||
.cloned()
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
},
|
||||
)
|
||||
.with_citations(citations),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
}
|
||||
api::message::Message::Summarization(summarization) => {
|
||||
let duration = summarization
|
||||
.finished_duration
|
||||
.map(|d| Duration::from_secs(d.seconds as u64));
|
||||
let (text, summarization_type, token_count) = match summarization.summary_type {
|
||||
Some(api::message::summarization::SummaryType::ConversationSummary(
|
||||
conv_summary,
|
||||
)) => {
|
||||
let token_count = if conv_summary.token_count > 0 {
|
||||
Some(conv_summary.token_count as u32)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let text = if !conv_summary.summary.is_empty() {
|
||||
AIAgentText {
|
||||
sections: parse_markdown_into_text_and_code_sections(
|
||||
&conv_summary.summary,
|
||||
),
|
||||
}
|
||||
} else {
|
||||
AIAgentText { sections: vec![] }
|
||||
};
|
||||
(text, SummarizationType::ConversationSummary, token_count)
|
||||
}
|
||||
Some(api::message::summarization::SummaryType::ToolCallResultSummary(_)) => (
|
||||
AIAgentText { sections: vec![] },
|
||||
SummarizationType::ToolCallResultSummary,
|
||||
None,
|
||||
),
|
||||
None => {
|
||||
// Default to ConversationSummary if not specified
|
||||
(
|
||||
AIAgentText { sections: vec![] },
|
||||
SummarizationType::ConversationSummary,
|
||||
None,
|
||||
)
|
||||
}
|
||||
};
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::summarization(
|
||||
MessageId::new(self.id),
|
||||
text,
|
||||
duration,
|
||||
summarization_type,
|
||||
token_count,
|
||||
),
|
||||
))
|
||||
}
|
||||
api::message::Message::UpdateReviewComments(update_comments) => {
|
||||
if let Some(operation) = update_comments.operation {
|
||||
match operation {
|
||||
api::message::update_review_comments::Operation::AddressReviewComments(
|
||||
address_comments,
|
||||
) => {
|
||||
if let Some(current_comments) = params.active_code_review {
|
||||
let addressed_comments = current_comments
|
||||
.addressed_comments
|
||||
.iter()
|
||||
.filter(|comment| {
|
||||
address_comments
|
||||
.comment_ids
|
||||
.iter()
|
||||
.any(|id| id == &comment.id.to_string())
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::comments_addressed(
|
||||
MessageId::new(self.id),
|
||||
addressed_comments,
|
||||
)
|
||||
.with_citations(citations),
|
||||
))
|
||||
} else {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
}
|
||||
api::message::Message::DebugOutput(debug_output) => {
|
||||
if ChannelState::enable_debug_features() {
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::debug_output(
|
||||
MessageId::new(self.id),
|
||||
debug_output.text,
|
||||
),
|
||||
))
|
||||
} else {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
}
|
||||
api::message::Message::ArtifactEvent(artifact_event) => match artifact_event.event {
|
||||
Some(api::message::artifact_event::Event::Created(artifact_created)) => {
|
||||
match artifact_created.artifact {
|
||||
Some(
|
||||
api::message::artifact_event::artifact_created::Artifact::PullRequest(
|
||||
pr,
|
||||
),
|
||||
) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::artifact_created(
|
||||
MessageId::new(self.id),
|
||||
ArtifactCreatedData::PullRequest {
|
||||
url: pr.url,
|
||||
branch: pr.branch,
|
||||
},
|
||||
)
|
||||
.with_citations(citations),
|
||||
)),
|
||||
Some(
|
||||
api::message::artifact_event::artifact_created::Artifact::Screenshot(
|
||||
screenshot,
|
||||
),
|
||||
) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::artifact_created(
|
||||
MessageId::new(self.id),
|
||||
ArtifactCreatedData::Screenshot {
|
||||
artifact_uid: screenshot.artifact_uid,
|
||||
mime_type: screenshot.mime_type,
|
||||
description: if screenshot.description.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(screenshot.description)
|
||||
},
|
||||
},
|
||||
)
|
||||
.with_citations(citations),
|
||||
)),
|
||||
Some(api::message::artifact_event::artifact_created::Artifact::File(
|
||||
file,
|
||||
)) => Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::artifact_created(
|
||||
MessageId::new(self.id),
|
||||
ArtifactCreatedData::File {
|
||||
artifact_uid: file.artifact_uid,
|
||||
filename: sanitized_basename(&file.filepath)
|
||||
.unwrap_or_else(|| file.filepath.clone()),
|
||||
filepath: file.filepath,
|
||||
mime_type: file.mime_type,
|
||||
description: if file.description.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(file.description)
|
||||
},
|
||||
size_bytes: file.size_bytes,
|
||||
},
|
||||
)
|
||||
.with_citations(citations),
|
||||
)),
|
||||
None => Ok(MaybeAIAgentOutputMessage::NoClientRepresentation),
|
||||
}
|
||||
}
|
||||
Some(api::message::artifact_event::Event::ForkArtifacts(_)) | None => {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
},
|
||||
api::message::Message::MessagesReceivedFromAgents(messages_received_from_agents) => {
|
||||
let messages = messages_received_from_agents
|
||||
.messages
|
||||
.into_iter()
|
||||
.map(|msg| crate::ai::agent::ReceivedMessageDisplay {
|
||||
message_id: msg.message_id,
|
||||
sender_agent_id: msg.sender_agent_id,
|
||||
addresses: msg.addresses,
|
||||
subject: msg.subject,
|
||||
message_body: msg.message_body,
|
||||
})
|
||||
.collect();
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::messages_received_from_agents(
|
||||
MessageId::new(self.id),
|
||||
messages,
|
||||
)
|
||||
.with_citations(citations),
|
||||
))
|
||||
}
|
||||
api::message::Message::EventsFromAgents(events) => {
|
||||
let event_ids = events
|
||||
.agent_events
|
||||
.iter()
|
||||
.map(|e| e.event_id.clone())
|
||||
.collect();
|
||||
Ok(MaybeAIAgentOutputMessage::Message(
|
||||
AIAgentOutputMessage::events_from_agents(MessageId::new(self.id), event_ids)
|
||||
.with_citations(citations),
|
||||
))
|
||||
}
|
||||
// These messages don't indicate an error but they don't translate to a client-side output message.
|
||||
api::message::Message::UserQuery(_)
|
||||
| api::message::Message::SystemQuery(_)
|
||||
| api::message::Message::ToolCallResult(_)
|
||||
| api::message::Message::CodeReview(_)
|
||||
| api::message::Message::ServerEvent(_)
|
||||
| api::message::Message::InvokeSkill(_)
|
||||
| api::message::Message::PassiveSuggestionResult(_) => {
|
||||
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<api::message::AgentOutput> for AIAgentText {
|
||||
fn from(value: api::message::AgentOutput) -> Self {
|
||||
AIAgentText {
|
||||
sections: parse_markdown_into_text_and_code_sections(value.text.as_str()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<api::message::AgentReasoning> for AIAgentText {
|
||||
fn from(value: api::message::AgentReasoning) -> Self {
|
||||
AIAgentText {
|
||||
sections: parse_markdown_into_text_and_code_sections(value.reasoning.as_str()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for converting an [`api::Message`] to an [`AIAgentOutputMessage`].
|
||||
trait ConvertAPIToolCallToAIAgentAction {
|
||||
fn to_action(
|
||||
self,
|
||||
params: ConversionParams,
|
||||
) -> Result<MaybeAIAgentAction, ToolToAIAgentActionError>;
|
||||
}
|
||||
|
||||
/// Trys to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
|
||||
///
|
||||
/// A [`Result::Error`] indicates an unexpected problem, while [`Ok(None)`]
|
||||
/// indicates a tool call that we aren't expected to parse.
|
||||
impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
|
||||
fn to_action(
|
||||
self,
|
||||
params: ConversionParams,
|
||||
) -> Result<MaybeAIAgentAction, ToolToAIAgentActionError> {
|
||||
let Some(tool) = self.tool else {
|
||||
return Err(ToolToAIAgentActionError::MissingTool);
|
||||
};
|
||||
|
||||
let create_standard_action = |action: AIAgentActionType| {
|
||||
Ok(MaybeAIAgentAction::Action(AIAgentAction {
|
||||
id: self.tool_call_id.clone().into(),
|
||||
task_id: params.task_id.clone(),
|
||||
action,
|
||||
requires_result: true,
|
||||
}))
|
||||
};
|
||||
|
||||
match tool {
|
||||
api::message::tool_call::Tool::RunShellCommand(run_shell_command) => {
|
||||
create_standard_action(run_shell_command.into())
|
||||
}
|
||||
api::message::tool_call::Tool::WriteToLongRunningShellCommand(
|
||||
write_to_long_running_shell_command,
|
||||
) => create_standard_action(write_to_long_running_shell_command.into()),
|
||||
api::message::tool_call::Tool::ReadFiles(read_files) => {
|
||||
create_standard_action(read_files.into())
|
||||
}
|
||||
api::message::tool_call::Tool::UploadFileArtifact(upload_file_artifact) => {
|
||||
create_standard_action(upload_file_artifact.try_into()?)
|
||||
}
|
||||
api::message::tool_call::Tool::SearchCodebase(search_codebase) => {
|
||||
create_standard_action(search_codebase.into())
|
||||
}
|
||||
api::message::tool_call::Tool::Grep(grep) => create_standard_action(grep.into()),
|
||||
#[allow(deprecated)]
|
||||
api::message::tool_call::Tool::FileGlob(glob) => create_standard_action(glob.into()),
|
||||
api::message::tool_call::Tool::FileGlobV2(glob) => create_standard_action(glob.into()),
|
||||
api::message::tool_call::Tool::ApplyFileDiffs(apply_file_diffs) => {
|
||||
create_standard_action(apply_file_diffs.into())
|
||||
}
|
||||
api::message::tool_call::Tool::ReadMcpResource(read_mcp_resource) => {
|
||||
create_standard_action(read_mcp_resource.into())
|
||||
}
|
||||
api::message::tool_call::Tool::CallMcpTool(call_mcp_tool) => {
|
||||
match call_mcp_tool.try_into() {
|
||||
Ok(call_mcp_tool_action) => create_standard_action(call_mcp_tool_action),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
api::message::tool_call::Tool::SuggestNewConversation(suggest_new_conversation) => {
|
||||
create_standard_action(suggest_new_conversation.into())
|
||||
}
|
||||
api::message::tool_call::Tool::SuggestPrompt(suggest_prompt) => {
|
||||
match suggest_prompt.try_into() {
|
||||
Ok(suggest_prompt_action) => create_standard_action(suggest_prompt_action),
|
||||
Err(_) => Ok(MaybeAIAgentAction::NoClientRepresentation),
|
||||
}
|
||||
}
|
||||
api::message::tool_call::Tool::OpenCodeReview(_) => {
|
||||
create_standard_action(AIAgentActionType::OpenCodeReview)
|
||||
}
|
||||
api::message::tool_call::Tool::InitProject(_) => {
|
||||
create_standard_action(AIAgentActionType::InitProject)
|
||||
}
|
||||
api::message::tool_call::Tool::ReadDocuments(read_documents) => {
|
||||
create_standard_action(read_documents.into())
|
||||
}
|
||||
api::message::tool_call::Tool::EditDocuments(edit_documents) => {
|
||||
create_standard_action(edit_documents.into())
|
||||
}
|
||||
api::message::tool_call::Tool::CreateDocuments(create_documents) => {
|
||||
create_standard_action(create_documents.into())
|
||||
}
|
||||
api::message::tool_call::Tool::ReadShellCommandOutput(read_shell_command_output) => {
|
||||
create_standard_action(read_shell_command_output.into())
|
||||
}
|
||||
api::message::tool_call::Tool::TransferShellCommandControlToUser(
|
||||
transfer_shell_command_control_to_user,
|
||||
) => create_standard_action(transfer_shell_command_control_to_user.into()),
|
||||
api::message::tool_call::Tool::UseComputer(use_computer) => {
|
||||
create_standard_action(use_computer.try_into()?)
|
||||
}
|
||||
api::message::tool_call::Tool::RequestComputerUse(request_computer_use) => {
|
||||
create_standard_action(request_computer_use.into())
|
||||
}
|
||||
api::message::tool_call::Tool::Subagent(subagent) => {
|
||||
use api::message::tool_call::subagent::Metadata;
|
||||
let subagent_type = match subagent.metadata {
|
||||
Some(Metadata::Cli(_)) => SubagentType::Cli,
|
||||
Some(Metadata::Research(_)) => SubagentType::Research,
|
||||
Some(Metadata::Advice(_)) => SubagentType::Advice,
|
||||
Some(Metadata::ComputerUse(_)) => SubagentType::ComputerUse,
|
||||
Some(Metadata::Summarization(_)) => SubagentType::Summarization,
|
||||
Some(Metadata::ConversationSearch(cs_meta)) => {
|
||||
let query = if cs_meta.query.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(cs_meta.query)
|
||||
};
|
||||
let conversation_id = if cs_meta.conversation_id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(cs_meta.conversation_id)
|
||||
};
|
||||
SubagentType::ConversationSearch {
|
||||
query,
|
||||
conversation_id,
|
||||
}
|
||||
}
|
||||
Some(Metadata::WarpDocumentationSearch(_)) => {
|
||||
SubagentType::WarpDocumentationSearch
|
||||
}
|
||||
None => SubagentType::Unknown,
|
||||
};
|
||||
Ok(MaybeAIAgentAction::Subagent(SubagentCall {
|
||||
task_id: subagent.task_id,
|
||||
subagent_type,
|
||||
}))
|
||||
}
|
||||
api::message::tool_call::Tool::StartAgent(start_agent) => {
|
||||
create_standard_action(AIAgentActionType::StartAgent {
|
||||
version: StartAgentVersion::V1,
|
||||
name: start_agent.name,
|
||||
prompt: start_agent.prompt,
|
||||
execution_mode: convert_start_agent_execution_mode(start_agent.execution_mode),
|
||||
lifecycle_subscription: start_agent.lifecycle_subscription.map(
|
||||
|subscription| {
|
||||
subscription
|
||||
.event_types
|
||||
.into_iter()
|
||||
.filter_map(convert_start_agent_lifecycle_event_type)
|
||||
.collect()
|
||||
},
|
||||
),
|
||||
})
|
||||
}
|
||||
api::message::tool_call::Tool::StartAgentV2(start_agent) => {
|
||||
create_standard_action(AIAgentActionType::StartAgent {
|
||||
version: StartAgentVersion::V2,
|
||||
name: start_agent.name,
|
||||
prompt: start_agent.prompt,
|
||||
execution_mode: convert_start_agent_v2_execution_mode(
|
||||
start_agent.execution_mode,
|
||||
),
|
||||
lifecycle_subscription: start_agent.lifecycle_subscription.map(
|
||||
|subscription| {
|
||||
subscription
|
||||
.event_types
|
||||
.into_iter()
|
||||
.filter_map(convert_start_agent_lifecycle_event_type)
|
||||
.collect()
|
||||
},
|
||||
),
|
||||
})
|
||||
}
|
||||
api::message::tool_call::Tool::SendMessageToAgent(send_message) => {
|
||||
create_standard_action(AIAgentActionType::SendMessageToAgent {
|
||||
addresses: send_message.addresses,
|
||||
subject: send_message.subject,
|
||||
message: send_message.message,
|
||||
})
|
||||
}
|
||||
api::message::tool_call::Tool::InsertReviewComments(insert_review_comments) => {
|
||||
create_standard_action(insert_review_comments.into())
|
||||
}
|
||||
api::message::tool_call::Tool::ReadSkill(read_skill) => {
|
||||
create_standard_action(read_skill.try_into()?)
|
||||
}
|
||||
api::message::tool_call::Tool::FetchConversation(fetch_conversation) => {
|
||||
create_standard_action(fetch_conversation.into())
|
||||
}
|
||||
api::message::tool_call::Tool::AskUserQuestion(ask) => {
|
||||
let questions = ask
|
||||
.questions
|
||||
.into_iter()
|
||||
.filter_map(convert_api_question)
|
||||
.collect();
|
||||
create_standard_action(AIAgentActionType::AskUserQuestion { questions })
|
||||
}
|
||||
// Clients do not need to know how to parse server tool-calls but receiving
|
||||
// them is not an error.
|
||||
api::message::tool_call::Tool::Server(_) => {
|
||||
Ok(MaybeAIAgentAction::NoClientRepresentation)
|
||||
}
|
||||
_ => Err(ToolToAIAgentActionError::UnexpectedTool),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<api::Suggestions> for Suggestions {
|
||||
fn from(api_suggestions: api::Suggestions) -> Self {
|
||||
Self {
|
||||
rules: api_suggestions
|
||||
.rules
|
||||
.into_iter()
|
||||
.map(|rule| SuggestedRule {
|
||||
name: rule.name,
|
||||
content: rule.content,
|
||||
logging_id: rule.logging_id.into(),
|
||||
})
|
||||
.collect(),
|
||||
agent_mode_workflows: api_suggestions
|
||||
.workflows
|
||||
.into_iter()
|
||||
.map(|workflow| SuggestedAgentModeWorkflow {
|
||||
name: workflow.name,
|
||||
prompt: workflow.prompt,
|
||||
logging_id: workflow.logging_id.into(),
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<api::TodoItem> for AIAgentTodo {
|
||||
fn from(value: api::TodoItem) -> Self {
|
||||
AIAgentTodo {
|
||||
id: value.id.into(),
|
||||
title: value.title,
|
||||
description: value.description,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reconstruct user inputs from the provided server messages
|
||||
/// (for use in shared agent exchanges where the input was not provided in this session)
|
||||
pub fn user_inputs_from_messages(messages: &[api::Message]) -> Vec<AIAgentInput> {
|
||||
let mut inputs = Vec::new();
|
||||
let mut document_versions: HashMap<AIDocumentId, AIDocumentVersion> = HashMap::new();
|
||||
for m in messages {
|
||||
let Some(inner) = &m.message else { continue };
|
||||
match inner {
|
||||
api::message::Message::UserQuery(uq) => {
|
||||
let context = convert_input_context(uq.context.as_ref());
|
||||
let referenced_attachments = uq
|
||||
.referenced_attachments
|
||||
.iter()
|
||||
.filter_map(|(key, attachment)| {
|
||||
AIAgentAttachment::try_from(attachment.clone())
|
||||
.ok()
|
||||
.map(|a| (key.clone(), a))
|
||||
})
|
||||
.collect();
|
||||
inputs.push(AIAgentInput::UserQuery {
|
||||
query: uq.query.clone(),
|
||||
context,
|
||||
static_query_type: None,
|
||||
referenced_attachments,
|
||||
user_query_mode: convert_user_query_mode(uq.mode.as_ref()),
|
||||
running_command: None,
|
||||
intended_agent: Some(uq.intended_agent()),
|
||||
});
|
||||
}
|
||||
api::message::Message::SystemQuery(sq) => {
|
||||
let ctx = convert_input_context(sq.context.as_ref());
|
||||
if let Some(t) = &sq.r#type {
|
||||
// These system queries appear as user inputs in ai blocks.
|
||||
match t {
|
||||
api::message::system_query::Type::CreateNewProject(p) => {
|
||||
inputs.push(AIAgentInput::CreateNewProject {
|
||||
query: p.query.clone(),
|
||||
context: ctx,
|
||||
});
|
||||
}
|
||||
api::message::system_query::Type::CloneRepository(p) => {
|
||||
inputs.push(AIAgentInput::CloneRepository {
|
||||
clone_repo_url: CloneRepositoryURL::new(p.url.clone()),
|
||||
context: ctx,
|
||||
});
|
||||
}
|
||||
api::message::system_query::Type::AutoCodeDiff(p) => {
|
||||
inputs.push(AIAgentInput::AutoCodeDiffQuery {
|
||||
query: p.query.clone(),
|
||||
context: ctx,
|
||||
});
|
||||
}
|
||||
api::message::system_query::Type::FetchReviewComments(fetch) => {
|
||||
inputs.push(AIAgentInput::FetchReviewComments {
|
||||
repo_path: fetch.repo_path.clone(),
|
||||
context: ctx,
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
api::message::Message::ToolCallResult(tcr) => {
|
||||
let task_id = TaskId::new(m.task_id.clone());
|
||||
if let Some(input) = convert_tool_call_result_to_input(
|
||||
&task_id,
|
||||
tcr,
|
||||
&HashMap::new(),
|
||||
&mut document_versions,
|
||||
) {
|
||||
inputs.push(input);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
inputs
|
||||
}
|
||||
|
||||
fn convert_api_question(
|
||||
q: api::ask_user_question::Question,
|
||||
) -> Option<ai::agent::action::AskUserQuestionItem> {
|
||||
let Some(QuestionType::MultipleChoice(mc)) = q.question_type else {
|
||||
return None;
|
||||
};
|
||||
|
||||
// Server sends -1 when there is no recommendation.
|
||||
let recommended_idx = usize::try_from(mc.recommended_option_index)
|
||||
.ok()
|
||||
.filter(|idx| *idx < mc.options.len());
|
||||
let options = mc
|
||||
.options
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, opt)| ai::agent::action::AskUserQuestionOption {
|
||||
label: opt.label.clone(),
|
||||
recommended: recommended_idx == Some(i),
|
||||
})
|
||||
.collect();
|
||||
Some(ai::agent::action::AskUserQuestionItem {
|
||||
question_id: q.question_id.clone(),
|
||||
question: q.question,
|
||||
question_type: ai::agent::action::AskUserQuestionType::MultipleChoice {
|
||||
is_multiselect: mc.is_multiselect,
|
||||
options,
|
||||
supports_other: mc.supports_other,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "convert_from_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,644 @@
|
||||
use super::{
|
||||
convert_api_question, ConversionParams, ConvertAPIMessageToClientOutputMessage,
|
||||
MaybeAIAgentOutputMessage,
|
||||
};
|
||||
use crate::ai::agent::task::TaskId;
|
||||
use crate::ai::agent::{
|
||||
AIAgentActionType, AIAgentOutputMessageType, LifecycleEventType, StartAgentExecutionMode,
|
||||
};
|
||||
use ai::agent::action::AskUserQuestionType;
|
||||
use ai::skills::SkillReference;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
fn start_agent_tool_call_message(
|
||||
name: &str,
|
||||
prompt: &str,
|
||||
execution_mode: Option<api::start_agent::ExecutionMode>,
|
||||
lifecycle_subscription_event_types: Option<Vec<i32>>,
|
||||
) -> api::Message {
|
||||
api::Message {
|
||||
id: "message-id".to_string(),
|
||||
task_id: "task-id".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: "tool-call-id".to_string(),
|
||||
tool: Some(api::message::tool_call::Tool::StartAgent(api::StartAgent {
|
||||
name: name.to_string(),
|
||||
prompt: prompt.to_string(),
|
||||
execution_mode,
|
||||
lifecycle_subscription: lifecycle_subscription_event_types
|
||||
.map(|event_types| api::start_agent::LifecycleSubscription { event_types }),
|
||||
})),
|
||||
})),
|
||||
request_id: "request-id".to_string(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn local_start_agent_v2_execution_mode(harness_type: &str) -> api::start_agent_v2::ExecutionMode {
|
||||
api::start_agent_v2::ExecutionMode {
|
||||
mode: Some(api::start_agent_v2::execution_mode::Mode::Local(
|
||||
api::start_agent_v2::execution_mode::Local {
|
||||
harness: Some(api::start_agent_v2::execution_mode::Harness {
|
||||
r#type: harness_type.to_string(),
|
||||
}),
|
||||
},
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn local_start_agent_v2_execution_mode_without_harness() -> api::start_agent_v2::ExecutionMode {
|
||||
api::start_agent_v2::ExecutionMode {
|
||||
mode: Some(api::start_agent_v2::execution_mode::Mode::Local(
|
||||
api::start_agent_v2::execution_mode::Local { harness: None },
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn start_agent_v2_tool_call_message(
|
||||
name: &str,
|
||||
prompt: &str,
|
||||
execution_mode: Option<api::start_agent_v2::ExecutionMode>,
|
||||
lifecycle_subscription_event_types: Option<Vec<i32>>,
|
||||
) -> api::Message {
|
||||
api::Message {
|
||||
id: "message-id".to_string(),
|
||||
task_id: "task-id".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: "tool-call-id".to_string(),
|
||||
tool: Some(api::message::tool_call::Tool::StartAgentV2(
|
||||
api::StartAgentV2 {
|
||||
name: name.to_string(),
|
||||
prompt: prompt.to_string(),
|
||||
execution_mode,
|
||||
lifecycle_subscription: lifecycle_subscription_event_types.map(|event_types| {
|
||||
api::start_agent_v2::LifecycleSubscription { event_types }
|
||||
}),
|
||||
},
|
||||
)),
|
||||
})),
|
||||
request_id: "request-id".to_string(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn upload_artifact_tool_call_message(path: &str, description: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: "message-id".to_string(),
|
||||
task_id: "task-id".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: "tool-call-id".to_string(),
|
||||
tool: Some(api::message::tool_call::Tool::UploadFileArtifact(
|
||||
api::UploadFileArtifact {
|
||||
file: Some(api::FilePathReference {
|
||||
file_path: path.to_string(),
|
||||
}),
|
||||
description: description.to_string(),
|
||||
},
|
||||
)),
|
||||
})),
|
||||
request_id: "request-id".to_string(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn remote_start_agent_v2_execution_mode(
|
||||
environment_id: &str,
|
||||
) -> api::start_agent_v2::ExecutionMode {
|
||||
api::start_agent_v2::ExecutionMode {
|
||||
mode: Some(api::start_agent_v2::execution_mode::Mode::Remote(
|
||||
api::start_agent_v2::execution_mode::Remote {
|
||||
environment_id: environment_id.to_string(),
|
||||
skills: vec![
|
||||
api::SkillRef {
|
||||
skill_reference: Some(api::skill_ref::SkillReference::Path(
|
||||
"/tmp/SKILL.md".to_string(),
|
||||
)),
|
||||
},
|
||||
api::SkillRef {
|
||||
skill_reference: Some(api::skill_ref::SkillReference::BundledSkillId(
|
||||
"review-comments".to_string(),
|
||||
)),
|
||||
},
|
||||
],
|
||||
model_id: "gpt-test".to_string(),
|
||||
computer_use_enabled: true,
|
||||
worker_host: "worker-host".to_string(),
|
||||
harness: Some(api::start_agent_v2::execution_mode::Harness {
|
||||
r#type: "claude-code".to_string(),
|
||||
}),
|
||||
title: "Remote child".to_string(),
|
||||
},
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn file_artifact_created_message(filepath: &str, description: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: "message-id".to_string(),
|
||||
task_id: "task-id".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ArtifactEvent(
|
||||
api::message::ArtifactEvent {
|
||||
event: Some(api::message::artifact_event::Event::Created(
|
||||
api::message::artifact_event::ArtifactCreated {
|
||||
artifact: Some(
|
||||
api::message::artifact_event::artifact_created::Artifact::File(
|
||||
api::message::artifact_event::FileArtifact {
|
||||
artifact_uid: "artifact-uid".to_string(),
|
||||
filepath: filepath.to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
size_bytes: 42,
|
||||
description: description.to_string(),
|
||||
},
|
||||
),
|
||||
),
|
||||
},
|
||||
)),
|
||||
},
|
||||
)),
|
||||
request_id: "request-id".to_string(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn remote_start_agent_execution_mode(environment_id: &str) -> api::start_agent::ExecutionMode {
|
||||
api::start_agent::ExecutionMode {
|
||||
mode: Some(api::start_agent::execution_mode::Mode::Remote(
|
||||
api::start_agent::execution_mode::Remote {
|
||||
environment_id: environment_id.to_string(),
|
||||
},
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_multiple_choice_question(
|
||||
recommended_option_index: i32,
|
||||
) -> api::ask_user_question::Question {
|
||||
api::ask_user_question::Question {
|
||||
question_id: "q1".to_string(),
|
||||
question: "Which option should we prefer?".to_string(),
|
||||
question_type: Some(
|
||||
api::ask_user_question::question::QuestionType::MultipleChoice(
|
||||
api::ask_user_question::MultipleChoice {
|
||||
is_multiselect: false,
|
||||
options: vec![
|
||||
api::ask_user_question::Option {
|
||||
label: "First".to_string(),
|
||||
},
|
||||
api::ask_user_question::Option {
|
||||
label: "Second".to_string(),
|
||||
},
|
||||
],
|
||||
recommended_option_index,
|
||||
supports_other: false,
|
||||
},
|
||||
),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_api_question_treats_negative_recommended_index_as_no_recommendation() {
|
||||
let converted = convert_api_question(build_multiple_choice_question(-1))
|
||||
.expect("multiple choice questions should convert");
|
||||
|
||||
let AskUserQuestionType::MultipleChoice { options, .. } = converted.question_type;
|
||||
assert_eq!(options.len(), 2);
|
||||
assert!(options.iter().all(|option| !option.recommended));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_api_question_uses_zero_based_recommended_index_when_present() {
|
||||
let converted = convert_api_question(build_multiple_choice_question(0))
|
||||
.expect("multiple choice questions should convert");
|
||||
|
||||
let AskUserQuestionType::MultipleChoice { options, .. } = converted.question_type;
|
||||
assert_eq!(options.len(), 2);
|
||||
assert!(options[0].recommended);
|
||||
assert!(!options[1].recommended);
|
||||
}
|
||||
|
||||
fn extract_start_agent_action(
|
||||
output: MaybeAIAgentOutputMessage,
|
||||
) -> (
|
||||
String,
|
||||
String,
|
||||
StartAgentExecutionMode,
|
||||
Option<Vec<LifecycleEventType>>,
|
||||
) {
|
||||
let MaybeAIAgentOutputMessage::Message(output_message) = output else {
|
||||
panic!("expected output message");
|
||||
};
|
||||
let AIAgentOutputMessageType::Action(action) = output_message.message else {
|
||||
panic!("expected action output message");
|
||||
};
|
||||
let AIAgentActionType::StartAgent {
|
||||
version: _,
|
||||
name,
|
||||
prompt,
|
||||
execution_mode,
|
||||
lifecycle_subscription,
|
||||
} = action.action
|
||||
else {
|
||||
panic!("expected StartAgent action");
|
||||
};
|
||||
(name, prompt, execution_mode, lifecycle_subscription)
|
||||
}
|
||||
|
||||
fn extract_upload_artifact_action(output: MaybeAIAgentOutputMessage) -> (String, Option<String>) {
|
||||
let MaybeAIAgentOutputMessage::Message(output_message) = output else {
|
||||
panic!("expected output message");
|
||||
};
|
||||
let AIAgentOutputMessageType::Action(action) = output_message.message else {
|
||||
panic!("expected action output message");
|
||||
};
|
||||
let AIAgentActionType::UploadArtifact(request) = action.action else {
|
||||
panic!("expected UploadArtifact action");
|
||||
};
|
||||
(request.file_path, request.description)
|
||||
}
|
||||
|
||||
fn extract_file_artifact_created(
|
||||
output: MaybeAIAgentOutputMessage,
|
||||
) -> (String, String, Option<String>, i64) {
|
||||
let MaybeAIAgentOutputMessage::Message(output_message) = output else {
|
||||
panic!("expected output message");
|
||||
};
|
||||
let AIAgentOutputMessageType::ArtifactCreated(artifact) = output_message.message else {
|
||||
panic!("expected artifact created output message");
|
||||
};
|
||||
let crate::ai::agent::ArtifactCreatedData::File {
|
||||
filepath,
|
||||
filename,
|
||||
description,
|
||||
size_bytes,
|
||||
..
|
||||
} = artifact
|
||||
else {
|
||||
panic!("expected file artifact created output message");
|
||||
};
|
||||
(filepath, filename, description, size_bytes)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_start_agent_tool_call_to_action_with_prompt() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message =
|
||||
start_agent_tool_call_message("Agent 1", "run tests and report failures", None, None);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
|
||||
assert_eq!(name, "Agent 1");
|
||||
assert_eq!(prompt, "run tests and report failures");
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
);
|
||||
assert_eq!(lifecycle_subscription, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_local_start_agent_v2_without_harness_type_to_defaults() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = start_agent_v2_tool_call_message(
|
||||
"Agent 7",
|
||||
"run in the default local harness",
|
||||
Some(local_start_agent_v2_execution_mode_without_harness()),
|
||||
None,
|
||||
);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
|
||||
assert_eq!(name, "Agent 7");
|
||||
assert_eq!(prompt, "run in the default local harness");
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
);
|
||||
assert_eq!(lifecycle_subscription, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_upload_artifact_tool_call_to_action() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = upload_artifact_tool_call_message(
|
||||
"/tmp/build/output.log",
|
||||
"Build output for the latest run",
|
||||
);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (file_path, description) = extract_upload_artifact_action(output);
|
||||
|
||||
assert_eq!(file_path, "/tmp/build/output.log");
|
||||
assert_eq!(
|
||||
description.as_deref(),
|
||||
Some("Build output for the latest run")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_file_artifact_created_message_with_filename() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message =
|
||||
file_artifact_created_message("outputs/report.txt", "Build output for the latest run");
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (filepath, filename, description, size_bytes) = extract_file_artifact_created(output);
|
||||
|
||||
assert_eq!(filepath, "outputs/report.txt");
|
||||
assert_eq!(filename, "report.txt");
|
||||
assert_eq!(
|
||||
description.as_deref(),
|
||||
Some("Build output for the latest run")
|
||||
);
|
||||
assert_eq!(size_bytes, 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_start_agent_tool_calls_with_different_prompt_lengths() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let partial_message = start_agent_tool_call_message("Agent 2", "run tests", None, None);
|
||||
let updated_message = start_agent_tool_call_message(
|
||||
"Agent 2",
|
||||
"run tests and then summarize failures",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
let partial_output = partial_message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("partial conversion should succeed");
|
||||
let updated_output = updated_message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("updated conversion should succeed");
|
||||
|
||||
let (_, partial_prompt, partial_execution_mode, _) = extract_start_agent_action(partial_output);
|
||||
let (_, updated_prompt, updated_execution_mode, _) = extract_start_agent_action(updated_output);
|
||||
|
||||
assert_eq!(partial_prompt, "run tests");
|
||||
assert_eq!(updated_prompt, "run tests and then summarize failures");
|
||||
assert_eq!(
|
||||
partial_execution_mode,
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
);
|
||||
assert_eq!(
|
||||
updated_execution_mode,
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_start_agent_with_explicit_empty_lifecycle_subscription() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = start_agent_tool_call_message("Agent 3", "run tests", None, Some(vec![]));
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (_, _, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
);
|
||||
assert_eq!(lifecycle_subscription, Some(vec![]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_start_agent_with_cancelled_and_blocked_lifecycle_subscription() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = start_agent_tool_call_message(
|
||||
"Agent 4",
|
||||
"wait for approval",
|
||||
None,
|
||||
Some(vec![
|
||||
api::LifecycleEventType::Cancelled as i32,
|
||||
api::LifecycleEventType::Blocked as i32,
|
||||
]),
|
||||
);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (_, _, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::local_with_defaults()
|
||||
);
|
||||
assert_eq!(
|
||||
lifecycle_subscription,
|
||||
Some(vec![
|
||||
LifecycleEventType::Cancelled,
|
||||
LifecycleEventType::Blocked
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_remote_start_agent_with_environment_id() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = start_agent_tool_call_message(
|
||||
"Agent 5",
|
||||
"run in the remote environment",
|
||||
Some(remote_start_agent_execution_mode("env-123")),
|
||||
None,
|
||||
);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
|
||||
assert_eq!(name, "Agent 5");
|
||||
assert_eq!(prompt, "run in the remote environment");
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::Remote {
|
||||
environment_id: "env-123".to_string(),
|
||||
skill_references: vec![],
|
||||
model_id: String::new(),
|
||||
computer_use_enabled: false,
|
||||
worker_host: String::new(),
|
||||
harness_type: String::new(),
|
||||
title: String::new(),
|
||||
}
|
||||
);
|
||||
assert_eq!(lifecycle_subscription, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_remote_start_agent_v2_with_skill_references() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = start_agent_v2_tool_call_message(
|
||||
"Agent 6",
|
||||
"run in the remote environment",
|
||||
Some(remote_start_agent_v2_execution_mode("env-123")),
|
||||
None,
|
||||
);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
|
||||
assert_eq!(name, "Agent 6");
|
||||
assert_eq!(prompt, "run in the remote environment");
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::Remote {
|
||||
environment_id: "env-123".to_string(),
|
||||
skill_references: vec![
|
||||
SkillReference::Path("/tmp/SKILL.md".into()),
|
||||
SkillReference::BundledSkillId("review-comments".to_string()),
|
||||
],
|
||||
model_id: "gpt-test".to_string(),
|
||||
computer_use_enabled: true,
|
||||
worker_host: "worker-host".to_string(),
|
||||
harness_type: "claude-code".to_string(),
|
||||
title: "Remote child".to_string(),
|
||||
}
|
||||
);
|
||||
assert_eq!(lifecycle_subscription, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_local_start_agent_v2_with_harness_type() {
|
||||
let task_id = TaskId::new("task-id".to_string());
|
||||
let message = start_agent_v2_tool_call_message(
|
||||
"Agent 6",
|
||||
"run in the local claude harness",
|
||||
Some(local_start_agent_v2_execution_mode("claude-code")),
|
||||
None,
|
||||
);
|
||||
|
||||
let output = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("conversion should succeed");
|
||||
|
||||
let (name, prompt, execution_mode, lifecycle_subscription) = extract_start_agent_action(output);
|
||||
|
||||
assert_eq!(name, "Agent 6");
|
||||
assert_eq!(prompt, "run in the local claude harness");
|
||||
assert_eq!(
|
||||
execution_mode,
|
||||
StartAgentExecutionMode::local_harness("claude-code".to_string())
|
||||
);
|
||||
assert_eq!(lifecycle_subscription, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transfer_control_tool_call_converts_to_action_message() {
|
||||
let task_id = TaskId::new("task".to_string());
|
||||
let reason = "Please finish the interactive flow".to_string();
|
||||
let message = api::Message {
|
||||
id: "message".to_string(),
|
||||
task_id: "task".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: "tool_call".to_string(),
|
||||
tool: Some(
|
||||
api::message::tool_call::Tool::TransferShellCommandControlToUser(
|
||||
api::message::tool_call::TransferShellCommandControlToUser {
|
||||
reason: reason.clone(),
|
||||
},
|
||||
),
|
||||
),
|
||||
})),
|
||||
request_id: "req".to_string(),
|
||||
timestamp: None,
|
||||
};
|
||||
|
||||
let converted = message
|
||||
.to_client_output_message(ConversionParams {
|
||||
task_id: &task_id,
|
||||
current_todo_list: None,
|
||||
active_code_review: None,
|
||||
})
|
||||
.expect("transfer-control conversion should succeed");
|
||||
|
||||
match converted {
|
||||
MaybeAIAgentOutputMessage::Message(output) => match output.message {
|
||||
AIAgentOutputMessageType::Action(action) => {
|
||||
assert_eq!(action.task_id, task_id);
|
||||
assert_eq!(
|
||||
action.action,
|
||||
AIAgentActionType::TransferShellCommandControlToUser { reason }
|
||||
);
|
||||
assert!(action.requires_result);
|
||||
}
|
||||
other => panic!("Expected action message, got {other:?}"),
|
||||
},
|
||||
MaybeAIAgentOutputMessage::NoClientRepresentation => {
|
||||
panic!("Expected transfer-control tool call to produce a client action")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,977 @@
|
||||
//! Conversions from application types to MAA API types.
|
||||
|
||||
use ai::agent::convert::ConvertToAPITypeError;
|
||||
use anyhow::anyhow;
|
||||
use chrono::{DateTime, Local, Timelike};
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::ai::{
|
||||
agent::{
|
||||
AIAgentActionResult, AIAgentActionResultType, AIAgentAttachment, AIAgentContext,
|
||||
AIAgentInput, DriveObjectPayload, MCPContext, PassiveSuggestionResultType,
|
||||
PassiveSuggestionTrigger, RunningCommand, StaticQueryType, Suggestions, UserQueryMode,
|
||||
},
|
||||
block_context::BlockContext,
|
||||
};
|
||||
|
||||
fn local_datetime_to_timestamp(timestamp: DateTime<Local>) -> prost_types::Timestamp {
|
||||
prost_types::Timestamp {
|
||||
seconds: timestamp.timestamp(),
|
||||
nanos: timestamp.timestamp_subsec_nanos() as i32,
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<StaticQueryType> for api::request::input::query_with_canned_response::Type {
|
||||
type Error = ConvertToAPITypeError;
|
||||
|
||||
fn try_from(value: StaticQueryType) -> Result<Self, Self::Error> {
|
||||
match value {
|
||||
StaticQueryType::Install => Ok(
|
||||
api::request::input::query_with_canned_response::Type::Install(
|
||||
api::request::input::query_with_canned_response::Install {},
|
||||
),
|
||||
),
|
||||
StaticQueryType::Code => {
|
||||
Ok(api::request::input::query_with_canned_response::Type::Code(
|
||||
api::request::input::query_with_canned_response::Code {},
|
||||
))
|
||||
}
|
||||
StaticQueryType::Deploy => Ok(
|
||||
api::request::input::query_with_canned_response::Type::Deploy(
|
||||
api::request::input::query_with_canned_response::Deploy {},
|
||||
),
|
||||
),
|
||||
StaticQueryType::SomethingElse => Ok(
|
||||
api::request::input::query_with_canned_response::Type::SomethingElse(
|
||||
api::request::input::query_with_canned_response::SomethingElse {},
|
||||
),
|
||||
),
|
||||
StaticQueryType::CustomOnboardingRequest => Ok(
|
||||
api::request::input::query_with_canned_response::Type::CustomOnboardingRequest(
|
||||
api::request::input::query_with_canned_response::CustomOnboardingRequest {},
|
||||
),
|
||||
),
|
||||
StaticQueryType::EvaluationSuite => {
|
||||
Err(anyhow::anyhow!("EvaluationSuite StaticQueryType not yet supported").into())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn convert_input(
|
||||
mut inputs: Vec<AIAgentInput>,
|
||||
) -> Result<api::request::Input, ConvertToAPITypeError> {
|
||||
if inputs.is_empty() {
|
||||
return Err(anyhow!("Attempted to send multi-agent request with no input").into());
|
||||
}
|
||||
let api_context = inputs
|
||||
.iter()
|
||||
.rev()
|
||||
.find_map(AIAgentInput::context)
|
||||
.map(convert_context);
|
||||
|
||||
let mut api_inputs = vec![];
|
||||
if inputs.len() == 1 {
|
||||
match inputs.pop().expect("Input exists.") {
|
||||
AIAgentInput::UserQuery {
|
||||
query,
|
||||
context,
|
||||
static_query_type: Some(query_type),
|
||||
..
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::QueryWithCannedResponse(
|
||||
api::request::input::QueryWithCannedResponse {
|
||||
query,
|
||||
r#type: Some(query_type.try_into()?),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::AutoCodeDiffQuery { query, context } => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::AutoCodeDiffQuery(
|
||||
api::request::input::AutoCodeDiffQuery { query },
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::ResumeConversation { context } => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::ResumeConversation(
|
||||
api::request::input::ResumeConversation {},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::InitProjectRules { context, .. } => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::InitProjectRules(
|
||||
api::request::input::InitProjectRules {},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::CreateEnvironment {
|
||||
context,
|
||||
repo_paths,
|
||||
..
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::CreateEnvironment(
|
||||
api::request::input::CreateEnvironment { repo_paths },
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::TriggerPassiveSuggestion {
|
||||
context,
|
||||
attachments,
|
||||
trigger,
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::GeneratePassiveSuggestions(
|
||||
api::request::input::GeneratePassiveSuggestions {
|
||||
attachments: attachments
|
||||
.into_iter()
|
||||
.map(|attachment| attachment.into())
|
||||
.collect(),
|
||||
trigger: Some(trigger.into()),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::CreateNewProject { query, context } => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::CreateNewProject(
|
||||
api::request::input::CreateNewProject { query },
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::CloneRepository {
|
||||
clone_repo_url,
|
||||
context,
|
||||
..
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::CloneRepository(
|
||||
api::request::input::CloneRepository {
|
||||
url: clone_repo_url.into_url(),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::CodeReview {
|
||||
context,
|
||||
review_comments,
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::CodeReview(
|
||||
api::request::input::CodeReview {
|
||||
operation: Some(
|
||||
api::request::input::code_review::Operation::InitialReviewComments(
|
||||
api::request::input::code_review::InitialReviewComments {
|
||||
review_comments: review_comments
|
||||
.comments
|
||||
.into_iter()
|
||||
.map(Into::into)
|
||||
.collect(),
|
||||
diff_set: Some(api::DiffSet {
|
||||
hunks: review_comments
|
||||
.diff_set
|
||||
.into_iter()
|
||||
.flat_map(|(file_path, hunks)| {
|
||||
hunks.into_iter().map(move |hunk| {
|
||||
hunk.convert_to_api(file_path.clone())
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
curr_ref: None,
|
||||
base_ref: None,
|
||||
}),
|
||||
},
|
||||
),
|
||||
),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::FetchReviewComments { repo_path, context } => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::FetchReviewComments(
|
||||
api::request::input::FetchReviewComments { repo_path },
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::SummarizeConversation { prompt } => {
|
||||
return Ok(api::request::Input {
|
||||
context: None,
|
||||
r#type: Some(api::request::input::Type::SummarizeConversation(
|
||||
api::request::input::SummarizeConversation {
|
||||
prompt: prompt.unwrap_or_default(),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::InvokeSkill {
|
||||
context,
|
||||
skill,
|
||||
user_query,
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::InvokeSkill(
|
||||
api::request::input::InvokeSkill {
|
||||
skill: Some(skill.into()),
|
||||
user_query: user_query.map(|user_query| {
|
||||
api::request::input::UserQuery {
|
||||
query: user_query.query,
|
||||
referenced_attachments: user_query
|
||||
.referenced_attachments
|
||||
.into_iter()
|
||||
.map(|(k, attachment)| (k, attachment.into()))
|
||||
.collect(),
|
||||
mode: None,
|
||||
intended_agent: Default::default(),
|
||||
}
|
||||
}),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
AIAgentInput::StartFromAmbientRunPrompt {
|
||||
ambient_run_id,
|
||||
context,
|
||||
runtime_skill,
|
||||
attachments_dir,
|
||||
} => {
|
||||
return Ok(api::request::Input {
|
||||
context: Some(convert_context(context.as_ref())),
|
||||
r#type: Some(api::request::input::Type::StartFromAmbientRunPrompt(
|
||||
api::request::input::StartFromAmbientRunPrompt {
|
||||
ambient_run_id,
|
||||
// Deprecated, we always resolve base_prompt from the stored task config.
|
||||
runtime_base_prompt: String::new(),
|
||||
|
||||
runtime_skill: runtime_skill.map(|skill| skill.into()),
|
||||
attachments_dir: attachments_dir.unwrap_or_default(),
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
other_input => match convert_input_to_user_input(other_input) {
|
||||
Ok(api_input) => api_inputs.push(api_input),
|
||||
Err(ConvertToAPITypeError::Ignore) => (),
|
||||
Err(e) => return Err(e),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
for input in inputs.into_iter() {
|
||||
match convert_input_to_user_input(input) {
|
||||
Ok(api_input) => api_inputs.push(api_input),
|
||||
Err(ConvertToAPITypeError::Ignore) => continue,
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(api::request::Input {
|
||||
context: api_context,
|
||||
r#type: Some(api::request::input::Type::UserInputs(
|
||||
api::request::input::UserInputs {
|
||||
inputs: api_inputs
|
||||
.into_iter()
|
||||
.map(|input| api::request::input::user_inputs::UserInput { input: Some(input) })
|
||||
.collect(),
|
||||
},
|
||||
)),
|
||||
})
|
||||
}
|
||||
|
||||
fn convert_input_to_user_input(
|
||||
input: AIAgentInput,
|
||||
) -> Result<api::request::input::user_inputs::user_input::Input, ConvertToAPITypeError> {
|
||||
match input {
|
||||
AIAgentInput::UserQuery {
|
||||
query,
|
||||
static_query_type: None,
|
||||
referenced_attachments,
|
||||
user_query_mode,
|
||||
running_command: None,
|
||||
intended_agent,
|
||||
..
|
||||
} => Ok(
|
||||
api::request::input::user_inputs::user_input::Input::UserQuery(
|
||||
api::request::input::UserQuery {
|
||||
query,
|
||||
referenced_attachments: referenced_attachments.into_iter().map(|(k, attachment)| (k, attachment.into())).collect(),
|
||||
mode: Some(user_query_mode.into()),
|
||||
intended_agent: intended_agent.map(|agent| agent.into()).unwrap_or_default(),
|
||||
},
|
||||
),
|
||||
),
|
||||
AIAgentInput::UserQuery {
|
||||
query,
|
||||
static_query_type: None,
|
||||
referenced_attachments,
|
||||
user_query_mode,
|
||||
running_command: Some(RunningCommand{
|
||||
command,
|
||||
block_id,
|
||||
grid_contents: output,
|
||||
cursor,
|
||||
requested_command_id,
|
||||
is_alt_screen_active,
|
||||
}),
|
||||
..
|
||||
} => {
|
||||
Ok(api::request::input::user_inputs::user_input::Input::CliAgentUserQuery(
|
||||
api::request::input::CliAgentUserQuery {
|
||||
user_query: Some(api::request::input::UserQuery {
|
||||
query,
|
||||
referenced_attachments: referenced_attachments.into_iter().map(|(k, attachment)| (k, attachment.into())).collect(),
|
||||
mode: Some(user_query_mode.into()),
|
||||
intended_agent: api::AgentType::Cli.into(),
|
||||
}),
|
||||
running_command: Some(api::RunningShellCommand{
|
||||
command,
|
||||
snapshot: Some(api::LongRunningShellCommandSnapshot {
|
||||
output,
|
||||
cursor,
|
||||
command_id: block_id.as_str().to_owned(),
|
||||
is_alt_screen_active,
|
||||
is_preempted: false,
|
||||
}),
|
||||
}),
|
||||
run_shell_command_tool_call_id: requested_command_id.map(|id| id.to_string()).unwrap_or_default(),
|
||||
}
|
||||
))
|
||||
}
|
||||
AIAgentInput::ActionResult { result, .. } => result.try_into(),
|
||||
AIAgentInput::MessagesReceivedFromAgents { messages } => Ok(
|
||||
api::request::input::user_inputs::user_input::Input::MessagesReceivedFromAgents(
|
||||
api::request::input::user_inputs::MessagesReceivedFromAgents {
|
||||
messages: messages
|
||||
.into_iter()
|
||||
.map(
|
||||
|msg| api::request::input::user_inputs::messages_received_from_agents::ReceivedMessage {
|
||||
message_id: msg.message_id,
|
||||
sender_agent_id: msg.sender_agent_id,
|
||||
addresses: msg.addresses,
|
||||
subject: msg.subject,
|
||||
message_body: msg.message_body,
|
||||
},
|
||||
)
|
||||
.collect(),
|
||||
},
|
||||
),
|
||||
),
|
||||
AIAgentInput::EventsFromAgents { events } => Ok(
|
||||
api::request::input::user_inputs::user_input::Input::EventsFromAgents(
|
||||
api::request::input::user_inputs::EventsFromAgents {
|
||||
agent_events: events,
|
||||
},
|
||||
),
|
||||
),
|
||||
AIAgentInput::PassiveSuggestionResult {
|
||||
trigger,
|
||||
suggestion,
|
||||
..
|
||||
} => {
|
||||
let api_trigger = match trigger {
|
||||
Some(PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger)) => Some(
|
||||
api::passive_suggestion_result_type::Trigger::ExecutedShellCommand(
|
||||
(*shell_trigger.executed_shell_command).into(),
|
||||
),
|
||||
),
|
||||
Some(PassiveSuggestionTrigger::AgentResponseCompleted { .. }) => Some(
|
||||
api::passive_suggestion_result_type::Trigger::AgentResponseCompleted(
|
||||
api::passive_suggestion_result_type::AgentResponseCompleted {},
|
||||
),
|
||||
),
|
||||
_ => None,
|
||||
};
|
||||
let api_suggestion = match suggestion {
|
||||
PassiveSuggestionResultType::Prompt { prompt } => Some(
|
||||
api::passive_suggestion_result_type::Suggestion::Prompt(
|
||||
api::passive_suggestion_result_type::Prompt { prompt },
|
||||
),
|
||||
),
|
||||
PassiveSuggestionResultType::CodeDiff {
|
||||
diffs,
|
||||
summary,
|
||||
accepted,
|
||||
} => Some(
|
||||
api::passive_suggestion_result_type::Suggestion::CodeDiff(
|
||||
api::passive_suggestion_result_type::CodeDiff {
|
||||
diffs: diffs
|
||||
.into_iter()
|
||||
.map(|d| api::passive_suggestion_result_type::code_diff::Diff {
|
||||
file_path: d.file_path,
|
||||
search: d.search,
|
||||
replace: d.replace,
|
||||
})
|
||||
.collect(),
|
||||
summary,
|
||||
accepted,
|
||||
},
|
||||
),
|
||||
),
|
||||
};
|
||||
Ok(
|
||||
api::request::input::user_inputs::user_input::Input::PassiveSuggestionResult(
|
||||
api::request::input::user_inputs::PassiveSuggestionResultInput {
|
||||
result: Some(api::PassiveSuggestionResultType {
|
||||
trigger: api_trigger,
|
||||
suggestion: api_suggestion,
|
||||
}),
|
||||
},
|
||||
),
|
||||
)
|
||||
}
|
||||
AIAgentInput::ResumeConversation { .. } => Err(ConvertToAPITypeError::Ignore),
|
||||
AIAgentInput::InitProjectRules { .. } => Err(ConvertToAPITypeError::Ignore),
|
||||
AIAgentInput::CodeReview { .. } => Err(ConvertToAPITypeError::Ignore),
|
||||
AIAgentInput::FetchReviewComments { .. } => Err(ConvertToAPITypeError::Ignore),
|
||||
AIAgentInput::CreateEnvironment { .. } => Err(ConvertToAPITypeError::Ignore),
|
||||
AIAgentInput::InvokeSkill { .. } => Err(ConvertToAPITypeError::Ignore),
|
||||
invalid_input => Err(anyhow!(
|
||||
"Cannot convert non user query or action result input into API UserInput: {invalid_input:?}"
|
||||
).into()),
|
||||
}
|
||||
}
|
||||
|
||||
impl From<PassiveSuggestionTrigger> for api::request::input::generate_passive_suggestions::Trigger {
|
||||
fn from(value: PassiveSuggestionTrigger) -> Self {
|
||||
match value {
|
||||
PassiveSuggestionTrigger::FilesChanged => {
|
||||
api::request::input::generate_passive_suggestions::Trigger::FilesChanged(())
|
||||
}
|
||||
PassiveSuggestionTrigger::CommandRun => {
|
||||
api::request::input::generate_passive_suggestions::Trigger::CommandRun(())
|
||||
}
|
||||
PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger) => {
|
||||
api::request::input::generate_passive_suggestions::Trigger::ShellCommandCompleted(
|
||||
api::request::input::generate_passive_suggestions::ShellCommandCompleted {
|
||||
executed_shell_command: Some(
|
||||
(*shell_trigger.executed_shell_command).into(),
|
||||
),
|
||||
relevant_files: shell_trigger
|
||||
.relevant_files
|
||||
.into_iter()
|
||||
.flat_map(|file| Vec::<api::AnyFileContent>::from(file).into_iter())
|
||||
.collect(),
|
||||
},
|
||||
)
|
||||
}
|
||||
PassiveSuggestionTrigger::AgentResponseCompleted { .. } => {
|
||||
api::request::input::generate_passive_suggestions::Trigger::AgentResponseCompleted(
|
||||
api::request::input::generate_passive_suggestions::AgentResponseCompleted {},
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<UserQueryMode> for warp_multi_agent_api::UserQueryMode {
|
||||
fn from(value: UserQueryMode) -> Self {
|
||||
match value {
|
||||
UserQueryMode::Normal => warp_multi_agent_api::UserQueryMode { r#type: None },
|
||||
UserQueryMode::Plan => warp_multi_agent_api::UserQueryMode {
|
||||
r#type: Some(warp_multi_agent_api::user_query_mode::Type::Plan(())),
|
||||
},
|
||||
UserQueryMode::Orchestrate => warp_multi_agent_api::UserQueryMode {
|
||||
r#type: Some(warp_multi_agent_api::user_query_mode::Type::Orchestrate(())),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AIAgentAttachment> for api::Attachment {
|
||||
fn from(attachment: AIAgentAttachment) -> Self {
|
||||
match attachment {
|
||||
AIAgentAttachment::PlainText(text) => api::Attachment {
|
||||
value: Some(api::attachment::Value::PlainText(text)),
|
||||
},
|
||||
AIAgentAttachment::Block(block) => api::Attachment {
|
||||
value: Some(api::attachment::Value::ExecutedShellCommand(block.into())),
|
||||
},
|
||||
AIAgentAttachment::DriveObject { uid, payload } => api::Attachment {
|
||||
value: Some(api::attachment::Value::DriveObject(api::DriveObject {
|
||||
uid,
|
||||
object_payload: payload.map(|p| match p {
|
||||
DriveObjectPayload::Workflow {
|
||||
name,
|
||||
description,
|
||||
command,
|
||||
} => api::drive_object::ObjectPayload::Workflow(api::Workflow {
|
||||
name,
|
||||
description,
|
||||
command,
|
||||
}),
|
||||
DriveObjectPayload::Notebook { title, content } => {
|
||||
api::drive_object::ObjectPayload::Notebook(api::Notebook {
|
||||
title,
|
||||
content,
|
||||
})
|
||||
}
|
||||
DriveObjectPayload::GenericStringObject {
|
||||
payload,
|
||||
object_type,
|
||||
} => api::drive_object::ObjectPayload::GenericStringObject(
|
||||
api::GenericStringObject {
|
||||
payload,
|
||||
object_type,
|
||||
},
|
||||
),
|
||||
}),
|
||||
})),
|
||||
},
|
||||
#[allow(deprecated)]
|
||||
AIAgentAttachment::DiffHunk {
|
||||
file_path,
|
||||
line_range,
|
||||
diff_content,
|
||||
lines_added,
|
||||
lines_removed,
|
||||
current,
|
||||
base,
|
||||
} => api::Attachment {
|
||||
value: Some(api::attachment::Value::DiffHunk(api::DiffHunk {
|
||||
file_path,
|
||||
line_range: Some(api::FileContentLineRange {
|
||||
start: line_range.start.as_usize() as u32,
|
||||
end: line_range.end.as_usize() as u32,
|
||||
}),
|
||||
diff_content,
|
||||
lines_added,
|
||||
lines_removed,
|
||||
current: current.map(Into::into),
|
||||
base: Some(base.into()),
|
||||
})),
|
||||
},
|
||||
AIAgentAttachment::DocumentContent {
|
||||
document_id,
|
||||
content,
|
||||
line_range,
|
||||
// TODO: Add attachment source to API
|
||||
..
|
||||
} => api::Attachment {
|
||||
value: Some(api::attachment::Value::DocumentContent(
|
||||
api::DocumentContent {
|
||||
document_id,
|
||||
content,
|
||||
line_range: line_range.map(|range| api::FileContentLineRange {
|
||||
start: range.start.as_usize() as u32,
|
||||
end: range.end.as_usize() as u32,
|
||||
}),
|
||||
},
|
||||
)),
|
||||
},
|
||||
AIAgentAttachment::DiffSet {
|
||||
file_diffs,
|
||||
current,
|
||||
base,
|
||||
} => api::Attachment {
|
||||
value: Some(api::attachment::Value::DiffSet(api::DiffSet {
|
||||
hunks: file_diffs
|
||||
.into_iter()
|
||||
.flat_map(|(file_path, hunks)| {
|
||||
hunks
|
||||
.into_iter()
|
||||
.map(move |hunk| hunk.convert_to_api(file_path.clone()))
|
||||
})
|
||||
.collect(),
|
||||
curr_ref: current.map(Into::into),
|
||||
base_ref: Some(base.into()),
|
||||
})),
|
||||
},
|
||||
AIAgentAttachment::FilePathReference { file_path, .. } => api::Attachment {
|
||||
value: Some(api::attachment::Value::FilePathReference(
|
||||
api::FilePathReference { file_path },
|
||||
)),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<AIAgentActionResult> for api::request::input::user_inputs::user_input::Input {
|
||||
type Error = ConvertToAPITypeError;
|
||||
|
||||
fn try_from(action_result: AIAgentActionResult) -> Result<Self, Self::Error> {
|
||||
let result = match action_result.result {
|
||||
AIAgentActionResultType::RequestCommandOutput(request_command_result) => {
|
||||
Some(request_command_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::WriteToLongRunningShellCommand(result) => {
|
||||
Some(result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::ReadFiles(read_files_result) => {
|
||||
Some(read_files_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::UploadArtifact(upload_artifact_result) => {
|
||||
Some(upload_artifact_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::SearchCodebase(search_codebase_result) => {
|
||||
Some(search_codebase_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::RequestFileEdits(request_file_edits_result) => {
|
||||
Some(request_file_edits_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::Grep(grep_result) => Some(grep_result.try_into()?),
|
||||
AIAgentActionResultType::FileGlob(file_glob_result) => {
|
||||
Some(file_glob_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::FileGlobV2(file_glob_result) => {
|
||||
Some(file_glob_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::ReadMCPResource(read_mcp_resource_result) => {
|
||||
Some(read_mcp_resource_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::CallMCPTool(call_mcp_tool_result) => {
|
||||
Some(call_mcp_tool_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::ReadSkill(read_skill_result) => {
|
||||
Some(read_skill_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::SuggestNewConversation(suggest_new_conversation_result) => {
|
||||
Some(suggest_new_conversation_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::SuggestPrompt(suggest_prompt_result) => {
|
||||
Some(suggest_prompt_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::OpenCodeReview => Some(
|
||||
warp_multi_agent_api::request::input::tool_call_result::Result::OpenCodeReview(
|
||||
warp_multi_agent_api::OpenCodeReviewResult {},
|
||||
),
|
||||
),
|
||||
AIAgentActionResultType::InsertReviewComments(insert_review_comments_result) => {
|
||||
Some(insert_review_comments_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::InitProject => Some(
|
||||
warp_multi_agent_api::request::input::tool_call_result::Result::InitProject(
|
||||
warp_multi_agent_api::InitProjectResult {},
|
||||
),
|
||||
),
|
||||
AIAgentActionResultType::ReadDocuments(read_documents_result) => {
|
||||
Some(read_documents_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::EditDocuments(edit_documents_result) => {
|
||||
Some(edit_documents_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::CreateDocuments(create_documents_result) => {
|
||||
Some(create_documents_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::ReadShellCommandOutput(read_shell_command_output_result) => {
|
||||
Some(read_shell_command_output_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::UseComputer(use_computer_result) => {
|
||||
Some(use_computer_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::RequestComputerUse(request_computer_use_result) => {
|
||||
Some(request_computer_use_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::FetchConversation(fetch_conversation_result) => {
|
||||
Some(fetch_conversation_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::StartAgent(start_agent_result) => {
|
||||
Some(start_agent_result.into())
|
||||
}
|
||||
AIAgentActionResultType::SendMessageToAgent(send_message_result) => {
|
||||
Some(send_message_result.into())
|
||||
}
|
||||
AIAgentActionResultType::TransferShellCommandControlToUser(transfer_control_result) => {
|
||||
Some(transfer_control_result.try_into()?)
|
||||
}
|
||||
AIAgentActionResultType::AskUserQuestion(ask_user_question_result) => {
|
||||
Some(ask_user_question_result.into())
|
||||
}
|
||||
};
|
||||
Ok(
|
||||
api::request::input::user_inputs::user_input::Input::ToolCallResult(
|
||||
api::request::input::ToolCallResult {
|
||||
tool_call_id: action_result.id.into(),
|
||||
result,
|
||||
},
|
||||
),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_context(context: &[AIAgentContext]) -> api::InputContext {
|
||||
let mut api_context = api::InputContext::default();
|
||||
for context in context.iter().cloned() {
|
||||
match context {
|
||||
AIAgentContext::Block(block) => {
|
||||
#[allow(deprecated)]
|
||||
api_context.executed_shell_commands.push((*block).into());
|
||||
}
|
||||
AIAgentContext::Directory {
|
||||
pwd,
|
||||
home_dir,
|
||||
are_file_symbols_indexed,
|
||||
} => {
|
||||
api_context.directory = Some(api::input_context::Directory {
|
||||
pwd: pwd.unwrap_or_default(),
|
||||
home: home_dir.unwrap_or_default(),
|
||||
pwd_file_symbols_indexed: are_file_symbols_indexed,
|
||||
});
|
||||
}
|
||||
AIAgentContext::SelectedText(text) => {
|
||||
api_context
|
||||
.selected_text
|
||||
.push(api::input_context::SelectedText { text });
|
||||
}
|
||||
AIAgentContext::ExecutionEnvironment(execution_ctx) => {
|
||||
api_context.shell = Some(api::input_context::Shell {
|
||||
name: execution_ctx.shell_name,
|
||||
version: execution_ctx.shell_version.unwrap_or_default(),
|
||||
});
|
||||
|
||||
if execution_ctx.os.category.is_none() && execution_ctx.os.distribution.is_none() {
|
||||
continue;
|
||||
}
|
||||
api_context.operating_system = Some(api::input_context::OperatingSystem {
|
||||
platform: execution_ctx.os.category.unwrap_or_default(),
|
||||
distribution: execution_ctx.os.distribution.unwrap_or_default(),
|
||||
});
|
||||
}
|
||||
AIAgentContext::CurrentTime { current_time } => {
|
||||
let utc_time = current_time.to_utc();
|
||||
api_context.current_time = Some(prost_types::Timestamp {
|
||||
seconds: utc_time.timestamp(),
|
||||
nanos: utc_time.nanosecond() as i32,
|
||||
});
|
||||
}
|
||||
AIAgentContext::Image(image_context) => {
|
||||
api_context.images.push(api::input_context::Image {
|
||||
data: image_context.data.into(),
|
||||
mime_type: image_context.mime_type,
|
||||
});
|
||||
}
|
||||
AIAgentContext::Codebase { path, name } => {
|
||||
api_context
|
||||
.codebases
|
||||
.push(api::input_context::Codebase { path, name });
|
||||
}
|
||||
AIAgentContext::ProjectRules {
|
||||
root_path,
|
||||
active_rules,
|
||||
additional_rule_paths,
|
||||
} => {
|
||||
api_context
|
||||
.project_rules
|
||||
.push(api::input_context::ProjectRules {
|
||||
root_path,
|
||||
active_rule_files: active_rules
|
||||
.into_iter()
|
||||
.flat_map(|rule| {
|
||||
let file_contents: Vec<api::FileContent> = rule.into();
|
||||
file_contents.into_iter()
|
||||
})
|
||||
.collect(),
|
||||
additional_rule_file_paths: additional_rule_paths,
|
||||
});
|
||||
}
|
||||
AIAgentContext::File(file_context) => {
|
||||
let contents: Vec<api::FileContent> = file_context.into();
|
||||
|
||||
for content in contents {
|
||||
api_context.files.push(api::input_context::File {
|
||||
content: Some(content),
|
||||
});
|
||||
}
|
||||
}
|
||||
AIAgentContext::Git { head, branch } => {
|
||||
api_context.git = Some(api::input_context::Git {
|
||||
head,
|
||||
branch: branch.unwrap_or_default(),
|
||||
});
|
||||
}
|
||||
AIAgentContext::Skills { skills } => {
|
||||
api_context.updated_skills_context = Some(api::input_context::SkillsContext {
|
||||
available_skills: skills
|
||||
.into_iter()
|
||||
.map(|skill| api::SkillDescriptor {
|
||||
skill_reference: Some(skill.reference.into()),
|
||||
name: skill.name,
|
||||
description: skill.description,
|
||||
provider: Some(skill.provider.into()),
|
||||
scope: Some(skill.scope.into()),
|
||||
})
|
||||
.collect(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
api_context
|
||||
}
|
||||
|
||||
impl From<Suggestions> for api::Suggestions {
|
||||
fn from(value: Suggestions) -> Self {
|
||||
Self {
|
||||
rules: value
|
||||
.rules
|
||||
.into_iter()
|
||||
.map(|rule| api::SuggestedRule {
|
||||
name: rule.name,
|
||||
content: rule.content,
|
||||
logging_id: rule.logging_id.to_string(),
|
||||
})
|
||||
.collect(),
|
||||
workflows: value
|
||||
.agent_mode_workflows
|
||||
.into_iter()
|
||||
.map(|workflow| api::SuggestedAgentModeWorkflow {
|
||||
name: workflow.name,
|
||||
prompt: workflow.prompt,
|
||||
logging_id: workflow.logging_id.to_string(),
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Convert rmcp resource to proto format.
|
||||
fn convert_mcp_resource(resource: rmcp::model::Resource) -> api::request::mcp_context::McpResource {
|
||||
let rmcp::model::RawResource {
|
||||
uri,
|
||||
name,
|
||||
description,
|
||||
mime_type,
|
||||
..
|
||||
} = resource.raw;
|
||||
api::request::mcp_context::McpResource {
|
||||
uri,
|
||||
name,
|
||||
description: description.unwrap_or_default(),
|
||||
mime_type: mime_type.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
// Convert rmcp tool to proto format, skipping tools with invalid schemas.
|
||||
fn convert_mcp_tool(tool: rmcp::model::Tool) -> Option<api::request::mcp_context::McpTool> {
|
||||
let Ok(prost_types::Value {
|
||||
kind: Some(prost_types::value::Kind::StructValue(input_schema)),
|
||||
}) = serde_json_to_prost(tool.input_schema.as_ref().clone().into())
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(api::request::mcp_context::McpTool {
|
||||
name: tool.name.to_string(),
|
||||
description: tool.description.map(|d| d.to_string()).unwrap_or_default(),
|
||||
input_schema: Some(input_schema),
|
||||
})
|
||||
}
|
||||
|
||||
impl From<MCPContext> for api::request::McpContext {
|
||||
#[allow(deprecated)]
|
||||
fn from(value: MCPContext) -> Self {
|
||||
// Check if we're using the old flat structure (no servers)
|
||||
// or the new grouped structure (servers populated)
|
||||
if value.servers.is_empty() {
|
||||
// Old behavior: use deprecated flat resources and tools lists
|
||||
api::request::McpContext {
|
||||
#[allow(deprecated)]
|
||||
resources: value
|
||||
.resources
|
||||
.into_iter()
|
||||
.map(convert_mcp_resource)
|
||||
.collect(),
|
||||
#[allow(deprecated)]
|
||||
tools: value
|
||||
.tools
|
||||
.into_iter()
|
||||
.filter_map(convert_mcp_tool)
|
||||
.collect(),
|
||||
servers: vec![], // Empty for old behavior
|
||||
}
|
||||
} else {
|
||||
// New behavior: group by server
|
||||
let servers: Vec<_> = value
|
||||
.servers
|
||||
.into_iter()
|
||||
.map(|server| api::request::mcp_context::McpServer {
|
||||
id: server.id,
|
||||
name: server.name,
|
||||
description: server.description,
|
||||
resources: server
|
||||
.resources
|
||||
.into_iter()
|
||||
.map(convert_mcp_resource)
|
||||
.collect(),
|
||||
tools: server
|
||||
.tools
|
||||
.into_iter()
|
||||
.filter_map(convert_mcp_tool)
|
||||
.collect(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
api::request::McpContext {
|
||||
#[allow(deprecated)]
|
||||
resources: vec![], // Empty - everything is grouped by server
|
||||
#[allow(deprecated)]
|
||||
tools: vec![], // Empty - everything is grouped by server
|
||||
servers,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BlockContext> for api::ExecutedShellCommand {
|
||||
fn from(block: BlockContext) -> Self {
|
||||
api::ExecutedShellCommand {
|
||||
command: block.command,
|
||||
output: block.output,
|
||||
exit_code: block.exit_code.value(),
|
||||
command_id: block.id.into(),
|
||||
is_auto_attached: block.is_auto_attached,
|
||||
started_ts: block.started_ts.map(local_datetime_to_timestamp),
|
||||
finished_ts: block.finished_ts.map(local_datetime_to_timestamp),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trys to convert a [`serde_json::Value`] to a [`prost_types::Value`].
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
fn serde_json_to_prost(value: serde_json::Value) -> Result<prost_types::Value, String> {
|
||||
use prost_types::value::Kind::*;
|
||||
use serde_json::Value::*;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
Ok(prost_types::Value {
|
||||
kind: Some(match value {
|
||||
Null => NullValue(0),
|
||||
Bool(v) => BoolValue(v),
|
||||
Number(n) => NumberValue(
|
||||
n.as_f64()
|
||||
.ok_or_else(|| format!("float {n} is not valid JSON number"))?,
|
||||
),
|
||||
String(s) => StringValue(s),
|
||||
Array(a) => ListValue(prost_types::ListValue {
|
||||
values: a
|
||||
.into_iter()
|
||||
.map(serde_json_to_prost)
|
||||
.collect::<Result<Vec<_>, std::string::String>>()?,
|
||||
}),
|
||||
Object(v) => StructValue(prost_types::Struct {
|
||||
fields: v
|
||||
.into_iter()
|
||||
.map(|(k, v)| serde_json_to_prost(v).map(|v| (k, v)))
|
||||
.collect::<Result<BTreeMap<_, _>, std::string::String>>()?,
|
||||
}),
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "convert_to_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,89 @@
|
||||
use crate::ai::agent::task::TaskId;
|
||||
use crate::ai::agent::{
|
||||
AIAgentActionResult, AIAgentActionResultType, TransferShellCommandControlToUserResult,
|
||||
};
|
||||
use crate::terminal::model::block::BlockId;
|
||||
use warp_core::command::ExitCode;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
#[test]
|
||||
fn transfer_control_snapshot_result_converts_to_tool_call_result_input() {
|
||||
let block_id = BlockId::default();
|
||||
let input =
|
||||
api::request::input::user_inputs::user_input::Input::try_from(AIAgentActionResult {
|
||||
id: "tool_call".to_string().into(),
|
||||
task_id: TaskId::new("task".to_string()),
|
||||
result: AIAgentActionResultType::TransferShellCommandControlToUser(
|
||||
TransferShellCommandControlToUserResult::Snapshot {
|
||||
block_id: block_id.clone(),
|
||||
grid_contents: "snapshot".to_string(),
|
||||
cursor: "<|cursor|>".to_string(),
|
||||
is_alt_screen_active: false,
|
||||
is_preempted: false,
|
||||
},
|
||||
),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
match input {
|
||||
api::request::input::user_inputs::user_input::Input::ToolCallResult(result) => {
|
||||
assert_eq!(result.tool_call_id, "tool_call");
|
||||
match result.result {
|
||||
Some(api::request::input::tool_call_result::Result::TransferShellCommandControlToUser(
|
||||
api_result,
|
||||
)) => match api_result.result {
|
||||
Some(
|
||||
api::transfer_shell_command_control_to_user_result::Result::LongRunningCommandSnapshot(snapshot),
|
||||
) => {
|
||||
assert_eq!(snapshot.command_id, block_id.to_string());
|
||||
assert_eq!(snapshot.output, "snapshot");
|
||||
assert_eq!(snapshot.cursor, "<|cursor|>");
|
||||
}
|
||||
other => panic!("Expected snapshot result, got {other:?}"),
|
||||
},
|
||||
other => panic!("Expected transfer-control tool call result, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("Expected tool-call-result input, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transfer_control_finished_result_converts_to_tool_call_result_input() {
|
||||
let block_id = BlockId::default();
|
||||
let input =
|
||||
api::request::input::user_inputs::user_input::Input::try_from(AIAgentActionResult {
|
||||
id: "tool_call".to_string().into(),
|
||||
task_id: TaskId::new("task".to_string()),
|
||||
result: AIAgentActionResultType::TransferShellCommandControlToUser(
|
||||
TransferShellCommandControlToUserResult::CommandFinished {
|
||||
block_id: block_id.clone(),
|
||||
output: "done".to_string(),
|
||||
exit_code: ExitCode::from(17),
|
||||
},
|
||||
),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
match input {
|
||||
api::request::input::user_inputs::user_input::Input::ToolCallResult(result) => {
|
||||
assert_eq!(result.tool_call_id, "tool_call");
|
||||
match result.result {
|
||||
Some(api::request::input::tool_call_result::Result::TransferShellCommandControlToUser(
|
||||
api_result,
|
||||
)) => match api_result.result {
|
||||
Some(
|
||||
api::transfer_shell_command_control_to_user_result::Result::CommandFinished(finished),
|
||||
) => {
|
||||
assert_eq!(finished.command_id, block_id.to_string());
|
||||
assert_eq!(finished.output, "done");
|
||||
assert_eq!(finished.exit_code, 17);
|
||||
}
|
||||
other => panic!("Expected command-finished result, got {other:?}"),
|
||||
},
|
||||
other => panic!("Expected transfer-control tool call result, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("Expected tool-call-result input, got {other:?}"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use crate::{ai::agent::redaction, terminal::model::session::SessionType};
|
||||
use futures_util::StreamExt;
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
use super::{convert_to::convert_input, ConvertToAPITypeError, RequestParams, ResponseStream};
|
||||
|
||||
pub async fn generate_multi_agent_output(
|
||||
server_api: Arc<ServerApi>,
|
||||
mut params: RequestParams,
|
||||
cancellation_rx: futures::channel::oneshot::Receiver<()>,
|
||||
) -> Result<ResponseStream, ConvertToAPITypeError> {
|
||||
let supported_tools = params
|
||||
.supported_tools_override
|
||||
.take()
|
||||
.unwrap_or_else(|| get_supported_tools(¶ms));
|
||||
let supported_cli_agent_tools = get_supported_cli_agent_tools(¶ms);
|
||||
let mut logging_metadata = HashMap::new();
|
||||
if let Some(metadata) = params.metadata {
|
||||
logging_metadata.insert(
|
||||
"is_autodetected_user_query".to_owned(),
|
||||
prost_types::Value {
|
||||
kind: Some(prost_types::value::Kind::BoolValue(
|
||||
metadata.is_autodetected_user_query,
|
||||
)),
|
||||
},
|
||||
);
|
||||
logging_metadata.insert(
|
||||
"entrypoint".to_owned(),
|
||||
prost_types::Value {
|
||||
kind: Some(prost_types::value::Kind::StringValue(
|
||||
metadata.entrypoint.entrypoint(),
|
||||
)),
|
||||
},
|
||||
);
|
||||
logging_metadata.insert(
|
||||
"is_auto_resume_after_error".to_owned(),
|
||||
prost_types::Value {
|
||||
kind: Some(prost_types::value::Kind::BoolValue(
|
||||
metadata.is_auto_resume_after_error,
|
||||
)),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
if params.should_redact_secrets {
|
||||
redaction::redact_inputs(&mut params.input);
|
||||
}
|
||||
|
||||
let mut api_keys = params.api_keys;
|
||||
if let Some(api_keys) = &mut api_keys {
|
||||
api_keys.allow_use_of_warp_credits = params.allow_use_of_warp_credits_with_byok;
|
||||
}
|
||||
|
||||
let request = api::Request {
|
||||
task_context: Some(api::request::TaskContext {
|
||||
tasks: params.tasks,
|
||||
}),
|
||||
input: Some(convert_input(params.input)?),
|
||||
settings: Some(api::request::Settings {
|
||||
model_config: Some(api::request::settings::ModelConfig {
|
||||
base: params.model.into(),
|
||||
cli_agent: params.cli_agent_model.into(),
|
||||
computer_use_agent: params.computer_use_model.into(),
|
||||
..Default::default()
|
||||
}),
|
||||
rules_enabled: params.is_memory_enabled,
|
||||
warp_drive_context_enabled: params.warp_drive_context_enabled,
|
||||
web_context_retrieval_enabled: true,
|
||||
supports_parallel_tool_calls: true,
|
||||
use_anthropic_text_editor_tools: false,
|
||||
planning_enabled: params.planning_enabled,
|
||||
supports_create_files: true,
|
||||
supported_tools: supported_tools.into_iter().map(Into::into).collect(),
|
||||
supports_long_running_commands: true,
|
||||
should_preserve_file_content_in_history: true,
|
||||
supports_todos_ui: true,
|
||||
supports_linked_code_blocks: FeatureFlag::LinkedCodeBlocks.is_enabled(),
|
||||
supports_started_child_task_message: true,
|
||||
supports_suggest_prompt: true,
|
||||
supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(),
|
||||
supports_reasoning_message: true,
|
||||
api_keys,
|
||||
autonomy_level: params.autonomy_level.into(),
|
||||
isolation_level: params.isolation_level.into(),
|
||||
web_search_enabled: params.web_search_enabled,
|
||||
supported_cli_agent_tools: supported_cli_agent_tools
|
||||
.into_iter()
|
||||
.map(Into::into)
|
||||
.collect(),
|
||||
supports_v4a_file_diffs: FeatureFlag::V4AFileDiffs.is_enabled(),
|
||||
supports_summarization_via_message_replacement:
|
||||
FeatureFlag::SummarizationViaMessageReplacement.is_enabled(),
|
||||
supports_bundled_skills: FeatureFlag::BundledSkills.is_enabled(),
|
||||
supports_research_agent: params.research_agent_enabled,
|
||||
supports_orchestration_v2: FeatureFlag::OrchestrationV2.is_enabled(),
|
||||
}),
|
||||
metadata: Some(api::request::Metadata {
|
||||
logging: logging_metadata,
|
||||
conversation_id: params
|
||||
.conversation_token
|
||||
.as_ref()
|
||||
.map(|token| token.as_str().to_string())
|
||||
.unwrap_or_default(),
|
||||
ambient_agent_task_id: params
|
||||
.ambient_agent_task_id
|
||||
.map(|id| id.to_string())
|
||||
.unwrap_or_default(),
|
||||
forked_from_conversation_id: if params.conversation_token.is_none() {
|
||||
// We only include this param on our initial request to the server
|
||||
// (when the forked conversation has not been asigned a new id yet).
|
||||
params
|
||||
.forked_from_conversation_token
|
||||
.map(|token| token.as_str().to_string())
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
String::new()
|
||||
},
|
||||
parent_agent_id: params.parent_agent_id.unwrap_or_default(),
|
||||
agent_name: params.agent_name.unwrap_or_default(),
|
||||
}),
|
||||
existing_suggestions: params
|
||||
.existing_suggestions
|
||||
.map(|suggestions| suggestions.into()),
|
||||
mcp_context: params.mcp_context.map(Into::into),
|
||||
};
|
||||
|
||||
let response_stream = server_api.generate_multi_agent_output(&request).await;
|
||||
match response_stream {
|
||||
Ok(stream) => {
|
||||
let output_stream = stream.take_until(cancellation_rx);
|
||||
Ok(Box::pin(output_stream))
|
||||
}
|
||||
Err(e) => {
|
||||
let (tx, rx) = async_channel::unbounded();
|
||||
let _ = tx.send(Err(e)).await;
|
||||
Ok(Box::pin(rx))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
|
||||
let mut supported_tools = vec![
|
||||
api::ToolType::Grep,
|
||||
api::ToolType::FileGlob,
|
||||
api::ToolType::FileGlobV2,
|
||||
api::ToolType::ReadMcpResource,
|
||||
api::ToolType::CallMcpTool,
|
||||
api::ToolType::InitProject,
|
||||
api::ToolType::OpenCodeReview,
|
||||
api::ToolType::RunShellCommand,
|
||||
api::ToolType::SuggestNewConversation,
|
||||
api::ToolType::Subagent,
|
||||
api::ToolType::WriteToLongRunningShellCommand,
|
||||
api::ToolType::ReadShellCommandOutput,
|
||||
api::ToolType::ReadDocuments,
|
||||
api::ToolType::CreateDocuments,
|
||||
api::ToolType::EditDocuments,
|
||||
api::ToolType::SuggestPrompt,
|
||||
];
|
||||
|
||||
if FeatureFlag::ConversationsAsContext.is_enabled() {
|
||||
supported_tools.push(api::ToolType::FetchConversation);
|
||||
}
|
||||
|
||||
match params.session_context.session_type() {
|
||||
None | Some(SessionType::Local) => {
|
||||
supported_tools.extend(&[
|
||||
api::ToolType::ReadFiles,
|
||||
api::ToolType::ApplyFileDiffs,
|
||||
api::ToolType::SearchCodebase,
|
||||
]);
|
||||
|
||||
if FeatureFlag::ArtifactCommand.is_enabled() {
|
||||
supported_tools.push(api::ToolType::UploadFileArtifact);
|
||||
}
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
|
||||
// Remote session with a known host — enable tools that route
|
||||
// through RemoteServerClient. The host_id is only populated
|
||||
// after a successful connection handshake, so its presence is a
|
||||
// sufficient proxy for client availability.
|
||||
// SearchCodebase remains disabled (follow-up work).
|
||||
supported_tools.extend(&[api::ToolType::ReadFiles, api::ToolType::ApplyFileDiffs]);
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: None }) => {
|
||||
// Feature flag off or not yet connected — no remote tools.
|
||||
}
|
||||
}
|
||||
|
||||
if FeatureFlag::AgentModeComputerUse.is_enabled() && params.computer_use_enabled {
|
||||
supported_tools.extend(&[api::ToolType::UseComputer]);
|
||||
supported_tools.extend(&[api::ToolType::RequestComputerUse])
|
||||
}
|
||||
|
||||
if FeatureFlag::PRCommentsSlashCommand.is_enabled() {
|
||||
supported_tools.push(api::ToolType::InsertReviewComments);
|
||||
}
|
||||
|
||||
if FeatureFlag::ListSkills.is_enabled() {
|
||||
supported_tools.push(api::ToolType::ReadSkill);
|
||||
}
|
||||
|
||||
if params.orchestration_enabled {
|
||||
supported_tools.push(if FeatureFlag::OrchestrationV2.is_enabled() {
|
||||
api::ToolType::StartAgentV2
|
||||
} else {
|
||||
api::ToolType::StartAgent
|
||||
});
|
||||
supported_tools.push(api::ToolType::SendMessageToAgent);
|
||||
}
|
||||
|
||||
if FeatureFlag::AskUserQuestion.is_enabled() && params.ask_user_question_enabled {
|
||||
supported_tools.push(api::ToolType::AskUserQuestion);
|
||||
}
|
||||
|
||||
supported_tools
|
||||
}
|
||||
|
||||
fn get_supported_cli_agent_tools(params: &RequestParams) -> Vec<api::ToolType> {
|
||||
let mut supported_cli_agent_tools = vec![
|
||||
api::ToolType::WriteToLongRunningShellCommand,
|
||||
api::ToolType::ReadShellCommandOutput,
|
||||
api::ToolType::Grep,
|
||||
api::ToolType::FileGlob,
|
||||
api::ToolType::FileGlobV2,
|
||||
];
|
||||
|
||||
if FeatureFlag::TransferControlTool.is_enabled() {
|
||||
supported_cli_agent_tools.push(api::ToolType::TransferShellCommandControlToUser);
|
||||
}
|
||||
|
||||
match params.session_context.session_type() {
|
||||
None | Some(SessionType::Local) => {
|
||||
supported_cli_agent_tools
|
||||
.extend(&[api::ToolType::ReadFiles, api::ToolType::SearchCodebase]);
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
|
||||
supported_cli_agent_tools.push(api::ToolType::ReadFiles);
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: None }) => {}
|
||||
}
|
||||
|
||||
supported_cli_agent_tools
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "impl_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,81 @@
|
||||
use crate::ai::agent::api::RequestParams;
|
||||
use crate::ai::blocklist::SessionContext;
|
||||
use crate::ai::llms::LLMId;
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use super::get_supported_tools;
|
||||
|
||||
fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool) -> RequestParams {
|
||||
let model = LLMId::from("test-model");
|
||||
|
||||
RequestParams {
|
||||
input: vec![],
|
||||
conversation_token: None,
|
||||
forked_from_conversation_token: None,
|
||||
ambient_agent_task_id: None,
|
||||
tasks: vec![],
|
||||
existing_suggestions: None,
|
||||
metadata: None,
|
||||
session_context: SessionContext::new_for_test(),
|
||||
model: model.clone(),
|
||||
coding_model: model.clone(),
|
||||
cli_agent_model: model.clone(),
|
||||
computer_use_model: model,
|
||||
is_memory_enabled: false,
|
||||
warp_drive_context_enabled: false,
|
||||
mcp_context: None,
|
||||
planning_enabled: true,
|
||||
should_redact_secrets: false,
|
||||
api_keys: None,
|
||||
allow_use_of_warp_credits_with_byok: false,
|
||||
autonomy_level: api::AutonomyLevel::Supervised,
|
||||
isolation_level: api::IsolationLevel::None,
|
||||
web_search_enabled: false,
|
||||
computer_use_enabled: false,
|
||||
ask_user_question_enabled,
|
||||
research_agent_enabled: false,
|
||||
orchestration_enabled: false,
|
||||
supported_tools_override: None,
|
||||
parent_agent_id: None,
|
||||
agent_name: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supported_tools_omits_ask_user_question_when_disabled() {
|
||||
let params = request_params_with_ask_user_question_enabled(false);
|
||||
let supported_tools = get_supported_tools(¶ms);
|
||||
|
||||
assert!(!supported_tools.contains(&api::ToolType::AskUserQuestion));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supported_tools_includes_ask_user_question_when_enabled_and_feature_flag_is_enabled() {
|
||||
if !FeatureFlag::AskUserQuestion.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let params = request_params_with_ask_user_question_enabled(true);
|
||||
let supported_tools = get_supported_tools(¶ms);
|
||||
|
||||
assert!(supported_tools.contains(&api::ToolType::AskUserQuestion));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supported_tools_include_upload_artifact_when_feature_flag_is_enabled() {
|
||||
let _flag = FeatureFlag::ArtifactCommand.override_enabled(true);
|
||||
let params = request_params_with_ask_user_question_enabled(false);
|
||||
let supported_tools = get_supported_tools(¶ms);
|
||||
|
||||
assert!(supported_tools.contains(&api::ToolType::UploadFileArtifact));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supported_tools_omit_upload_artifact_when_feature_flag_is_disabled() {
|
||||
let _flag = FeatureFlag::ArtifactCommand.override_enabled(false);
|
||||
let params = request_params_with_ask_user_question_enabled(false);
|
||||
let supported_tools = get_supported_tools(¶ms);
|
||||
|
||||
assert!(!supported_tools.contains(&api::ToolType::UploadFileArtifact));
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
use crate::code_review::comments::CommentId;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// The current state of a code review.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CodeReview {
|
||||
/// Comments that are currently pending (have yet to be addressed).
|
||||
pub pending_comments: Vec<ReviewComment>,
|
||||
/// Comments that have been addressed.
|
||||
pub addressed_comments: Vec<ReviewComment>,
|
||||
}
|
||||
|
||||
impl CodeReview {
|
||||
pub fn new_with_pending_comments(pending_comments: Vec<ReviewComment>) -> Self {
|
||||
Self {
|
||||
pending_comments,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub struct ReviewComment {
|
||||
pub id: CommentId,
|
||||
pub content: String,
|
||||
pub diff: ReviewDiff,
|
||||
pub head_title: Option<String>,
|
||||
}
|
||||
|
||||
impl ReviewComment {
|
||||
pub fn title(&self) -> String {
|
||||
match (&self.diff.file_path, self.diff.line_number) {
|
||||
(Some(file_path), Some(line_number)) => {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.unwrap_or("Invalid File Name");
|
||||
let display_line = line_number + 1;
|
||||
format!("{file_name}:{display_line}")
|
||||
}
|
||||
(Some(file_path), None) => {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.unwrap_or("Invalid File Name");
|
||||
file_name.to_string()
|
||||
}
|
||||
(None, _) => self
|
||||
.head_title
|
||||
.as_ref()
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "Review Comment".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crate::code_review::comments::AttachedReviewComment> for ReviewComment {
|
||||
fn from(comment: crate::code_review::comments::AttachedReviewComment) -> Self {
|
||||
let head_title = comment.head().map(|head| head.title());
|
||||
|
||||
ReviewComment {
|
||||
id: comment.id,
|
||||
content: comment.content,
|
||||
diff: comment.target.into(),
|
||||
head_title,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crate::code_review::comments::AttachedReviewCommentTarget> for ReviewDiff {
|
||||
fn from(val: crate::code_review::comments::AttachedReviewCommentTarget) -> Self {
|
||||
// Convert from the server format of a line number (which is zero indexed)
|
||||
// to one that is one-indexed to display within the blocklist.
|
||||
match val {
|
||||
crate::code_review::comments::AttachedReviewCommentTarget::Line {
|
||||
absolute_file_path,
|
||||
line,
|
||||
content: _,
|
||||
} => {
|
||||
let line_number = line
|
||||
.line_number()
|
||||
.map(|line_number| line_number.as_usize() + 1);
|
||||
Self {
|
||||
file_path: Some(absolute_file_path),
|
||||
line_number,
|
||||
}
|
||||
}
|
||||
crate::code_review::comments::AttachedReviewCommentTarget::File {
|
||||
absolute_file_path,
|
||||
} => Self {
|
||||
file_path: Some(absolute_file_path),
|
||||
line_number: None,
|
||||
},
|
||||
crate::code_review::comments::AttachedReviewCommentTarget::General => Self {
|
||||
file_path: None,
|
||||
line_number: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub struct ReviewDiff {
|
||||
pub file_path: Option<PathBuf>,
|
||||
pub line_number: Option<usize>,
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,199 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::{
|
||||
artifact_from_fork_proto, AIConversation, AIConversationAutoexecuteMode, AIConversationId,
|
||||
};
|
||||
use crate::ai::artifacts::Artifact;
|
||||
use crate::persistence::model::AgentConversationData;
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
fn restored_conversation(conversation_data: Option<AgentConversationData>) -> AIConversation {
|
||||
AIConversation::new_restored(
|
||||
AIConversationId::new(),
|
||||
vec![api::Task {
|
||||
id: "root-task".to_string(),
|
||||
messages: vec![],
|
||||
dependencies: None,
|
||||
description: String::new(),
|
||||
summary: String::new(),
|
||||
server_data: String::new(),
|
||||
}],
|
||||
conversation_data,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn user_query_message(id: &str, request_id: &str, query: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: "root-task".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::UserQuery(api::message::UserQuery {
|
||||
query: query.to_string(),
|
||||
context: None,
|
||||
referenced_attachments: HashMap::new(),
|
||||
mode: None,
|
||||
intended_agent: Default::default(),
|
||||
})),
|
||||
request_id: request_id.to_string(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn agent_output_message(id: &str, request_id: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: "root-task".to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::AgentOutput(
|
||||
api::message::AgentOutput {
|
||||
text: "Done".to_string(),
|
||||
},
|
||||
)),
|
||||
request_id: request_id.to_string(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn restored_conversation_with_queries(queries: &[&str]) -> AIConversation {
|
||||
let messages = queries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.flat_map(|(index, query)| {
|
||||
let request_id = format!("request-{index}");
|
||||
[
|
||||
user_query_message(&format!("user-{index}"), &request_id, query),
|
||||
agent_output_message(&format!("agent-{index}"), &request_id),
|
||||
]
|
||||
})
|
||||
.collect();
|
||||
|
||||
AIConversation::new_restored(
|
||||
AIConversationId::new(),
|
||||
vec![api::Task {
|
||||
id: "root-task".to_string(),
|
||||
messages,
|
||||
dependencies: None,
|
||||
description: String::new(),
|
||||
summary: String::new(),
|
||||
server_data: String::new(),
|
||||
}],
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn latest_user_query_returns_latest_non_empty_user_query() {
|
||||
let conversation =
|
||||
restored_conversation_with_queries(&["write unit tests", "fix the failing test"]);
|
||||
|
||||
assert_eq!(
|
||||
conversation.latest_user_query(),
|
||||
Some("fix the failing test".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn latest_user_query_trims_and_skips_empty_queries() {
|
||||
let conversation = restored_conversation_with_queries(&[" write unit tests ", " "]);
|
||||
|
||||
assert_eq!(
|
||||
conversation.latest_user_query(),
|
||||
Some("write unit tests".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restored_conversation_defaults_autoexecute_override_when_not_persisted() {
|
||||
let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
|
||||
let conversation_data: AgentConversationData =
|
||||
serde_json::from_str(r#"{"server_conversation_token":null}"#).unwrap();
|
||||
|
||||
let conversation = restored_conversation(Some(conversation_data));
|
||||
|
||||
assert_eq!(
|
||||
conversation.autoexecute_override(),
|
||||
AIConversationAutoexecuteMode::RespectUserSettings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restored_conversation_defaults_unknown_persisted_autoexecute_override() {
|
||||
let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
|
||||
let conversation_data: AgentConversationData = serde_json::from_str(
|
||||
r#"{"server_conversation_token":null,"autoexecute_override":"UnexpectedValue"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conversation = restored_conversation(Some(conversation_data));
|
||||
|
||||
assert_eq!(
|
||||
conversation.autoexecute_override(),
|
||||
AIConversationAutoexecuteMode::RespectUserSettings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restored_conversation_uses_persisted_autoexecute_override_when_enabled() {
|
||||
let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
|
||||
let conversation_data: AgentConversationData = serde_json::from_str(
|
||||
r#"{"server_conversation_token":null,"autoexecute_override":"RunToCompletion"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conversation = restored_conversation(Some(conversation_data));
|
||||
|
||||
assert_eq!(
|
||||
conversation.autoexecute_override(),
|
||||
AIConversationAutoexecuteMode::RunToCompletion
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restored_conversation_ignores_persisted_autoexecute_override_when_disabled() {
|
||||
let _flag = FeatureFlag::RememberFastForwardState.override_enabled(false);
|
||||
let conversation_data: AgentConversationData = serde_json::from_str(
|
||||
r#"{"server_conversation_token":null,"autoexecute_override":"RunToCompletion"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conversation = restored_conversation(Some(conversation_data));
|
||||
|
||||
assert_eq!(
|
||||
conversation.autoexecute_override(),
|
||||
AIConversationAutoexecuteMode::RespectUserSettings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fork_artifacts_adds_file_artifacts_to_conversation() {
|
||||
let proto_artifact = api::message::artifact_event::ConversationArtifact {
|
||||
artifact: Some(
|
||||
api::message::artifact_event::conversation_artifact::Artifact::File(
|
||||
api::message::artifact_event::FileArtifact {
|
||||
artifact_uid: "artifact-file-1".to_string(),
|
||||
filepath: "outputs/report.txt".to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
size_bytes: 42,
|
||||
description: "Daily summary".to_string(),
|
||||
},
|
||||
),
|
||||
),
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
artifact_from_fork_proto(&proto_artifact),
|
||||
Some(Artifact::File {
|
||||
artifact_uid: "artifact-file-1".to_string(),
|
||||
filepath: "outputs/report.txt".to_string(),
|
||||
filename: "report.txt".to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
description: Some("Daily summary".to_string()),
|
||||
size_bytes: Some(42),
|
||||
})
|
||||
);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,451 @@
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::test_util::ai_agent_tasks::{
|
||||
create_api_subtask, create_api_task, create_message, create_subagent_tool_call_message,
|
||||
};
|
||||
|
||||
use super::{base_dir, materialize_tasks_to_yaml};
|
||||
|
||||
/// Lists filenames (not full paths) in a directory, sorted.
|
||||
fn list_dir_sorted(dir: &Path) -> Vec<String> {
|
||||
let mut entries: Vec<String> = fs::read_dir(dir)
|
||||
.unwrap()
|
||||
.filter_map(|e| e.ok())
|
||||
.map(|e| e.file_name().to_string_lossy().into_owned())
|
||||
.collect();
|
||||
entries.sort();
|
||||
entries
|
||||
}
|
||||
|
||||
fn make_user_query_message(id: &str, task_id: &str, query: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::UserQuery(api::message::UserQuery {
|
||||
query: query.to_string(),
|
||||
context: None,
|
||||
mode: None,
|
||||
referenced_attachments: Default::default(),
|
||||
intended_agent: Default::default(),
|
||||
})),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_tool_call_message(
|
||||
id: &str,
|
||||
task_id: &str,
|
||||
tool_call_id: &str,
|
||||
tool: api::message::tool_call::Tool,
|
||||
) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
tool: Some(tool),
|
||||
})),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_tool_call_result_message(
|
||||
id: &str,
|
||||
task_id: &str,
|
||||
tool_call_id: &str,
|
||||
result: api::message::tool_call_result::Result,
|
||||
) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCallResult(
|
||||
api::message::ToolCallResult {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
result: Some(result),
|
||||
context: None,
|
||||
},
|
||||
)),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn cleanup_dir(path: &str) {
|
||||
let _ = fs::remove_dir_all(path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_message_types_produce_sequentially_indexed_files() {
|
||||
let task_id = "root";
|
||||
let tasks = vec![create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
make_user_query_message("m1", task_id, "hello"),
|
||||
// AgentOutput via create_message helper
|
||||
create_message("m2", task_id),
|
||||
make_tool_call_message(
|
||||
"m3",
|
||||
task_id,
|
||||
"tc1",
|
||||
api::message::tool_call::Tool::Grep(api::message::tool_call::Grep {
|
||||
queries: vec!["foo".into()],
|
||||
path: "/src".into(),
|
||||
}),
|
||||
),
|
||||
],
|
||||
)];
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&tasks).unwrap();
|
||||
assert!(
|
||||
Path::new(&dir).starts_with(base_dir()),
|
||||
"returned path should be under temp_dir(), got: {dir}",
|
||||
);
|
||||
// Verify no mixed separators: on Windows the path should use only '\',
|
||||
// on Unix only '/'. This catches the original bug where tempdir_in
|
||||
// joined a forward-slash parent with a native backslash separator.
|
||||
assert!(
|
||||
!dir.contains('/') || !dir.contains('\\'),
|
||||
"returned path has mixed separators: {dir}",
|
||||
);
|
||||
let files = list_dir_sorted(Path::new(&dir));
|
||||
|
||||
assert_eq!(files.len(), 3);
|
||||
assert!(files[0].starts_with("000.m1.user_query"));
|
||||
assert!(files[1].starts_with("001.m2.agent_output"));
|
||||
assert!(files[2].starts_with("002.m3.tool_call.tc1.grep"));
|
||||
|
||||
// Verify user_query content is searchable.
|
||||
let content = fs::read_to_string(Path::new(&dir).join(&files[0])).unwrap();
|
||||
assert!(content.contains("type: user_query"));
|
||||
assert!(content.contains("hello"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subagent_file_and_subdirectory_share_same_index() {
|
||||
let root_id = "root";
|
||||
let subtask_id = "subtask1";
|
||||
|
||||
let root_task = create_api_task(
|
||||
root_id,
|
||||
vec![
|
||||
make_user_query_message("m1", root_id, "search my conversation"),
|
||||
create_subagent_tool_call_message(
|
||||
"m2",
|
||||
root_id,
|
||||
subtask_id,
|
||||
Some(
|
||||
api::message::tool_call::subagent::Metadata::ConversationSearch(
|
||||
Default::default(),
|
||||
),
|
||||
),
|
||||
),
|
||||
],
|
||||
);
|
||||
let subtask = create_api_subtask(
|
||||
subtask_id,
|
||||
root_id,
|
||||
vec![create_message("sub_m1", subtask_id)],
|
||||
);
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&[root_task, subtask]).unwrap();
|
||||
let entries = list_dir_sorted(Path::new(&dir));
|
||||
|
||||
// Should have: 000.m1.user_query.yaml, 001.m2.subagent.*.yaml, 001.subtask1/ (directory)
|
||||
assert_eq!(entries.len(), 3);
|
||||
|
||||
// The subagent YAML file and its subdirectory must share the same "001" prefix.
|
||||
let subagent_file = entries
|
||||
.iter()
|
||||
.find(|e| e.contains("subagent") && e.ends_with(".yaml"))
|
||||
.expect("should have subagent yaml file");
|
||||
let subdir = entries
|
||||
.iter()
|
||||
.find(|e| e.contains(subtask_id) && !e.ends_with(".yaml"))
|
||||
.expect("should have subtask directory");
|
||||
|
||||
let file_prefix: String = subagent_file.chars().take(3).collect();
|
||||
let dir_prefix: String = subdir.chars().take(3).collect();
|
||||
assert_eq!(
|
||||
file_prefix, dir_prefix,
|
||||
"subagent file ({subagent_file}) and directory ({subdir}) must share the same index prefix"
|
||||
);
|
||||
assert_eq!(file_prefix, "001");
|
||||
|
||||
// Verify subtask directory contains the subtask's messages.
|
||||
let sub_entries = list_dir_sorted(&Path::new(&dir).join(subdir));
|
||||
assert_eq!(sub_entries.len(), 1);
|
||||
assert!(sub_entries[0].contains("sub_m1"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_subtask_in_task_map_produces_file_but_no_directory() {
|
||||
let root_id = "root";
|
||||
|
||||
// Subagent references subtask "missing_task" which is not in the task list.
|
||||
let root_task = create_api_task(
|
||||
root_id,
|
||||
vec![create_subagent_tool_call_message(
|
||||
"m1",
|
||||
root_id,
|
||||
"missing_task",
|
||||
Some(api::message::tool_call::subagent::Metadata::Cli(
|
||||
Default::default(),
|
||||
)),
|
||||
)],
|
||||
);
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&[root_task]).unwrap();
|
||||
let entries = list_dir_sorted(Path::new(&dir));
|
||||
|
||||
// Should have just the YAML file, no subdirectory since the subtask is missing.
|
||||
assert_eq!(entries.len(), 1);
|
||||
assert!(entries[0].ends_with(".yaml"));
|
||||
assert!(entries[0].contains("subagent"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_task_list_returns_error() {
|
||||
let result = materialize_tasks_to_yaml(&[]);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("No root task found"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_call_result_resolves_tool_name_from_matching_call() {
|
||||
let task_id = "root";
|
||||
let tasks = vec![create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
make_tool_call_message(
|
||||
"m1",
|
||||
task_id,
|
||||
"tc1",
|
||||
api::message::tool_call::Tool::Grep(api::message::tool_call::Grep {
|
||||
queries: vec!["pattern".into()],
|
||||
path: "/src".into(),
|
||||
}),
|
||||
),
|
||||
make_tool_call_result_message(
|
||||
"m2",
|
||||
task_id,
|
||||
"tc1",
|
||||
api::message::tool_call_result::Result::Grep(api::GrepResult {
|
||||
result: Some(api::grep_result::Result::Success(
|
||||
api::grep_result::Success {
|
||||
matched_files: vec![api::grep_result::success::GrepFileMatch {
|
||||
file_path: "foo.rs".into(),
|
||||
matched_lines: vec![
|
||||
api::grep_result::success::grep_file_match::GrepLineMatch {
|
||||
line_number: 42,
|
||||
},
|
||||
],
|
||||
}],
|
||||
},
|
||||
)),
|
||||
}),
|
||||
),
|
||||
],
|
||||
)];
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&tasks).unwrap();
|
||||
let files = list_dir_sorted(Path::new(&dir));
|
||||
|
||||
assert_eq!(files.len(), 2);
|
||||
// The result file should contain "grep" in its name, resolved from the tool call.
|
||||
assert!(
|
||||
files[1].contains("grep"),
|
||||
"result filename should contain tool name 'grep', got: {}",
|
||||
files[1]
|
||||
);
|
||||
|
||||
// Verify line numbers are serialized.
|
||||
let content = fs::read_to_string(Path::new(&dir).join(&files[1])).unwrap();
|
||||
assert!(content.contains("foo.rs"), "should contain file path");
|
||||
assert!(content.contains("42"), "should contain line number");
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_tool_calls_are_skipped() {
|
||||
let task_id = "root";
|
||||
let tasks = vec![create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
make_user_query_message("m1", task_id, "hello"),
|
||||
make_tool_call_message(
|
||||
"m2",
|
||||
task_id,
|
||||
"tc_server",
|
||||
api::message::tool_call::Tool::Server(api::message::tool_call::Server {
|
||||
payload: String::new(),
|
||||
}),
|
||||
),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
)];
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&tasks).unwrap();
|
||||
let files = list_dir_sorted(Path::new(&dir));
|
||||
|
||||
// Server tool call should be skipped; only user_query and agent_output.
|
||||
assert_eq!(files.len(), 2);
|
||||
assert!(files[0].contains("user_query"));
|
||||
assert!(files[1].contains("agent_output"));
|
||||
// Index should still be sequential (000, 001) since server call was skipped.
|
||||
assert!(files[0].starts_with("000"));
|
||||
assert!(files[1].starts_with("001"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn start_agent_v2_tool_call_serializes_name_and_prompt() {
|
||||
let task_id = "root";
|
||||
let tasks = vec![create_api_task(
|
||||
task_id,
|
||||
vec![make_tool_call_message(
|
||||
"m1",
|
||||
task_id,
|
||||
"tc_start_agent_v2",
|
||||
api::message::tool_call::Tool::StartAgentV2(api::StartAgentV2 {
|
||||
name: "Remote child".to_string(),
|
||||
prompt: "Investigate the build failure".to_string(),
|
||||
execution_mode: None,
|
||||
lifecycle_subscription: None,
|
||||
}),
|
||||
)],
|
||||
)];
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&tasks).unwrap();
|
||||
let files = list_dir_sorted(Path::new(&dir));
|
||||
let content = fs::read_to_string(Path::new(&dir).join(&files[0])).unwrap();
|
||||
|
||||
assert!(content.contains("tool_name: start_agent"));
|
||||
assert!(content.contains("name: \"Remote child\""));
|
||||
assert!(content.contains("prompt: |"));
|
||||
assert!(content.contains("Investigate the build failure"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn start_agent_v2_tool_call_result_serializes_agent_id_and_error() {
|
||||
let task_id = "root";
|
||||
let tasks = vec![create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
make_tool_call_message(
|
||||
"m1",
|
||||
task_id,
|
||||
"tc_start_agent_v2",
|
||||
api::message::tool_call::Tool::StartAgentV2(api::StartAgentV2 {
|
||||
name: "Remote child".to_string(),
|
||||
prompt: "Investigate the build failure".to_string(),
|
||||
execution_mode: None,
|
||||
lifecycle_subscription: None,
|
||||
}),
|
||||
),
|
||||
make_tool_call_result_message(
|
||||
"m2",
|
||||
task_id,
|
||||
"tc_start_agent_v2",
|
||||
api::message::tool_call_result::Result::StartAgentV2(api::StartAgentV2Result {
|
||||
result: Some(api::start_agent_v2_result::Result::Success(
|
||||
api::start_agent_v2_result::Success {
|
||||
agent_id: "agent-123".to_string(),
|
||||
},
|
||||
)),
|
||||
}),
|
||||
),
|
||||
make_tool_call_result_message(
|
||||
"m3",
|
||||
task_id,
|
||||
"tc_start_agent_v2",
|
||||
api::message::tool_call_result::Result::StartAgentV2(api::StartAgentV2Result {
|
||||
result: Some(api::start_agent_v2_result::Result::Error(
|
||||
api::start_agent_v2_result::Error {
|
||||
error: "child failed".to_string(),
|
||||
},
|
||||
)),
|
||||
}),
|
||||
),
|
||||
],
|
||||
)];
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&tasks).unwrap();
|
||||
let files = list_dir_sorted(Path::new(&dir));
|
||||
let success_content = fs::read_to_string(Path::new(&dir).join(&files[1])).unwrap();
|
||||
let error_content = fs::read_to_string(Path::new(&dir).join(&files[2])).unwrap();
|
||||
|
||||
assert!(success_content.contains("agent_id: agent-123"));
|
||||
assert!(error_content.contains("error: child failed"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upload_file_artifact_tool_call_result_serializes_only_supported_success_fields() {
|
||||
let task_id = "root";
|
||||
let tasks = vec![create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
make_tool_call_message(
|
||||
"m1",
|
||||
task_id,
|
||||
"tc_upload_file_artifact",
|
||||
api::message::tool_call::Tool::UploadFileArtifact(api::UploadFileArtifact {
|
||||
file: Some(api::FilePathReference {
|
||||
file_path: "outputs/report.txt".to_string(),
|
||||
}),
|
||||
description: "Daily summary".to_string(),
|
||||
}),
|
||||
),
|
||||
make_tool_call_result_message(
|
||||
"m2",
|
||||
task_id,
|
||||
"tc_upload_file_artifact",
|
||||
api::message::tool_call_result::Result::UploadFileArtifact(
|
||||
api::UploadFileArtifactResult {
|
||||
result: Some(api::upload_file_artifact_result::Result::Success(
|
||||
api::upload_file_artifact_result::Success {
|
||||
artifact_uid: "artifact-123".to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
size_bytes: 42,
|
||||
},
|
||||
)),
|
||||
},
|
||||
),
|
||||
),
|
||||
],
|
||||
)];
|
||||
|
||||
let dir = materialize_tasks_to_yaml(&tasks).unwrap();
|
||||
let files = list_dir_sorted(Path::new(&dir));
|
||||
let success_content = fs::read_to_string(Path::new(&dir).join(&files[1])).unwrap();
|
||||
|
||||
assert!(success_content.contains("artifact_uid: artifact-123"));
|
||||
assert!(success_content.contains("mime_type: text/plain"));
|
||||
assert!(success_content.contains("size_bytes: 42"));
|
||||
assert!(!success_content.contains("filepath:"));
|
||||
assert!(!success_content.contains("description:"));
|
||||
|
||||
cleanup_dir(&dir);
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
use warp_core::ui::{appearance::Appearance, theme::AnsiColorIdentifier};
|
||||
|
||||
use crate::ui_components::{blended_colors, icons::Icon};
|
||||
|
||||
pub fn todo_list_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::BulletedListBlock.into(),
|
||||
blended_colors::neutral_7(appearance.theme()),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn pending_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::Queued.into(),
|
||||
blended_colors::neutral_5(appearance.theme()),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn in_progress_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::Circle.into(),
|
||||
AnsiColorIdentifier::Magenta.to_ansi_color(&appearance.theme().terminal_colors().normal),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn succeeded_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::Check.into(),
|
||||
AnsiColorIdentifier::Green.to_ansi_color(&appearance.theme().terminal_colors().normal),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn addressed_comment_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::AddressedComment.into(),
|
||||
AnsiColorIdentifier::Green.to_ansi_color(&appearance.theme().terminal_colors().normal),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn failed_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::Triangle.into(),
|
||||
AnsiColorIdentifier::Red.to_ansi_color(&appearance.theme().terminal_colors().normal),
|
||||
)
|
||||
}
|
||||
|
||||
/// Not running, does not need user's attention
|
||||
pub fn gray_stop_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::StopFilled.into(),
|
||||
blended_colors::neutral_5(appearance.theme()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Agent is waiting for user to follow-up with next prompt.
|
||||
pub fn gray_clock_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::ClockSnooze.into(),
|
||||
blended_colors::neutral_5(appearance.theme()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Loading but not actionable yet.
|
||||
pub fn gray_circle_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::Circle.into(),
|
||||
blended_colors::neutral_5(appearance.theme()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Not running, requires user's attention
|
||||
pub fn yellow_stop_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::StopFilled.into(),
|
||||
AnsiColorIdentifier::Yellow.to_ansi_color(&appearance.theme().terminal_colors().normal),
|
||||
)
|
||||
}
|
||||
|
||||
/// To be used for actions (like running commands/reading files) that are long-running and executing.
|
||||
pub fn yellow_running_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(
|
||||
Icon::Circle.into(),
|
||||
AnsiColorIdentifier::Yellow.to_ansi_color(&appearance.theme().terminal_colors().normal),
|
||||
)
|
||||
}
|
||||
|
||||
/// Used for buttons that stop the current task
|
||||
pub fn red_stop_icon(appearance: &Appearance) -> warpui::elements::Icon {
|
||||
warpui::elements::Icon::new(Icon::StopFilled.into(), appearance.theme().ansi_fg_red())
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
//! Linearization utilities for task messages.
|
||||
//!
|
||||
//! This module provides pure functions for linearizing task messages in a conversation,
|
||||
//! following a DFS traversal that interleaves subtask messages at subagent tool calls.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::ai::agent::task::helper::TaskExt as _;
|
||||
|
||||
/// Computes the set of "active" task IDs in a task tree.
|
||||
///
|
||||
/// An active task is one that is still in progress. The algorithm:
|
||||
/// 1. Start with a queue containing the root task ID
|
||||
/// 2. For each task in the queue, walk through its messages:
|
||||
/// - When encountering a subagent ToolCall, add the subtask to the queue
|
||||
/// - When encountering a ToolCallResult matching a subagent call, remove from queue
|
||||
/// 3. After processing all messages, add the task to the active set
|
||||
/// 4. Repeat until the queue is empty
|
||||
pub fn compute_active_task_ids<'a>(
|
||||
root_task_id: &str,
|
||||
tasks: &HashMap<&str, &'a api::Task>,
|
||||
) -> HashSet<&'a str> {
|
||||
let mut active_tasks = HashSet::new();
|
||||
let mut visited = HashSet::new();
|
||||
let mut queue = vec![root_task_id];
|
||||
|
||||
while let Some(task_id) = queue.pop() {
|
||||
// Cycle protection: skip tasks we've already processed.
|
||||
if !visited.insert(task_id) {
|
||||
log::error!("Cycle detected in active task computation at task {task_id}");
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(task) = tasks.get(task_id) else {
|
||||
// Task not found - skip it.
|
||||
continue;
|
||||
};
|
||||
|
||||
// Track subagent tool calls: tool_call_id -> subtask_id.
|
||||
let mut pending_subagents: HashMap<&str, &str> = HashMap::new();
|
||||
|
||||
for message in &task.messages {
|
||||
match &message.message {
|
||||
Some(api::message::Message::ToolCall(tool_call)) => {
|
||||
// Check if this is a subagent call.
|
||||
if let Some(api::message::tool_call::Tool::Subagent(subagent)) = &tool_call.tool
|
||||
{
|
||||
if !subagent.task_id.is_empty() {
|
||||
// Add subtask to the queue.
|
||||
queue.push(subagent.task_id.as_str());
|
||||
// Track this subagent call so we can remove it when we see the result.
|
||||
pending_subagents
|
||||
.insert(tool_call.tool_call_id.as_str(), subagent.task_id.as_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(api::message::Message::ToolCallResult(result)) => {
|
||||
// If this result matches a pending subagent call, remove from queue.
|
||||
if let Some(subtask_id) = pending_subagents.remove(result.tool_call_id.as_str())
|
||||
{
|
||||
queue.retain(|id| *id != subtask_id);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// After processing all messages, add this task to the active set.
|
||||
active_tasks.insert(task.id.as_str());
|
||||
}
|
||||
|
||||
active_tasks
|
||||
}
|
||||
|
||||
/// Computes the depth (distance from root) for each task in the map.
|
||||
///
|
||||
/// Tasks with no parent have depth 0. Tasks whose parent chain contains a cycle or leads to a
|
||||
/// missing task are assigned depth 0.
|
||||
pub fn compute_task_depths(tasks: &HashMap<String, api::Task>) -> HashMap<&str, usize> {
|
||||
let mut depths = HashMap::new();
|
||||
for (task_id, _) in tasks.iter() {
|
||||
let mut depth = 0;
|
||||
let mut current_id: &str = task_id;
|
||||
let mut visited = HashSet::new();
|
||||
while let Some(task) = tasks.get(current_id) {
|
||||
if !visited.insert(current_id) {
|
||||
// Cycle detected; treat as depth 0.
|
||||
log::error!("Cycle detected in task parent chain starting from task {task_id}");
|
||||
depth = 0;
|
||||
break;
|
||||
}
|
||||
if let Some(parent_id) = task.parent_id() {
|
||||
depth += 1;
|
||||
current_id = parent_id;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
depths.insert(task_id.as_str(), depth);
|
||||
}
|
||||
depths
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "linearization_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,427 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::*;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
// Helper function to create a basic message
|
||||
fn create_message(id: &str, task_id: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: "server_data".to_string(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::AgentOutput(
|
||||
api::message::AgentOutput {
|
||||
text: format!("Message content for {id}"),
|
||||
},
|
||||
)),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn create_subagent_tool_call_message(id: &str, task_id: &str, subtask_id: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: "server_data".to_string(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: format!("{id}_tool_call"),
|
||||
tool: Some(api::message::tool_call::Tool::Subagent(
|
||||
api::message::tool_call::Subagent {
|
||||
task_id: subtask_id.to_string(),
|
||||
payload: String::new(),
|
||||
metadata: None,
|
||||
},
|
||||
)),
|
||||
})),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create a tool call result message.
|
||||
fn create_tool_call_result_message(id: &str, task_id: &str, tool_call_id: &str) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: "server_data".to_string(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCallResult(
|
||||
api::message::ToolCallResult {
|
||||
tool_call_id: tool_call_id.to_string(),
|
||||
context: None,
|
||||
result: None,
|
||||
},
|
||||
)),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create a task with dependencies
|
||||
fn create_task(id: &str, messages: Vec<api::Message>, parent_task_id: Option<String>) -> api::Task {
|
||||
let dependencies = parent_task_id.map(|parent_id| api::task::Dependencies {
|
||||
parent_task_id: parent_id,
|
||||
});
|
||||
|
||||
api::Task {
|
||||
id: id.to_string(),
|
||||
messages,
|
||||
dependencies,
|
||||
description: format!("Task {id}"),
|
||||
summary: format!("Summary for task {id}"),
|
||||
server_data: "".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create a root task (no parent)
|
||||
fn create_root_task(id: &str, messages: Vec<api::Message>) -> api::Task {
|
||||
create_task(id, messages, None)
|
||||
}
|
||||
|
||||
// Helper function to create a child task
|
||||
fn create_child_task(id: &str, messages: Vec<api::Message>, parent_id: &str) -> api::Task {
|
||||
create_task(id, messages, Some(parent_id.to_string()))
|
||||
}
|
||||
|
||||
// Helper to build a task map from a slice of tasks.
|
||||
fn make_task_map(tasks: &[api::Task]) -> HashMap<&str, &api::Task> {
|
||||
tasks.iter().map(|t| (t.id.as_str(), t)).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_single_root() {
|
||||
let root = create_root_task("root", vec![create_message("m1", "root")]);
|
||||
let tasks = vec![root];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
assert_eq!(active.len(), 1);
|
||||
assert!(active.contains("root"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_subagent_in_progress() {
|
||||
// Root calls a subagent but has not received the result yet.
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "child"),
|
||||
];
|
||||
let child_messages = vec![create_message("child_m1", "child")];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let child = create_child_task("child", child_messages, "root");
|
||||
let tasks = vec![root, child];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// Both root and child are active.
|
||||
assert_eq!(active.len(), 2);
|
||||
assert!(active.contains("root"));
|
||||
assert!(active.contains("child"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_subagent_completed() {
|
||||
// Root calls a subagent and has received the result.
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "child"),
|
||||
create_tool_call_result_message("result1", "root", "call1_tool_call"),
|
||||
create_message("m2", "root"),
|
||||
];
|
||||
let child_messages = vec![create_message("child_m1", "child")];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let child = create_child_task("child", child_messages, "root");
|
||||
let tasks = vec![root, child];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// Only root is active - child completed.
|
||||
assert_eq!(active.len(), 1);
|
||||
assert!(active.contains("root"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_nested_subagents() {
|
||||
// root -> child (in progress) -> grandchild (in progress)
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "child"),
|
||||
];
|
||||
let child_messages = vec![
|
||||
create_message("child_m1", "child"),
|
||||
create_subagent_tool_call_message("call2", "child", "grandchild"),
|
||||
];
|
||||
let grandchild_messages = vec![create_message("gc_m1", "grandchild")];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let child = create_child_task("child", child_messages, "root");
|
||||
let grandchild = create_child_task("grandchild", grandchild_messages, "child");
|
||||
let tasks = vec![root, child, grandchild];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// All three are active.
|
||||
assert_eq!(active.len(), 3);
|
||||
assert!(active.contains("root"));
|
||||
assert!(active.contains("child"));
|
||||
assert!(active.contains("grandchild"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_nested_subagents_partial_completion() {
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "child"),
|
||||
];
|
||||
let child_messages = vec![
|
||||
create_message("child_m1", "child"),
|
||||
create_subagent_tool_call_message("call2", "child", "grandchild"),
|
||||
create_tool_call_result_message("result2", "child", "call2_tool_call"),
|
||||
];
|
||||
let grandchild_messages = vec![create_message("gc_m1", "grandchild")];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let child = create_child_task("child", child_messages, "root");
|
||||
let grandchild = create_child_task("grandchild", grandchild_messages, "child");
|
||||
let tasks = vec![root, child, grandchild];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// Grandchild completed, but root and its child are still running.
|
||||
assert_eq!(active.len(), 2);
|
||||
assert!(active.contains("root"));
|
||||
assert!(active.contains("child"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_multiple_parallel_subagents() {
|
||||
// Root calls two subagents, one completed and one in progress.
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "child1"),
|
||||
create_subagent_tool_call_message("call2", "root", "child2"),
|
||||
create_tool_call_result_message("result1", "root", "call1_tool_call"),
|
||||
];
|
||||
let child1_messages = vec![create_message("c1_m1", "child1")];
|
||||
let child2_messages = vec![create_message("c2_m1", "child2")];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let child1 = create_child_task("child1", child1_messages, "root");
|
||||
let child2 = create_child_task("child2", child2_messages, "root");
|
||||
let tasks = vec![root, child1, child2];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// Root and child2 are active, child1 is completed.
|
||||
assert_eq!(active.len(), 2);
|
||||
assert!(active.contains("root"));
|
||||
assert!(active.contains("child2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_missing_subtask() {
|
||||
// Root calls a subagent that doesn't exist in the task list.
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "nonexistent"),
|
||||
];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let tasks = vec![root];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// Only root is active - the missing subtask is skipped.
|
||||
assert_eq!(active.len(), 1);
|
||||
assert!(active.contains("root"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_missing_root() {
|
||||
// Root task ID doesn't exist in the map.
|
||||
let tasks: HashMap<&str, &api::Task> = HashMap::new();
|
||||
|
||||
let active = compute_active_task_ids("nonexistent", &tasks);
|
||||
|
||||
assert!(active.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_active_task_ids_cycle_protection() {
|
||||
// Create a cycle: root -> child -> root (via subagent call).
|
||||
// This should not cause an infinite loop.
|
||||
let root_messages = vec![
|
||||
create_message("m1", "root"),
|
||||
create_subagent_tool_call_message("call1", "root", "child"),
|
||||
];
|
||||
let child_messages = vec![
|
||||
create_message("child_m1", "child"),
|
||||
// Child calls back to root, creating a cycle.
|
||||
create_subagent_tool_call_message("call2", "child", "root"),
|
||||
];
|
||||
|
||||
let root = create_root_task("root", root_messages);
|
||||
let child = create_child_task("child", child_messages, "root");
|
||||
let tasks = vec![root, child];
|
||||
let tasks = make_task_map(&tasks);
|
||||
|
||||
// This should complete without infinite looping.
|
||||
let active = compute_active_task_ids("root", &tasks);
|
||||
|
||||
// Both root and child are active (cycle is broken by visited check).
|
||||
assert_eq!(active.len(), 2);
|
||||
assert!(active.contains("root"));
|
||||
assert!(active.contains("child"));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// compute_task_depths tests
|
||||
// ============================================================================
|
||||
|
||||
/// Creates a task with the given ID and optional parent, without any messages.
|
||||
fn create_task_for_depth(id: &str, parent_id: Option<&str>) -> api::Task {
|
||||
create_task(id, vec![], parent_id.map(str::to_string))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_empty() {
|
||||
let tasks = HashMap::new();
|
||||
let depths = compute_task_depths(&tasks);
|
||||
assert!(depths.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_single_root() {
|
||||
let tasks: HashMap<String, _> =
|
||||
[("root".to_string(), create_task_for_depth("root", None))].into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
assert_eq!(depths.get("root"), Some(&0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_linear_chain() {
|
||||
// root -> child -> grandchild
|
||||
let tasks: HashMap<String, _> = [
|
||||
("root".to_string(), create_task_for_depth("root", None)),
|
||||
(
|
||||
"child".to_string(),
|
||||
create_task_for_depth("child", Some("root")),
|
||||
),
|
||||
(
|
||||
"grandchild".to_string(),
|
||||
create_task_for_depth("grandchild", Some("child")),
|
||||
),
|
||||
]
|
||||
.into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
assert_eq!(depths.get("root"), Some(&0));
|
||||
assert_eq!(depths.get("child"), Some(&1));
|
||||
assert_eq!(depths.get("grandchild"), Some(&2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_tree_structure() {
|
||||
// root -> child1 -> grandchild1
|
||||
// -> child2
|
||||
let tasks: HashMap<String, _> = [
|
||||
("root".to_string(), create_task_for_depth("root", None)),
|
||||
(
|
||||
"child1".to_string(),
|
||||
create_task_for_depth("child1", Some("root")),
|
||||
),
|
||||
(
|
||||
"child2".to_string(),
|
||||
create_task_for_depth("child2", Some("root")),
|
||||
),
|
||||
(
|
||||
"grandchild1".to_string(),
|
||||
create_task_for_depth("grandchild1", Some("child1")),
|
||||
),
|
||||
]
|
||||
.into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
assert_eq!(depths.get("root"), Some(&0));
|
||||
assert_eq!(depths.get("child1"), Some(&1));
|
||||
assert_eq!(depths.get("child2"), Some(&1));
|
||||
assert_eq!(depths.get("grandchild1"), Some(&2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_orphan() {
|
||||
// Task with parent that doesn't exist in the map.
|
||||
let tasks: HashMap<String, _> = [(
|
||||
"orphan".to_string(),
|
||||
create_task_for_depth("orphan", Some("missing_parent")),
|
||||
)]
|
||||
.into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
// Orphan's parent is missing, so the chain breaks immediately after computing depth 1.
|
||||
assert_eq!(depths.get("orphan"), Some(&1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_cycle_two_tasks() {
|
||||
// a -> b -> a (cycle)
|
||||
let tasks: HashMap<String, _> = [
|
||||
("a".to_string(), create_task_for_depth("a", Some("b"))),
|
||||
("b".to_string(), create_task_for_depth("b", Some("a"))),
|
||||
]
|
||||
.into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
// Both tasks are in a cycle, so they should get depth 0.
|
||||
assert_eq!(depths.get("a"), Some(&0));
|
||||
assert_eq!(depths.get("b"), Some(&0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_self_referential_cycle() {
|
||||
// a -> a (self-referential)
|
||||
let tasks: HashMap<String, _> =
|
||||
[("a".to_string(), create_task_for_depth("a", Some("a")))].into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
// Self-referential task should get depth 0.
|
||||
assert_eq!(depths.get("a"), Some(&0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_task_depths_cycle_with_tail() {
|
||||
// Task c points to a cycle: c -> a -> b -> a
|
||||
let tasks: HashMap<String, _> = [
|
||||
("a".to_string(), create_task_for_depth("a", Some("b"))),
|
||||
("b".to_string(), create_task_for_depth("b", Some("a"))),
|
||||
("c".to_string(), create_task_for_depth("c", Some("a"))),
|
||||
]
|
||||
.into();
|
||||
|
||||
let depths = compute_task_depths(&tasks);
|
||||
|
||||
// a and b are in a cycle, so depth 0.
|
||||
assert_eq!(depths.get("a"), Some(&0));
|
||||
assert_eq!(depths.get("b"), Some(&0));
|
||||
// c's parent chain leads into a cycle, so it also gets depth 0.
|
||||
assert_eq!(depths.get("c"), Some(&0));
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,191 @@
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
|
||||
use warp_multi_agent_api::{FileContent, FileContentLineRange};
|
||||
|
||||
use crate::ai::agent::{
|
||||
AIAgentOutput, AIAgentOutputMessage, AIAgentOutputMessageType, AIAgentText, AIAgentTextSection,
|
||||
AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram, AnyFileContent,
|
||||
FileContext, FormattedTextWrapper, MessageId, ProgrammingLanguage,
|
||||
};
|
||||
use crate::terminal::shell::ShellType;
|
||||
use markdown_parser::{FormattedText, FormattedTextFragment, FormattedTextLine};
|
||||
|
||||
fn to_range(range: Range<u32>) -> Option<FileContentLineRange> {
|
||||
Some(FileContentLineRange {
|
||||
start: range.start,
|
||||
end: range.end,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_text_wrapper_shares_arc_across_calls() {
|
||||
let text = FormattedText::new([FormattedTextLine::Line(vec![
|
||||
FormattedTextFragment::plain_text("hello world"),
|
||||
])]);
|
||||
let wrapper = FormattedTextWrapper::from(text);
|
||||
let arc1 = wrapper.formatted_text_arc();
|
||||
let arc2 = wrapper.formatted_text_arc();
|
||||
// Both calls must return the same allocation — not independent deep copies.
|
||||
assert!(Arc::ptr_eq(&arc1, &arc2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_text_wrapper_preserves_content() {
|
||||
let text = FormattedText::new([
|
||||
FormattedTextLine::Line(vec![FormattedTextFragment::plain_text("line one")]),
|
||||
FormattedTextLine::Line(vec![FormattedTextFragment::plain_text("line two")]),
|
||||
]);
|
||||
let wrapper = FormattedTextWrapper::from(text);
|
||||
// lines() metadata matches the cached Arc
|
||||
assert_eq!(wrapper.lines().len(), 2);
|
||||
assert_eq!(wrapper.lines()[0].raw_text(), "line one\n");
|
||||
assert_eq!(wrapper.lines()[1].raw_text(), "line two\n");
|
||||
// Arc contains the same lines
|
||||
let ft = wrapper.formatted_text_arc();
|
||||
assert_eq!(ft.lines.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_files() {
|
||||
let a = FileContext::new(
|
||||
"a.txt".to_string(),
|
||||
AnyFileContent::StringContent("hey\nyou".to_string()),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
Into::<Vec<FileContent>>::into(a),
|
||||
vec![FileContent {
|
||||
file_path: "a.txt".to_string(),
|
||||
content: "hey\nyou".to_string(),
|
||||
line_range: None,
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_files_range() {
|
||||
// Content is pre-sliced to match the line range.
|
||||
let a = FileContext::new(
|
||||
"a.txt".to_string(),
|
||||
AnyFileContent::StringContent("hey\nyou".to_string()),
|
||||
Some(1..2),
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
Into::<Vec<FileContent>>::into(a),
|
||||
vec![FileContent {
|
||||
file_path: "a.txt".to_string(),
|
||||
content: "hey\nyou".to_string(),
|
||||
line_range: to_range(1..2),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_files_range_out_of_bounds() {
|
||||
// Even with an out-of-bounds range, content is passed through as-is.
|
||||
let a = FileContext::new(
|
||||
"a.txt".to_string(),
|
||||
AnyFileContent::StringContent(String::new()),
|
||||
Some(10..20),
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
Into::<Vec<FileContent>>::into(a),
|
||||
vec![FileContent {
|
||||
file_path: "a.txt".to_string(),
|
||||
content: String::new(),
|
||||
line_range: to_range(10..20),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_programming_language_from_string() {
|
||||
// Shell language specifiers should produce Shell variants
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("bash".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::Bash)
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("shell".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::Bash)
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("sh".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::Bash)
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("zsh".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::Zsh)
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("fish".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::Fish)
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("powershell".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::PowerShell)
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("pwsh".to_string()),
|
||||
ProgrammingLanguage::Shell(ShellType::PowerShell)
|
||||
);
|
||||
|
||||
// Non-shell languages should produce Other variants
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("python".to_string()),
|
||||
ProgrammingLanguage::Other("python".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("rust".to_string()),
|
||||
ProgrammingLanguage::Other("rust".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
ProgrammingLanguage::from("javascript".to_string()),
|
||||
ProgrammingLanguage::Other("javascript".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_for_copy_preserves_visual_markdown_sections() {
|
||||
let output = AIAgentOutput {
|
||||
messages: vec".to_string(),
|
||||
layout: AgentOutputImageLayout::Block,
|
||||
},
|
||||
},
|
||||
AIAgentTextSection::MermaidDiagram {
|
||||
diagram: AgentOutputMermaidDiagram {
|
||||
source: "graph TD\nA --> B".to_string(),
|
||||
markdown_source: "```mermaid\ngraph TD\nA --> B\n```".to_string(),
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
citations: Vec::new(),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
output.format_for_copy(None),
|
||||
"Intro\n\n```mermaid\ngraph TD\nA --> B\n```"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,409 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ai::agent::{
|
||||
AIAgentActionResultType, AIAgentAttachment, AIAgentContext, AIAgentInput, AnyFileContent,
|
||||
AskUserQuestionAnswerItem, AskUserQuestionResult, BlockContext, PassiveSuggestionResultType,
|
||||
PassiveSuggestionTrigger, RequestCommandOutputResult, TransferShellCommandControlToUserResult,
|
||||
};
|
||||
|
||||
use super::super::blocklist::block::secret_redaction::{
|
||||
find_secrets_in_text, SECRET_REDACTION_REPLACEMENT_CHARACTER,
|
||||
};
|
||||
|
||||
/// Redact all detected secrets in-place within the given string.
|
||||
pub(crate) fn redact_secrets(input: &mut String) {
|
||||
let mut secrets: Vec<_> = find_secrets_in_text(input)
|
||||
.into_iter()
|
||||
.map(|r| r.byte_range)
|
||||
.collect();
|
||||
// Replace from the end to preserve indices
|
||||
secrets.sort_by_key(|range| range.start);
|
||||
for range in secrets.into_iter().rev() {
|
||||
let replacement =
|
||||
SECRET_REDACTION_REPLACEMENT_CHARACTER.repeat(range.end.saturating_sub(range.start));
|
||||
input.replace_range(range.start..range.end, &replacement);
|
||||
}
|
||||
}
|
||||
|
||||
/// Redact secrets in-place for all user-provided text fields inside the inputs that will be
|
||||
/// sent to the server.
|
||||
pub(crate) fn redact_inputs(inputs: &mut [AIAgentInput]) {
|
||||
for input in inputs.iter_mut() {
|
||||
match input {
|
||||
AIAgentInput::UserQuery {
|
||||
query,
|
||||
context,
|
||||
referenced_attachments,
|
||||
..
|
||||
} => {
|
||||
redact_secrets(query);
|
||||
redact_context(Arc::make_mut(context));
|
||||
referenced_attachments
|
||||
.values_mut()
|
||||
.for_each(redact_attachment);
|
||||
}
|
||||
AIAgentInput::AutoCodeDiffQuery { query, context, .. } => {
|
||||
redact_secrets(query);
|
||||
redact_context(Arc::make_mut(context));
|
||||
}
|
||||
AIAgentInput::CreateNewProject { context, .. }
|
||||
| AIAgentInput::CloneRepository { context, .. }
|
||||
| AIAgentInput::ResumeConversation { context }
|
||||
| AIAgentInput::InitProjectRules { context, .. }
|
||||
| AIAgentInput::StartFromAmbientRunPrompt { context, .. } => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
}
|
||||
AIAgentInput::SummarizeConversation { prompt } => {
|
||||
if let Some(p) = prompt {
|
||||
redact_secrets(p);
|
||||
}
|
||||
}
|
||||
AIAgentInput::CreateEnvironment { context, .. } => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
}
|
||||
AIAgentInput::TriggerPassiveSuggestion {
|
||||
context,
|
||||
attachments,
|
||||
trigger,
|
||||
} => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
attachments.iter_mut().for_each(redact_attachment);
|
||||
if let PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger) = trigger {
|
||||
redact_secrets(&mut shell_trigger.executed_shell_command.command);
|
||||
redact_secrets(&mut shell_trigger.executed_shell_command.output);
|
||||
for file in shell_trigger.relevant_files.iter_mut() {
|
||||
if let AnyFileContent::StringContent(content) = &mut file.content {
|
||||
redact_secrets(content);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentInput::CodeReview {
|
||||
context,
|
||||
review_comments,
|
||||
} => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
for comment in review_comments.comments.iter_mut() {
|
||||
redact_secrets(&mut comment.content);
|
||||
match &mut comment.target {
|
||||
crate::code_review::comments::AttachedReviewCommentTarget::Line {
|
||||
content,
|
||||
..
|
||||
} => {
|
||||
redact_secrets(&mut content.content);
|
||||
}
|
||||
crate::code_review::comments::AttachedReviewCommentTarget::File {
|
||||
..
|
||||
}
|
||||
| crate::code_review::comments::AttachedReviewCommentTarget::General => {}
|
||||
}
|
||||
}
|
||||
|
||||
for diff in review_comments.diff_set.values_mut().flatten() {
|
||||
redact_secrets(&mut diff.diff_content);
|
||||
}
|
||||
}
|
||||
// No user-provided text to redact in inter-agent relay inputs.
|
||||
AIAgentInput::MessagesReceivedFromAgents { .. }
|
||||
| AIAgentInput::EventsFromAgents { .. } => {}
|
||||
AIAgentInput::ActionResult { result, context } => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
match &mut result.result {
|
||||
AIAgentActionResultType::RequestCommandOutput(output) => {
|
||||
if let RequestCommandOutputResult::Completed { output, .. } = output {
|
||||
redact_secrets(output);
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::WriteToLongRunningShellCommand(result) => {
|
||||
use crate::ai::agent::WriteToLongRunningShellCommandResult::*;
|
||||
match result {
|
||||
Snapshot { grid_contents, .. } => redact_secrets(grid_contents),
|
||||
CommandFinished { output, .. } => redact_secrets(output),
|
||||
Error(_) | Cancelled => {}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::ReadShellCommandOutput(result) => {
|
||||
use crate::ai::agent::ReadShellCommandOutputResult::*;
|
||||
match result {
|
||||
CommandFinished { output, .. } => redact_secrets(output),
|
||||
LongRunningCommandSnapshot { grid_contents, .. } => {
|
||||
redact_secrets(grid_contents)
|
||||
}
|
||||
Error(_) | Cancelled => {}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::ReadFiles(read_files_result) => {
|
||||
if let crate::ai::agent::ReadFilesResult::Success { files } =
|
||||
read_files_result
|
||||
{
|
||||
for file in files {
|
||||
if let AnyFileContent::StringContent(content) = &mut file.content {
|
||||
redact_secrets(content);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::UploadArtifact(upload_result) => {
|
||||
use crate::ai::agent::UploadArtifactResult;
|
||||
match upload_result {
|
||||
UploadArtifactResult::Success {
|
||||
filepath,
|
||||
description,
|
||||
..
|
||||
} => {
|
||||
if let Some(filepath) = filepath {
|
||||
redact_secrets(filepath);
|
||||
}
|
||||
if let Some(description) = description {
|
||||
redact_secrets(description);
|
||||
}
|
||||
}
|
||||
UploadArtifactResult::Error(error) => redact_secrets(error),
|
||||
UploadArtifactResult::Cancelled => {}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::SearchCodebase(search_codebase_result) => {
|
||||
if let crate::ai::agent::SearchCodebaseResult::Success { files } =
|
||||
search_codebase_result
|
||||
{
|
||||
for file in files {
|
||||
if let AnyFileContent::StringContent(content) = &mut file.content {
|
||||
redact_secrets(content);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::RequestFileEdits(request_file_edits_result) => {
|
||||
if let crate::ai::agent::RequestFileEditsResult::Success {
|
||||
diff,
|
||||
updated_files,
|
||||
deleted_files,
|
||||
..
|
||||
} = request_file_edits_result
|
||||
{
|
||||
redact_secrets(diff);
|
||||
for file in updated_files {
|
||||
if let AnyFileContent::StringContent(content) =
|
||||
&mut file.file_context.content
|
||||
{
|
||||
redact_secrets(content);
|
||||
}
|
||||
}
|
||||
for file_path in deleted_files {
|
||||
redact_secrets(file_path);
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::InsertReviewComments(result) => {
|
||||
use crate::ai::agent::InsertReviewCommentsResult::*;
|
||||
match result {
|
||||
Success { repo_path } => redact_secrets(repo_path),
|
||||
Error { repo_path, message } => {
|
||||
redact_secrets(repo_path);
|
||||
redact_secrets(message);
|
||||
}
|
||||
Cancelled => {}
|
||||
}
|
||||
}
|
||||
|
||||
// These are effectively flow control and don't contain secrets
|
||||
AIAgentActionResultType::SuggestNewConversation { .. }
|
||||
| AIAgentActionResultType::OpenCodeReview
|
||||
| AIAgentActionResultType::InitProject => {}
|
||||
|
||||
// Contains only file path/line number information
|
||||
AIAgentActionResultType::Grep(_)
|
||||
| AIAgentActionResultType::FileGlob(_)
|
||||
| AIAgentActionResultType::FileGlobV2(_) => {}
|
||||
|
||||
// TODO: Redact MCP-related results
|
||||
AIAgentActionResultType::CallMCPTool { .. }
|
||||
| AIAgentActionResultType::ReadSkill { .. }
|
||||
| AIAgentActionResultType::ReadMCPResource { .. }
|
||||
| AIAgentActionResultType::SuggestPrompt { .. }
|
||||
| AIAgentActionResultType::ReadDocuments(_)
|
||||
| AIAgentActionResultType::EditDocuments(_)
|
||||
| AIAgentActionResultType::CreateDocuments(_) => {}
|
||||
|
||||
// TODO(AGENT-2282): figure out whether there's any reasonable way to
|
||||
// do redaction here (probably not).
|
||||
AIAgentActionResultType::UseComputer(_) => {}
|
||||
|
||||
// Request computer use just contains screen dimensions, no secrets
|
||||
AIAgentActionResultType::RequestComputerUse(_) => {}
|
||||
|
||||
// FetchConversation results contain tasks returned from the server,
|
||||
// which were already redacted before being sent as client inputs.
|
||||
// (client inputs -> redaction -> server request -> task messages)
|
||||
AIAgentActionResultType::FetchConversation(_) => {}
|
||||
|
||||
// StartAgent results contain only an agent ID string, no secrets
|
||||
AIAgentActionResultType::StartAgent(_) => {}
|
||||
|
||||
// SendMessageToAgent results contain only a message ID or error string, no secrets
|
||||
AIAgentActionResultType::SendMessageToAgent(_) => {}
|
||||
// TransferShellCommandControlToUser result - similar to WriteToLongRunningShellCommand
|
||||
AIAgentActionResultType::TransferShellCommandControlToUser(result) => {
|
||||
match result {
|
||||
TransferShellCommandControlToUserResult::Snapshot {
|
||||
grid_contents,
|
||||
..
|
||||
} => redact_secrets(grid_contents),
|
||||
TransferShellCommandControlToUserResult::CommandFinished {
|
||||
output,
|
||||
..
|
||||
} => redact_secrets(output),
|
||||
TransferShellCommandControlToUserResult::Error(_)
|
||||
| TransferShellCommandControlToUserResult::Cancelled => {}
|
||||
}
|
||||
}
|
||||
AIAgentActionResultType::AskUserQuestion(result) => {
|
||||
redact_ask_user_question_result(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentInput::FetchReviewComments { repo_path, context } => {
|
||||
redact_secrets(repo_path);
|
||||
redact_context(Arc::make_mut(context));
|
||||
}
|
||||
AIAgentInput::InvokeSkill {
|
||||
context,
|
||||
skill,
|
||||
user_query,
|
||||
} => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
redact_secrets(&mut skill.content);
|
||||
if let Some(user_query) = user_query {
|
||||
redact_secrets(&mut user_query.query);
|
||||
for attachment in user_query.referenced_attachments.values_mut() {
|
||||
redact_attachment(attachment);
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentInput::PassiveSuggestionResult {
|
||||
trigger,
|
||||
suggestion,
|
||||
context,
|
||||
} => {
|
||||
redact_context(Arc::make_mut(context));
|
||||
match suggestion {
|
||||
PassiveSuggestionResultType::Prompt { prompt } => redact_secrets(prompt),
|
||||
PassiveSuggestionResultType::CodeDiff { diffs, .. } => {
|
||||
for diff in diffs {
|
||||
redact_secrets(&mut diff.file_path);
|
||||
redact_secrets(&mut diff.search);
|
||||
redact_secrets(&mut diff.replace);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(PassiveSuggestionTrigger::ShellCommandCompleted(shell_trigger)) =
|
||||
trigger
|
||||
{
|
||||
redact_secrets(&mut shell_trigger.executed_shell_command.command);
|
||||
redact_secrets(&mut shell_trigger.executed_shell_command.output);
|
||||
for file in shell_trigger.relevant_files.iter_mut() {
|
||||
if let AnyFileContent::StringContent(content) = &mut file.content {
|
||||
redact_secrets(content);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn redact_ask_user_question_result(result: &mut AskUserQuestionResult) {
|
||||
match result {
|
||||
AskUserQuestionResult::Success { answers } => {
|
||||
for answer in answers {
|
||||
if let AskUserQuestionAnswerItem::Answered { other_text, .. } = answer {
|
||||
redact_secrets(other_text);
|
||||
}
|
||||
}
|
||||
}
|
||||
AskUserQuestionResult::SkippedByAutoApprove { .. } => {}
|
||||
AskUserQuestionResult::Error(message) => redact_secrets(message),
|
||||
AskUserQuestionResult::Cancelled => {}
|
||||
}
|
||||
}
|
||||
fn redact_context(context: &mut [AIAgentContext]) {
|
||||
for context_item in context {
|
||||
match context_item {
|
||||
AIAgentContext::Block(context) => {
|
||||
redact_secrets(&mut context.command);
|
||||
redact_secrets(&mut context.output);
|
||||
}
|
||||
AIAgentContext::SelectedText(text) => {
|
||||
redact_secrets(text);
|
||||
}
|
||||
// Other context types don't contain user-provided text that needs redaction
|
||||
AIAgentContext::Directory { .. }
|
||||
| AIAgentContext::ExecutionEnvironment(_)
|
||||
| AIAgentContext::CurrentTime { .. }
|
||||
| AIAgentContext::Image(_)
|
||||
| AIAgentContext::Codebase { .. }
|
||||
| AIAgentContext::ProjectRules { .. }
|
||||
| AIAgentContext::Git { .. }
|
||||
| AIAgentContext::File(_)
|
||||
| AIAgentContext::Skills { .. } => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn redact_attachment(attachment: &mut AIAgentAttachment) {
|
||||
match attachment {
|
||||
AIAgentAttachment::PlainText(text) => {
|
||||
redact_secrets(text);
|
||||
}
|
||||
AIAgentAttachment::Block(BlockContext {
|
||||
command, output, ..
|
||||
}) => {
|
||||
redact_secrets(command);
|
||||
redact_secrets(output);
|
||||
}
|
||||
AIAgentAttachment::DriveObject { payload, .. } => {
|
||||
if let Some(drive_payload) = payload {
|
||||
match drive_payload {
|
||||
crate::ai::agent::DriveObjectPayload::Workflow {
|
||||
name,
|
||||
description,
|
||||
command,
|
||||
} => {
|
||||
redact_secrets(name);
|
||||
redact_secrets(description);
|
||||
redact_secrets(command);
|
||||
}
|
||||
crate::ai::agent::DriveObjectPayload::Notebook { title, content } => {
|
||||
redact_secrets(title);
|
||||
redact_secrets(content);
|
||||
}
|
||||
crate::ai::agent::DriveObjectPayload::GenericStringObject {
|
||||
payload, ..
|
||||
} => {
|
||||
redact_secrets(payload);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentAttachment::DiffHunk {
|
||||
file_path,
|
||||
diff_content,
|
||||
..
|
||||
} => {
|
||||
redact_secrets(file_path);
|
||||
redact_secrets(diff_content);
|
||||
}
|
||||
AIAgentAttachment::DiffSet { file_diffs, .. } => {
|
||||
for hunks in file_diffs.values_mut() {
|
||||
for hunk in hunks {
|
||||
redact_secrets(&mut hunk.diff_content);
|
||||
}
|
||||
}
|
||||
}
|
||||
AIAgentAttachment::DocumentContent { content, .. } => {
|
||||
redact_secrets(content);
|
||||
}
|
||||
// FilePathReference only contains a file ID and filename, no user secrets.
|
||||
AIAgentAttachment::FilePathReference { .. } => {}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
use crate::ai::agent::{
|
||||
SuggestedAgentModeWorkflow, SuggestedLoggingId, SuggestedRule, Suggestions,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_extend_suggestions() {
|
||||
// Create base suggestions
|
||||
let mut base_suggestions = Suggestions {
|
||||
rules: vec![
|
||||
SuggestedRule {
|
||||
name: "rule1".into(),
|
||||
content: "content1".into(),
|
||||
logging_id: SuggestedLoggingId::from("id1".to_string()),
|
||||
},
|
||||
SuggestedRule {
|
||||
name: "rule2".into(),
|
||||
content: "content2".into(),
|
||||
logging_id: SuggestedLoggingId::from("id2".to_string()),
|
||||
},
|
||||
],
|
||||
agent_mode_workflows: vec![
|
||||
SuggestedAgentModeWorkflow {
|
||||
name: "workflow1".into(),
|
||||
prompt: "prompt1".into(),
|
||||
logging_id: SuggestedLoggingId::from("wid1".to_string()),
|
||||
},
|
||||
SuggestedAgentModeWorkflow {
|
||||
name: "workflow2".into(),
|
||||
prompt: "prompt2".into(),
|
||||
logging_id: SuggestedLoggingId::from("wid2".to_string()),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
// Create additional suggestions with both unique and duplicate logging_ids
|
||||
let additional_suggestions = Suggestions {
|
||||
rules: vec![
|
||||
// Duplicate logging_id but different name/content
|
||||
SuggestedRule {
|
||||
name: "rule1_modified".into(),
|
||||
content: "content1_modified".into(),
|
||||
logging_id: SuggestedLoggingId::from("id1".to_string()),
|
||||
},
|
||||
// New unique rule
|
||||
SuggestedRule {
|
||||
name: "rule3".into(),
|
||||
content: "content3".into(),
|
||||
logging_id: SuggestedLoggingId::from("id3".to_string()),
|
||||
},
|
||||
// Another new unique rule
|
||||
SuggestedRule {
|
||||
name: "rule4".into(),
|
||||
content: "content4".into(),
|
||||
logging_id: SuggestedLoggingId::from("id4".to_string()),
|
||||
},
|
||||
],
|
||||
agent_mode_workflows: vec![
|
||||
// Duplicate workflow logging_id but different name/prompt
|
||||
SuggestedAgentModeWorkflow {
|
||||
name: "workflow1_modified".into(),
|
||||
prompt: "prompt1_modified".into(),
|
||||
logging_id: SuggestedLoggingId::from("wid1".to_string()),
|
||||
},
|
||||
// New unique workflow
|
||||
SuggestedAgentModeWorkflow {
|
||||
name: "workflow3".into(),
|
||||
prompt: "prompt3".into(),
|
||||
logging_id: SuggestedLoggingId::from("wid3".to_string()),
|
||||
},
|
||||
// Another new unique workflow
|
||||
SuggestedAgentModeWorkflow {
|
||||
name: "workflow4".into(),
|
||||
prompt: "prompt4".into(),
|
||||
logging_id: SuggestedLoggingId::from("wid4".to_string()),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
// Extend base suggestions with additional ones
|
||||
base_suggestions.extend(&additional_suggestions);
|
||||
|
||||
// Verify rules
|
||||
|
||||
// Verify the length (should be 4 because one was a duplicate)
|
||||
assert_eq!(base_suggestions.rules.len(), 4);
|
||||
|
||||
// Verify that original rules with id1 and id2 are still present and unchanged
|
||||
assert!(base_suggestions
|
||||
.rules
|
||||
.iter()
|
||||
.any(|r| r.logging_id.to_string() == "id1"
|
||||
&& r.name == "rule1"
|
||||
&& r.content == "content1"));
|
||||
assert!(base_suggestions
|
||||
.rules
|
||||
.iter()
|
||||
.any(|r| r.logging_id.to_string() == "id2"
|
||||
&& r.name == "rule2"
|
||||
&& r.content == "content2"));
|
||||
|
||||
// Verify that new unique rules (id3 and id4) were added
|
||||
assert!(base_suggestions
|
||||
.rules
|
||||
.iter()
|
||||
.any(|r| r.logging_id.to_string() == "id3"
|
||||
&& r.name == "rule3"
|
||||
&& r.content == "content3"));
|
||||
assert!(base_suggestions
|
||||
.rules
|
||||
.iter()
|
||||
.any(|r| r.logging_id.to_string() == "id4"
|
||||
&& r.name == "rule4"
|
||||
&& r.content == "content4"));
|
||||
|
||||
// Verify that the modified version of id1 was not added (deduplication worked)
|
||||
assert!(!base_suggestions
|
||||
.rules
|
||||
.iter()
|
||||
.any(|r| r.logging_id.to_string() == "id1" && r.name == "rule1_modified"));
|
||||
|
||||
// Verify workflows
|
||||
|
||||
// Verify the length (should be 4 because one was a duplicate)
|
||||
assert_eq!(base_suggestions.agent_mode_workflows.len(), 4);
|
||||
|
||||
// Verify that original workflows with wid1 and wid2 are still present and unchanged
|
||||
assert!(base_suggestions
|
||||
.agent_mode_workflows
|
||||
.iter()
|
||||
.any(|w| w.logging_id.to_string() == "wid1"
|
||||
&& w.name == "workflow1"
|
||||
&& w.prompt == "prompt1"));
|
||||
assert!(base_suggestions
|
||||
.agent_mode_workflows
|
||||
.iter()
|
||||
.any(|w| w.logging_id.to_string() == "wid2"
|
||||
&& w.name == "workflow2"
|
||||
&& w.prompt == "prompt2"));
|
||||
|
||||
// Verify that new unique workflows (wid3 and wid4) were added
|
||||
assert!(base_suggestions
|
||||
.agent_mode_workflows
|
||||
.iter()
|
||||
.any(|w| w.logging_id.to_string() == "wid3"
|
||||
&& w.name == "workflow3"
|
||||
&& w.prompt == "prompt3"));
|
||||
assert!(base_suggestions
|
||||
.agent_mode_workflows
|
||||
.iter()
|
||||
.any(|w| w.logging_id.to_string() == "wid4"
|
||||
&& w.name == "workflow4"
|
||||
&& w.prompt == "prompt4"));
|
||||
|
||||
// Verify that the modified version of wid1 was not added (deduplication worked)
|
||||
assert!(!base_suggestions
|
||||
.agent_mode_workflows
|
||||
.iter()
|
||||
.any(|w| w.logging_id.to_string() == "wid1" && w.name == "workflow1_modified"));
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,214 @@
|
||||
//! This module contains traits and trait implementations for exposing helper methods for accessing
|
||||
//! proto fields.
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
pub trait TaskExt {
|
||||
fn parent_id(&self) -> Option<&str>;
|
||||
}
|
||||
|
||||
impl TaskExt for api::Task {
|
||||
fn parent_id(&self) -> Option<&str> {
|
||||
self.dependencies
|
||||
.as_ref()
|
||||
.map(|deps| deps.parent_task_id.as_str())
|
||||
.filter(|id| !id.is_empty())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait MessageExt {
|
||||
fn todos_op(&self) -> Option<&api::message::update_todos::Operation>;
|
||||
fn tool_call(&self) -> Option<&api::message::ToolCall>;
|
||||
fn tool_call_mut(&mut self) -> Option<&mut api::message::ToolCall>;
|
||||
fn tool_call_result(&self) -> Option<&api::message::ToolCallResult>;
|
||||
}
|
||||
|
||||
pub trait ToolCallExt {
|
||||
fn subagent(&self) -> Option<&api::message::tool_call::Subagent>;
|
||||
fn subagent_mut(&mut self) -> Option<&mut api::message::tool_call::Subagent>;
|
||||
}
|
||||
|
||||
pub trait ToolExt {
|
||||
fn name(&self) -> &'static str;
|
||||
}
|
||||
|
||||
pub trait SubagentExt {
|
||||
fn is_cli(&self) -> bool;
|
||||
fn is_advice(&self) -> bool;
|
||||
fn is_computer_use(&self) -> bool;
|
||||
fn is_summarization(&self) -> bool;
|
||||
fn is_conversation_search(&self) -> bool;
|
||||
fn is_warp_documentation_search(&self) -> bool;
|
||||
fn type_name(&self) -> &'static str;
|
||||
}
|
||||
|
||||
impl MessageExt for api::Message {
|
||||
fn todos_op(&self) -> Option<&api::message::update_todos::Operation> {
|
||||
self.message.as_ref().and_then(|message| {
|
||||
if let api::message::Message::UpdateTodos(update) = message {
|
||||
update.operation.as_ref()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn tool_call(&self) -> Option<&api::message::ToolCall> {
|
||||
self.message.as_ref().and_then(|message| {
|
||||
if let api::message::Message::ToolCall(tool_call) = message {
|
||||
Some(tool_call)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn tool_call_mut(&mut self) -> Option<&mut api::message::ToolCall> {
|
||||
self.message.as_mut().and_then(|message| {
|
||||
if let api::message::Message::ToolCall(tool_call) = message {
|
||||
Some(tool_call)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn tool_call_result(&self) -> Option<&api::message::ToolCallResult> {
|
||||
self.message.as_ref().and_then(|message| {
|
||||
if let api::message::Message::ToolCallResult(result) = message {
|
||||
Some(result)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolCallExt for api::message::ToolCall {
|
||||
fn subagent(&self) -> Option<&api::message::tool_call::Subagent> {
|
||||
match self.tool.as_ref() {
|
||||
Some(api::message::tool_call::Tool::Subagent(subagent)) => Some(subagent),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn subagent_mut(&mut self) -> Option<&mut api::message::tool_call::Subagent> {
|
||||
match self.tool.as_mut() {
|
||||
Some(api::message::tool_call::Tool::Subagent(subagent)) => Some(subagent),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolExt for api::message::tool_call::Tool {
|
||||
fn name(&self) -> &'static str {
|
||||
use api::message::tool_call::Tool;
|
||||
match self {
|
||||
Tool::RunShellCommand(_) => "run_shell_command",
|
||||
Tool::SearchCodebase(_) => "search_codebase",
|
||||
Tool::ReadFiles(_) => "read_files",
|
||||
Tool::UploadFileArtifact(_) => "upload_artifact",
|
||||
Tool::ApplyFileDiffs(_) => "apply_file_diffs",
|
||||
Tool::Grep(_) => "grep",
|
||||
#[allow(deprecated)]
|
||||
Tool::FileGlob(_) => "file_glob",
|
||||
Tool::FileGlobV2(_) => "file_glob_v2",
|
||||
Tool::ReadMcpResource(_) => "read_mcp_resource",
|
||||
Tool::CallMcpTool(_) => "call_mcp_tool",
|
||||
Tool::WriteToLongRunningShellCommand(_) => "write_to_lrc",
|
||||
Tool::ReadDocuments(_) => "read_documents",
|
||||
Tool::EditDocuments(_) => "edit_documents",
|
||||
Tool::CreateDocuments(_) => "create_documents",
|
||||
Tool::ReadShellCommandOutput(_) => "read_shell_command_output",
|
||||
Tool::UseComputer(_) => "use_computer",
|
||||
Tool::RequestComputerUse(_) => "request_computer_use",
|
||||
Tool::FetchConversation(_) => "fetch_conversation",
|
||||
Tool::InsertReviewComments(_) => "insert_review_comments",
|
||||
Tool::ReadSkill(_) => "read_skill",
|
||||
Tool::SuggestPlan(_) => "suggest_plan",
|
||||
Tool::SuggestCreatePlan(_) => "suggest_create_plan",
|
||||
Tool::SuggestNewConversation(_) => "suggest_new_conversation",
|
||||
Tool::SuggestPrompt(_) => "suggest_prompt",
|
||||
Tool::OpenCodeReview(_) => "open_code_review",
|
||||
Tool::InitProject(_) => "init_project",
|
||||
Tool::StartAgent(_) => "start_agent",
|
||||
// Keep the logical tool name stable across the v1/v2 schema split so analytics,
|
||||
// history, and UI handling continue to treat both as the same tool.
|
||||
Tool::StartAgentV2(_) => "start_agent",
|
||||
Tool::Server(_) => "server",
|
||||
Tool::Subagent(_) => "subagent",
|
||||
Tool::AskUserQuestion(_) => "ask_user_question",
|
||||
Tool::SendMessageToAgent(_) => "send_message_to_agent",
|
||||
Tool::TransferShellCommandControlToUser(_) => "transfer_shell_command_control",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SubagentExt for api::message::tool_call::Subagent {
|
||||
fn is_cli(&self) -> bool {
|
||||
self.metadata.as_ref().is_some_and(|metadata| {
|
||||
matches!(
|
||||
metadata,
|
||||
api::message::tool_call::subagent::Metadata::Cli(_)
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_advice(&self) -> bool {
|
||||
self.metadata.as_ref().is_some_and(|metadata| {
|
||||
matches!(
|
||||
metadata,
|
||||
api::message::tool_call::subagent::Metadata::Advice(_)
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_computer_use(&self) -> bool {
|
||||
self.metadata.as_ref().is_some_and(|metadata| {
|
||||
matches!(
|
||||
metadata,
|
||||
api::message::tool_call::subagent::Metadata::ComputerUse(_)
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_summarization(&self) -> bool {
|
||||
self.metadata.as_ref().is_some_and(|metadata| {
|
||||
matches!(
|
||||
metadata,
|
||||
api::message::tool_call::subagent::Metadata::Summarization(_)
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_conversation_search(&self) -> bool {
|
||||
self.metadata.as_ref().is_some_and(|metadata| {
|
||||
matches!(
|
||||
metadata,
|
||||
api::message::tool_call::subagent::Metadata::ConversationSearch(_)
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_warp_documentation_search(&self) -> bool {
|
||||
self.metadata.as_ref().is_some_and(|metadata| {
|
||||
matches!(
|
||||
metadata,
|
||||
api::message::tool_call::subagent::Metadata::WarpDocumentationSearch(_)
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn type_name(&self) -> &'static str {
|
||||
use api::message::tool_call::subagent::Metadata;
|
||||
match &self.metadata {
|
||||
Some(Metadata::Cli(_)) => "cli",
|
||||
Some(Metadata::Research(_)) => "research",
|
||||
Some(Metadata::Advice(_)) => "advice",
|
||||
Some(Metadata::ComputerUse(_)) => "computer_use",
|
||||
Some(Metadata::Summarization(_)) => "summarization",
|
||||
Some(Metadata::ConversationSearch(_)) => "conversation_search",
|
||||
Some(Metadata::WarpDocumentationSearch(_)) => "warp_documentation_search",
|
||||
None => "unknown",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::ai::agent::task::TaskId;
|
||||
|
||||
use super::Task;
|
||||
|
||||
/// Keeps track of the state of tasks before they are modified.
|
||||
/// Messages are assumed to be only updated during the same transaction
|
||||
/// in which they were added, so we can clean up message by simply
|
||||
/// deleting them.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Transaction {
|
||||
saved_tasks: HashMap<TaskId, SavedTask>,
|
||||
}
|
||||
|
||||
/// Saves state for either a newly added task or a pre-existing task
|
||||
/// modified during a transaction.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum SavedTask {
|
||||
New(TaskId),
|
||||
Existing(Box<Task>),
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
saved_tasks: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// A map of the tasks modified in this transaction.
|
||||
pub fn saved_tasks(self) -> HashMap<TaskId, SavedTask> {
|
||||
self.saved_tasks
|
||||
}
|
||||
|
||||
/// Saves a SavedTask::New to the transaction, representing a newly added task.
|
||||
pub fn checkpoint_new_task(&mut self, task_id: &TaskId) {
|
||||
if !self.saved_tasks.contains_key(task_id) {
|
||||
let task = SavedTask::New(task_id.clone());
|
||||
self.saved_tasks.insert(task_id.clone(), task);
|
||||
}
|
||||
}
|
||||
|
||||
/// Saves a SavedTask::Existing to the transaction, representing an existing
|
||||
/// task which is being modified.
|
||||
pub fn checkpoint_task(&mut self, task: &Task) {
|
||||
if !self.saved_tasks.contains_key(task.id()) {
|
||||
self.saved_tasks.insert(
|
||||
task.id().clone(),
|
||||
SavedTask::Existing(Box::new(task.clone())),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::ai::{
|
||||
agent::{AIAgentContext, AIAgentInput},
|
||||
skills::SkillDescriptor,
|
||||
};
|
||||
|
||||
use super::{
|
||||
task::{
|
||||
helper::{MessageExt, ToolCallExt},
|
||||
Task, TaskId,
|
||||
},
|
||||
AIAgentExchange, AIAgentExchangeId, AIAgentOutputMessageType,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ExchangeRef {
|
||||
task_id: TaskId,
|
||||
exchange_index: usize,
|
||||
}
|
||||
|
||||
/// Task storage with a linearized exchange index for O(1) first/last access.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaskStore {
|
||||
root_task_id: TaskId,
|
||||
tasks: HashMap<TaskId, Task>,
|
||||
linearized_refs: Vec<ExchangeRef>,
|
||||
}
|
||||
|
||||
impl TaskStore {
|
||||
pub fn with_root_task(root_task: Task) -> Self {
|
||||
let root_task_id = root_task.id().clone();
|
||||
let mut store = Self {
|
||||
tasks: HashMap::new(),
|
||||
linearized_refs: Vec::new(),
|
||||
root_task_id: root_task_id.clone(),
|
||||
};
|
||||
store.tasks.insert(root_task_id, root_task);
|
||||
store.rebuild_linearized_refs_index();
|
||||
store
|
||||
}
|
||||
|
||||
/// Creates a TaskStore from an existing HashMap of tasks.
|
||||
/// Rebuilds the linearized index after construction.
|
||||
pub fn from_tasks(tasks: HashMap<TaskId, Task>, root_task_id: TaskId) -> Self {
|
||||
let mut store = Self {
|
||||
tasks,
|
||||
linearized_refs: Vec::new(),
|
||||
root_task_id,
|
||||
};
|
||||
store.rebuild_linearized_refs_index();
|
||||
store
|
||||
}
|
||||
|
||||
pub fn root_task_id(&self) -> &TaskId {
|
||||
&self.root_task_id
|
||||
}
|
||||
|
||||
pub fn get(&self, task_id: &TaskId) -> Option<&Task> {
|
||||
self.tasks.get(task_id)
|
||||
}
|
||||
|
||||
pub fn tasks(&self) -> impl Iterator<Item = &Task> {
|
||||
self.tasks.values()
|
||||
}
|
||||
|
||||
pub fn task_count(&self) -> usize {
|
||||
self.tasks.len()
|
||||
}
|
||||
|
||||
/// Appends an exchange to a task and rebuilds the index.
|
||||
/// Returns true if the task was found and the exchange was appended.
|
||||
pub fn append_exchange(&mut self, task_id: &TaskId, exchange: AIAgentExchange) -> bool {
|
||||
let Some(task) = self.tasks.get_mut(task_id) else {
|
||||
return false;
|
||||
};
|
||||
task.append_exchange(exchange);
|
||||
self.rebuild_linearized_refs_index();
|
||||
true
|
||||
}
|
||||
|
||||
/// Removes an exchange from a task and rebuilds the index.
|
||||
/// Returns the removed exchange if found.
|
||||
pub fn remove_task_exchange(
|
||||
&mut self,
|
||||
task_id: &TaskId,
|
||||
exchange_id: AIAgentExchangeId,
|
||||
) -> Option<AIAgentExchange> {
|
||||
let task = self.tasks.get_mut(task_id)?;
|
||||
let exchange = task.remove_exchange(exchange_id)?;
|
||||
self.rebuild_linearized_refs_index();
|
||||
Some(exchange)
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to an exchange by its ID, searching all tasks.
|
||||
pub fn exchange_mut(&mut self, exchange_id: AIAgentExchangeId) -> Option<&mut AIAgentExchange> {
|
||||
for task in self.tasks.values_mut() {
|
||||
if let Some(exchange) = task.exchange_mut(exchange_id) {
|
||||
return Some(exchange);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Modifies a task via the provided closure and rebuilds the exchange index
|
||||
/// if exchanges changed.
|
||||
pub fn modify_task<R>(
|
||||
&mut self,
|
||||
task_id: &TaskId,
|
||||
f: impl FnOnce(&mut Task) -> R,
|
||||
) -> Option<R> {
|
||||
let exchange_count_before = self.tasks.get(task_id)?.exchanges_len();
|
||||
let task = self.tasks.get_mut(task_id)?;
|
||||
let result = f(task);
|
||||
let exchange_count_after = self
|
||||
.tasks
|
||||
.get(task_id)
|
||||
.map(|t| t.exchanges_len())
|
||||
.unwrap_or(0);
|
||||
if exchange_count_before != exchange_count_after {
|
||||
self.rebuild_linearized_refs_index();
|
||||
}
|
||||
Some(result)
|
||||
}
|
||||
|
||||
/// Modifies the root task via the provided closure and rebuilds the exchange index if exchanges changed.
|
||||
pub fn modify_root_task<R>(&mut self, f: impl FnOnce(&mut Task) -> R) -> Option<R> {
|
||||
let root_task_id = self.root_task_id.clone();
|
||||
self.modify_task(&root_task_id, f)
|
||||
}
|
||||
|
||||
pub fn root_task(&self) -> Option<&Task> {
|
||||
self.tasks.get(&self.root_task_id)
|
||||
}
|
||||
|
||||
/// Sets or replaces the root task, removing any previous root if it exists.
|
||||
pub fn set_root_task(&mut self, root_task: Task) {
|
||||
// Remove the old root task and its exchange refs
|
||||
let old_root_id = self.root_task_id.clone();
|
||||
self.remove(&old_root_id);
|
||||
|
||||
let new_root_id = root_task.id().clone();
|
||||
self.root_task_id = new_root_id;
|
||||
self.insert(root_task);
|
||||
}
|
||||
|
||||
pub fn first_exchange(&self) -> Option<&AIAgentExchange> {
|
||||
self.linearized_refs
|
||||
.first()
|
||||
.and_then(|r| self.lookup_exchange(r))
|
||||
}
|
||||
|
||||
pub fn latest_exchange(&self) -> Option<&AIAgentExchange> {
|
||||
self.linearized_refs
|
||||
.last()
|
||||
.and_then(|r| self.lookup_exchange(r))
|
||||
}
|
||||
|
||||
pub fn exchange_count(&self) -> usize {
|
||||
self.linearized_refs.len()
|
||||
}
|
||||
|
||||
pub fn all_exchanges(&self) -> impl Iterator<Item = &AIAgentExchange> {
|
||||
self.linearized_refs
|
||||
.iter()
|
||||
.filter_map(|r| self.lookup_exchange(r))
|
||||
}
|
||||
|
||||
pub fn all_exchanges_rev(&self) -> impl Iterator<Item = &AIAgentExchange> {
|
||||
self.linearized_refs
|
||||
.iter()
|
||||
.rev()
|
||||
.filter_map(|r| self.lookup_exchange(r))
|
||||
}
|
||||
|
||||
pub fn all_exchanges_by_task(&self) -> Vec<(TaskId, Vec<&AIAgentExchange>)> {
|
||||
let mut result: Vec<(TaskId, Vec<&AIAgentExchange>)> = Vec::new();
|
||||
|
||||
for exchange_ref in &self.linearized_refs {
|
||||
let Some(exchange) = self.lookup_exchange(exchange_ref) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// Check if we should append to the last group or start a new one
|
||||
if let Some((last_task_id, exchanges)) = result.last_mut() {
|
||||
if last_task_id == &exchange_ref.task_id {
|
||||
exchanges.push(exchange);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Start a new group
|
||||
result.push((exchange_ref.task_id.clone(), vec![exchange]));
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub fn latest_skills(&self) -> Option<Vec<SkillDescriptor>> {
|
||||
self.linearized_refs.iter().rev().find_map(|exchange_ref| {
|
||||
let exchange = self.lookup_exchange(exchange_ref);
|
||||
|
||||
if let Some(exchange) = exchange {
|
||||
let skills = exchange.input.iter().find_map(|input| {
|
||||
let context = match input {
|
||||
AIAgentInput::UserQuery { context, .. } => Some(context),
|
||||
AIAgentInput::ResumeConversation { context, .. } => Some(context),
|
||||
AIAgentInput::ActionResult { context, .. } => Some(context),
|
||||
AIAgentInput::TriggerPassiveSuggestion { context, .. } => Some(context),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
context.and_then(|ctx| {
|
||||
ctx.iter().find_map(|context| {
|
||||
if let AIAgentContext::Skills { skills } = context {
|
||||
Some(skills)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
});
|
||||
|
||||
skills.cloned()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns all messages in linearized DFS order, interleaving subtask messages
|
||||
/// immediately after their parent subagent call messages.
|
||||
pub fn all_linearized_messages(&self) -> Vec<&api::Message> {
|
||||
fn collect_messages_dfs<'a>(
|
||||
me: &'a TaskStore,
|
||||
messages: &mut Vec<&'a api::Message>,
|
||||
task: &'a Task,
|
||||
) {
|
||||
for message in task.messages() {
|
||||
messages.push(message);
|
||||
// If this message is a subagent call, recursively add subtask messages
|
||||
if let Some(subagent_call) = message
|
||||
.tool_call()
|
||||
.and_then(|tc: &api::message::ToolCall| tc.subagent())
|
||||
{
|
||||
if let Some(subtask) = me.get(&TaskId::new(subagent_call.task_id.clone())) {
|
||||
collect_messages_dfs(me, messages, subtask);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut messages = Vec::new();
|
||||
if let Some(root_task) = self.root_task() {
|
||||
collect_messages_dfs(self, &mut messages, root_task);
|
||||
}
|
||||
messages
|
||||
}
|
||||
|
||||
pub fn insert(&mut self, task: Task) {
|
||||
self.tasks.insert(task.id().clone(), task);
|
||||
self.rebuild_linearized_refs_index();
|
||||
}
|
||||
|
||||
pub fn remove(&mut self, task_id: &TaskId) -> Option<Task> {
|
||||
let task = self.tasks.remove(task_id)?;
|
||||
self.linearized_refs.retain(|r| &r.task_id != task_id);
|
||||
Some(task)
|
||||
}
|
||||
|
||||
fn lookup_exchange(&self, r: &ExchangeRef) -> Option<&AIAgentExchange> {
|
||||
self.tasks
|
||||
.get(&r.task_id)?
|
||||
.exchanges()
|
||||
.nth(r.exchange_index)
|
||||
}
|
||||
|
||||
/// Rebuilds the linearized index from scratch using DFS traversal.
|
||||
fn rebuild_linearized_refs_index(&mut self) {
|
||||
self.linearized_refs = Self::build_linearized_refs(&self.tasks, &self.root_task_id);
|
||||
}
|
||||
|
||||
/// Builds linearized exchange refs via DFS traversal without mutating self.
|
||||
/// This allows us to borrow `tasks` immutably throughout the traversal.
|
||||
fn build_linearized_refs(
|
||||
tasks: &HashMap<TaskId, Task>,
|
||||
root_task_id: &TaskId,
|
||||
) -> Vec<ExchangeRef> {
|
||||
let mut refs = Vec::new();
|
||||
|
||||
fn append_refs_for_task(
|
||||
tasks: &HashMap<TaskId, Task>,
|
||||
refs: &mut Vec<ExchangeRef>,
|
||||
task: &Task,
|
||||
) {
|
||||
let task_id = task.id().clone();
|
||||
|
||||
for (exchange_index, exchange) in task.exchanges().enumerate() {
|
||||
refs.push(ExchangeRef {
|
||||
task_id: task_id.clone(),
|
||||
exchange_index,
|
||||
});
|
||||
|
||||
// Check for subagent calls in the exchange output.
|
||||
if let Some(output) = exchange.output_status.output() {
|
||||
for output_message in output.get().messages.iter() {
|
||||
if let AIAgentOutputMessageType::Subagent(subagent_call) =
|
||||
&output_message.message
|
||||
{
|
||||
if let Some(subtask) =
|
||||
tasks.get(&TaskId::new(subagent_call.task_id.clone()))
|
||||
{
|
||||
append_refs_for_task(tasks, refs, subtask);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(root_task) = tasks.get(root_task_id) {
|
||||
append_refs_for_task(tasks, &mut refs, root_task);
|
||||
}
|
||||
|
||||
refs
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod testing {
|
||||
use crate::ai::agent::task::TaskId;
|
||||
|
||||
use super::TaskStore;
|
||||
|
||||
impl TaskStore {
|
||||
pub fn contains(&self, task_id: &TaskId) -> bool {
|
||||
self.tasks.contains_key(task_id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "task_store_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,625 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use chrono::Local;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai::{
|
||||
agent::{
|
||||
task::{Task, TaskId},
|
||||
AIAgentExchange, AIAgentExchangeId, AIAgentOutput, AIAgentOutputMessage,
|
||||
AIAgentOutputMessageType, AIAgentOutputStatus, FinishedAIAgentOutput, MessageId, Shared,
|
||||
SubagentCall,
|
||||
},
|
||||
llms::LLMId,
|
||||
};
|
||||
|
||||
use super::TaskStore;
|
||||
|
||||
fn create_test_exchange() -> AIAgentExchange {
|
||||
AIAgentExchange {
|
||||
id: AIAgentExchangeId::new(),
|
||||
input: vec![],
|
||||
output_status: AIAgentOutputStatus::Streaming { output: None },
|
||||
added_message_ids: HashSet::new(),
|
||||
start_time: Local::now(),
|
||||
finish_time: None,
|
||||
time_to_first_token_ms: None,
|
||||
working_directory: None,
|
||||
model_id: LLMId::from(""),
|
||||
request_cost: None,
|
||||
coding_model_id: LLMId::from(""),
|
||||
cli_agent_model_id: LLMId::from(""),
|
||||
computer_use_model_id: LLMId::from(""),
|
||||
response_initiator: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn create_test_task_with_exchanges(exchange_count: usize) -> Task {
|
||||
let mut task = Task::new_optimistic_root();
|
||||
for _ in 0..exchange_count {
|
||||
task.append_exchange(create_test_exchange());
|
||||
}
|
||||
task
|
||||
}
|
||||
|
||||
fn create_test_subtask_with_exchanges(exchange_count: usize) -> Task {
|
||||
use crate::terminal::model::block::BlockId;
|
||||
let mut task = Task::new_optimistic_cli_agent_subtask(BlockId::new());
|
||||
for _ in 0..exchange_count {
|
||||
task.append_exchange(create_test_exchange());
|
||||
}
|
||||
task
|
||||
}
|
||||
|
||||
/// Creates an exchange with a finished output containing a subagent call to the given task_id.
|
||||
fn create_exchange_with_subagent_call(subtask_id: &TaskId) -> AIAgentExchange {
|
||||
let output = AIAgentOutput {
|
||||
messages: vec![AIAgentOutputMessage {
|
||||
id: MessageId::new(Uuid::new_v4().to_string()),
|
||||
message: AIAgentOutputMessageType::Subagent(SubagentCall {
|
||||
task_id: subtask_id.to_string(),
|
||||
subagent_type: crate::ai::agent::SubagentType::Unknown,
|
||||
}),
|
||||
citations: vec![],
|
||||
}],
|
||||
citations: vec![],
|
||||
server_output_id: None,
|
||||
api_metadata_bytes: None,
|
||||
suggestions: None,
|
||||
telemetry_events: vec![],
|
||||
model_info: None,
|
||||
request_cost: None,
|
||||
};
|
||||
|
||||
AIAgentExchange {
|
||||
id: AIAgentExchangeId::new(),
|
||||
input: vec![],
|
||||
output_status: AIAgentOutputStatus::Finished {
|
||||
finished_output: FinishedAIAgentOutput::Success {
|
||||
output: Shared::new(output),
|
||||
},
|
||||
},
|
||||
added_message_ids: HashSet::new(),
|
||||
start_time: Local::now(),
|
||||
finish_time: None,
|
||||
time_to_first_token_ms: None,
|
||||
working_directory: None,
|
||||
model_id: LLMId::from(""),
|
||||
request_cost: None,
|
||||
coding_model_id: LLMId::from(""),
|
||||
cli_agent_model_id: LLMId::from(""),
|
||||
computer_use_model_id: LLMId::from(""),
|
||||
response_initiator: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_root_task() {
|
||||
let task = create_test_task_with_exchanges(3);
|
||||
let task_id = task.id().clone();
|
||||
let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
|
||||
let store = TaskStore::with_root_task(task);
|
||||
|
||||
assert_eq!(store.task_count(), 1);
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
assert_eq!(store.root_task_id(), &task_id);
|
||||
assert_eq!(store.root_task().expect("task exists").id(), &task_id);
|
||||
assert_eq!(store.first_exchange().map(|e| e.id), Some(exchange_ids[0]));
|
||||
assert_eq!(store.latest_exchange().map(|e| e.id), Some(exchange_ids[2]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_first_and_latest_exchange_o1() {
|
||||
let task = create_test_task_with_exchanges(5);
|
||||
let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
let store = TaskStore::with_root_task(task);
|
||||
|
||||
// These should be O(1) operations
|
||||
let first = store.first_exchange().expect("has exchanges");
|
||||
let latest = store.latest_exchange().expect("has exchanges");
|
||||
|
||||
assert_eq!(first.id, exchange_ids[0]);
|
||||
assert_eq!(latest.id, exchange_ids[4]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_subtask() {
|
||||
// Create root task with 1 exchange, then we'll add a subagent call exchange
|
||||
let root_task = create_test_task_with_exchanges(1);
|
||||
let root_task_id = root_task.id().clone();
|
||||
let mut store = TaskStore::with_root_task(root_task);
|
||||
|
||||
// Create subtask with 1 exchange
|
||||
let subtask = create_test_subtask_with_exchanges(1);
|
||||
let subtask_id = subtask.id().clone();
|
||||
|
||||
// Add exchange with subagent call to root task BEFORE inserting subtask
|
||||
let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
|
||||
store.append_exchange(&root_task_id, subagent_exchange);
|
||||
|
||||
// Now insert the subtask - its exchanges should be included via the subagent call
|
||||
store.insert(subtask);
|
||||
|
||||
assert_eq!(store.task_count(), 2);
|
||||
// 1 root exchange + 1 subagent call exchange + 1 subtask exchange = 3
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
assert!(store.get(&root_task_id).is_some());
|
||||
assert!(store.get(&subtask_id).is_some());
|
||||
assert!(store.contains(&subtask_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_task() {
|
||||
let task = create_test_task_with_exchanges(3);
|
||||
let task_id = task.id().clone();
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
assert_eq!(store.task_count(), 1);
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
|
||||
let removed = store.remove(&task_id);
|
||||
assert!(removed.is_some());
|
||||
assert_eq!(store.task_count(), 0);
|
||||
assert_eq!(store.exchange_count(), 0);
|
||||
assert!(store.get(&task_id).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_nonexistent_task() {
|
||||
let task = create_test_task_with_exchanges(2);
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
let nonexistent_id = TaskId::new(Uuid::new_v4().to_string());
|
||||
let removed = store.remove(&nonexistent_id);
|
||||
assert!(removed.is_none());
|
||||
assert_eq!(store.task_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_exchanges_iteration() {
|
||||
let task = create_test_task_with_exchanges(4);
|
||||
let expected_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
let store = TaskStore::with_root_task(task);
|
||||
|
||||
let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
|
||||
assert_eq!(actual_ids, expected_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_exchanges_by_task() {
|
||||
let task = create_test_task_with_exchanges(3);
|
||||
let task_id = task.id().clone();
|
||||
let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
let store = TaskStore::with_root_task(task);
|
||||
|
||||
let by_task = store.all_exchanges_by_task();
|
||||
assert_eq!(by_task.len(), 1);
|
||||
assert_eq!(by_task[0].0, task_id);
|
||||
assert_eq!(by_task[0].1.len(), 3);
|
||||
|
||||
let actual_ids: Vec<_> = by_task[0].1.iter().map(|e| e.id).collect();
|
||||
assert_eq!(actual_ids, exchange_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_set_root_task_replaces_old() {
|
||||
let task1 = create_test_task_with_exchanges(2);
|
||||
let task1_id = task1.id().clone();
|
||||
let mut store = TaskStore::with_root_task(task1);
|
||||
|
||||
let task2 = create_test_task_with_exchanges(3);
|
||||
let task2_id = task2.id().clone();
|
||||
let task2_exchange_ids: Vec<_> = task2.exchanges().map(|e| e.id).collect();
|
||||
store.set_root_task(task2);
|
||||
|
||||
assert_eq!(store.task_count(), 1);
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
assert!(store.get(&task1_id).is_none());
|
||||
assert!(store.get(&task2_id).is_some());
|
||||
assert_eq!(store.root_task_id(), &task2_id);
|
||||
|
||||
let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(actual_ids, task2_exchange_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_append_exchange() {
|
||||
let task = create_test_task_with_exchanges(2);
|
||||
let task_id = task.id().clone();
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
let new_exchange = create_test_exchange();
|
||||
let new_exchange_id = new_exchange.id;
|
||||
|
||||
let result = store.append_exchange(&task_id, new_exchange);
|
||||
assert!(result);
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
|
||||
// Verify the new exchange is accessible
|
||||
assert_eq!(store.latest_exchange().map(|e| e.id), Some(new_exchange_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_task_exchange() {
|
||||
let task = create_test_task_with_exchanges(3);
|
||||
let task_id = task.id().clone();
|
||||
let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
// First verify initial state
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
|
||||
// Remove the middle exchange
|
||||
let removed = store.remove_task_exchange(&task_id, exchange_ids[1]);
|
||||
assert!(removed.is_some());
|
||||
assert_eq!(removed.unwrap().id, exchange_ids[1]);
|
||||
|
||||
// After removal, we should have 2 exchanges
|
||||
assert_eq!(store.exchange_count(), 2);
|
||||
|
||||
// Verify the remaining exchanges are correct
|
||||
let remaining_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(remaining_ids, vec![exchange_ids[0], exchange_ids[2]]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_append_exchange_to_nonexistent_task() {
|
||||
let task = create_test_task_with_exchanges(1);
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
let nonexistent_id = TaskId::new(Uuid::new_v4().to_string());
|
||||
let result = store.append_exchange(&nonexistent_id, create_test_exchange());
|
||||
assert!(!result);
|
||||
assert_eq!(store.exchange_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_exchange_from_nonexistent_task() {
|
||||
let task = create_test_task_with_exchanges(1);
|
||||
let exchange_id = task.exchanges().next().unwrap().id;
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
let nonexistent_task_id = TaskId::new(Uuid::new_v4().to_string());
|
||||
let result = store.remove_task_exchange(&nonexistent_task_id, exchange_id);
|
||||
assert!(result.is_none());
|
||||
assert_eq!(store.exchange_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tasks_iteration() {
|
||||
let task1 = create_test_task_with_exchanges(2);
|
||||
let task1_id = task1.id().clone();
|
||||
let mut store = TaskStore::with_root_task(task1);
|
||||
|
||||
let task2 = create_test_subtask_with_exchanges(1);
|
||||
let task2_id = task2.id().clone();
|
||||
store.insert(task2);
|
||||
|
||||
let task_ids: HashSet<_> = store.tasks().map(|t| t.id().clone()).collect();
|
||||
assert_eq!(task_ids.len(), 2);
|
||||
assert!(task_ids.contains(&task1_id));
|
||||
assert!(task_ids.contains(&task2_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_tasks_exchange_order() {
|
||||
// Create root task with 1 exchange, then add subagent call exchange
|
||||
let root_task = create_test_task_with_exchanges(1);
|
||||
let root_task_id = root_task.id().clone();
|
||||
let first_root_exchange_id = root_task.exchanges().next().unwrap().id;
|
||||
let mut store = TaskStore::with_root_task(root_task);
|
||||
|
||||
// Create subtask with 2 exchanges
|
||||
let subtask = create_test_subtask_with_exchanges(2);
|
||||
let subtask_id = subtask.id().clone();
|
||||
let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();
|
||||
|
||||
// Add subagent call exchange to root
|
||||
let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
|
||||
let subagent_exchange_id = subagent_exchange.id;
|
||||
store.append_exchange(&root_task_id, subagent_exchange);
|
||||
|
||||
// Insert subtask - now linked via subagent call
|
||||
store.insert(subtask);
|
||||
|
||||
// 1 root + 1 subagent call + 2 subtask = 4 exchanges
|
||||
assert_eq!(store.exchange_count(), 4);
|
||||
|
||||
// First/last should still work
|
||||
assert_eq!(
|
||||
store.first_exchange().map(|e| e.id),
|
||||
Some(first_root_exchange_id)
|
||||
);
|
||||
assert_eq!(
|
||||
store.latest_exchange().map(|e| e.id),
|
||||
Some(subtask_exchange_ids[1])
|
||||
);
|
||||
|
||||
// Order: root[0], root[subagent_call], subtask[0], subtask[1]
|
||||
let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(all_ids.len(), 4);
|
||||
assert_eq!(all_ids[0], first_root_exchange_id);
|
||||
assert_eq!(all_ids[1], subagent_exchange_id);
|
||||
assert_eq!(all_ids[2], subtask_exchange_ids[0]);
|
||||
assert_eq!(all_ids[3], subtask_exchange_ids[1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_task_handling() {
|
||||
let task = create_test_task_with_exchanges(0);
|
||||
let task_id = task.id().clone();
|
||||
let store = TaskStore::with_root_task(task);
|
||||
|
||||
assert_eq!(store.task_count(), 1);
|
||||
assert_eq!(store.exchange_count(), 0);
|
||||
assert!(store.first_exchange().is_none());
|
||||
assert!(store.latest_exchange().is_none());
|
||||
assert!(store.get(&task_id).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_tasks() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
let task = create_test_task_with_exchanges(3);
|
||||
let task_id = task.id().clone();
|
||||
let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
|
||||
let mut tasks = HashMap::new();
|
||||
tasks.insert(task_id.clone(), task);
|
||||
|
||||
let store = TaskStore::from_tasks(tasks, task_id.clone());
|
||||
|
||||
assert_eq!(store.task_count(), 1);
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
assert_eq!(store.root_task_id(), &task_id);
|
||||
|
||||
let actual_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(actual_ids, exchange_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exchange_mut() {
|
||||
let task = create_test_task_with_exchanges(2);
|
||||
let exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
// Can find existing exchange
|
||||
let exchange = store.exchange_mut(exchange_ids[0]);
|
||||
assert!(exchange.is_some());
|
||||
assert_eq!(exchange.unwrap().id, exchange_ids[0]);
|
||||
|
||||
// Returns None for non-existent exchange
|
||||
let fake_id = AIAgentExchangeId::new();
|
||||
assert!(store.exchange_mut(fake_id).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_modify_task_conditional_rebuild() {
|
||||
let task = create_test_task_with_exchanges(2);
|
||||
let task_id = task.id().clone();
|
||||
let original_exchange_ids: Vec<_> = task.exchanges().map(|e| e.id).collect();
|
||||
let mut store = TaskStore::with_root_task(task);
|
||||
|
||||
// Verify initial state
|
||||
assert_eq!(store.exchange_count(), 2);
|
||||
|
||||
// modify_task with no exchange change should still work
|
||||
let result = store.modify_task(&task_id, |task| {
|
||||
assert_eq!(task.exchanges_len(), 2);
|
||||
"no change"
|
||||
});
|
||||
assert_eq!(result, Some("no change"));
|
||||
assert_eq!(store.exchange_count(), 2);
|
||||
|
||||
// modify_task that adds an exchange should update the index
|
||||
let new_exchange = create_test_exchange();
|
||||
let new_exchange_id = new_exchange.id;
|
||||
store.modify_task(&task_id, |task| {
|
||||
task.append_exchange(new_exchange);
|
||||
});
|
||||
assert_eq!(store.exchange_count(), 3);
|
||||
assert_eq!(store.latest_exchange().map(|e| e.id), Some(new_exchange_id));
|
||||
|
||||
// Verify all exchanges are in the index
|
||||
let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(all_ids.len(), 3);
|
||||
assert_eq!(all_ids[0], original_exchange_ids[0]);
|
||||
assert_eq!(all_ids[1], original_exchange_ids[1]);
|
||||
assert_eq!(all_ids[2], new_exchange_id);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Subtask Linearization Tests
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_linearization_parent_with_one_subtask() {
|
||||
// Create a subtask first so we have its ID
|
||||
let subtask = create_test_subtask_with_exchanges(2);
|
||||
let subtask_id = subtask.id().clone();
|
||||
let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();
|
||||
|
||||
// Create root task: exchange1, exchange_with_subagent_call, exchange3
|
||||
let mut root_task = Task::new_optimistic_root();
|
||||
let exchange1 = create_test_exchange();
|
||||
let exchange1_id = exchange1.id;
|
||||
root_task.append_exchange(exchange1);
|
||||
|
||||
let exchange_with_call = create_exchange_with_subagent_call(&subtask_id);
|
||||
let exchange_with_call_id = exchange_with_call.id;
|
||||
root_task.append_exchange(exchange_with_call);
|
||||
|
||||
let exchange3 = create_test_exchange();
|
||||
let exchange3_id = exchange3.id;
|
||||
root_task.append_exchange(exchange3);
|
||||
|
||||
// Build the store
|
||||
let mut store = TaskStore::with_root_task(root_task);
|
||||
store.insert(subtask);
|
||||
|
||||
// Total exchanges: 3 root + 2 subtask = 5
|
||||
assert_eq!(store.exchange_count(), 5);
|
||||
|
||||
// Expected order: root[0], root[1] (with subagent call), subtask[0], subtask[1], root[2]
|
||||
let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(all_ids.len(), 5);
|
||||
assert_eq!(all_ids[0], exchange1_id);
|
||||
assert_eq!(all_ids[1], exchange_with_call_id);
|
||||
assert_eq!(all_ids[2], subtask_exchange_ids[0]);
|
||||
assert_eq!(all_ids[3], subtask_exchange_ids[1]);
|
||||
assert_eq!(all_ids[4], exchange3_id);
|
||||
|
||||
// Verify first and last
|
||||
assert_eq!(store.first_exchange().map(|e| e.id), Some(exchange1_id));
|
||||
assert_eq!(store.latest_exchange().map(|e| e.id), Some(exchange3_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_linearization_nested_subtasks() {
|
||||
// Create nested subtask (grandchild) first
|
||||
let grandchild_subtask = create_test_subtask_with_exchanges(1);
|
||||
let grandchild_id = grandchild_subtask.id().clone();
|
||||
let grandchild_exchange_id = grandchild_subtask.exchanges().next().unwrap().id;
|
||||
|
||||
// Create child subtask with a call to grandchild
|
||||
use crate::terminal::model::block::BlockId;
|
||||
let mut child_subtask = Task::new_optimistic_cli_agent_subtask(BlockId::new());
|
||||
let child_id = child_subtask.id().clone();
|
||||
|
||||
let child_exchange1 = create_test_exchange();
|
||||
let child_exchange1_id = child_exchange1.id;
|
||||
child_subtask.append_exchange(child_exchange1);
|
||||
|
||||
let child_call_to_grandchild = create_exchange_with_subagent_call(&grandchild_id);
|
||||
let child_call_exchange_id = child_call_to_grandchild.id;
|
||||
child_subtask.append_exchange(child_call_to_grandchild);
|
||||
|
||||
// Create root task with a call to child
|
||||
let mut root_task = Task::new_optimistic_root();
|
||||
|
||||
let root_exchange1 = create_test_exchange();
|
||||
let root_exchange1_id = root_exchange1.id;
|
||||
root_task.append_exchange(root_exchange1);
|
||||
|
||||
let root_call_to_child = create_exchange_with_subagent_call(&child_id);
|
||||
let root_call_exchange_id = root_call_to_child.id;
|
||||
root_task.append_exchange(root_call_to_child);
|
||||
|
||||
// Build the store
|
||||
let mut store = TaskStore::with_root_task(root_task);
|
||||
store.insert(child_subtask);
|
||||
store.insert(grandchild_subtask);
|
||||
|
||||
// Total: 2 root + 2 child + 1 grandchild = 5
|
||||
assert_eq!(store.exchange_count(), 5);
|
||||
|
||||
// Expected DFS order:
|
||||
// root[0], root[1] (calls child) -> child[0], child[1] (calls grandchild) -> grandchild[0]
|
||||
let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(all_ids.len(), 5);
|
||||
assert_eq!(all_ids[0], root_exchange1_id);
|
||||
assert_eq!(all_ids[1], root_call_exchange_id);
|
||||
assert_eq!(all_ids[2], child_exchange1_id);
|
||||
assert_eq!(all_ids[3], child_call_exchange_id);
|
||||
assert_eq!(all_ids[4], grandchild_exchange_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_linearization_multiple_subtasks_same_parent() {
|
||||
// Create two subtasks
|
||||
let subtask1 = create_test_subtask_with_exchanges(1);
|
||||
let subtask1_id = subtask1.id().clone();
|
||||
let subtask1_exchange_id = subtask1.exchanges().next().unwrap().id;
|
||||
|
||||
let subtask2 = create_test_subtask_with_exchanges(2);
|
||||
let subtask2_id = subtask2.id().clone();
|
||||
let subtask2_exchange_ids: Vec<_> = subtask2.exchanges().map(|e| e.id).collect();
|
||||
|
||||
// Create root task with calls to both subtasks in separate exchanges
|
||||
let mut root_task = Task::new_optimistic_root();
|
||||
|
||||
let call_to_subtask1 = create_exchange_with_subagent_call(&subtask1_id);
|
||||
let call_to_subtask1_id = call_to_subtask1.id;
|
||||
root_task.append_exchange(call_to_subtask1);
|
||||
|
||||
let middle_exchange = create_test_exchange();
|
||||
let middle_exchange_id = middle_exchange.id;
|
||||
root_task.append_exchange(middle_exchange);
|
||||
|
||||
let call_to_subtask2 = create_exchange_with_subagent_call(&subtask2_id);
|
||||
let call_to_subtask2_id = call_to_subtask2.id;
|
||||
root_task.append_exchange(call_to_subtask2);
|
||||
|
||||
// Build the store
|
||||
let mut store = TaskStore::with_root_task(root_task);
|
||||
store.insert(subtask1);
|
||||
store.insert(subtask2);
|
||||
|
||||
// Total: 3 root + 1 subtask1 + 2 subtask2 = 6
|
||||
assert_eq!(store.exchange_count(), 6);
|
||||
|
||||
// Expected order:
|
||||
// root[0] (calls subtask1) -> subtask1[0], root[1], root[2] (calls subtask2) -> subtask2[0], subtask2[1]
|
||||
let all_ids: Vec<_> = store.all_exchanges().map(|e| e.id).collect();
|
||||
assert_eq!(all_ids.len(), 6);
|
||||
assert_eq!(all_ids[0], call_to_subtask1_id);
|
||||
assert_eq!(all_ids[1], subtask1_exchange_id);
|
||||
assert_eq!(all_ids[2], middle_exchange_id);
|
||||
assert_eq!(all_ids[3], call_to_subtask2_id);
|
||||
assert_eq!(all_ids[4], subtask2_exchange_ids[0]);
|
||||
assert_eq!(all_ids[5], subtask2_exchange_ids[1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_exchanges_by_task_with_subtasks() {
|
||||
// Create a subtask
|
||||
let subtask = create_test_subtask_with_exchanges(2);
|
||||
let subtask_id = subtask.id().clone();
|
||||
let subtask_exchange_ids: Vec<_> = subtask.exchanges().map(|e| e.id).collect();
|
||||
|
||||
// Create root task with a call to subtask
|
||||
let mut root_task = Task::new_optimistic_root();
|
||||
let root_id = root_task.id().clone();
|
||||
|
||||
let root_exchange1 = create_test_exchange();
|
||||
let root_exchange1_id = root_exchange1.id;
|
||||
root_task.append_exchange(root_exchange1);
|
||||
|
||||
let call_to_subtask = create_exchange_with_subagent_call(&subtask_id);
|
||||
let call_exchange_id = call_to_subtask.id;
|
||||
root_task.append_exchange(call_to_subtask);
|
||||
|
||||
let root_exchange3 = create_test_exchange();
|
||||
let root_exchange3_id = root_exchange3.id;
|
||||
root_task.append_exchange(root_exchange3);
|
||||
|
||||
// Build the store
|
||||
let mut store = TaskStore::with_root_task(root_task);
|
||||
store.insert(subtask);
|
||||
|
||||
// Check all_exchanges_by_task grouping
|
||||
let by_task = store.all_exchanges_by_task();
|
||||
|
||||
// Should have 3 groups: root[0-1], subtask[0-1], root[2]
|
||||
assert_eq!(by_task.len(), 3);
|
||||
|
||||
// First group: root task's first two exchanges
|
||||
assert_eq!(by_task[0].0, root_id);
|
||||
assert_eq!(by_task[0].1.len(), 2);
|
||||
assert_eq!(by_task[0].1[0].id, root_exchange1_id);
|
||||
assert_eq!(by_task[0].1[1].id, call_exchange_id);
|
||||
|
||||
// Second group: subtask's exchanges
|
||||
assert_eq!(by_task[1].0, subtask_id);
|
||||
assert_eq!(by_task[1].1.len(), 2);
|
||||
assert_eq!(by_task[1].1[0].id, subtask_exchange_ids[0]);
|
||||
assert_eq!(by_task[1].1[1].id, subtask_exchange_ids[1]);
|
||||
|
||||
// Third group: root task's last exchange
|
||||
assert_eq!(by_task[2].0, root_id);
|
||||
assert_eq!(by_task[2].1.len(), 1);
|
||||
assert_eq!(by_task[2].1[0].id, root_exchange3_id);
|
||||
}
|
||||
@@ -0,0 +1,586 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use crate::ai::agent::{
|
||||
AIAgentActionType, AIAgentExchange, AIAgentOutput, AIAgentOutputMessageType,
|
||||
AIAgentOutputStatus, MessageId, Shared,
|
||||
};
|
||||
use crate::ai::llms::LLMId;
|
||||
use crate::test_util::ai_agent_tasks::{
|
||||
create_api_subtask, create_api_task, create_message, create_subagent_tool_call_message,
|
||||
};
|
||||
use chrono::Local;
|
||||
use prost_types::FieldMask;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use super::{ExtractMessagesError, Task};
|
||||
|
||||
/// Creates a Task backed by server data from the given api::Task.
|
||||
fn create_server_task(api_task: api::Task) -> Task {
|
||||
Task::new_restored_root(api_task, std::iter::empty())
|
||||
}
|
||||
|
||||
fn create_streaming_exchange_with_output() -> AIAgentExchange {
|
||||
AIAgentExchange {
|
||||
id: Default::default(),
|
||||
input: vec![],
|
||||
output_status: AIAgentOutputStatus::Streaming {
|
||||
output: Some(Shared::new(AIAgentOutput::default())),
|
||||
},
|
||||
added_message_ids: HashSet::new(),
|
||||
start_time: Local::now(),
|
||||
finish_time: None,
|
||||
time_to_first_token_ms: None,
|
||||
working_directory: None,
|
||||
model_id: LLMId::from(""),
|
||||
request_cost: None,
|
||||
coding_model_id: LLMId::from(""),
|
||||
cli_agent_model_id: LLMId::from(""),
|
||||
computer_use_model_id: LLMId::from(""),
|
||||
response_initiator: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn create_start_agent_tool_call_message(
|
||||
id: &str,
|
||||
task_id: &str,
|
||||
name: &str,
|
||||
prompt: &str,
|
||||
) -> api::Message {
|
||||
api::Message {
|
||||
id: id.to_string(),
|
||||
task_id: task_id.to_string(),
|
||||
server_message_data: String::new(),
|
||||
citations: vec![],
|
||||
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
||||
tool_call_id: format!("{id}_tool_call"),
|
||||
tool: Some(api::message::tool_call::Tool::StartAgent(api::StartAgent {
|
||||
name: name.to_string(),
|
||||
prompt: prompt.to_string(),
|
||||
execution_mode: None,
|
||||
lifecycle_subscription: None,
|
||||
})),
|
||||
})),
|
||||
request_id: String::new(),
|
||||
timestamp: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_start_agent_prompt(
|
||||
task: &Task,
|
||||
exchange_id: crate::ai::agent::AIAgentExchangeId,
|
||||
prompt: &str,
|
||||
) {
|
||||
let exchange = task.exchange(exchange_id).expect("exchange should exist");
|
||||
let output = exchange
|
||||
.output_status
|
||||
.output()
|
||||
.expect("output should be initialized");
|
||||
let output = output.get();
|
||||
let output_message = output
|
||||
.messages
|
||||
.iter()
|
||||
.find(|message| message.id == MessageId::new("start_agent_message".to_string()))
|
||||
.expect("start agent output message should exist");
|
||||
|
||||
let AIAgentOutputMessageType::Action(action) = &output_message.message else {
|
||||
panic!("expected action output message");
|
||||
};
|
||||
let AIAgentActionType::StartAgent {
|
||||
prompt: current_prompt,
|
||||
..
|
||||
} = &action.action
|
||||
else {
|
||||
panic!("expected StartAgent action");
|
||||
};
|
||||
|
||||
assert_eq!(current_prompt, prompt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_upsert_message_adds_start_agent_prompt_to_output() {
|
||||
let task_id = "task1";
|
||||
let mut task = create_server_task(create_api_task(task_id, vec![]));
|
||||
|
||||
let exchange = create_streaming_exchange_with_output();
|
||||
let exchange_id = exchange.id;
|
||||
task.append_exchange(exchange);
|
||||
|
||||
task.upsert_message(
|
||||
create_start_agent_tool_call_message(
|
||||
"start_agent_message",
|
||||
task_id,
|
||||
"Agent 1",
|
||||
"run tests",
|
||||
),
|
||||
exchange_id,
|
||||
None,
|
||||
None,
|
||||
FieldMask {
|
||||
paths: vec!["message.tool_call".to_string()],
|
||||
},
|
||||
false,
|
||||
)
|
||||
.expect("initial upsert should succeed");
|
||||
assert_start_agent_prompt(&task, exchange_id, "run tests");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests for Task::splice_messages()
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_happy_path() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
create_message("m4", task_id),
|
||||
create_message("m5", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract m2, m3, m4 (middle 3 messages).
|
||||
let replacement = vec![create_message("replacement", task_id)];
|
||||
let result = task.splice_messages("m2", "m4", 3, replacement);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let extracted = result.unwrap();
|
||||
assert_eq!(extracted.len(), 3);
|
||||
assert_eq!(extracted[0].id, "m2");
|
||||
assert_eq!(extracted[1].id, "m3");
|
||||
assert_eq!(extracted[2].id, "m4");
|
||||
|
||||
// Verify the task now has: m1, replacement, m5.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["m1", "replacement", "m5"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_single_message() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract just m2.
|
||||
let replacement = vec![create_message("replacement", task_id)];
|
||||
let result = task.splice_messages("m2", "m2", 1, replacement);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let extracted = result.unwrap();
|
||||
assert_eq!(extracted.len(), 1);
|
||||
assert_eq!(extracted[0].id, "m2");
|
||||
|
||||
// Verify the task now has: m1, replacement, m3.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["m1", "replacement", "m3"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_all_messages() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract all messages.
|
||||
let replacement = vec![create_message("replacement", task_id)];
|
||||
let result = task.splice_messages("m1", "m3", 3, replacement);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let extracted = result.unwrap();
|
||||
assert_eq!(extracted.len(), 3);
|
||||
|
||||
// Verify the task now only has the replacement.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["replacement"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_empty_replacement() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract m2 with no replacement (pure deletion).
|
||||
let result = task.splice_messages("m2", "m2", 1, vec![]);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let extracted = result.unwrap();
|
||||
assert_eq!(extracted.len(), 1);
|
||||
assert_eq!(extracted[0].id, "m2");
|
||||
|
||||
// Verify the task now has: m1, m3.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["m1", "m3"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_multiple_replacements() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract m2 and replace with two messages.
|
||||
let replacement = vec![create_message("r1", task_id), create_message("r2", task_id)];
|
||||
let result = task.splice_messages("m2", "m2", 1, replacement);
|
||||
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Verify the task now has: m1, r1, r2, m3.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["m1", "r1", "r2", "m3"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_first_message_not_found() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![create_message("m1", task_id), create_message("m2", task_id)],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
let result = task.splice_messages("nonexistent", "m2", 1, vec![]);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ExtractMessagesError::FirstMessageNotFound(id)) if id == "nonexistent"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_last_message_not_found() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![create_message("m1", task_id), create_message("m2", task_id)],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
let result = task.splice_messages("m1", "nonexistent", 1, vec![]);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ExtractMessagesError::LastMessageNotFound(id)) if id == "nonexistent"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_invalid_range() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// first_message_id appears after last_message_id.
|
||||
let result = task.splice_messages("m3", "m1", 3, vec![]);
|
||||
|
||||
assert!(matches!(result, Err(ExtractMessagesError::InvalidRange)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_checksum_mismatch_too_few() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Claim there are 5 messages when there are only 3 in the range.
|
||||
let result = task.splice_messages("m1", "m3", 5, vec![]);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ExtractMessagesError::ChecksumMismatch {
|
||||
expected: 5,
|
||||
actual: 3
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_checksum_mismatch_too_many() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Claim there is 1 message when there are 3 in the range.
|
||||
let result = task.splice_messages("m1", "m3", 1, vec![]);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ExtractMessagesError::ChecksumMismatch {
|
||||
expected: 1,
|
||||
actual: 3
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_optimistic_task_not_initialized() {
|
||||
let mut task = Task::new_optimistic_root();
|
||||
|
||||
let result = task.splice_messages("m1", "m2", 2, vec![]);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ExtractMessagesError::TaskNotInitialized)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_from_beginning() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
create_message("m4", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract from the beginning.
|
||||
let replacement = vec![create_message("replacement", task_id)];
|
||||
let result = task.splice_messages("m1", "m2", 2, replacement);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let extracted = result.unwrap();
|
||||
assert_eq!(extracted.len(), 2);
|
||||
|
||||
// Verify the task now has: replacement, m3, m4.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["replacement", "m3", "m4"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_splice_messages_from_end() {
|
||||
let task_id = "task1";
|
||||
let api_task = create_api_task(
|
||||
task_id,
|
||||
vec![
|
||||
create_message("m1", task_id),
|
||||
create_message("m2", task_id),
|
||||
create_message("m3", task_id),
|
||||
create_message("m4", task_id),
|
||||
],
|
||||
);
|
||||
let mut task = create_server_task(api_task);
|
||||
|
||||
// Extract from the end.
|
||||
let replacement = vec![create_message("replacement", task_id)];
|
||||
let result = task.splice_messages("m3", "m4", 2, replacement);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let extracted = result.unwrap();
|
||||
assert_eq!(extracted.len(), 2);
|
||||
|
||||
// Verify the task now has: m1, m2, replacement.
|
||||
let remaining_ids: Vec<_> = task.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(remaining_ids, vec!["m1", "m2", "replacement"]);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests for Task::new_moved_messages_subtask()
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_new_moved_messages_subtask_basic() {
|
||||
let parent_id = "parent";
|
||||
let subtask_id = "subtask";
|
||||
|
||||
// Create parent task with a subagent call referencing the subtask.
|
||||
let parent_api_task = create_api_task(
|
||||
parent_id,
|
||||
vec![
|
||||
create_message("m1", parent_id),
|
||||
create_subagent_tool_call_message("subagent_call", parent_id, subtask_id, None),
|
||||
create_message("m2", parent_id),
|
||||
],
|
||||
);
|
||||
|
||||
// Create the subtask api::Task with some messages.
|
||||
let subtask_api_task = create_api_task(
|
||||
subtask_id,
|
||||
vec![
|
||||
create_message("s1", subtask_id),
|
||||
create_message("s2", subtask_id),
|
||||
],
|
||||
);
|
||||
|
||||
let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
|
||||
|
||||
assert_eq!(subtask.id().to_string(), subtask_id);
|
||||
assert!(subtask.exchanges().next().is_none()); // No exchanges.
|
||||
assert_eq!(subtask.messages().count(), 2);
|
||||
|
||||
// Should have subagent_params extracted from parent.
|
||||
let subagent_params = subtask.subagent_params();
|
||||
assert!(subagent_params.is_some());
|
||||
assert_eq!(
|
||||
subagent_params.unwrap().tool_call_id,
|
||||
"subagent_call_tool_call"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_moved_messages_subtask_with_summarization_metadata() {
|
||||
let parent_id = "parent";
|
||||
let subtask_id = "subtask";
|
||||
|
||||
// Create parent task with a summarization subagent call.
|
||||
let parent_api_task = create_api_task(
|
||||
parent_id,
|
||||
vec![create_subagent_tool_call_message(
|
||||
"summary_call",
|
||||
parent_id,
|
||||
subtask_id,
|
||||
Some(api::message::tool_call::subagent::Metadata::Summarization(
|
||||
(),
|
||||
)),
|
||||
)],
|
||||
);
|
||||
|
||||
let subtask_api_task = create_api_task(subtask_id, vec![create_message("s1", subtask_id)]);
|
||||
|
||||
let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
|
||||
|
||||
// Check that subagent_params has the summarization metadata.
|
||||
let subagent_params = subtask.subagent_params();
|
||||
assert!(subagent_params.is_some());
|
||||
|
||||
let call = &subagent_params.unwrap().call;
|
||||
assert!(matches!(
|
||||
call.metadata,
|
||||
Some(api::message::tool_call::subagent::Metadata::Summarization(
|
||||
_
|
||||
))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_moved_messages_subtask_no_matching_subagent_call() {
|
||||
let parent_id = "parent";
|
||||
let subtask_id = "subtask";
|
||||
|
||||
// Parent task has no subagent call to this subtask.
|
||||
let parent_api_task = create_api_task(
|
||||
parent_id,
|
||||
vec![
|
||||
create_message("m1", parent_id),
|
||||
// Subagent call references a different task.
|
||||
create_subagent_tool_call_message("other_call", parent_id, "other_task", None),
|
||||
],
|
||||
);
|
||||
|
||||
let subtask_api_task = create_api_task(subtask_id, vec![create_message("s1", subtask_id)]);
|
||||
|
||||
let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
|
||||
|
||||
// No subagent_params since no matching call was found.
|
||||
assert!(subtask.subagent_params().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_moved_messages_subtask_preserves_messages() {
|
||||
let parent_id = "parent";
|
||||
let subtask_id = "subtask";
|
||||
|
||||
let parent_api_task = create_api_task(
|
||||
parent_id,
|
||||
vec![create_subagent_tool_call_message(
|
||||
"call", parent_id, subtask_id, None,
|
||||
)],
|
||||
);
|
||||
|
||||
// Subtask with multiple messages.
|
||||
let subtask_api_task = create_api_task(
|
||||
subtask_id,
|
||||
vec![
|
||||
create_message("s1", subtask_id),
|
||||
create_message("s2", subtask_id),
|
||||
create_message("s3", subtask_id),
|
||||
],
|
||||
);
|
||||
|
||||
let subtask = Task::new_moved_messages_subtask(subtask_api_task, &parent_api_task);
|
||||
|
||||
// All messages should be preserved.
|
||||
let message_ids: Vec<_> = subtask.messages().map(|m| m.id.as_str()).collect();
|
||||
assert_eq!(message_ids, vec!["s1", "s2", "s3"]);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests for Warp docs subagent classification
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_is_warp_documentation_search_subagent() {
|
||||
let parent_id = "parent";
|
||||
let subtask_id = "subtask";
|
||||
let parent_api_task = create_api_task(
|
||||
parent_id,
|
||||
vec![create_subagent_tool_call_message(
|
||||
"docs_call",
|
||||
parent_id,
|
||||
subtask_id,
|
||||
Some(api::message::tool_call::subagent::Metadata::WarpDocumentationSearch(())),
|
||||
)],
|
||||
);
|
||||
let subtask_api_task = create_api_subtask(subtask_id, parent_id, vec![]);
|
||||
let subtask = Task::new_restored_subtask(subtask_api_task, &parent_api_task, vec![]);
|
||||
|
||||
assert!(subtask.is_warp_documentation_search_subagent());
|
||||
assert!(!subtask.is_conversation_search_subagent());
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
use serde::Serialize;
|
||||
use warpui::{AppContext, SingletonEntity};
|
||||
|
||||
use crate::ai::llms::LLMId;
|
||||
use crate::CloudModel;
|
||||
use crate::{
|
||||
server::telemetry::AgentModeCitation as CitationForTelemetry,
|
||||
terminal::view::block_onboarding::onboarding_agentic_suggestions_block::OnboardingChipType,
|
||||
};
|
||||
|
||||
use super::conversation::AIConversationId;
|
||||
use super::{
|
||||
AIAgentCitation, AIAgentExchangeId, EntrypointType, PassiveSuggestionTriggerType,
|
||||
ServerOutputId,
|
||||
};
|
||||
|
||||
pub trait ForTelemetry {
|
||||
type Output;
|
||||
|
||||
fn for_telemetry(&self, ctx: &AppContext) -> Option<Self::Output>;
|
||||
}
|
||||
|
||||
impl ForTelemetry for AIAgentCitation {
|
||||
type Output = CitationForTelemetry;
|
||||
|
||||
fn for_telemetry(&self, ctx: &AppContext) -> Option<Self::Output> {
|
||||
match self {
|
||||
Self::WarpDriveObject { uid } => {
|
||||
CloudModel::as_ref(ctx).get_by_uid(uid).map(|object| {
|
||||
CitationForTelemetry::WarpDriveObject {
|
||||
object_type: object.object_type(),
|
||||
uid: object.uid(),
|
||||
}
|
||||
})
|
||||
}
|
||||
Self::WarpDocumentation { path } => {
|
||||
Some(CitationForTelemetry::WarpDocs { page: path.clone() })
|
||||
}
|
||||
Self::WebPage { url } => Some(CitationForTelemetry::WebPage { url: url.clone() }),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl EntrypointType {
|
||||
pub fn entrypoint(&self) -> String {
|
||||
match self {
|
||||
Self::Onboarding { chip_type } => {
|
||||
format!(
|
||||
"ONBOARDING.{}",
|
||||
match chip_type {
|
||||
OnboardingChipType::FixAnIssue => "FIX_AN_ISSUE",
|
||||
OnboardingChipType::PullCloudLogs => "PULL_CLOUD_LOGS",
|
||||
OnboardingChipType::StartAFeature => "START_A_FEATURE",
|
||||
OnboardingChipType::PythonSnakeGame => "PYTHON_SNAKE_GAME",
|
||||
OnboardingChipType::ExploreGitHistory => "EXPLORE_GIT_HISTORY",
|
||||
OnboardingChipType::MatrixThemePicker => "MATRIX_THEME_PICKER",
|
||||
OnboardingChipType::Other => "OTHER",
|
||||
}
|
||||
)
|
||||
}
|
||||
Self::PromptSuggestion {
|
||||
is_static,
|
||||
is_coding,
|
||||
} => match (is_static, is_coding) {
|
||||
(true, true) => "PROMPT_SUGGESTION.CODING_STATIC".to_string(),
|
||||
(true, false) => "PROMPT_SUGGESTION.STATIC".to_string(),
|
||||
(false, true) => "PROMPT_SUGGESTION.CODING".to_string(),
|
||||
(false, false) => "PROMPT_SUGGESTION.SIMPLE".to_string(),
|
||||
},
|
||||
Self::ZeroStateAgentModePromptSuggestion => {
|
||||
"ZERO_STATE_AGENT_MODE_PROMPT_SUGGESTION".to_string()
|
||||
}
|
||||
Self::InitProjectRules => "INIT_PROJECT_RULES".to_string(),
|
||||
Self::UserInitiated => "USER_INITIATED".to_string(),
|
||||
Self::AgentInitiated => "AGENT_INITIATED".to_string(),
|
||||
Self::TriggerPassiveSuggestion { trigger } => {
|
||||
let trigger_name = match trigger {
|
||||
Some(PassiveSuggestionTriggerType::FilesChanged) => "FILES_CHANGED",
|
||||
Some(PassiveSuggestionTriggerType::CommandRun) => "COMMAND_RUN",
|
||||
Some(PassiveSuggestionTriggerType::ShellCommandCompleted) => {
|
||||
"SHELL_COMMAND_COMPLETED"
|
||||
}
|
||||
Some(PassiveSuggestionTriggerType::AgentResponseCompleted) => {
|
||||
"AGENT_RESPONSE_COMPLETED"
|
||||
}
|
||||
None => "NONE",
|
||||
};
|
||||
format!("TRIGGER_SUGGEST_PROMPT.{trigger_name}")
|
||||
}
|
||||
Self::CloneRepository => "CLONE_REPOSITORY".to_string(),
|
||||
Self::SharedSession => "SHARED_SESSION".to_string(),
|
||||
Self::ResumeConversation => "RESUME_CONVERSATION".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, Debug, Serialize)]
|
||||
pub struct AIIdentifiers {
|
||||
/// Useful for joining to client-side telemetry.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub client_conversation_id: Option<AIConversationId>,
|
||||
/// A stable ID to relate failures for the same underlying request.
|
||||
#[serde(rename = "exchange_id", skip_serializing_if = "Option::is_none")]
|
||||
pub client_exchange_id: Option<AIAgentExchangeId>,
|
||||
/// Unique ID for this output coming from the AI API. Generated by the server. Only passed in
|
||||
/// the initial response chunk.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub server_output_id: Option<ServerOutputId>,
|
||||
/// The conversation ID included in the response chunk.
|
||||
///
|
||||
/// Once this is set, it is never updated. That shouldn't be an issue because the conversation
|
||||
/// ID is only expected to be passed in the initial response chunk.
|
||||
///
|
||||
/// Note that this conversation ID is server-scoped; it is _not_ related to the
|
||||
/// `AIConversationId`, which is entirely a client-side abstraction.
|
||||
///
|
||||
/// This is mainly used for server-side logging and analytics.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub server_conversation_id: Option<String>,
|
||||
/// The ID of the model actually used to generate the output. This may differ from the requested model.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub model_id: Option<LLMId>,
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
use crate::ai::agent::AIAgentTodo;
|
||||
|
||||
use super::AIAgentTodoId;
|
||||
pub(crate) mod popup;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct AIAgentTodoList {
|
||||
completed_items: Vec<AIAgentTodo>,
|
||||
pending_items: Vec<AIAgentTodo>,
|
||||
}
|
||||
|
||||
impl AIAgentTodoList {
|
||||
pub fn with_pending_items(mut self, pending_items: Vec<AIAgentTodo>) -> Self {
|
||||
self.pending_items = pending_items;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_completed_items(mut self, completed_items: Vec<AIAgentTodo>) -> Self {
|
||||
self.completed_items = completed_items;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn update_pending_items(&mut self, pending_items: Vec<AIAgentTodo>) {
|
||||
self.pending_items = pending_items;
|
||||
}
|
||||
|
||||
pub fn clear_pending_items(&mut self) {
|
||||
self.pending_items.clear();
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.pending_items.len() + self.completed_items.len()
|
||||
}
|
||||
|
||||
pub fn is_finished(&self) -> bool {
|
||||
self.pending_items.is_empty() && !self.completed_items.is_empty()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.pending_items.is_empty() && self.completed_items.is_empty()
|
||||
}
|
||||
|
||||
pub fn in_progress_item(&self) -> Option<&AIAgentTodo> {
|
||||
self.pending_items.first()
|
||||
}
|
||||
|
||||
pub fn pending_items(&self) -> &[AIAgentTodo] {
|
||||
&self.pending_items
|
||||
}
|
||||
|
||||
pub fn completed_items(&self) -> &[AIAgentTodo] {
|
||||
&self.completed_items
|
||||
}
|
||||
|
||||
pub fn is_pending(&self, todo_id: &AIAgentTodoId) -> bool {
|
||||
self.pending_items.iter().any(|item| &item.id == todo_id)
|
||||
}
|
||||
|
||||
pub fn is_completed(&self, todo_id: &AIAgentTodoId) -> bool {
|
||||
self.completed_items.iter().any(|item| &item.id == todo_id)
|
||||
}
|
||||
|
||||
pub fn get_item(&self, todo_id: &AIAgentTodoId) -> Option<&AIAgentTodo> {
|
||||
self.items().find(|item| &item.id == todo_id)
|
||||
}
|
||||
|
||||
pub fn get_item_index(&self, todo_id: &AIAgentTodoId) -> Option<usize> {
|
||||
self.items().position(|item| &item.id == todo_id)
|
||||
}
|
||||
|
||||
fn items(&self) -> impl Iterator<Item = &AIAgentTodo> {
|
||||
self.completed_items.iter().chain(self.pending_items.iter())
|
||||
}
|
||||
|
||||
pub fn update_pending_todos(&mut self, todos: Vec<AIAgentTodo>) {
|
||||
self.pending_items = todos;
|
||||
}
|
||||
|
||||
pub fn mark_todos_complete(&mut self, completed_todo_ids: Vec<String>) {
|
||||
for completed_todo_id in completed_todo_ids.into_iter() {
|
||||
if let Some(item) = self
|
||||
.pending_items
|
||||
.iter()
|
||||
.position(|item| item.id == completed_todo_id.clone().into())
|
||||
.map(|i| self.pending_items.remove(i))
|
||||
{
|
||||
self.completed_items.push(item);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
use crate::ai::blocklist::{BlocklistAIContextEvent, BlocklistAIContextModel};
|
||||
use pathfinder_color::ColorU;
|
||||
use warp_core::ui::appearance::Appearance;
|
||||
use warp_core::ui::theme::Fill;
|
||||
use warpui::elements::{
|
||||
ClippedScrollStateHandle, ClippedScrollable, Dismiss, Empty, Expanded, ParentElement,
|
||||
SavePosition, ScrollTarget, ScrollToPositionMode, ScrollbarWidth, Shrinkable,
|
||||
};
|
||||
use warpui::fonts::FamilyId;
|
||||
use warpui::ModelHandle;
|
||||
use warpui::SingletonEntity;
|
||||
use warpui::{
|
||||
elements::{
|
||||
Border, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, DropShadow, Flex,
|
||||
MainAxisSize, Radius, Text,
|
||||
},
|
||||
fonts::{Properties, Weight},
|
||||
keymap::FixedBinding,
|
||||
AppContext, Element, Entity, EntityId, TypedActionView, View, ViewContext,
|
||||
};
|
||||
|
||||
use crate::ai::agent::icons::{in_progress_icon, pending_icon, succeeded_icon};
|
||||
use crate::ai::agent::todos::AIAgentTodoList;
|
||||
use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel};
|
||||
use crate::ui_components::blended_colors;
|
||||
|
||||
pub struct AgentTodosPopupView {
|
||||
terminal_view_id: EntityId,
|
||||
ai_context_model: ModelHandle<BlocklistAIContextModel>,
|
||||
scroll_state: ClippedScrollStateHandle,
|
||||
}
|
||||
|
||||
const IN_PROGRESS_POSITION_ID: &str = "AgentTodosPopup-in-progress";
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum AgentTodosPopupAction {
|
||||
ClosePopup,
|
||||
}
|
||||
|
||||
pub enum AgentTodosPopupEvent {
|
||||
Close,
|
||||
}
|
||||
|
||||
struct Styles {
|
||||
ui_font_family: FamilyId,
|
||||
background: Fill,
|
||||
main_text_color: ColorU,
|
||||
sub_text_color: ColorU,
|
||||
detail_font_size: f32,
|
||||
}
|
||||
|
||||
pub fn init(app: &mut AppContext) {
|
||||
use warpui::keymap::macros::*;
|
||||
|
||||
app.register_fixed_bindings([FixedBinding::new(
|
||||
"escape",
|
||||
AgentTodosPopupAction::ClosePopup,
|
||||
id!(AgentTodosPopupView::ui_name()),
|
||||
)]);
|
||||
}
|
||||
|
||||
impl AgentTodosPopupView {
|
||||
pub fn new(
|
||||
terminal_view_id: EntityId,
|
||||
ai_context_model: ModelHandle<BlocklistAIContextModel>,
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) -> Self {
|
||||
let blocklist_history_model = BlocklistAIHistoryModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&blocklist_history_model, move |me, _, event, ctx| {
|
||||
me.handle_blocklist_history_event(event, ctx);
|
||||
});
|
||||
ctx.subscribe_to_model(&ai_context_model, move |_, _, event, ctx| {
|
||||
if let BlocklistAIContextEvent::PendingQueryStateUpdated = event {
|
||||
ctx.notify();
|
||||
}
|
||||
});
|
||||
Self {
|
||||
terminal_view_id,
|
||||
ai_context_model,
|
||||
scroll_state: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_blocklist_history_event(
|
||||
&mut self,
|
||||
event: &BlocklistAIHistoryEvent,
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) {
|
||||
if let BlocklistAIHistoryEvent::UpdatedTodoList { terminal_view_id } = event {
|
||||
if *terminal_view_id == self.terminal_view_id {
|
||||
ctx.notify();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Scroll to the in-progress item, if not currently visible.
|
||||
pub fn scroll_to_in_progress_item(&self) {
|
||||
self.scroll_state.scroll_to_position(ScrollTarget {
|
||||
position_id: IN_PROGRESS_POSITION_ID.to_string(),
|
||||
mode: ScrollToPositionMode::FullyIntoView,
|
||||
});
|
||||
}
|
||||
|
||||
fn close(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
ctx.emit(AgentTodosPopupEvent::Close);
|
||||
}
|
||||
|
||||
fn styles(&self, appearance: &Appearance) -> Styles {
|
||||
let theme = appearance.theme();
|
||||
let background = theme.surface_1();
|
||||
let main_text_color = blended_colors::text_main(theme, background);
|
||||
let sub_text_color = blended_colors::text_sub(theme, background);
|
||||
let detail_font_size = appearance.ui_font_size();
|
||||
let ui_font_family = appearance.ui_font_family();
|
||||
|
||||
Styles {
|
||||
ui_font_family,
|
||||
background,
|
||||
main_text_color,
|
||||
sub_text_color,
|
||||
detail_font_size,
|
||||
}
|
||||
}
|
||||
|
||||
fn render_header(
|
||||
&self,
|
||||
app: &warpui::AppContext,
|
||||
todo_list: &AIAgentTodoList,
|
||||
) -> Box<dyn warpui::Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let styles = self.styles(appearance);
|
||||
let theme = appearance.theme();
|
||||
|
||||
let completed_count = todo_list.completed_items().len();
|
||||
let total_count = todo_list.pending_items().len() + completed_count;
|
||||
|
||||
let mut header_row = Flex::row().with_cross_axis_alignment(CrossAxisAlignment::Center);
|
||||
let mut header = Text::new(
|
||||
"Tasks".to_string(),
|
||||
appearance.header_font_family(),
|
||||
styles.detail_font_size + 2.,
|
||||
)
|
||||
.with_color(styles.main_text_color)
|
||||
.with_style(Properties::default().weight(Weight::Semibold));
|
||||
|
||||
header.add_text_with_highlights(
|
||||
format!(" {completed_count}/{total_count}"),
|
||||
theme.sub_text_color(theme.surface_1()).into(),
|
||||
Properties::default().weight(Weight::Semibold),
|
||||
);
|
||||
|
||||
header_row.add_child(header.finish());
|
||||
header_row.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl View for AgentTodosPopupView {
|
||||
fn ui_name() -> &'static str {
|
||||
"AgentTodosPopup"
|
||||
}
|
||||
|
||||
fn render(&self, app: &warpui::AppContext) -> Box<dyn warpui::Element> {
|
||||
let Some(todo_list) = self
|
||||
.ai_context_model
|
||||
.as_ref(app)
|
||||
.selected_conversation_todolist(app)
|
||||
else {
|
||||
// We don't have an empty state.
|
||||
// Assume the popup will only be shown if there are todos.
|
||||
return Empty::new().finish();
|
||||
};
|
||||
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let styles = self.styles(appearance);
|
||||
let theme = appearance.theme();
|
||||
|
||||
let background = styles.background;
|
||||
let main_text_color = styles.main_text_color;
|
||||
let sub_text_color = styles.sub_text_color;
|
||||
let detail_font_size = styles.detail_font_size;
|
||||
let ui_font_family = styles.ui_font_family;
|
||||
|
||||
let mut list_col = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_spacing(12.);
|
||||
|
||||
let items_with_icons = todo_list
|
||||
.completed_items()
|
||||
.iter()
|
||||
.map(|item| (item, succeeded_icon(appearance)))
|
||||
.chain(
|
||||
todo_list
|
||||
.pending_items()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, item)| {
|
||||
(
|
||||
item,
|
||||
if i == 0 {
|
||||
in_progress_icon(appearance)
|
||||
} else {
|
||||
pending_icon(appearance)
|
||||
},
|
||||
)
|
||||
}),
|
||||
);
|
||||
|
||||
for (item, status_icon) in items_with_icons {
|
||||
let mut row = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_size(MainAxisSize::Max);
|
||||
|
||||
// Status icon
|
||||
row.add_child(
|
||||
Container::new(
|
||||
ConstrainedBox::new(status_icon.finish())
|
||||
.with_width(16.)
|
||||
.with_height(16.)
|
||||
.finish(),
|
||||
)
|
||||
.with_margin_right(8.)
|
||||
.finish(),
|
||||
);
|
||||
|
||||
let is_in_progress = todo_list
|
||||
.in_progress_item()
|
||||
.map(|t| t.id.clone())
|
||||
.as_ref()
|
||||
.map(|id| &item.id == id)
|
||||
.unwrap_or(false);
|
||||
|
||||
// Title and status
|
||||
let mut text_col =
|
||||
Flex::column().with_cross_axis_alignment(CrossAxisAlignment::Stretch);
|
||||
|
||||
let text_color = if is_in_progress {
|
||||
main_text_color
|
||||
} else {
|
||||
sub_text_color
|
||||
};
|
||||
|
||||
text_col.add_child(
|
||||
Text::new(item.title.clone(), ui_font_family, detail_font_size)
|
||||
.with_color(text_color)
|
||||
.finish(),
|
||||
);
|
||||
|
||||
row.add_child(Expanded::new(1.0, text_col.finish()).finish());
|
||||
|
||||
let row = if is_in_progress {
|
||||
SavePosition::new(row.finish(), IN_PROGRESS_POSITION_ID).finish()
|
||||
} else {
|
||||
row.finish()
|
||||
};
|
||||
|
||||
list_col.add_child(row);
|
||||
}
|
||||
|
||||
let header = Container::new(self.render_header(app, todo_list))
|
||||
.with_padding_top(16.)
|
||||
.with_horizontal_padding(16.)
|
||||
.with_padding_bottom(8.)
|
||||
.finish();
|
||||
|
||||
let scrollable_body = ClippedScrollable::vertical(
|
||||
self.scroll_state.clone(),
|
||||
Container::new(list_col.finish())
|
||||
.with_horizontal_padding(16.)
|
||||
.with_padding_bottom(16.)
|
||||
.finish(),
|
||||
ScrollbarWidth::Auto,
|
||||
theme.nonactive_ui_detail().into(),
|
||||
theme.active_ui_detail().into(),
|
||||
warpui::elements::Fill::None,
|
||||
)
|
||||
.with_overlayed_scrollbar()
|
||||
.finish();
|
||||
|
||||
let panel_col = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_child(header)
|
||||
.with_child(Shrinkable::new(1.0, scrollable_body).finish());
|
||||
|
||||
Dismiss::new(
|
||||
ConstrainedBox::new(
|
||||
Container::new(panel_col.finish())
|
||||
.with_background(background)
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(8.)))
|
||||
.with_border(Border::all(1.).with_border_fill(theme.outline()))
|
||||
.with_drop_shadow(DropShadow::default())
|
||||
.finish(),
|
||||
)
|
||||
.with_width(300.)
|
||||
.with_max_height(420.)
|
||||
.finish(),
|
||||
)
|
||||
.prevent_interaction_with_other_elements()
|
||||
.on_dismiss(|ctx, _app| {
|
||||
ctx.dispatch_typed_action(AgentTodosPopupAction::ClosePopup);
|
||||
})
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for AgentTodosPopupView {
|
||||
type Event = AgentTodosPopupEvent;
|
||||
}
|
||||
|
||||
impl TypedActionView for AgentTodosPopupView {
|
||||
type Action = AgentTodosPopupAction;
|
||||
|
||||
fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
|
||||
match action {
|
||||
AgentTodosPopupAction::ClosePopup => {
|
||||
self.close(ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
use super::{
|
||||
AIAgentTextSection, AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram,
|
||||
AgentOutputTable, ProgrammingLanguage,
|
||||
};
|
||||
use crate::code::editor_management::CodeSource;
|
||||
use crate::features::FeatureFlag;
|
||||
use ai::gfm_table::{format_gfm_table, maybe_collect_gfm_table_lines};
|
||||
use itertools::Itertools;
|
||||
use lazy_static::lazy_static;
|
||||
use markdown_parser::{
|
||||
parse_image_run_line, parse_markdown_with_gfm_tables, FormattedImage, FormattedTextLine,
|
||||
};
|
||||
use mermaid_to_svg::is_mermaid_diagram;
|
||||
use regex::Regex;
|
||||
use std::{collections::HashMap, path::PathBuf};
|
||||
use warp_util::path::LineAndColumnArg;
|
||||
|
||||
lazy_static! {
|
||||
/// Markdown prefix for code blocks. Matches on triple backticks followed by a language.
|
||||
/// Importantly, parameters for linked code blocks are captured into their own group.
|
||||
static ref CODE_START_REGEX: Regex = Regex::new(r"^\s*```([\w+-]*)(.*)$").expect("Regex is valid");
|
||||
|
||||
/// Markdown suffix for code blocks.
|
||||
static ref CODE_END_REGEX: Regex = Regex::new(r"^\s*```\s*$").expect("Regex is valid");
|
||||
|
||||
/// Extracts key-value parameters from text in the format: key=value, used for code block metadata.
|
||||
/// Expects to match on text with format path=/path/to/file start=<line_number>
|
||||
static ref CODE_PARAMS_REGEX: Regex = Regex::new(r"(\w+)=([^\s]+)").expect("Regex is valid");
|
||||
}
|
||||
|
||||
/// Converts the given `markdown_text` into corresponding `Text` and `Code` `AIAgentOutputStep`s.
|
||||
pub(super) fn parse_markdown_into_text_and_code_sections(
|
||||
markdown_text: &str,
|
||||
) -> Vec<AIAgentTextSection> {
|
||||
let mut sections = vec![];
|
||||
let mut current_section = CurrentSection::PlainText(String::new());
|
||||
|
||||
let mut lines = markdown_text.lines().peekable();
|
||||
while let Some(line) = lines.next() {
|
||||
match &mut current_section {
|
||||
CurrentSection::PlainText(text) => {
|
||||
// Detect tables and render them as formatted table sections.
|
||||
if let Some(table_lines) = maybe_collect_gfm_table_lines(line, &mut lines, |l| {
|
||||
CODE_START_REGEX.is_match(l)
|
||||
}) {
|
||||
let markdown_source = table_lines.join("\n");
|
||||
let table_section = if FeatureFlag::BlocklistMarkdownTableRendering.is_enabled()
|
||||
{
|
||||
parse_agent_output_table(&markdown_source)
|
||||
} else {
|
||||
Some(AgentOutputTable::legacy(format_gfm_table(&table_lines)))
|
||||
};
|
||||
if let Some(table_section) = table_section {
|
||||
if !text.is_empty() {
|
||||
flush_plain_text_sections(text, &mut sections);
|
||||
text.clear();
|
||||
}
|
||||
sections.push(AIAgentTextSection::Table {
|
||||
table: table_section,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(&markdown_source);
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some((_, [language, param_str])) = CODE_START_REGEX
|
||||
.captures(line)
|
||||
.map(|capture_group| capture_group.extract())
|
||||
{
|
||||
if !text.is_empty() {
|
||||
flush_plain_text_sections(text, &mut sections);
|
||||
}
|
||||
|
||||
let source = {
|
||||
let mut params = HashMap::new();
|
||||
for (_, [key, value]) in CODE_PARAMS_REGEX
|
||||
.captures_iter(param_str)
|
||||
.map(|c| c.extract())
|
||||
{
|
||||
params.insert(key, value);
|
||||
}
|
||||
match (params.get("path"), params.get("start")) {
|
||||
(Some(path), Some(start)) => {
|
||||
start
|
||||
.parse::<usize>()
|
||||
.ok()
|
||||
.map(|line_num| CodeSource::Link {
|
||||
path: PathBuf::from(path),
|
||||
range_start: Some(LineAndColumnArg {
|
||||
line_num,
|
||||
column_num: None,
|
||||
}),
|
||||
range_end: None,
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
};
|
||||
current_section = CurrentSection::Code {
|
||||
code: String::new(),
|
||||
language_token: Some(language.to_owned()).filter(|l| !l.is_empty()),
|
||||
language: Some(language)
|
||||
.filter(|l| !l.is_empty())
|
||||
.map(|l| l.to_owned().into()),
|
||||
source,
|
||||
};
|
||||
} else {
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(line);
|
||||
}
|
||||
}
|
||||
CurrentSection::Code {
|
||||
code,
|
||||
language,
|
||||
language_token,
|
||||
source,
|
||||
} => {
|
||||
if CODE_END_REGEX.is_match(line) {
|
||||
if !code.is_empty() {
|
||||
if let Some(CodeSource::Link {
|
||||
range_start: Some(start),
|
||||
range_end,
|
||||
..
|
||||
}) = source.as_mut()
|
||||
{
|
||||
*range_end = Some(LineAndColumnArg {
|
||||
line_num: start.line_num + code.lines().count() - 1,
|
||||
column_num: None,
|
||||
});
|
||||
}
|
||||
push_code_or_mermaid_section(
|
||||
std::mem::take(code),
|
||||
language.clone(),
|
||||
language_token.as_deref(),
|
||||
source.take(),
|
||||
&mut sections,
|
||||
);
|
||||
}
|
||||
current_section = CurrentSection::PlainText(String::new());
|
||||
} else {
|
||||
if !code.is_empty() {
|
||||
code.push('\n');
|
||||
}
|
||||
code.push_str(line);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match current_section {
|
||||
CurrentSection::PlainText(text) => {
|
||||
flush_plain_text_sections(&text, &mut sections);
|
||||
}
|
||||
CurrentSection::Code {
|
||||
code,
|
||||
language,
|
||||
language_token,
|
||||
source,
|
||||
} => {
|
||||
push_code_or_mermaid_section(
|
||||
code,
|
||||
language,
|
||||
language_token.as_deref(),
|
||||
source,
|
||||
&mut sections,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if sections.is_empty() {
|
||||
sections.push(AIAgentTextSection::PlainText {
|
||||
text: String::new().into(),
|
||||
});
|
||||
}
|
||||
sections
|
||||
}
|
||||
|
||||
fn parse_agent_output_table(markdown_source: &str) -> Option<AgentOutputTable> {
|
||||
let formatted_text = parse_markdown_with_gfm_tables(markdown_source).ok()?;
|
||||
let table = formatted_text
|
||||
.lines
|
||||
.into_iter()
|
||||
.exactly_one()
|
||||
.ok()
|
||||
.and_then(|line| match line {
|
||||
FormattedTextLine::Table(table) => Some(table),
|
||||
_ => None,
|
||||
})?;
|
||||
Some(AgentOutputTable::structured(
|
||||
markdown_source.to_owned(),
|
||||
table,
|
||||
))
|
||||
}
|
||||
|
||||
enum CurrentSection {
|
||||
PlainText(String),
|
||||
Code {
|
||||
code: String,
|
||||
language_token: Option<String>,
|
||||
language: Option<ProgrammingLanguage>,
|
||||
source: Option<CodeSource>,
|
||||
},
|
||||
}
|
||||
|
||||
fn flush_plain_text_sections(markdown_text: &str, sections: &mut Vec<AIAgentTextSection>) {
|
||||
if markdown_text.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut plain_text = String::new();
|
||||
for line in markdown_text.split_inclusive('\n') {
|
||||
if let Some(images) = parse_image_run_line(line) {
|
||||
if !plain_text.is_empty() {
|
||||
sections.push(AIAgentTextSection::PlainText {
|
||||
text: std::mem::take(&mut plain_text).into(),
|
||||
});
|
||||
}
|
||||
|
||||
if images.len() == 1 {
|
||||
if let Some(image) = images.into_iter().next() {
|
||||
sections.push(image_section(image, AgentOutputImageLayout::Block));
|
||||
}
|
||||
} else {
|
||||
sections.extend(
|
||||
images
|
||||
.into_iter()
|
||||
.map(|image| image_section(image, AgentOutputImageLayout::Inline)),
|
||||
);
|
||||
}
|
||||
} else {
|
||||
plain_text.push_str(line);
|
||||
}
|
||||
}
|
||||
|
||||
if !plain_text.is_empty() {
|
||||
sections.push(AIAgentTextSection::PlainText {
|
||||
text: plain_text.into(),
|
||||
});
|
||||
}
|
||||
}
|
||||
fn image_section(image: FormattedImage, layout: AgentOutputImageLayout) -> AIAgentTextSection {
|
||||
AIAgentTextSection::Image {
|
||||
image: AgentOutputImage {
|
||||
markdown_source: markdown_source_for_image(&image),
|
||||
alt_text: image.alt_text,
|
||||
source: image.source,
|
||||
title: image.title,
|
||||
layout,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn markdown_source_for_image(image: &FormattedImage) -> String {
|
||||
warp_editor::content::text::format_image_markdown(
|
||||
&image.alt_text,
|
||||
&image.source,
|
||||
image.title.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn markdown_source_for_mermaid(source: &str) -> String {
|
||||
format!("```mermaid\n{source}\n```")
|
||||
}
|
||||
|
||||
fn push_code_or_mermaid_section(
|
||||
code: String,
|
||||
language: Option<ProgrammingLanguage>,
|
||||
language_token: Option<&str>,
|
||||
source: Option<CodeSource>,
|
||||
sections: &mut Vec<AIAgentTextSection>,
|
||||
) {
|
||||
if code.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
if language_token.is_some_and(is_mermaid_diagram) {
|
||||
sections.push(AIAgentTextSection::MermaidDiagram {
|
||||
diagram: AgentOutputMermaidDiagram {
|
||||
markdown_source: markdown_source_for_mermaid(&code),
|
||||
source: code,
|
||||
},
|
||||
});
|
||||
} else {
|
||||
sections.push(AIAgentTextSection::Code {
|
||||
code,
|
||||
language,
|
||||
source,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "util_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,280 @@
|
||||
use super::parse_markdown_into_text_and_code_sections;
|
||||
use crate::ai::agent::{AIAgentTextSection, AgentOutputImageLayout, AgentOutputTableRendering};
|
||||
use crate::features::FeatureFlag;
|
||||
|
||||
#[test]
|
||||
fn extracts_gfm_pipe_table_into_table_section() {
|
||||
let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(true);
|
||||
let input = "Intro\n\n| A | B |\n| --- | --- |\n| 1 | 2 |\n\nOutro";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 3);
|
||||
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::PlainText { text } => {
|
||||
assert!(text.text().contains("Intro"));
|
||||
}
|
||||
_ => panic!("expected first section to be PlainText"),
|
||||
}
|
||||
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::Table { table } => {
|
||||
assert_eq!(table.markdown_source, "| A | B |\n| --- | --- |\n| 1 | 2 |");
|
||||
match &table.rendering {
|
||||
AgentOutputTableRendering::Legacy { .. } => {
|
||||
panic!("expected structured table rendering")
|
||||
}
|
||||
AgentOutputTableRendering::Structured { table } => {
|
||||
assert_eq!(table.headers.len(), 2);
|
||||
assert_eq!(table.rows.len(), 1);
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
table.rendered_lines(),
|
||||
vec!["A\tB".to_string(), "1\t2".to_string()]
|
||||
);
|
||||
}
|
||||
_ => panic!("expected second section to be Table"),
|
||||
}
|
||||
|
||||
match §ions[2] {
|
||||
AIAgentTextSection::PlainText { text } => {
|
||||
assert!(text.text().contains("Outro"));
|
||||
}
|
||||
_ => panic!("expected third section to be PlainText"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_extract_pipe_text_without_separator_row() {
|
||||
let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(true);
|
||||
let input = "a | b\nc | d";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 1);
|
||||
assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn table_can_be_followed_immediately_by_text() {
|
||||
let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(true);
|
||||
let input = "| A | B |\n|---|---|\n| 1 | 2 |\nAfter";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 2);
|
||||
assert!(matches!(sections[0], AIAgentTextSection::Table { .. }));
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::PlainText { text } => {
|
||||
assert!(text.text().contains("After"));
|
||||
}
|
||||
_ => panic!("expected second section to be PlainText"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_gfm_pipe_table_into_legacy_table_section_when_flag_disabled() {
|
||||
let _flag = FeatureFlag::BlocklistMarkdownTableRendering.override_enabled(false);
|
||||
let input = "Intro\n\n| A | B |\n|---|---|\n| 1 | 2 |\n\nOutro";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 3);
|
||||
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::Table { table } => {
|
||||
assert_eq!(
|
||||
table.markdown_source,
|
||||
"| A | B |\n| --- | --- |\n| 1 | 2 |"
|
||||
);
|
||||
assert_eq!(
|
||||
table.rendered_lines(),
|
||||
vec![
|
||||
"| A | B |".to_string(),
|
||||
"| --- | --- |".to_string(),
|
||||
"| 1 | 2 |".to_string()
|
||||
]
|
||||
);
|
||||
assert!(matches!(
|
||||
&table.rendering,
|
||||
AgentOutputTableRendering::Legacy { .. }
|
||||
));
|
||||
}
|
||||
_ => panic!("expected second section to be Table"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_markdown_image_into_image_section() {
|
||||
let input = "Intro\n\n\n\nOutro";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 3);
|
||||
assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.alt_text, "Diagram");
|
||||
assert_eq!(image.source, "./diagram.png");
|
||||
assert_eq!(image.markdown_source, "");
|
||||
assert_eq!(image.layout, AgentOutputImageLayout::Block);
|
||||
}
|
||||
_ => panic!("expected second section to be Image"),
|
||||
}
|
||||
assert!(matches!(sections[2], AIAgentTextSection::PlainText { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_multiple_markdown_images_in_order() {
|
||||
let input = "\n\n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 2);
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.markdown_source, "");
|
||||
assert_eq!(image.layout, AgentOutputImageLayout::Block);
|
||||
}
|
||||
_ => panic!("expected first section to be Image"),
|
||||
}
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.markdown_source, "");
|
||||
assert_eq!(image.layout, AgentOutputImageLayout::Block);
|
||||
}
|
||||
_ => panic!("expected second section to be Image"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_same_line_markdown_images_into_inline_image_sections() {
|
||||
let input = " \n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 2);
|
||||
for (section, expected_markdown) in sections.iter().zip(["", ""])
|
||||
{
|
||||
match section {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.markdown_source, expected_markdown);
|
||||
assert_eq!(image.layout, AgentOutputImageLayout::Inline);
|
||||
}
|
||||
_ => panic!("expected inline image section"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_extract_inline_image_run_from_mixed_text_line() {
|
||||
let input = "Intro  \n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 1);
|
||||
assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_block_image_with_commonmark_title() {
|
||||
let input = "![Rex](./rex.png \"My dog Rex\")\n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 1);
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.alt_text, "Rex");
|
||||
assert_eq!(image.source, "./rex.png");
|
||||
assert_eq!(image.title.as_deref(), Some("My dog Rex"));
|
||||
// Right-click copy uses `markdown_source`, so it must round-trip
|
||||
// the authored title (product invariant 9).
|
||||
assert_eq!(image.markdown_source, "![Rex](./rex.png \"My dog Rex\")");
|
||||
}
|
||||
_ => panic!("expected image section"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_inline_image_run_with_partial_title() {
|
||||
// The inline run contains two images where only the second carries a title.
|
||||
let input = " ![Two](two.png \"caption\")\n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 2);
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.title, None);
|
||||
assert_eq!(image.markdown_source, "");
|
||||
assert_eq!(image.layout, AgentOutputImageLayout::Inline);
|
||||
}
|
||||
_ => panic!("expected first inline image"),
|
||||
}
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.title.as_deref(), Some("caption"));
|
||||
assert_eq!(image.markdown_source, "![Two](two.png \"caption\")");
|
||||
assert_eq!(image.layout, AgentOutputImageLayout::Inline);
|
||||
}
|
||||
_ => panic!("expected second inline image"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn block_image_with_empty_title_normalizes_to_none() {
|
||||
let input = "![Alt](image.png \"\")\n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 1);
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::Image { image } => {
|
||||
assert_eq!(image.title, None);
|
||||
// Empty titles normalize away, so `markdown_source` is the
|
||||
// canonical untitled form, not the original source text.
|
||||
assert_eq!(image.markdown_source, "");
|
||||
}
|
||||
_ => panic!("expected image section"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn block_image_with_unclosed_title_falls_back_to_plain_text() {
|
||||
let input = "![Alt](image.png \"unterminated)\n";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 1);
|
||||
// Unclosed titles cause the whole image to render as plain text.
|
||||
assert!(matches!(sections[0], AIAgentTextSection::PlainText { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_mermaid_code_block_into_mermaid_section() {
|
||||
let input = "```mermaid\ngraph TD\nA[Start] --> B[Finish]\n```";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 1);
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::MermaidDiagram { diagram } => {
|
||||
assert_eq!(diagram.source, "graph TD\nA[Start] --> B[Finish]");
|
||||
assert_eq!(
|
||||
diagram.markdown_source,
|
||||
"```mermaid\ngraph TD\nA[Start] --> B[Finish]\n```"
|
||||
);
|
||||
}
|
||||
_ => panic!("expected mermaid diagram section"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_multiple_mermaid_code_blocks_in_order() {
|
||||
let input = "```mermaid\ngraph TD\nA --> B\n```\n\n```mermaid\ngraph TD\nB --> C\n```";
|
||||
let sections = parse_markdown_into_text_and_code_sections(input);
|
||||
|
||||
assert_eq!(sections.len(), 2);
|
||||
match §ions[0] {
|
||||
AIAgentTextSection::MermaidDiagram { diagram } => {
|
||||
assert_eq!(diagram.source, "graph TD\nA --> B");
|
||||
}
|
||||
_ => panic!("expected first section to be MermaidDiagram"),
|
||||
}
|
||||
match §ions[1] {
|
||||
AIAgentTextSection::MermaidDiagram { diagram } => {
|
||||
assert_eq!(diagram.source, "graph TD\nB --> C");
|
||||
}
|
||||
_ => panic!("expected second section to be MermaidDiagram"),
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,366 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use futures::future::Either;
|
||||
use futures::StreamExt;
|
||||
use instant::Instant;
|
||||
use warpui::r#async::Timer;
|
||||
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_PROACTIVE_RECONNECT: Duration = Duration::from_secs(14 * 60);
|
||||
pub(crate) const DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG: usize = 5;
|
||||
|
||||
/// Configuration for the shared agent-event stream driver.
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct AgentEventDriverConfig {
|
||||
/// Run IDs whose events should be multiplexed into a single stream.
|
||||
pub run_ids: Vec<String>,
|
||||
/// Last fully handled event sequence. Events at or below this cursor are
|
||||
/// ignored on reconnect so the consumer only sees new work.
|
||||
pub since_sequence: i64,
|
||||
/// Exponential-ish reconnect delays, in seconds, used after stream open
|
||||
/// failures, stream errors, and clean stream termination.
|
||||
pub reconnect_backoff_steps: &'static [u64],
|
||||
/// Optional deadline for proactively recycling an otherwise healthy stream
|
||||
/// before upstream infrastructure times it out (for example, before Cloud
|
||||
/// Run's 20-minute streaming timeout).
|
||||
pub proactive_reconnect_after: Option<Duration>,
|
||||
/// Failure count at which reconnect logging is escalated from debug to warn.
|
||||
/// This only affects log severity; retry behavior stays the same.
|
||||
pub failures_before_error_log: usize,
|
||||
}
|
||||
|
||||
impl AgentEventDriverConfig {
|
||||
/// Build the production reconnecting configuration used by long-lived
|
||||
/// orchestration and harness listeners.
|
||||
pub(crate) fn retry_forever(run_ids: Vec<String>, since_sequence: i64) -> Self {
|
||||
Self {
|
||||
run_ids,
|
||||
since_sequence,
|
||||
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: Some(DEFAULT_AGENT_EVENT_PROACTIVE_RECONNECT),
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tells the shared driver whether to continue or stop after a handled event.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum AgentEventConsumerControlFlow {
|
||||
Continue,
|
||||
#[cfg_attr(not(test), allow(dead_code))]
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// High-level connection state updates emitted by the shared driver.
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum AgentEventDriverState {
|
||||
Connected,
|
||||
RetryScheduled {
|
||||
/// Number of consecutive failed reconnect cycles since the last
|
||||
/// successful stream open or event delivery.
|
||||
consecutive_failures: usize,
|
||||
/// Delay before the next reconnect attempt.
|
||||
backoff: Duration,
|
||||
/// Whether the retry happened before the stream ever connected.
|
||||
is_initial_connect: bool,
|
||||
},
|
||||
/// A healthy stream was intentionally recycled after
|
||||
/// `proactive_reconnect_after`.
|
||||
ProactiveReconnect,
|
||||
}
|
||||
|
||||
/// Parsed items emitted by an [`AgentEventSource`].
|
||||
pub(crate) enum AgentEventSourceItem {
|
||||
Open,
|
||||
Event(AgentRunEvent),
|
||||
}
|
||||
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(target_family = "wasm")] {
|
||||
type AgentEventSourceStream =
|
||||
futures::stream::LocalBoxStream<'static, Result<AgentEventSourceItem>>;
|
||||
} else {
|
||||
type AgentEventSourceStream =
|
||||
futures::stream::BoxStream<'static, Result<AgentEventSourceItem>>;
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens a stream of parsed agent events for one or more run IDs.
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
pub(crate) trait AgentEventSource: Send + Sync {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
run_ids: &[String],
|
||||
since_sequence: i64,
|
||||
) -> Result<AgentEventSourceStream>;
|
||||
}
|
||||
|
||||
/// [`AgentEventSource`] backed by [`ServerApi::stream_agent_events`].
|
||||
pub(crate) struct ServerApiAgentEventSource {
|
||||
server_api: Arc<ServerApi>,
|
||||
}
|
||||
|
||||
impl ServerApiAgentEventSource {
|
||||
pub(crate) fn new(server_api: Arc<ServerApi>) -> Self {
|
||||
Self { server_api }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
impl AgentEventSource for ServerApiAgentEventSource {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
run_ids: &[String],
|
||||
since_sequence: i64,
|
||||
) -> Result<AgentEventSourceStream> {
|
||||
let stream = self
|
||||
.server_api
|
||||
.stream_agent_events(run_ids, since_sequence)
|
||||
.await?;
|
||||
|
||||
let stream = stream.filter_map(|event_result| async move {
|
||||
match event_result {
|
||||
Ok(reqwest_eventsource::Event::Open) => Some(Ok(AgentEventSourceItem::Open)),
|
||||
Ok(reqwest_eventsource::Event::Message(message)) => {
|
||||
match serde_json::from_str::<AgentRunEvent>(&message.data) {
|
||||
Ok(event) => Some(Ok(AgentEventSourceItem::Event(event))),
|
||||
Err(err) => {
|
||||
log::warn!("Skipping malformed agent event from SSE stream: {err}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => Some(Err(anyhow!("SSE stream error: {err:?}"))),
|
||||
}
|
||||
});
|
||||
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(target_family = "wasm")] {
|
||||
Ok(stream.boxed_local())
|
||||
} else {
|
||||
Ok(stream.boxed())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Consumes events produced by [`run_agent_event_driver`].
|
||||
///
|
||||
/// Errors from [`on_event`](Self::on_event) are fatal because the event could not
|
||||
/// be safely processed. [`persist_cursor`](Self::persist_cursor) and
|
||||
/// [`on_driver_state`](Self::on_driver_state) are treated as best-effort hooks:
|
||||
/// their errors are logged and the driver continues.
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
pub(crate) trait AgentEventConsumer: Send {
|
||||
async fn on_event(&mut self, event: AgentRunEvent) -> Result<AgentEventConsumerControlFlow>;
|
||||
|
||||
async fn persist_cursor(&mut self, _sequence: i64) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn on_driver_state(&mut self, _state: AgentEventDriverState) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs a reconnecting agent-event stream until the consumer stops it or a
|
||||
/// fatal event-processing error occurs.
|
||||
pub(crate) async fn run_agent_event_driver<S, C>(
|
||||
source: S,
|
||||
config: AgentEventDriverConfig,
|
||||
consumer: &mut C,
|
||||
) -> Result<()>
|
||||
where
|
||||
S: AgentEventSource,
|
||||
C: AgentEventConsumer,
|
||||
{
|
||||
let mut since_sequence = config.since_sequence;
|
||||
let mut failures = 0usize;
|
||||
let mut has_connected_once = false;
|
||||
|
||||
loop {
|
||||
let mut stream = match source.open_stream(&config.run_ids, since_sequence).await {
|
||||
Ok(stream) => {
|
||||
failures = 0;
|
||||
has_connected_once = true;
|
||||
notify_driver_state(consumer, AgentEventDriverState::Connected).await;
|
||||
stream
|
||||
}
|
||||
Err(err) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
log_stream_failure(
|
||||
&config.run_ids,
|
||||
failures,
|
||||
backoff,
|
||||
&err,
|
||||
config.failures_before_error_log,
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures: failures,
|
||||
backoff,
|
||||
is_initial_connect: !has_connected_once,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
Timer::after(backoff).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let proactive_reconnect_deadline = config
|
||||
.proactive_reconnect_after
|
||||
.map(|duration| Instant::now() + duration);
|
||||
|
||||
loop {
|
||||
let next_item = if let Some(deadline) = proactive_reconnect_deadline {
|
||||
let remaining = deadline.saturating_duration_since(Instant::now());
|
||||
if remaining.is_zero() {
|
||||
NextDriverItem::ProactiveReconnect
|
||||
} else {
|
||||
let next_stream_item = stream.next();
|
||||
let reconnect_timer = Timer::after(remaining);
|
||||
futures::pin_mut!(next_stream_item);
|
||||
futures::pin_mut!(reconnect_timer);
|
||||
match futures::future::select(next_stream_item, reconnect_timer).await {
|
||||
Either::Left((stream_item, _)) => NextDriverItem::StreamItem(stream_item),
|
||||
Either::Right(_) => NextDriverItem::ProactiveReconnect,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
NextDriverItem::StreamItem(stream.next().await)
|
||||
};
|
||||
|
||||
match next_item {
|
||||
NextDriverItem::ProactiveReconnect => {
|
||||
notify_driver_state(consumer, AgentEventDriverState::ProactiveReconnect).await;
|
||||
break;
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Open))) => {
|
||||
failures = 0;
|
||||
log::info!("Agent event stream opened for {:?}", config.run_ids);
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Event(event)))) => {
|
||||
failures = 0;
|
||||
if event.sequence <= since_sequence {
|
||||
continue;
|
||||
}
|
||||
|
||||
let event_sequence = event.sequence;
|
||||
let control_flow = consumer.on_event(event).await?;
|
||||
since_sequence = event_sequence;
|
||||
|
||||
if let Err(err) = consumer.persist_cursor(since_sequence).await {
|
||||
log::warn!(
|
||||
"Ignoring agent event cursor persistence failure at sequence {since_sequence}: {err:#}"
|
||||
);
|
||||
}
|
||||
|
||||
if matches!(control_flow, AgentEventConsumerControlFlow::Stop) {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
NextDriverItem::StreamItem(Some(Err(err))) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
log_stream_failure(
|
||||
&config.run_ids,
|
||||
failures,
|
||||
backoff,
|
||||
&err,
|
||||
config.failures_before_error_log,
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures: failures,
|
||||
backoff,
|
||||
is_initial_connect: false,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
Timer::after(backoff).await;
|
||||
break;
|
||||
}
|
||||
NextDriverItem::StreamItem(None) => {
|
||||
failures += 1;
|
||||
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
|
||||
log::warn!(
|
||||
"Agent event stream closed for {:?}, reconnecting in {backoff:?}",
|
||||
config.run_ids
|
||||
);
|
||||
notify_driver_state(
|
||||
consumer,
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures: failures,
|
||||
backoff,
|
||||
is_initial_connect: false,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
Timer::after(backoff).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum NextDriverItem {
|
||||
StreamItem(Option<Result<AgentEventSourceItem>>),
|
||||
ProactiveReconnect,
|
||||
}
|
||||
|
||||
async fn notify_driver_state<C: AgentEventConsumer>(
|
||||
consumer: &mut C,
|
||||
state: AgentEventDriverState,
|
||||
) {
|
||||
if let Err(err) = consumer.on_driver_state(state.clone()).await {
|
||||
log::warn!("Ignoring agent event driver state callback error for {state:?}: {err:#}");
|
||||
}
|
||||
}
|
||||
|
||||
fn log_stream_failure(
|
||||
run_ids: &[String],
|
||||
failures: usize,
|
||||
backoff: Duration,
|
||||
err: &anyhow::Error,
|
||||
failures_before_error_log: usize,
|
||||
) {
|
||||
if agent_event_failures_exceeded_threshold(failures, failures_before_error_log) {
|
||||
log::error!(
|
||||
"Agent event stream failed {failures} consecutive times for {:?}, retrying in {backoff:?}: {err:#}",
|
||||
run_ids
|
||||
);
|
||||
} else {
|
||||
log::warn!(
|
||||
"Agent event stream failed for {:?}, retrying in {backoff:?}: {err:#}",
|
||||
run_ids
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn agent_event_backoff(failures: usize, backoff_steps: &[u64]) -> Duration {
|
||||
let safe_steps = if backoff_steps.is_empty() {
|
||||
DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS
|
||||
} else {
|
||||
backoff_steps
|
||||
};
|
||||
let index = failures.saturating_sub(1).min(safe_steps.len() - 1);
|
||||
Duration::from_secs(safe_steps[index])
|
||||
}
|
||||
|
||||
pub(crate) fn agent_event_failures_exceeded_threshold(failures: usize, threshold: usize) -> bool {
|
||||
failures >= threshold
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::anyhow;
|
||||
use async_trait::async_trait;
|
||||
use futures::stream::{self, BoxStream};
|
||||
use futures::StreamExt;
|
||||
|
||||
use super::*;
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
|
||||
const ZERO_BACKOFF_STEPS: &[u64] = &[0];
|
||||
|
||||
struct FakeAgentEventSource {
|
||||
responses: Mutex<VecDeque<anyhow::Result<Vec<anyhow::Result<AgentEventSourceItem>>>>>,
|
||||
}
|
||||
|
||||
impl FakeAgentEventSource {
|
||||
fn new(responses: Vec<anyhow::Result<Vec<anyhow::Result<AgentEventSourceItem>>>>) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(VecDeque::from(responses)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentEventSource for FakeAgentEventSource {
|
||||
async fn open_stream(
|
||||
&self,
|
||||
_run_ids: &[String],
|
||||
_since_sequence: i64,
|
||||
) -> anyhow::Result<BoxStream<'static, anyhow::Result<AgentEventSourceItem>>> {
|
||||
let response = self
|
||||
.responses
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.expect("fake response missing");
|
||||
let stream = response?;
|
||||
Ok(stream::iter(stream).boxed())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingConsumer {
|
||||
handled_sequences: Vec<i64>,
|
||||
persisted_sequences: Vec<i64>,
|
||||
driver_states: Vec<AgentEventDriverState>,
|
||||
stop_after: usize,
|
||||
fail_persist_cursor: bool,
|
||||
fail_driver_state: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentEventConsumer for RecordingConsumer {
|
||||
async fn on_event(
|
||||
&mut self,
|
||||
event: AgentRunEvent,
|
||||
) -> anyhow::Result<AgentEventConsumerControlFlow> {
|
||||
self.handled_sequences.push(event.sequence);
|
||||
if self.handled_sequences.len() >= self.stop_after {
|
||||
Ok(AgentEventConsumerControlFlow::Stop)
|
||||
} else {
|
||||
Ok(AgentEventConsumerControlFlow::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_cursor(&mut self, sequence: i64) -> anyhow::Result<()> {
|
||||
self.persisted_sequences.push(sequence);
|
||||
if self.fail_persist_cursor {
|
||||
Err(anyhow!("persist failed"))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_driver_state(&mut self, state: AgentEventDriverState) -> anyhow::Result<()> {
|
||||
self.driver_states.push(state);
|
||||
if self.fail_driver_state {
|
||||
Err(anyhow!("state callback failed"))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn make_run_event(
|
||||
sequence: i64,
|
||||
event_type: &str,
|
||||
run_id: &str,
|
||||
ref_id: Option<&str>,
|
||||
) -> AgentRunEvent {
|
||||
AgentRunEvent {
|
||||
event_type: event_type.to_string(),
|
||||
run_id: run_id.to_string(),
|
||||
ref_id: ref_id.map(|value| value.to_string()),
|
||||
execution_id: None,
|
||||
occurred_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
sequence,
|
||||
}
|
||||
}
|
||||
|
||||
fn ok_stream(
|
||||
items: Vec<anyhow::Result<AgentEventSourceItem>>,
|
||||
) -> anyhow::Result<Vec<anyhow::Result<AgentEventSourceItem>>> {
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_skips_duplicate_sequences_and_persists_new_cursor() {
|
||||
let source = FakeAgentEventSource::new(vec![ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
2,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-2"),
|
||||
))),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
3,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-3"),
|
||||
))),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
4,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-4"),
|
||||
))),
|
||||
])]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 2,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 2,
|
||||
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(consumer.handled_sequences, vec![3, 4]);
|
||||
assert_eq!(consumer.persisted_sequences, vec![3, 4]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_resets_failures_after_successful_event_delivery() {
|
||||
let source = FakeAgentEventSource::new(vec![
|
||||
ok_stream(vec![Ok(AgentEventSourceItem::Open), Err(anyhow!("boom-1"))]),
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
Err(anyhow!("boom-2")),
|
||||
]),
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
2,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-2"),
|
||||
))),
|
||||
]),
|
||||
]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 2,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let retry_failures = consumer
|
||||
.driver_states
|
||||
.into_iter()
|
||||
.filter_map(|state| match state {
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures,
|
||||
..
|
||||
} => Some(consecutive_failures),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(retry_failures, vec![1, 1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_ignores_persist_cursor_errors() {
|
||||
let source = FakeAgentEventSource::new(vec![ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
])]);
|
||||
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 1,
|
||||
fail_persist_cursor: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(consumer.handled_sequences, vec![1]);
|
||||
assert_eq!(consumer.persisted_sequences, vec![1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_ignores_driver_state_errors() {
|
||||
let source = FakeAgentEventSource::new(vec![ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
])]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 1,
|
||||
fail_driver_state: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(consumer.handled_sequences, vec![1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn driver_retries_initial_connection_until_stream_opens() {
|
||||
let source = FakeAgentEventSource::new(vec![
|
||||
Err(anyhow!("boom-1")),
|
||||
Err(anyhow!("boom-2")),
|
||||
ok_stream(vec![
|
||||
Ok(AgentEventSourceItem::Open),
|
||||
Ok(AgentEventSourceItem::Event(make_run_event(
|
||||
1,
|
||||
"new_message",
|
||||
"child-run",
|
||||
Some("msg-1"),
|
||||
))),
|
||||
]),
|
||||
]);
|
||||
let mut consumer = RecordingConsumer {
|
||||
stop_after: 1,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = AgentEventDriverConfig {
|
||||
run_ids: vec!["child-run".to_string()],
|
||||
since_sequence: 0,
|
||||
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
|
||||
proactive_reconnect_after: None,
|
||||
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
};
|
||||
|
||||
run_agent_event_driver(source, config, &mut consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(consumer.handled_sequences, vec![1]);
|
||||
let retry_failures = consumer
|
||||
.driver_states
|
||||
.into_iter()
|
||||
.filter_map(|state| match state {
|
||||
AgentEventDriverState::RetryScheduled {
|
||||
consecutive_failures,
|
||||
is_initial_connect,
|
||||
..
|
||||
} if is_initial_connect => Some(consecutive_failures),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(retry_failures, vec![1, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backoff_escalates_then_caps() {
|
||||
assert_eq!(
|
||||
agent_event_backoff(1, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(1)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(2, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(2)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(3, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(5)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(4, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(10)
|
||||
);
|
||||
assert_eq!(
|
||||
agent_event_backoff(100, DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS),
|
||||
Duration::from_secs(10)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_threshold_is_reached_at_and_above_limit() {
|
||||
assert!(!agent_event_failures_exceeded_threshold(4, 5));
|
||||
assert!(agent_event_failures_exceeded_threshold(5, 5));
|
||||
assert!(agent_event_failures_exceeded_threshold(6, 5));
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use futures::future::Either;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use warpui::r#async::Timer;
|
||||
|
||||
use crate::ai::agent::ReceivedMessageInput;
|
||||
use crate::server::server_api::ai::{AIClient, AgentRunEvent, ReadAgentMessageResponse};
|
||||
|
||||
pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Hydrates `new_message` agent events into full message payloads and delivery
|
||||
/// acknowledgements.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MessageHydrator {
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
fetch_timeout: Duration,
|
||||
}
|
||||
|
||||
impl MessageHydrator {
|
||||
pub(crate) fn new(ai_client: Arc<dyn AIClient>) -> Self {
|
||||
Self::with_fetch_timeout(ai_client, DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT)
|
||||
}
|
||||
|
||||
pub(crate) fn with_fetch_timeout(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
fetch_timeout: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
ai_client,
|
||||
fetch_timeout,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn hydrate_event_for_recipient(
|
||||
&self,
|
||||
event: &AgentRunEvent,
|
||||
recipient_run_id: &str,
|
||||
) -> Option<ReceivedMessageInput> {
|
||||
if event.event_type != "new_message" || event.run_id != recipient_run_id {
|
||||
return None;
|
||||
}
|
||||
|
||||
let message = match self.read_message_from_event_with_timeout(event).await {
|
||||
Ok(message) => message,
|
||||
Err(err) => {
|
||||
log::warn!(
|
||||
"Failed to hydrate agent message for event ref_id={:?}: {err:#}",
|
||||
event.ref_id
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some(ReceivedMessageInput {
|
||||
message_id: message.message_id,
|
||||
sender_agent_id: message.sender_run_id,
|
||||
addresses: vec![recipient_run_id.to_string()],
|
||||
subject: message.subject,
|
||||
message_body: message.body,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub(crate) async fn read_message_with_timeout(
|
||||
&self,
|
||||
message_id: &str,
|
||||
) -> Result<ReadAgentMessageResponse> {
|
||||
let read_message = self.ai_client.read_agent_message(message_id);
|
||||
let timeout = Timer::after(self.fetch_timeout);
|
||||
futures::pin_mut!(read_message);
|
||||
futures::pin_mut!(timeout);
|
||||
|
||||
match futures::future::select(read_message, timeout).await {
|
||||
Either::Left((result, _)) => {
|
||||
result.with_context(|| format!("Failed to read agent message {message_id}"))
|
||||
}
|
||||
Either::Right(_) => Err(anyhow!("Timed out reading agent message {message_id}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_family = "wasm")]
|
||||
pub(crate) async fn read_message_with_timeout(
|
||||
&self,
|
||||
message_id: &str,
|
||||
) -> Result<ReadAgentMessageResponse> {
|
||||
self.ai_client
|
||||
.read_agent_message(message_id)
|
||||
.await
|
||||
.with_context(|| format!("Failed to read agent message {message_id}"))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_message_from_event_with_timeout(
|
||||
&self,
|
||||
event: &AgentRunEvent,
|
||||
) -> Result<ReadAgentMessageResponse> {
|
||||
let Some(message_id) = event.ref_id.as_deref() else {
|
||||
return Err(anyhow!("Agent event is missing ref_id"));
|
||||
};
|
||||
self.read_message_with_timeout(message_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_message_delivered(&self, message_id: &str) -> Result<()> {
|
||||
self.ai_client
|
||||
.mark_message_delivered(message_id)
|
||||
.await
|
||||
.with_context(|| format!("Failed to mark agent message {message_id} as delivered"))
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_messages_delivered_best_effort<'a, I>(
|
||||
&self,
|
||||
message_ids: I,
|
||||
) -> Vec<(String, anyhow::Error)>
|
||||
where
|
||||
I: IntoIterator<Item = &'a str>,
|
||||
{
|
||||
let mut failures = Vec::new();
|
||||
// TODO(REMOTE-1266): Parallelize delivery acknowledgements for bursty
|
||||
// batches once the parent-bridge restore path is hardened enough to
|
||||
// tolerate a concurrent FuturesUnordered/join_all flow here.
|
||||
for message_id in message_ids {
|
||||
if let Err(err) = self.mark_message_delivered(message_id).await {
|
||||
failures.push((message_id.to_string(), err));
|
||||
}
|
||||
}
|
||||
failures
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use mockall::predicate::eq;
|
||||
|
||||
use super::*;
|
||||
use crate::server::server_api::ai::{
|
||||
AIClient, AgentRunEvent, MockAIClient, ReadAgentMessageResponse,
|
||||
};
|
||||
|
||||
fn make_run_event(
|
||||
sequence: i64,
|
||||
event_type: &str,
|
||||
run_id: &str,
|
||||
ref_id: Option<&str>,
|
||||
) -> AgentRunEvent {
|
||||
AgentRunEvent {
|
||||
event_type: event_type.to_string(),
|
||||
run_id: run_id.to_string(),
|
||||
ref_id: ref_id.map(|value| value.to_string()),
|
||||
execution_id: None,
|
||||
occurred_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
sequence,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hydrator_reads_new_message_for_matching_run() {
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.times(1)
|
||||
.returning(|_| {
|
||||
Ok(ReadAgentMessageResponse {
|
||||
message_id: "msg-123".to_string(),
|
||||
sender_run_id: "parent-run".to_string(),
|
||||
subject: "Need a redirect".to_string(),
|
||||
body: "Switch to the failing test first.".to_string(),
|
||||
sent_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
delivered_at: None,
|
||||
read_at: Some("2026-01-01T00:00:01Z".to_string()),
|
||||
})
|
||||
});
|
||||
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
|
||||
let hydrator = MessageHydrator::new(ai_client);
|
||||
let event = make_run_event(7, "new_message", "child-run", Some("msg-123"));
|
||||
|
||||
let hydrated = hydrator
|
||||
.hydrate_event_for_recipient(&event, "child-run")
|
||||
.await
|
||||
.expect("expected hydrated message");
|
||||
|
||||
assert_eq!(hydrated.message_id, "msg-123");
|
||||
assert_eq!(hydrated.sender_agent_id, "parent-run");
|
||||
assert_eq!(hydrated.subject, "Need a redirect");
|
||||
assert_eq!(hydrated.message_body, "Switch to the failing test first.");
|
||||
assert_eq!(hydrated.addresses, vec!["child-run".to_string()]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hydrator_ignores_events_for_other_runs() {
|
||||
let ai_client: Arc<dyn AIClient> = Arc::new(MockAIClient::new());
|
||||
let hydrator = MessageHydrator::new(ai_client);
|
||||
let event = make_run_event(7, "new_message", "other-run", Some("msg-123"));
|
||||
|
||||
assert!(hydrator
|
||||
.hydrate_event_for_recipient(&event, "child-run")
|
||||
.await
|
||||
.is_none());
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
//! Shared agent-event stream utilities used by orchestration consumers and
|
||||
//! third-party harness bridges.
|
||||
|
||||
mod driver;
|
||||
mod message_hydrator;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) use driver::{
|
||||
agent_event_backoff, agent_event_failures_exceeded_threshold, AgentEventDriverState,
|
||||
AgentEventSource, AgentEventSourceItem, DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
|
||||
DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
|
||||
};
|
||||
pub(crate) use driver::{
|
||||
run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow,
|
||||
AgentEventDriverConfig, ServerApiAgentEventSource,
|
||||
};
|
||||
pub(crate) use message_hydrator::MessageHydrator;
|
||||
|
||||
#[cfg(test)]
|
||||
mod driver_tests;
|
||||
#[cfg(test)]
|
||||
mod message_hydrator_tests;
|
||||
@@ -0,0 +1,536 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warpui::{AppContext, Entity, EntityId, ModelContext, SingletonEntity, WindowId};
|
||||
|
||||
use crate::settings::AISettings;
|
||||
|
||||
use crate::ai::active_agent_views_model::{ActiveAgentViewsEvent, ActiveAgentViewsModel};
|
||||
use crate::ai::agent::conversation::{AIConversationId, ConversationStatus};
|
||||
use crate::ai::agent_management::notifications::{
|
||||
NotificationCategory, NotificationId, NotificationItem, NotificationItems, NotificationOrigin,
|
||||
NotificationSourceAgent,
|
||||
};
|
||||
use crate::ai::artifacts::Artifact;
|
||||
use crate::ai::blocklist::BlocklistAIHistoryEvent;
|
||||
use crate::server::telemetry::TelemetryEvent;
|
||||
use crate::terminal::cli_agent_sessions::{
|
||||
CLIAgentSessionStatus, CLIAgentSessionsModel, CLIAgentSessionsModelEvent,
|
||||
};
|
||||
use crate::terminal::CLIAgent;
|
||||
use crate::workspace::util::is_terminal_view_in_same_tab;
|
||||
use crate::workspace::{Workspace, WorkspaceRegistry};
|
||||
use crate::BlocklistAIHistoryModel;
|
||||
use warp_core::send_telemetry_from_ctx;
|
||||
|
||||
/// Singleton model responsible for triggering in-app notifications on blocking conversation
|
||||
/// status updates and tracking/storing these notifications for the notifications mailbox.
|
||||
/// Tracks and stores notifications for both warp agent conversations and other supported
|
||||
/// cli agent sessions.
|
||||
pub struct AgentNotificationsModel {
|
||||
notifications: NotificationItems,
|
||||
/// Artifacts accumulated during the current turn for each conversation.
|
||||
/// Drained into the notification when a terminal state fires, cleared on InProgress.
|
||||
pub(crate) pending_artifacts: HashMap<AIConversationId, Vec<Artifact>>,
|
||||
}
|
||||
|
||||
impl Entity for AgentNotificationsModel {
|
||||
type Event = AgentManagementEvent;
|
||||
}
|
||||
|
||||
impl SingletonEntity for AgentNotificationsModel {}
|
||||
|
||||
impl AgentNotificationsModel {
|
||||
pub(crate) fn new(ctx: &mut ModelContext<Self>) -> Self {
|
||||
let history_model = BlocklistAIHistoryModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&history_model, move |me, event, ctx| {
|
||||
me.handle_history_event(event, ctx);
|
||||
});
|
||||
|
||||
let cli_sessions_model = CLIAgentSessionsModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&cli_sessions_model, |me, event, ctx| {
|
||||
me.handle_cli_agent_session_event(event, ctx);
|
||||
});
|
||||
|
||||
let active_views_model = ActiveAgentViewsModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&active_views_model, |me, event, ctx| {
|
||||
me.handle_active_agent_views_changed(event, ctx);
|
||||
});
|
||||
|
||||
Self {
|
||||
notifications: NotificationItems::default(),
|
||||
pending_artifacts: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn notifications(&self) -> &NotificationItems {
|
||||
&self.notifications
|
||||
}
|
||||
|
||||
pub(crate) fn mark_item_read(&mut self, id: NotificationId, ctx: &mut ModelContext<Self>) {
|
||||
if self.notifications.mark_item_read(id) {
|
||||
ctx.emit(AgentManagementEvent::NotificationUpdated);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn mark_all_items_read(&mut self, ctx: &mut ModelContext<Self>) {
|
||||
if self.notifications.mark_all_items_read() {
|
||||
ctx.emit(AgentManagementEvent::AllNotificationsMarkedRead);
|
||||
}
|
||||
}
|
||||
|
||||
/// Marks all notifications from the given terminal view as read.
|
||||
pub(crate) fn mark_items_from_terminal_view_read(
|
||||
&mut self,
|
||||
terminal_view_id: EntityId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if !FeatureFlag::HOANotifications.is_enabled() {
|
||||
return;
|
||||
}
|
||||
if self
|
||||
.notifications
|
||||
.mark_all_terminal_view_items_as_read(terminal_view_id)
|
||||
{
|
||||
ctx.emit(AgentManagementEvent::NotificationUpdated);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_active_agent_views_changed(
|
||||
&mut self,
|
||||
event: &ActiveAgentViewsEvent,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if !FeatureFlag::HOANotifications.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
match event {
|
||||
ActiveAgentViewsEvent::ConversationClosed { conversation_id } => {
|
||||
// When a conversation is closed, clean up its notifications
|
||||
// (as there's no conversation to navigate to when you click said notifications).
|
||||
if self
|
||||
.notifications
|
||||
.remove_by_origin(NotificationOrigin::Conversation(*conversation_id))
|
||||
{
|
||||
ctx.emit(AgentManagementEvent::NotificationUpdated);
|
||||
}
|
||||
}
|
||||
ActiveAgentViewsEvent::TerminalViewFocused
|
||||
| ActiveAgentViewsEvent::WindowClosed
|
||||
| ActiveAgentViewsEvent::AmbientSessionOpened { .. }
|
||||
| ActiveAgentViewsEvent::AmbientSessionClosed { .. } => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_cli_agent_session_event(
|
||||
&mut self,
|
||||
event: &CLIAgentSessionsModelEvent,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if !FeatureFlag::HOANotifications.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
match event {
|
||||
CLIAgentSessionsModelEvent::Ended {
|
||||
terminal_view_id, ..
|
||||
} => {
|
||||
self.remove_notification_by_source(
|
||||
NotificationOrigin::CLISession(*terminal_view_id),
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
CLIAgentSessionsModelEvent::Started { .. }
|
||||
| CLIAgentSessionsModelEvent::InputSessionChanged { .. }
|
||||
| CLIAgentSessionsModelEvent::SessionUpdated { .. } => {}
|
||||
CLIAgentSessionsModelEvent::StatusChanged {
|
||||
terminal_view_id,
|
||||
agent,
|
||||
status,
|
||||
session_context,
|
||||
} => match status {
|
||||
// When the agent resumes its work we can assume that the previous notification is stale.
|
||||
CLIAgentSessionStatus::InProgress => {
|
||||
self.remove_notification_by_source(
|
||||
NotificationOrigin::CLISession(*terminal_view_id),
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
CLIAgentSessionStatus::Success => {
|
||||
let title = session_context
|
||||
.display_title()
|
||||
.unwrap_or_else(|| format!("{} completed", agent.display_name()));
|
||||
let message = match agent {
|
||||
CLIAgent::Codex => "Notification from Codex",
|
||||
_ => "Task completed.",
|
||||
};
|
||||
self.add_notification(
|
||||
title,
|
||||
message.to_owned(),
|
||||
NotificationCategory::Complete,
|
||||
NotificationSourceAgent::CLI(*agent),
|
||||
NotificationOrigin::CLISession(*terminal_view_id),
|
||||
*terminal_view_id,
|
||||
vec![],
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
CLIAgentSessionStatus::Blocked { message } => {
|
||||
let title = session_context
|
||||
.display_title()
|
||||
.unwrap_or_else(|| format!("{} needs attention", agent.display_name()));
|
||||
self.add_notification(
|
||||
title,
|
||||
message
|
||||
.clone()
|
||||
.unwrap_or_else(|| "Waiting for input.".to_owned()),
|
||||
NotificationCategory::Request,
|
||||
NotificationSourceAgent::CLI(*agent),
|
||||
NotificationOrigin::CLISession(*terminal_view_id),
|
||||
*terminal_view_id,
|
||||
vec![],
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_history_event(
|
||||
&mut self,
|
||||
event: &BlocklistAIHistoryEvent,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
// When a conversation is deleted or removed, clean up its notification and pending artifacts.
|
||||
if let BlocklistAIHistoryEvent::DeletedConversation {
|
||||
conversation_id, ..
|
||||
}
|
||||
| BlocklistAIHistoryEvent::RemoveConversation {
|
||||
conversation_id, ..
|
||||
} = event
|
||||
{
|
||||
if FeatureFlag::HOANotifications.is_enabled() {
|
||||
self.pending_artifacts.remove(conversation_id);
|
||||
self.remove_notification_by_source(
|
||||
NotificationOrigin::Conversation(*conversation_id),
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Accumulate artifacts as they arrive during the conversation.
|
||||
if let BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
conversation_id,
|
||||
artifact,
|
||||
..
|
||||
} = event
|
||||
{
|
||||
if FeatureFlag::HOANotifications.is_enabled() {
|
||||
self.pending_artifacts
|
||||
.entry(*conversation_id)
|
||||
.or_default()
|
||||
.push(artifact.clone());
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
let BlocklistAIHistoryEvent::UpdatedConversationStatus {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
// We shouldn't trigger toasts when restoring conversations on startup.
|
||||
is_restored: false,
|
||||
} = event
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let ai_history_model = BlocklistAIHistoryModel::as_ref(ctx);
|
||||
let Some(updated_conversation) = ai_history_model.conversation(conversation_id) else {
|
||||
return;
|
||||
};
|
||||
|
||||
if updated_conversation.should_exclude_from_navigation() {
|
||||
return;
|
||||
}
|
||||
|
||||
let status = updated_conversation.status().clone();
|
||||
let latest_query = updated_conversation.latest_user_query();
|
||||
if FeatureFlag::HOANotifications.is_enabled() {
|
||||
self.handle_history_event_for_mailbox(
|
||||
&status,
|
||||
*conversation_id,
|
||||
latest_query,
|
||||
*terminal_view_id,
|
||||
ctx,
|
||||
);
|
||||
// The new mailbox path handled the event — skip the legacy toast path below.
|
||||
return;
|
||||
}
|
||||
|
||||
if !status.should_trigger_notification() {
|
||||
return;
|
||||
}
|
||||
|
||||
if is_terminal_view_visible(*terminal_view_id, ctx) {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some((window_id, tab_index)) =
|
||||
window_and_tab_idx_id_for_conversation(*conversation_id, ctx)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
ctx.emit(AgentManagementEvent::ConversationNeedsAttention {
|
||||
window_id,
|
||||
tab_index,
|
||||
terminal_view_id: *terminal_view_id,
|
||||
conversation_id: *conversation_id,
|
||||
});
|
||||
}
|
||||
|
||||
fn handle_history_event_for_mailbox(
|
||||
&mut self,
|
||||
status: &ConversationStatus,
|
||||
conversation_id: AIConversationId,
|
||||
latest_query: Option<String>,
|
||||
terminal_view_id: EntityId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let origin = NotificationOrigin::Conversation(conversation_id);
|
||||
|
||||
// If the conversation view is no longer open, don't create notifications for it
|
||||
// (there's nothing to navigate to when clicking them).
|
||||
if !ActiveAgentViewsModel::as_ref(ctx).is_conversation_open(conversation_id, ctx) {
|
||||
self.pending_artifacts.remove(&conversation_id);
|
||||
self.remove_notification_by_source(origin, ctx);
|
||||
return;
|
||||
}
|
||||
|
||||
let title = latest_query.unwrap_or_else(|| "Agent task".to_owned());
|
||||
|
||||
match status {
|
||||
// When the agent resumes its work, clear stale notifications.
|
||||
ConversationStatus::InProgress => {
|
||||
self.remove_notification_by_source(origin, ctx);
|
||||
}
|
||||
ConversationStatus::Success => {
|
||||
let artifacts = self.flush_pending_artifacts(conversation_id);
|
||||
self.add_notification(
|
||||
title,
|
||||
"Task completed.".to_owned(),
|
||||
NotificationCategory::Complete,
|
||||
NotificationSourceAgent::Oz,
|
||||
origin,
|
||||
terminal_view_id,
|
||||
artifacts,
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
ConversationStatus::Cancelled => {
|
||||
let artifacts = self.flush_pending_artifacts(conversation_id);
|
||||
self.add_notification(
|
||||
title,
|
||||
"Task was cancelled.".to_owned(),
|
||||
NotificationCategory::Complete,
|
||||
NotificationSourceAgent::Oz,
|
||||
origin,
|
||||
terminal_view_id,
|
||||
artifacts,
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
ConversationStatus::Blocked { blocked_action } => {
|
||||
self.add_notification(
|
||||
title,
|
||||
blocked_action.clone(),
|
||||
NotificationCategory::Request,
|
||||
NotificationSourceAgent::Oz,
|
||||
origin,
|
||||
terminal_view_id,
|
||||
vec![],
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
ConversationStatus::Error => {
|
||||
let artifacts = self.flush_pending_artifacts(conversation_id);
|
||||
self.add_notification(
|
||||
title,
|
||||
"Something went wrong.".to_owned(),
|
||||
NotificationCategory::Error,
|
||||
NotificationSourceAgent::Oz,
|
||||
origin,
|
||||
terminal_view_id,
|
||||
artifacts,
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Removes the existing notification for the given source (if any) and emits an update event.
|
||||
fn remove_notification_by_source(
|
||||
&mut self,
|
||||
origin: NotificationOrigin,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if self.notifications.remove_by_origin(origin) {
|
||||
ctx.emit(AgentManagementEvent::NotificationUpdated);
|
||||
}
|
||||
}
|
||||
|
||||
/// Drains and returns the pending artifacts for a conversation.
|
||||
pub(crate) fn flush_pending_artifacts(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
) -> Vec<Artifact> {
|
||||
self.pending_artifacts
|
||||
.remove(&conversation_id)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn add_notification(
|
||||
&mut self,
|
||||
title: String,
|
||||
message: String,
|
||||
category: NotificationCategory,
|
||||
agent: NotificationSourceAgent,
|
||||
origin: NotificationOrigin,
|
||||
terminal_view_id: EntityId,
|
||||
artifacts: Vec<Artifact>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if !*AISettings::as_ref(ctx).show_agent_notifications {
|
||||
return;
|
||||
}
|
||||
|
||||
let is_visible = is_terminal_view_visible(terminal_view_id, ctx);
|
||||
let branch = resolve_git_branch_for_terminal_view(terminal_view_id, ctx);
|
||||
let item = NotificationItem::new(
|
||||
title,
|
||||
message,
|
||||
category,
|
||||
agent,
|
||||
origin,
|
||||
is_visible,
|
||||
terminal_view_id,
|
||||
artifacts,
|
||||
branch,
|
||||
);
|
||||
send_telemetry_from_ctx!(
|
||||
TelemetryEvent::AgentNotificationShown {
|
||||
agent_variant: agent.into(),
|
||||
},
|
||||
ctx
|
||||
);
|
||||
|
||||
let id = item.id;
|
||||
self.notifications.push(item);
|
||||
ctx.emit(AgentManagementEvent::NotificationAdded { id });
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum AgentManagementEvent {
|
||||
/// A Warp-native conversation needs attention and is not visible in the current window/tab.
|
||||
ConversationNeedsAttention {
|
||||
window_id: WindowId,
|
||||
tab_index: usize,
|
||||
terminal_view_id: EntityId,
|
||||
conversation_id: AIConversationId,
|
||||
},
|
||||
/// A new notification was added to the persistent notification center.
|
||||
NotificationAdded { id: NotificationId },
|
||||
/// A notification's read state changed.
|
||||
NotificationUpdated,
|
||||
/// All notifications were marked as read.
|
||||
AllNotificationsMarkedRead,
|
||||
}
|
||||
|
||||
impl ConversationStatus {
|
||||
/// Returns true if the updating the conversation with this status should trigger some
|
||||
/// notification to the user.
|
||||
pub fn should_trigger_notification(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
ConversationStatus::Success
|
||||
| ConversationStatus::Blocked { .. }
|
||||
| ConversationStatus::Error
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn is_terminal_view_visible(terminal_view_id: EntityId, app: &AppContext) -> bool {
|
||||
let Some(active_id) = active_focused_terminal_id(app) else {
|
||||
return false;
|
||||
};
|
||||
active_id == terminal_view_id
|
||||
|| is_terminal_view_in_same_tab(&active_id, &terminal_view_id, app)
|
||||
}
|
||||
|
||||
fn window_and_tab_idx_id_for_conversation(
|
||||
conversation_id: AIConversationId,
|
||||
app: &AppContext,
|
||||
) -> Option<(WindowId, usize)> {
|
||||
WorkspaceRegistry::as_ref(app)
|
||||
.all_workspaces(app)
|
||||
.iter()
|
||||
.find_map(|(window_id, workspace_handle)| {
|
||||
workspace_handle
|
||||
.as_ref(app)
|
||||
.tab_views()
|
||||
.enumerate()
|
||||
.find_map(|(tab_idx, pane_group)| {
|
||||
pane_group
|
||||
.as_ref(app)
|
||||
.terminal_pane_ids()
|
||||
.filter_map(|pane_id| {
|
||||
pane_group
|
||||
.as_ref(app)
|
||||
.terminal_view_from_pane_id(pane_id, app)
|
||||
})
|
||||
.find_map(|terminal_view| {
|
||||
let terminal_view_conversation_id =
|
||||
terminal_view.as_ref(app).active_conversation_id(app)?;
|
||||
(terminal_view_conversation_id == conversation_id)
|
||||
.then_some((*window_id, tab_idx))
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_git_branch_for_terminal_view(
|
||||
terminal_view_id: EntityId,
|
||||
app: &AppContext,
|
||||
) -> Option<String> {
|
||||
for (_, workspace_handle) in WorkspaceRegistry::as_ref(app).all_workspaces(app) {
|
||||
for pane_group in workspace_handle.as_ref(app).tab_views() {
|
||||
let pane_group = pane_group.as_ref(app);
|
||||
for pane_id in pane_group.terminal_pane_ids() {
|
||||
if let Some(terminal_view) = pane_group.terminal_view_from_pane_id(pane_id, app) {
|
||||
if terminal_view.id() == terminal_view_id {
|
||||
return terminal_view.as_ref(app).current_git_branch(app);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn active_focused_terminal_id(app: &AppContext) -> Option<EntityId> {
|
||||
let active_window = app.windows().active_window()?;
|
||||
let workspace = app
|
||||
.views_of_type::<Workspace>(active_window)
|
||||
.and_then(|views| views.first().cloned())?;
|
||||
|
||||
let workspace = workspace.as_ref(app);
|
||||
workspace.active_terminal_id(app)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "agent_management_model_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,211 @@
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warpui::{App, EntityId, ModelHandle};
|
||||
|
||||
use crate::ai::active_agent_views_model::ActiveAgentViewsModel;
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::artifacts::Artifact;
|
||||
use crate::ai::blocklist::BlocklistAIHistoryEvent;
|
||||
use crate::terminal::cli_agent_sessions::CLIAgentSessionsModel;
|
||||
use crate::BlocklistAIHistoryModel;
|
||||
|
||||
use super::AgentNotificationsModel;
|
||||
|
||||
fn setup_app(
|
||||
app: &mut App,
|
||||
) -> (
|
||||
ModelHandle<BlocklistAIHistoryModel>,
|
||||
ModelHandle<AgentNotificationsModel>,
|
||||
) {
|
||||
let history = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
|
||||
app.add_singleton_model(|_| CLIAgentSessionsModel::new());
|
||||
app.add_singleton_model(|_| ActiveAgentViewsModel::new());
|
||||
let notifications = app.add_singleton_model(AgentNotificationsModel::new);
|
||||
(history, notifications)
|
||||
}
|
||||
|
||||
fn make_pr_artifact(url: &str, branch: &str) -> Artifact {
|
||||
Artifact::PullRequest {
|
||||
url: url.to_string(),
|
||||
branch: branch.to_string(),
|
||||
repo: None,
|
||||
number: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_plan_artifact(doc_uid: &str, title: &str) -> Artifact {
|
||||
Artifact::Plan {
|
||||
document_uid: doc_uid.to_string(),
|
||||
notebook_uid: None,
|
||||
title: Some(title.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn artifact_event_accumulates_into_pending() {
|
||||
App::test((), |mut app| async move {
|
||||
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
|
||||
let (history, notifications) = setup_app(&mut app);
|
||||
|
||||
let conversation_id = AIConversationId::new();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
artifact: make_pr_artifact("https://github.com/org/repo/pull/42", "feature-branch"),
|
||||
});
|
||||
});
|
||||
|
||||
notifications.read(&app, |model, _| {
|
||||
let pending = model.pending_artifacts.get(&conversation_id).unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert!(matches!(&pending[0], Artifact::PullRequest { branch, .. } if branch == "feature-branch"));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_artifacts_accumulated_across_turns() {
|
||||
App::test((), |mut app| async move {
|
||||
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
|
||||
let (history, notifications) = setup_app(&mut app);
|
||||
|
||||
let conversation_id = AIConversationId::new();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
artifact: make_plan_artifact("doc-1", "My Plan"),
|
||||
});
|
||||
});
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "main"),
|
||||
});
|
||||
});
|
||||
|
||||
notifications.read(&app, |model, _| {
|
||||
let pending = model.pending_artifacts.get(&conversation_id).unwrap();
|
||||
assert_eq!(pending.len(), 2);
|
||||
assert!(matches!(&pending[0], Artifact::Plan { title: Some(t), .. } if t == "My Plan"));
|
||||
assert!(matches!(&pending[1], Artifact::PullRequest { .. }));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_drains_pending_artifacts() {
|
||||
App::test((), |mut app| async move {
|
||||
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
|
||||
let (history, notifications) = setup_app(&mut app);
|
||||
|
||||
let conversation_id = AIConversationId::new();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "branch-1"),
|
||||
});
|
||||
});
|
||||
|
||||
notifications.update(&mut app, |model, _| {
|
||||
let artifacts = model.flush_pending_artifacts(conversation_id);
|
||||
assert_eq!(artifacts.len(), 1);
|
||||
assert!(matches!(&artifacts[0], Artifact::PullRequest { branch, .. } if branch == "branch-1"));
|
||||
});
|
||||
|
||||
notifications.read(&app, |model, _| {
|
||||
assert!(!model.pending_artifacts.contains_key(&conversation_id));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_returns_empty_vec_when_no_artifacts() {
|
||||
App::test((), |mut app| async move {
|
||||
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
|
||||
let (_history, notifications) = setup_app(&mut app);
|
||||
|
||||
let conversation_id = AIConversationId::new();
|
||||
|
||||
notifications.update(&mut app, |model, _| {
|
||||
let artifacts = model.flush_pending_artifacts(conversation_id);
|
||||
assert!(artifacts.is_empty());
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deletion_cleans_up_pending_artifacts() {
|
||||
App::test((), |mut app| async move {
|
||||
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
|
||||
let (history, notifications) = setup_app(&mut app);
|
||||
|
||||
let conversation_id = AIConversationId::new();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "branch-1"),
|
||||
});
|
||||
});
|
||||
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::DeletedConversation {
|
||||
terminal_view_id,
|
||||
conversation_id,
|
||||
conversation_title: None,
|
||||
});
|
||||
});
|
||||
|
||||
notifications.read(&app, |model, _| {
|
||||
assert!(!model.pending_artifacts.contains_key(&conversation_id));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn separate_conversations_have_independent_pending_artifacts() {
|
||||
App::test((), |mut app| async move {
|
||||
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
|
||||
let (history, notifications) = setup_app(&mut app);
|
||||
|
||||
let conv_a = AIConversationId::new();
|
||||
let conv_b = AIConversationId::new();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id: conv_a,
|
||||
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "branch-a"),
|
||||
});
|
||||
});
|
||||
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
|
||||
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
|
||||
terminal_view_id,
|
||||
conversation_id: conv_b,
|
||||
artifact: make_plan_artifact("doc-b", "Plan B"),
|
||||
});
|
||||
});
|
||||
|
||||
notifications.update(&mut app, |model, _| {
|
||||
let a = model.flush_pending_artifacts(conv_a);
|
||||
assert_eq!(a.len(), 1);
|
||||
assert!(matches!(&a[0], Artifact::PullRequest { branch, .. } if branch == "branch-a"));
|
||||
|
||||
let b = model.flush_pending_artifacts(conv_b);
|
||||
assert_eq!(b.len(), 1);
|
||||
assert!(matches!(&b[0], Artifact::Plan { title: Some(t), .. } if t == "Plan B"));
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,474 @@
|
||||
//! Agent type selector modal.
|
||||
//!
|
||||
//! This modal is displayed when users click "New agent" to choose between
|
||||
//! cloud and local agent modes.
|
||||
|
||||
use crate::appearance::Appearance;
|
||||
use crate::ui_components::icons::Icon;
|
||||
use pathfinder_color::ColorU;
|
||||
use pathfinder_geometry::vector::vec2f;
|
||||
use warp_core::ui::color::blend::Blend;
|
||||
use warp_core::ui::theme::color::internal_colors;
|
||||
use warpui::elements::{
|
||||
Align, Border, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, Dismiss,
|
||||
DropShadow, Element, Flex, Hoverable, MainAxisAlignment, MainAxisSize, MouseStateHandle,
|
||||
ParentElement, Radius, Shrinkable, Text,
|
||||
};
|
||||
use warpui::fonts::{Properties, Weight};
|
||||
use warpui::keymap::{FixedBinding, Keystroke};
|
||||
use warpui::platform::Cursor;
|
||||
use warpui::ui_components::components::{Coords, UiComponent, UiComponentStyles};
|
||||
use warpui::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext};
|
||||
|
||||
// Modal dimensions based on Figma design.
|
||||
const MODAL_WIDTH: f32 = 440.;
|
||||
const DIALOG_CORNER_RADIUS: f32 = 8.;
|
||||
|
||||
const HEADER_PADDING_TOP: f32 = 24.;
|
||||
const HEADER_PADDING_BOTTOM: f32 = 12.;
|
||||
const HEADER_PADDING_HORIZONTAL: f32 = 24.;
|
||||
|
||||
const BODY_PADDING_VERTICAL: f32 = 16.;
|
||||
const BODY_PADDING_HORIZONTAL: f32 = 20.;
|
||||
|
||||
const OPTION_PADDING_VERTICAL: f32 = 8.;
|
||||
const OPTION_PADDING_HORIZONTAL: f32 = 12.;
|
||||
const OPTION_CORNER_RADIUS: f32 = 4.;
|
||||
const OPTION_GAP: f32 = 12.;
|
||||
const OPTIONS_VERTICAL_GAP: f32 = 8.;
|
||||
|
||||
const AVATAR_SIZE: f32 = 48.;
|
||||
const AVATAR_ICON_SIZE: f32 = 24.;
|
||||
|
||||
const TITLE_FONT_SIZE: f32 = 16.;
|
||||
const OPTION_TITLE_FONT_SIZE: f32 = 14.;
|
||||
const OPTION_DESC_FONT_SIZE: f32 = 12.;
|
||||
|
||||
/// The type of agent selected by the user.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AgentType {
|
||||
/// Cloud agent - runs autonomously in a cloud environment.
|
||||
Cloud,
|
||||
/// Local agent - runs on the user's machine.
|
||||
Local,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AgentTypeSelectorAction {
|
||||
SelectCloudAgent,
|
||||
SelectLocalAgent,
|
||||
Dismiss,
|
||||
HoveredIn(usize),
|
||||
ArrowUp,
|
||||
ArrowDown,
|
||||
Enter,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AgentTypeSelectorEvent {
|
||||
Selected(AgentType),
|
||||
Dismissed,
|
||||
}
|
||||
|
||||
pub struct AgentTypeSelector {
|
||||
close_button_mouse_state: MouseStateHandle,
|
||||
cloud_agent_mouse_state: MouseStateHandle,
|
||||
local_agent_mouse_state: MouseStateHandle,
|
||||
dialog_mouse_state: MouseStateHandle,
|
||||
selected_option_index: usize,
|
||||
}
|
||||
|
||||
pub fn init(app: &mut AppContext) {
|
||||
use warpui::keymap::macros::*;
|
||||
|
||||
app.register_fixed_bindings(vec![
|
||||
FixedBinding::new(
|
||||
"escape",
|
||||
AgentTypeSelectorAction::Dismiss,
|
||||
id!("AgentTypeSelector"),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"enter",
|
||||
AgentTypeSelectorAction::Enter,
|
||||
id!("AgentTypeSelector"),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"numpadenter",
|
||||
AgentTypeSelectorAction::Enter,
|
||||
id!("AgentTypeSelector"),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"up",
|
||||
AgentTypeSelectorAction::ArrowUp,
|
||||
id!("AgentTypeSelector"),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"down",
|
||||
AgentTypeSelectorAction::ArrowDown,
|
||||
id!("AgentTypeSelector"),
|
||||
),
|
||||
]);
|
||||
}
|
||||
|
||||
impl AgentTypeSelector {
|
||||
pub fn new(_ctx: &mut ViewContext<Self>) -> Self {
|
||||
Self {
|
||||
close_button_mouse_state: MouseStateHandle::default(),
|
||||
cloud_agent_mouse_state: MouseStateHandle::default(),
|
||||
local_agent_mouse_state: MouseStateHandle::default(),
|
||||
dialog_mouse_state: MouseStateHandle::default(),
|
||||
// Cloud agent is selected by default (index 0).
|
||||
selected_option_index: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn render_header(&self, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let title = Text::new(
|
||||
"Choose your agent".to_string(),
|
||||
appearance.ui_font_family(),
|
||||
TITLE_FONT_SIZE,
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Bold))
|
||||
.with_color(theme.active_ui_text_color().into())
|
||||
.soft_wrap(true)
|
||||
.finish();
|
||||
|
||||
let close_button = appearance
|
||||
.ui_builder()
|
||||
.close_button(16., self.close_button_mouse_state.clone())
|
||||
.build()
|
||||
.on_click(|ctx, _, _| {
|
||||
ctx.dispatch_typed_action(AgentTypeSelectorAction::Dismiss);
|
||||
})
|
||||
.finish();
|
||||
|
||||
let esc_keystroke = Keystroke::parse("escape").expect("escape keystroke parses");
|
||||
let esc_pill = appearance
|
||||
.ui_builder()
|
||||
.keyboard_shortcut(&esc_keystroke)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(OPTION_DESC_FONT_SIZE),
|
||||
font_color: Some(theme.nonactive_ui_text_color().into_solid()),
|
||||
border_radius: Some(CornerRadius::with_all(Radius::Pixels(4.))),
|
||||
padding: Some(Coords {
|
||||
top: 0.,
|
||||
bottom: 0.,
|
||||
left: 3.,
|
||||
right: 3.,
|
||||
}),
|
||||
height: Some(16.),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish();
|
||||
|
||||
let right_controls = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_child(close_button)
|
||||
.with_child(Container::new(esc_pill).with_margin_left(4.).finish())
|
||||
.finish();
|
||||
|
||||
Flex::row()
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Start)
|
||||
.with_spacing(12.)
|
||||
.with_child(Shrinkable::new(1., title).finish())
|
||||
.with_child(right_controls)
|
||||
.finish()
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn render_option(
|
||||
&self,
|
||||
index: usize,
|
||||
icon: Icon,
|
||||
title: &'static str,
|
||||
description: &'static str,
|
||||
is_suggested: bool,
|
||||
mouse_state: MouseStateHandle,
|
||||
action: AgentTypeSelectorAction,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let font_family = appearance.ui_font_family();
|
||||
let active_text = theme.active_ui_text_color();
|
||||
let nonactive_text = theme.nonactive_ui_text_color();
|
||||
|
||||
let base_background = theme.surface_2();
|
||||
let hover_background = base_background.blend(&internal_colors::accent_overlay_1(theme));
|
||||
|
||||
let base_border = internal_colors::neutral_4(theme);
|
||||
let hover_border = theme.accent().into_solid();
|
||||
|
||||
let avatar_background = internal_colors::neutral_2(theme);
|
||||
let avatar_border = internal_colors::neutral_3(theme);
|
||||
|
||||
let badge_background = internal_colors::neutral_2(theme);
|
||||
let badge_border = internal_colors::neutral_3(theme);
|
||||
let badge_text_color = internal_colors::neutral_5(theme);
|
||||
|
||||
let icon_color = nonactive_text;
|
||||
|
||||
let is_selected = self.selected_option_index == index;
|
||||
let action = action.clone();
|
||||
Hoverable::new(mouse_state, move |state| {
|
||||
let is_hovered = state.is_hovered() || state.is_clicked();
|
||||
let (background, border_color) = if is_hovered || is_selected {
|
||||
(hover_background, hover_border)
|
||||
} else {
|
||||
(base_background, base_border)
|
||||
};
|
||||
|
||||
let avatar_icon = ConstrainedBox::new(icon.to_warpui_icon(icon_color).finish())
|
||||
.with_width(AVATAR_ICON_SIZE)
|
||||
.with_height(AVATAR_ICON_SIZE)
|
||||
.finish();
|
||||
|
||||
let avatar_contents = ConstrainedBox::new(
|
||||
Flex::row()
|
||||
.with_main_axis_alignment(MainAxisAlignment::Center)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(avatar_icon)
|
||||
.finish(),
|
||||
)
|
||||
.with_width(AVATAR_SIZE)
|
||||
.with_height(AVATAR_SIZE)
|
||||
.finish();
|
||||
|
||||
let avatar = Container::new(avatar_contents)
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Percentage(50.)))
|
||||
.with_background(avatar_background)
|
||||
.with_border(Border::all(1.).with_border_color(avatar_border))
|
||||
.finish();
|
||||
|
||||
let title_text = Text::new(title.to_string(), font_family, OPTION_TITLE_FONT_SIZE)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(active_text.into())
|
||||
.finish();
|
||||
|
||||
let mut title_row = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_spacing(4.)
|
||||
.with_child(title_text);
|
||||
|
||||
if is_suggested {
|
||||
let suggested_text =
|
||||
Text::new("Suggested".to_string(), font_family, OPTION_DESC_FONT_SIZE)
|
||||
.with_style(Properties::default().weight(Weight::Medium))
|
||||
.with_color(badge_text_color)
|
||||
.finish();
|
||||
|
||||
let suggested = Container::new(suggested_text)
|
||||
.with_horizontal_padding(8.)
|
||||
.with_vertical_padding(2.)
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Percentage(50.)))
|
||||
.with_background(badge_background)
|
||||
.with_border(Border::all(1.).with_border_color(badge_border))
|
||||
.finish();
|
||||
|
||||
title_row.add_child(Container::new(suggested).with_margin_left(8.).finish());
|
||||
}
|
||||
|
||||
let description_text =
|
||||
Text::new(description.to_string(), font_family, OPTION_DESC_FONT_SIZE)
|
||||
.with_style(Properties::default().weight(Weight::Normal))
|
||||
.with_color(nonactive_text.into())
|
||||
.soft_wrap(true)
|
||||
.finish();
|
||||
|
||||
let text_content = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Start)
|
||||
.with_child(title_row.finish())
|
||||
.with_child(
|
||||
Container::new(description_text)
|
||||
.with_margin_top(4.)
|
||||
.finish(),
|
||||
)
|
||||
.finish();
|
||||
|
||||
Container::new(
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_spacing(OPTION_GAP)
|
||||
.with_child(avatar)
|
||||
.with_child(Shrinkable::new(1., text_content).finish())
|
||||
.finish(),
|
||||
)
|
||||
.with_padding_left(OPTION_PADDING_HORIZONTAL)
|
||||
.with_padding_right(OPTION_PADDING_HORIZONTAL)
|
||||
.with_padding_top(OPTION_PADDING_VERTICAL)
|
||||
.with_padding_bottom(OPTION_PADDING_VERTICAL)
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(OPTION_CORNER_RADIUS)))
|
||||
.with_border(Border::all(1.).with_border_color(border_color))
|
||||
.with_background(background)
|
||||
.finish()
|
||||
})
|
||||
.with_cursor(Cursor::PointingHand)
|
||||
.additional_on_hover(move |is_hovered, ctx, _app, _pos| {
|
||||
if is_hovered {
|
||||
ctx.dispatch_typed_action(AgentTypeSelectorAction::HoveredIn(index));
|
||||
}
|
||||
})
|
||||
.on_click(move |ctx, _, _| {
|
||||
ctx.dispatch_typed_action(action.clone());
|
||||
})
|
||||
.finish()
|
||||
}
|
||||
|
||||
fn render_modal(&self, appearance: &Appearance, app: &AppContext) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let header = Container::new(self.render_header(appearance))
|
||||
.with_padding_top(HEADER_PADDING_TOP)
|
||||
.with_padding_bottom(HEADER_PADDING_BOTTOM)
|
||||
.with_padding_left(HEADER_PADDING_HORIZONTAL)
|
||||
.with_padding_right(HEADER_PADDING_HORIZONTAL)
|
||||
.finish();
|
||||
|
||||
let cloud_agent_option = self.render_option(
|
||||
0,
|
||||
Icon::OzCloud,
|
||||
"Cloud agent",
|
||||
"Runs autonomously in a cloud environment you choose. Best for parallel or long-running work.",
|
||||
true,
|
||||
self.cloud_agent_mouse_state.clone(),
|
||||
AgentTypeSelectorAction::SelectCloudAgent,
|
||||
appearance,
|
||||
);
|
||||
|
||||
let local_agent_option = self.render_option(
|
||||
1,
|
||||
Icon::Oz,
|
||||
"Local agent",
|
||||
"Runs on your machine and requires supervision. Best for quick, interactive tasks.",
|
||||
false,
|
||||
self.local_agent_mouse_state.clone(),
|
||||
AgentTypeSelectorAction::SelectLocalAgent,
|
||||
appearance,
|
||||
);
|
||||
|
||||
let options = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_spacing(OPTIONS_VERTICAL_GAP)
|
||||
.with_child(cloud_agent_option)
|
||||
.with_child(local_agent_option)
|
||||
.finish();
|
||||
|
||||
let body = Container::new(options)
|
||||
.with_padding_top(BODY_PADDING_VERTICAL)
|
||||
.with_padding_bottom(BODY_PADDING_VERTICAL)
|
||||
.with_padding_left(BODY_PADDING_HORIZONTAL)
|
||||
.with_padding_right(BODY_PADDING_HORIZONTAL)
|
||||
.finish();
|
||||
|
||||
let dialog_contents = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_child(header)
|
||||
.with_child(body)
|
||||
.finish();
|
||||
|
||||
let dialog_background = theme
|
||||
.surface_1()
|
||||
.blend(&internal_colors::fg_overlay_1(theme));
|
||||
let dialog_border = internal_colors::neutral_4(theme);
|
||||
|
||||
let dialog = Container::new(dialog_contents)
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(DIALOG_CORNER_RADIUS)))
|
||||
.with_border(Border::all(1.).with_border_color(dialog_border))
|
||||
.with_background(dialog_background)
|
||||
.with_drop_shadow(DropShadow {
|
||||
color: ColorU::new(0, 0, 0, 77),
|
||||
offset: vec2f(0., 7.),
|
||||
blur_radius: 7.,
|
||||
spread_radius: 0.,
|
||||
})
|
||||
.finish();
|
||||
|
||||
// Wrap dialog in a Hoverable to consume click events within the modal,
|
||||
// preventing the Dismiss from triggering when clicking inside the dialog.
|
||||
let clickable_dialog =
|
||||
Hoverable::new(self.dialog_mouse_state.clone(), |_| dialog).on_click(|_, _, _| {});
|
||||
|
||||
let constrained_dialog = ConstrainedBox::new(clickable_dialog.finish())
|
||||
.with_width(MODAL_WIDTH)
|
||||
.finish();
|
||||
|
||||
let dismiss_dialog = Dismiss::new(constrained_dialog)
|
||||
.prevent_interaction_with_other_elements()
|
||||
.on_dismiss(|ctx, _app| {
|
||||
ctx.dispatch_typed_action(AgentTypeSelectorAction::Dismiss);
|
||||
})
|
||||
.finish();
|
||||
|
||||
Container::new(Align::new(dismiss_dialog).finish())
|
||||
.with_background_color(ColorU::new(0, 0, 0, 179))
|
||||
.with_corner_radius(app.windows().window_corner_radius())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for AgentTypeSelector {
|
||||
type Event = AgentTypeSelectorEvent;
|
||||
}
|
||||
|
||||
impl TypedActionView for AgentTypeSelector {
|
||||
type Action = AgentTypeSelectorAction;
|
||||
|
||||
fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
|
||||
match action {
|
||||
AgentTypeSelectorAction::SelectCloudAgent => {
|
||||
ctx.emit(AgentTypeSelectorEvent::Selected(AgentType::Cloud));
|
||||
ctx.notify();
|
||||
}
|
||||
AgentTypeSelectorAction::SelectLocalAgent => {
|
||||
ctx.emit(AgentTypeSelectorEvent::Selected(AgentType::Local));
|
||||
ctx.notify();
|
||||
}
|
||||
AgentTypeSelectorAction::Dismiss => {
|
||||
ctx.emit(AgentTypeSelectorEvent::Dismissed);
|
||||
ctx.notify();
|
||||
}
|
||||
AgentTypeSelectorAction::HoveredIn(index) => {
|
||||
self.selected_option_index = *index;
|
||||
ctx.notify();
|
||||
}
|
||||
AgentTypeSelectorAction::ArrowUp | AgentTypeSelectorAction::ArrowDown => {
|
||||
// Toggle between the two options.
|
||||
self.selected_option_index = if self.selected_option_index == 0 {
|
||||
1
|
||||
} else {
|
||||
0
|
||||
};
|
||||
ctx.notify();
|
||||
}
|
||||
AgentTypeSelectorAction::Enter => {
|
||||
match self.selected_option_index {
|
||||
0 => {
|
||||
ctx.emit(AgentTypeSelectorEvent::Selected(AgentType::Cloud));
|
||||
}
|
||||
1 => {
|
||||
ctx.emit(AgentTypeSelectorEvent::Selected(AgentType::Local));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
ctx.notify();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl View for AgentTypeSelector {
|
||||
fn ui_name() -> &'static str {
|
||||
"AgentTypeSelector"
|
||||
}
|
||||
|
||||
fn render(&self, app: &AppContext) -> Box<dyn Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
self.render_modal(appearance, app)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,735 @@
|
||||
use crate::ai::agent_management::telemetry::{AgentManagementTelemetryEvent, SetupGuideStep};
|
||||
use crate::ai::blocklist::code_block::{
|
||||
render_code_block_plain, CodeBlockOptions, CodeSnippetButtonHandles,
|
||||
};
|
||||
use crate::appearance::Appearance;
|
||||
use crate::completer::SessionAgnosticContext;
|
||||
use crate::send_telemetry_from_ctx;
|
||||
use crate::view_components::action_button::{ActionButton, SecondaryTheme};
|
||||
use crate::workflows::workflow::{Argument, ArgumentType, Workflow};
|
||||
use crate::workflows::WorkflowType;
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use string_offset::CharCounter;
|
||||
use warp_completer::signatures::CommandRegistry;
|
||||
use warp_completer::{util::parse_current_commands_and_tokens, ParsedTokensSnapshot};
|
||||
use warp_core::report_error;
|
||||
use warp_core::ui::theme::{AnsiColorIdentifier, AnsiColors};
|
||||
use warpui::clipboard::ClipboardContent;
|
||||
use warpui::elements::{
|
||||
new_scrollable::{ClippedAxisConfiguration, DualAxisConfig, NewScrollable},
|
||||
Align, Border, ClippedScrollStateHandle, ConstrainedBox, Container, CornerRadius,
|
||||
CrossAxisAlignment, Element, Empty, Expanded, Flex, Highlight, HighlightedRange,
|
||||
MainAxisAlignment, MainAxisSize, MouseStateHandle, ParentElement, Radius, Text,
|
||||
};
|
||||
use warpui::fonts::{Properties, Weight};
|
||||
use warpui::prelude::ChildView;
|
||||
use warpui::text_layout::TextStyle;
|
||||
use warpui::ui_components::components::{UiComponent, UiComponentStyles};
|
||||
use warpui::ViewHandle;
|
||||
use warpui::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext};
|
||||
|
||||
const DOCS_URL: &str = "https://docs.warp.dev/agent-platform/cloud-agents/overview";
|
||||
const ENV_DOCS_URL: &str =
|
||||
"https://docs.warp.dev/reference/cli/integration-setup#creating-an-environment";
|
||||
const OZ_URL: &str = "https://oz.warp.dev";
|
||||
|
||||
const CONTENT_MAX_WIDTH: f32 = 720.;
|
||||
|
||||
const CREATE_ENV_SLASH_CMD: &str = "/create-environment";
|
||||
const CREATE_ENV_CLI_CMD: &str =
|
||||
"oz environment create [OPTIONS] --name <NAME> --docker-image <DOCKER_IMAGE>";
|
||||
const CREATE_SLACK_INTEGRATION_CMD: &str =
|
||||
"oz integration create slack --environment {{environment_id}}";
|
||||
const CREATE_LINEAR_INTEGRATION_CMD: &str =
|
||||
"oz integration create linear --environment {{environment_id}}";
|
||||
|
||||
pub struct CloudSetupGuideView {
|
||||
create_env_code_handles: CodeSnippetButtonHandles,
|
||||
create_env_cli_code_handles: CodeSnippetButtonHandles,
|
||||
create_slack_integration_code_handles: CodeSnippetButtonHandles,
|
||||
create_linear_integration_code_handles: CodeSnippetButtonHandles,
|
||||
docs_link_mouse_state: MouseStateHandle,
|
||||
env_docs_link_mouse_state: MouseStateHandle,
|
||||
integration_docs_link_mouse_state: MouseStateHandle,
|
||||
visit_oz_button: ViewHandle<ActionButton>,
|
||||
parsed_tokens: HashMap<&'static str, ParsedTokensSnapshot>,
|
||||
vertical_scroll_state: ClippedScrollStateHandle,
|
||||
horizontal_scroll_state: ClippedScrollStateHandle,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum CloudSetupGuideAction {
|
||||
CopyCode {
|
||||
code: String,
|
||||
step: SetupGuideStep,
|
||||
},
|
||||
RunWorkflow {
|
||||
workflow: Box<WorkflowType>,
|
||||
step: SetupGuideStep,
|
||||
},
|
||||
VisitOz,
|
||||
OpenDocs {
|
||||
docs: SetupGuideDocs,
|
||||
},
|
||||
}
|
||||
|
||||
/// Which URL the user clicked in the setup guide (also used in telemetry)
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
pub enum SetupGuideDocs {
|
||||
Main,
|
||||
Environment,
|
||||
Integration,
|
||||
}
|
||||
|
||||
pub enum CloudSetupGuideEvent {
|
||||
OpenNewTabAndInsertWorkflow(WorkflowType),
|
||||
}
|
||||
|
||||
impl CloudSetupGuideView {
|
||||
pub fn new(ctx: &mut ViewContext<Self>) -> Self {
|
||||
let code_snippets: Vec<&'static str> = vec![
|
||||
CREATE_ENV_SLASH_CMD,
|
||||
CREATE_ENV_CLI_CMD,
|
||||
CREATE_SLACK_INTEGRATION_CMD,
|
||||
CREATE_LINEAR_INTEGRATION_CMD,
|
||||
];
|
||||
|
||||
// We spawn a background task to parse the commands in the code blocks,
|
||||
// and re-render the blocks once it's done. Colors are computed at render time
|
||||
// so they respond to theme changes.
|
||||
let context = SessionAgnosticContext::new(CommandRegistry::global_instance());
|
||||
ctx.spawn(
|
||||
async move {
|
||||
let mut results = HashMap::new();
|
||||
for snippet in code_snippets {
|
||||
let parsed =
|
||||
parse_current_commands_and_tokens(snippet.to_string(), &context).await;
|
||||
results.insert(snippet, parsed);
|
||||
}
|
||||
results
|
||||
},
|
||||
|view, results, ctx| {
|
||||
view.parsed_tokens = results;
|
||||
ctx.notify();
|
||||
},
|
||||
);
|
||||
|
||||
let visit_oz_button = ctx.add_typed_action_view(|_| {
|
||||
ActionButton::new("Visit Oz", SecondaryTheme)
|
||||
.on_click(|ctx| ctx.dispatch_typed_action(CloudSetupGuideAction::VisitOz))
|
||||
});
|
||||
|
||||
Self {
|
||||
create_env_code_handles: CodeSnippetButtonHandles::default(),
|
||||
create_env_cli_code_handles: CodeSnippetButtonHandles::default(),
|
||||
create_slack_integration_code_handles: CodeSnippetButtonHandles::default(),
|
||||
create_linear_integration_code_handles: CodeSnippetButtonHandles::default(),
|
||||
docs_link_mouse_state: MouseStateHandle::default(),
|
||||
env_docs_link_mouse_state: MouseStateHandle::default(),
|
||||
integration_docs_link_mouse_state: MouseStateHandle::default(),
|
||||
visit_oz_button,
|
||||
parsed_tokens: HashMap::new(),
|
||||
vertical_scroll_state: ClippedScrollStateHandle::default(),
|
||||
horizontal_scroll_state: ClippedScrollStateHandle::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Render the main header for the setup guide.
|
||||
fn render_header(&self, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let title_font_size = 24.;
|
||||
let subtitle_font_size = 16.;
|
||||
|
||||
let mut header_container = Flex::column().with_spacing(8.);
|
||||
|
||||
let title = Text::new(
|
||||
"Getting started with Oz cloud agents",
|
||||
appearance.ui_font_family(),
|
||||
title_font_size,
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(theme.active_ui_text_color().into_solid())
|
||||
.finish();
|
||||
header_container.add_child(title);
|
||||
|
||||
let subtitle = Text::new(
|
||||
"Start Oz cloud agents directly in Warp from an integration (Linear, Slack), with an event (GitHub, built-in schedule), or programmatically with the Oz SDK or CLI.",
|
||||
appearance.ui_font_family(),
|
||||
subtitle_font_size,
|
||||
)
|
||||
.with_color(theme.nonactive_ui_text_color().into_solid())
|
||||
.finish();
|
||||
header_container.add_child(subtitle);
|
||||
|
||||
// Documentation link line.
|
||||
let docs_line = Flex::row()
|
||||
.with_child(
|
||||
Text::new_inline(
|
||||
"Check out the ",
|
||||
appearance.ui_font_family(),
|
||||
subtitle_font_size,
|
||||
)
|
||||
.with_color(theme.nonactive_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
)
|
||||
.with_child(
|
||||
appearance
|
||||
.ui_builder()
|
||||
.link(
|
||||
"Oz documentation".to_string(),
|
||||
None,
|
||||
Some(Box::new(|ctx| {
|
||||
ctx.dispatch_typed_action(CloudSetupGuideAction::OpenDocs {
|
||||
docs: SetupGuideDocs::Main,
|
||||
});
|
||||
})),
|
||||
self.docs_link_mouse_state.clone(),
|
||||
)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(subtitle_font_size),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish(),
|
||||
)
|
||||
.with_child(
|
||||
Text::new_inline(
|
||||
" to learn more.",
|
||||
appearance.ui_font_family(),
|
||||
subtitle_font_size,
|
||||
)
|
||||
.with_color(theme.nonactive_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
);
|
||||
header_container.add_child(docs_line.finish());
|
||||
|
||||
header_container.finish()
|
||||
}
|
||||
|
||||
/// Render the quick start banner with link to oz.warp.dev.
|
||||
fn render_quick_start_banner(&self, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let font_size = 16.;
|
||||
|
||||
let text = Text::new_inline(
|
||||
"Quick start: Visit oz.warp.dev for a UI-based setup experience.",
|
||||
appearance.ui_font_family(),
|
||||
font_size,
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(theme.active_ui_text_color().into_solid())
|
||||
.finish();
|
||||
|
||||
// Use cyan overlay for the blue border per Figma spec.
|
||||
let border_color = theme.ansi_overlay_2(theme.terminal_colors().normal.cyan);
|
||||
|
||||
Container::new(
|
||||
Flex::row()
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_child(text)
|
||||
.with_child(ChildView::new(&self.visit_oz_button).finish())
|
||||
.finish(),
|
||||
)
|
||||
.with_background(theme.surface_overlay_1())
|
||||
.with_border(Border::all(1.).with_border_fill(border_color))
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(8.)))
|
||||
.with_horizontal_padding(16.)
|
||||
.with_vertical_padding(12.)
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Render the manual setup section header.
|
||||
fn render_manual_setup_header(&self, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let font_size = 16.;
|
||||
|
||||
Text::new(
|
||||
"Manual setup: Create a Slack or Linear integration with the Oz CLI",
|
||||
appearance.ui_font_family(),
|
||||
font_size,
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(theme.active_ui_text_color().into_solid())
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Render a styled number to be displayed with each step.
|
||||
fn render_step_number(number: u32, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let number_text = Text::new(
|
||||
number.to_string(),
|
||||
appearance.ui_font_family(),
|
||||
appearance.ui_font_size(),
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(theme.active_ui_text_color().into_solid())
|
||||
.finish();
|
||||
|
||||
let centered_number = Flex::row()
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_main_axis_alignment(MainAxisAlignment::Center)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_child(number_text)
|
||||
.finish();
|
||||
|
||||
Container::new(
|
||||
ConstrainedBox::new(centered_number)
|
||||
.with_width(28.)
|
||||
.with_height(28.)
|
||||
.finish(),
|
||||
)
|
||||
.with_background(theme.surface_1())
|
||||
.with_border(Border::all(1.).with_border_fill(theme.outline()))
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Percentage(50.)))
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Render a description that includes a link at the end
|
||||
/// (e.g. "Use warp's environment setup command to have an agent help you through it. LINK[Visit docs]")
|
||||
fn render_description_with_link(
|
||||
prefix: &'static str,
|
||||
link_text: &'static str,
|
||||
link_mouse_state: MouseStateHandle,
|
||||
telemetry_url: SetupGuideDocs,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let step_desc_font_size = 14.;
|
||||
let link = appearance
|
||||
.ui_builder()
|
||||
.link(
|
||||
link_text.to_string(),
|
||||
None,
|
||||
Some(Box::new(move |ctx| {
|
||||
ctx.dispatch_typed_action(CloudSetupGuideAction::OpenDocs {
|
||||
docs: telemetry_url,
|
||||
});
|
||||
})),
|
||||
link_mouse_state,
|
||||
)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(step_desc_font_size),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish();
|
||||
|
||||
Flex::row()
|
||||
.with_child(
|
||||
Text::new_inline(prefix, appearance.ui_font_family(), step_desc_font_size)
|
||||
.with_color(appearance.theme().nonactive_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
)
|
||||
.with_child(link)
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Render a code block with buttons to copy and run the code.
|
||||
fn render_code_block(
|
||||
&self,
|
||||
code: &'static str,
|
||||
handles: CodeSnippetButtonHandles,
|
||||
app: &AppContext,
|
||||
) -> Box<dyn Element> {
|
||||
let terminal_colors = Appearance::as_ref(app).theme().terminal_colors().normal;
|
||||
let highlights = self
|
||||
.parsed_tokens
|
||||
.get(code)
|
||||
.map(|parsed| tokens_to_highlight_ranges(parsed, &terminal_colors))
|
||||
.unwrap_or_default();
|
||||
|
||||
// Match command to formatted workflow with correct args.
|
||||
let Some((workflow, setup_step)) = (match code {
|
||||
CREATE_ENV_SLASH_CMD => Some((
|
||||
WorkflowType::Local(
|
||||
Workflow::new("Create Environment", CREATE_ENV_SLASH_CMD).with_arguments(vec![
|
||||
Argument::new("github link or local filepath", ArgumentType::Text)
|
||||
.with_description("GitHub link or local filepath to the repository"),
|
||||
]),
|
||||
),
|
||||
SetupGuideStep::CreateEnvironment,
|
||||
)),
|
||||
CREATE_ENV_CLI_CMD => Some((
|
||||
WorkflowType::Local(
|
||||
Workflow::new("Create Environment (CLI)", CREATE_ENV_CLI_CMD).with_arguments(
|
||||
vec![
|
||||
Argument::new("NAME", ArgumentType::Text)
|
||||
.with_description("Name for the environment"),
|
||||
Argument::new("DOCKER_IMAGE", ArgumentType::Text)
|
||||
.with_description("Docker image to use for the environment"),
|
||||
],
|
||||
),
|
||||
),
|
||||
SetupGuideStep::CreateEnvironmentCli,
|
||||
)),
|
||||
CREATE_SLACK_INTEGRATION_CMD => Some((
|
||||
WorkflowType::Local(
|
||||
Workflow::new("Create Slack Integration", CREATE_SLACK_INTEGRATION_CMD)
|
||||
.with_arguments(vec![Argument::new("environment_id", ArgumentType::Text)
|
||||
.with_description("ID of the environment to integrate with")]),
|
||||
),
|
||||
SetupGuideStep::CreateSlackIntegration,
|
||||
)),
|
||||
CREATE_LINEAR_INTEGRATION_CMD => Some((
|
||||
WorkflowType::Local(
|
||||
Workflow::new("Create Linear Integration", CREATE_LINEAR_INTEGRATION_CMD)
|
||||
.with_arguments(vec![Argument::new("environment_id", ArgumentType::Text)
|
||||
.with_description("ID of the environment to integrate with")]),
|
||||
),
|
||||
SetupGuideStep::CreateLinearIntegration,
|
||||
)),
|
||||
_ => None,
|
||||
}) else {
|
||||
report_error!(anyhow::anyhow!(
|
||||
"Received unknown code in render_code_block: {}",
|
||||
code
|
||||
));
|
||||
return Empty::new().finish();
|
||||
};
|
||||
|
||||
render_code_block_plain(
|
||||
code,
|
||||
highlights.into_iter(),
|
||||
CodeBlockOptions {
|
||||
on_open: None,
|
||||
on_execute: Some(Box::new(move |_code, ctx| {
|
||||
ctx.dispatch_typed_action(CloudSetupGuideAction::RunWorkflow {
|
||||
workflow: Box::new(workflow.clone()),
|
||||
step: setup_step,
|
||||
});
|
||||
})),
|
||||
on_copy: Some(Box::new(move |_code, ctx| {
|
||||
ctx.dispatch_typed_action(CloudSetupGuideAction::CopyCode {
|
||||
code: code.to_string().clone(),
|
||||
step: setup_step,
|
||||
});
|
||||
})),
|
||||
on_insert: None,
|
||||
footer_element: None,
|
||||
mouse_handles: Some(handles),
|
||||
file_path: None,
|
||||
},
|
||||
app,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
/// Render step 1: Create an environment.
|
||||
fn render_step_1(&self, appearance: &Appearance, app: &AppContext) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let step_title_font_size = 14.;
|
||||
let step_desc_font_size = 14.;
|
||||
|
||||
let title_row = Flex::row()
|
||||
.with_spacing(16.)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_child(Self::render_step_number(1, appearance))
|
||||
.with_child(
|
||||
Text::new(
|
||||
"Create an environment",
|
||||
appearance.ui_font_family(),
|
||||
step_title_font_size,
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(theme.active_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
)
|
||||
.finish();
|
||||
|
||||
let description = Container::new(
|
||||
Text::new(
|
||||
"First, set up an environment to create an integration.",
|
||||
appearance.ui_font_family(),
|
||||
step_desc_font_size,
|
||||
)
|
||||
.with_color(theme.nonactive_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
)
|
||||
.with_padding_left(46.)
|
||||
.finish();
|
||||
|
||||
let sub_description = Container::new(Self::render_description_with_link(
|
||||
"Use Warp's environment setup command to have an agent help you through it. ",
|
||||
"Visit docs",
|
||||
self.env_docs_link_mouse_state.clone(),
|
||||
SetupGuideDocs::Environment,
|
||||
appearance,
|
||||
))
|
||||
.with_padding_left(46.)
|
||||
.with_padding_bottom(8.)
|
||||
.finish();
|
||||
|
||||
let slash_cmd_code_block = Container::new(self.render_code_block(
|
||||
CREATE_ENV_SLASH_CMD,
|
||||
self.create_env_code_handles.clone(),
|
||||
app,
|
||||
))
|
||||
.with_padding_left(46.)
|
||||
.finish();
|
||||
|
||||
let or_text = Container::new(
|
||||
Text::new(
|
||||
"Or, supply your own existing docker image.",
|
||||
appearance.ui_font_family(),
|
||||
step_desc_font_size,
|
||||
)
|
||||
.with_color(theme.nonactive_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
)
|
||||
.with_padding_left(46.)
|
||||
.with_padding_top(8.)
|
||||
.with_padding_bottom(8.)
|
||||
.finish();
|
||||
|
||||
let cli_code_block = Container::new(self.render_code_block(
|
||||
CREATE_ENV_CLI_CMD,
|
||||
self.create_env_cli_code_handles.clone(),
|
||||
app,
|
||||
))
|
||||
.with_padding_left(46.)
|
||||
.finish();
|
||||
|
||||
Flex::column()
|
||||
.with_spacing(8.)
|
||||
.with_child(title_row)
|
||||
.with_child(description)
|
||||
.with_child(sub_description)
|
||||
.with_child(slash_cmd_code_block)
|
||||
.with_child(or_text)
|
||||
.with_child(cli_code_block)
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Render step 2: Create an integration.
|
||||
fn render_step_2(&self, appearance: &Appearance, app: &AppContext) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let step_title_font_size = 14.;
|
||||
|
||||
let title_row = Flex::row()
|
||||
.with_spacing(16.)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_child(Self::render_step_number(2, appearance))
|
||||
.with_child(
|
||||
Text::new(
|
||||
"Create an integration",
|
||||
appearance.ui_font_family(),
|
||||
step_title_font_size,
|
||||
)
|
||||
.with_style(Properties::default().weight(Weight::Semibold))
|
||||
.with_color(theme.active_ui_text_color().into_solid())
|
||||
.finish(),
|
||||
)
|
||||
.finish();
|
||||
|
||||
let sub_description = Container::new(Self::render_description_with_link(
|
||||
"Integrate Slack or Linear to assign Warp's Agent tasks with @Warp. ",
|
||||
"Visit docs",
|
||||
self.integration_docs_link_mouse_state.clone(),
|
||||
SetupGuideDocs::Integration,
|
||||
appearance,
|
||||
))
|
||||
.with_padding_left(46.)
|
||||
.with_padding_bottom(8.)
|
||||
.finish();
|
||||
|
||||
let code_block_slack = Container::new(self.render_code_block(
|
||||
CREATE_SLACK_INTEGRATION_CMD,
|
||||
self.create_slack_integration_code_handles.clone(),
|
||||
app,
|
||||
))
|
||||
.with_padding_left(46.)
|
||||
.finish();
|
||||
|
||||
let code_block_linear = Container::new(self.render_code_block(
|
||||
CREATE_LINEAR_INTEGRATION_CMD,
|
||||
self.create_linear_integration_code_handles.clone(),
|
||||
app,
|
||||
))
|
||||
.with_padding_left(46.)
|
||||
.finish();
|
||||
|
||||
Flex::column()
|
||||
.with_spacing(8.)
|
||||
.with_child(title_row)
|
||||
.with_child(sub_description)
|
||||
.with_child(code_block_slack)
|
||||
.with_child(code_block_linear)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for CloudSetupGuideView {
|
||||
type Event = CloudSetupGuideEvent;
|
||||
}
|
||||
|
||||
impl View for CloudSetupGuideView {
|
||||
fn ui_name() -> &'static str {
|
||||
"AgentManagementHelpPageView"
|
||||
}
|
||||
|
||||
fn render(&self, app: &AppContext) -> Box<dyn Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let theme = appearance.theme();
|
||||
|
||||
let steps = Flex::column()
|
||||
.with_spacing(24.)
|
||||
.with_child(self.render_step_1(appearance, app))
|
||||
.with_child(self.render_step_2(appearance, app))
|
||||
.finish();
|
||||
|
||||
let mut content = Flex::column()
|
||||
.with_spacing(24.)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch);
|
||||
|
||||
content.add_child(self.render_header(appearance));
|
||||
content.add_child(self.render_quick_start_banner(appearance));
|
||||
content.add_child(self.render_manual_setup_header(appearance));
|
||||
content.add_child(steps);
|
||||
|
||||
let content = content.finish();
|
||||
|
||||
let scrollable = NewScrollable::horizontal_and_vertical(
|
||||
DualAxisConfig::Clipped {
|
||||
horizontal: ClippedAxisConfiguration {
|
||||
handle: self.horizontal_scroll_state.clone(),
|
||||
max_size: None,
|
||||
stretch_child: true,
|
||||
},
|
||||
vertical: ClippedAxisConfiguration {
|
||||
handle: self.vertical_scroll_state.clone(),
|
||||
max_size: None,
|
||||
stretch_child: false,
|
||||
},
|
||||
child: Align::new(
|
||||
Container::new(
|
||||
ConstrainedBox::new(content)
|
||||
.with_max_width(CONTENT_MAX_WIDTH)
|
||||
.finish(),
|
||||
)
|
||||
.with_uniform_padding(24.)
|
||||
.finish(),
|
||||
)
|
||||
.top_center()
|
||||
.finish(),
|
||||
},
|
||||
theme.nonactive_ui_detail().into(),
|
||||
theme.active_ui_detail().into(),
|
||||
warpui::elements::Fill::None,
|
||||
)
|
||||
.finish();
|
||||
|
||||
Flex::column()
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_child(Expanded::new(1., scrollable).finish())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl TypedActionView for CloudSetupGuideView {
|
||||
type Action = CloudSetupGuideAction;
|
||||
|
||||
fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
|
||||
match action {
|
||||
CloudSetupGuideAction::CopyCode { code, step } => {
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::SetupGuideStepCopy { step: *step },
|
||||
ctx
|
||||
);
|
||||
ctx.clipboard()
|
||||
.write(ClipboardContent::plain_text(code.clone()));
|
||||
}
|
||||
CloudSetupGuideAction::RunWorkflow { workflow, step } => {
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::SetupGuideStepRun { step: *step },
|
||||
ctx
|
||||
);
|
||||
ctx.emit(CloudSetupGuideEvent::OpenNewTabAndInsertWorkflow(
|
||||
(**workflow).clone(),
|
||||
));
|
||||
}
|
||||
CloudSetupGuideAction::VisitOz => {
|
||||
ctx.open_url(OZ_URL);
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::SetupGuideStepRun {
|
||||
step: SetupGuideStep::VisitOz
|
||||
},
|
||||
ctx
|
||||
);
|
||||
}
|
||||
CloudSetupGuideAction::OpenDocs { docs } => {
|
||||
let url = match docs {
|
||||
SetupGuideDocs::Main => DOCS_URL,
|
||||
SetupGuideDocs::Environment => ENV_DOCS_URL,
|
||||
SetupGuideDocs::Integration => DOCS_URL,
|
||||
};
|
||||
ctx.open_url(url);
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::SetupGuideDocsLink { docs: *docs },
|
||||
ctx
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Highlight the commands in command blocks correctly (including the command prefix).
|
||||
/// It's unfortunate that we have to do this manually, but the alternative is inserting a custom code editor into this component
|
||||
/// and that would be a lot of bloat for not much benefit.
|
||||
fn tokens_to_highlight_ranges(
|
||||
parsed_tokens: &ParsedTokensSnapshot,
|
||||
terminal_colors: &AnsiColors,
|
||||
) -> Vec<HighlightedRange> {
|
||||
let code = &parsed_tokens.buffer_text;
|
||||
let mut highlights = Vec::new();
|
||||
|
||||
// Handle slash commands: if code starts with '/', highlight the command prefix in magenta
|
||||
if code.starts_with('/') {
|
||||
if let Some(space_idx) = code.find(' ') {
|
||||
let color = AnsiColorIdentifier::Magenta.to_ansi_color(terminal_colors);
|
||||
highlights.push(HighlightedRange {
|
||||
highlight: Highlight::new()
|
||||
.with_text_style(TextStyle::new().with_foreground_color(color.into())),
|
||||
highlight_indices: (0..space_idx).collect(),
|
||||
});
|
||||
return highlights;
|
||||
}
|
||||
}
|
||||
|
||||
// Highlight commands in the code block (converting bytes to char indexes as we go).
|
||||
let mut char_counter = CharCounter::new(code);
|
||||
for token_data in &parsed_tokens.parsed_tokens {
|
||||
let Some(description) = &token_data.token_description else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let byte_start = token_data.token.span.start();
|
||||
let byte_end = token_data.token.span.end();
|
||||
|
||||
let Some(char_start) = char_counter.char_offset(byte_start) else {
|
||||
continue;
|
||||
};
|
||||
let Some(char_end) = char_counter.char_offset(byte_end) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let char_indices: Vec<usize> = (char_start.as_usize()..char_end.as_usize()).collect();
|
||||
if char_indices.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let color_id: AnsiColorIdentifier = description.suggestion_type.to_name().into();
|
||||
let color = color_id.to_ansi_color(terminal_colors);
|
||||
|
||||
highlights.push(HighlightedRange {
|
||||
highlight: Highlight::new()
|
||||
.with_text_style(TextStyle::new().with_foreground_color(color.into())),
|
||||
highlight_indices: char_indices,
|
||||
});
|
||||
}
|
||||
|
||||
highlights
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
//! Action buttons row for conversation details panel.
|
||||
|
||||
use warp_core::ui::theme::AnsiColorIdentifier;
|
||||
use warpui::elements::{ChildView, CrossAxisAlignment, Empty, Flex, ParentElement};
|
||||
use warpui::{AppContext, Element, Entity, TypedActionView, View, ViewContext, ViewHandle};
|
||||
|
||||
use crate::view_components::copyable_text_field::COPY_FEEDBACK_DURATION;
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::agent_conversations_model::AgentRunDisplayStatus;
|
||||
use crate::ai::agent_management::view::ManagementCardItemId;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::ui_components::icons::Icon;
|
||||
use crate::view_components::action_button::{ActionButton, ButtonSize, SecondaryTheme};
|
||||
use crate::workspace::WorkspaceAction;
|
||||
|
||||
const BUTTON_SPACING: f32 = 4.;
|
||||
|
||||
/// Per-button config for the action buttons row.
|
||||
/// Each field controls one button independently.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ActionButtonsConfig {
|
||||
pub open_action: Option<WorkspaceAction>,
|
||||
pub cancel_task_id: Option<AmbientAgentTaskId>,
|
||||
pub fork_conversation_id: Option<AIConversationId>,
|
||||
/// Shows an info button for viewing more details.
|
||||
/// Only used in management view hover toolbelt.
|
||||
pub view_details_item_id: Option<ManagementCardItemId>,
|
||||
/// Conversation link URL (either to the transcript or live session) for copy link button.
|
||||
pub copy_link_url: Option<String>,
|
||||
}
|
||||
|
||||
impl ActionButtonsConfig {
|
||||
/// Returns true if no buttons will be rendered.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.open_action.is_none()
|
||||
&& self.cancel_task_id.is_none()
|
||||
&& self.fork_conversation_id.is_none()
|
||||
&& self.view_details_item_id.is_none()
|
||||
&& self.copy_link_url.is_none()
|
||||
}
|
||||
|
||||
/// Create config for a task.
|
||||
/// - `display_status`: used to determine if cancel button should show.
|
||||
/// - `open_action`: pass `Some(action)` to show open button, `None` to hide
|
||||
/// - `copy_link_url`: conversation link URL, or `None` to hide
|
||||
pub fn for_task(
|
||||
task_id: AmbientAgentTaskId,
|
||||
display_status: &AgentRunDisplayStatus,
|
||||
open_action: Option<WorkspaceAction>,
|
||||
copy_link_url: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
open_action,
|
||||
cancel_task_id: if display_status.is_cancellable() {
|
||||
Some(task_id)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
fork_conversation_id: None,
|
||||
view_details_item_id: None,
|
||||
copy_link_url,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create config for a conversation.
|
||||
/// - `open_action`: pass `Some(action)` to show open button, `None` to hide
|
||||
/// - `copy_link_url`: conversation link URL, or `None` to hide
|
||||
pub fn for_conversation(
|
||||
conversation_id: AIConversationId,
|
||||
open_action: Option<WorkspaceAction>,
|
||||
copy_link_url: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
open_action,
|
||||
cancel_task_id: None,
|
||||
fork_conversation_id: Some(conversation_id),
|
||||
view_details_item_id: None,
|
||||
copy_link_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Events emitted by the action buttons.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AgentDetailsButtonEvent {
|
||||
Open,
|
||||
CancelTask { task_id: AmbientAgentTaskId },
|
||||
ForkConversation { conversation_id: AIConversationId },
|
||||
ViewDetails { item_id: ManagementCardItemId },
|
||||
CopyLink { link: String },
|
||||
}
|
||||
|
||||
/// Actions dispatched by button clicks (internal).
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AgentDetailsAction {
|
||||
Open,
|
||||
CancelTask,
|
||||
ForkConversation,
|
||||
ViewDetails,
|
||||
CopyLink,
|
||||
}
|
||||
|
||||
/// Reusable action buttons row for details panel.
|
||||
pub struct ConversationActionButtonsRow {
|
||||
config: ActionButtonsConfig,
|
||||
open_button: ViewHandle<ActionButton>,
|
||||
cancel_task_button: ViewHandle<ActionButton>,
|
||||
fork_conversation_button: ViewHandle<ActionButton>,
|
||||
view_details_button: ViewHandle<ActionButton>,
|
||||
copy_link_button: ViewHandle<ActionButton>,
|
||||
}
|
||||
|
||||
impl ConversationActionButtonsRow {
|
||||
pub fn new(ctx: &mut ViewContext<Self>) -> Self {
|
||||
let open_button = ctx.add_typed_action_view(|_| {
|
||||
Self::make_action_button(
|
||||
Icon::LinkExternal,
|
||||
"Open conversation",
|
||||
None,
|
||||
AgentDetailsAction::Open,
|
||||
)
|
||||
});
|
||||
|
||||
let cancel_task_button = ctx.add_typed_action_view(|_| {
|
||||
Self::make_action_button(
|
||||
Icon::StopFilled,
|
||||
"Cancel task",
|
||||
Some(AnsiColorIdentifier::Red),
|
||||
AgentDetailsAction::CancelTask,
|
||||
)
|
||||
});
|
||||
|
||||
let fork_conversation_button = ctx.add_typed_action_view(|_| {
|
||||
Self::make_action_button(
|
||||
Icon::ArrowSplit,
|
||||
"Fork conversation",
|
||||
None,
|
||||
AgentDetailsAction::ForkConversation,
|
||||
)
|
||||
});
|
||||
|
||||
let view_details_button = ctx.add_typed_action_view(|_| {
|
||||
Self::make_action_button(
|
||||
Icon::Info,
|
||||
"View details",
|
||||
None,
|
||||
AgentDetailsAction::ViewDetails,
|
||||
)
|
||||
});
|
||||
|
||||
let copy_link_button = ctx.add_typed_action_view(|_| {
|
||||
Self::make_action_button(
|
||||
Icon::Link,
|
||||
"Copy link to run",
|
||||
None,
|
||||
AgentDetailsAction::CopyLink,
|
||||
)
|
||||
});
|
||||
|
||||
Self {
|
||||
config: ActionButtonsConfig::default(),
|
||||
open_button,
|
||||
cancel_task_button,
|
||||
fork_conversation_button,
|
||||
view_details_button,
|
||||
copy_link_button,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the config and rerender.
|
||||
pub fn set_config(&mut self, config: ActionButtonsConfig, ctx: &mut ViewContext<Self>) {
|
||||
self.config = config;
|
||||
ctx.notify();
|
||||
}
|
||||
|
||||
/// Returns true if no buttons will be rendered.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.config.is_empty()
|
||||
}
|
||||
|
||||
fn make_action_button(
|
||||
icon: Icon,
|
||||
tooltip: &str,
|
||||
icon_color: Option<AnsiColorIdentifier>,
|
||||
action: AgentDetailsAction,
|
||||
) -> ActionButton {
|
||||
let mut button = ActionButton::new("", SecondaryTheme)
|
||||
.with_icon(icon)
|
||||
.with_size(ButtonSize::Small)
|
||||
.with_tooltip(tooltip)
|
||||
.on_click(move |ctx| {
|
||||
ctx.dispatch_typed_action(action.clone());
|
||||
});
|
||||
if let Some(color) = icon_color {
|
||||
button = button.with_icon_ansi_color(color);
|
||||
}
|
||||
button
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for ConversationActionButtonsRow {
|
||||
type Event = AgentDetailsButtonEvent;
|
||||
}
|
||||
|
||||
impl View for ConversationActionButtonsRow {
|
||||
fn ui_name() -> &'static str {
|
||||
"ConversationActionButtonsRow"
|
||||
}
|
||||
|
||||
fn render(&self, _app: &AppContext) -> Box<dyn Element> {
|
||||
if self.config.is_empty() {
|
||||
return Empty::new().finish();
|
||||
}
|
||||
|
||||
let mut row = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_spacing(BUTTON_SPACING);
|
||||
|
||||
if self.config.copy_link_url.is_some() {
|
||||
row.add_child(ChildView::new(&self.copy_link_button).finish());
|
||||
}
|
||||
|
||||
if self.config.open_action.is_some() {
|
||||
row.add_child(ChildView::new(&self.open_button).finish());
|
||||
}
|
||||
if self.config.cancel_task_id.is_some() {
|
||||
row.add_child(ChildView::new(&self.cancel_task_button).finish());
|
||||
}
|
||||
if self.config.fork_conversation_id.is_some() && !cfg!(target_family = "wasm") {
|
||||
row.add_child(ChildView::new(&self.fork_conversation_button).finish());
|
||||
}
|
||||
if self.config.view_details_item_id.is_some() {
|
||||
row.add_child(ChildView::new(&self.view_details_button).finish());
|
||||
}
|
||||
|
||||
row.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl TypedActionView for ConversationActionButtonsRow {
|
||||
type Action = AgentDetailsAction;
|
||||
|
||||
fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
|
||||
match action {
|
||||
AgentDetailsAction::Open => {
|
||||
if self.config.open_action.is_some() {
|
||||
ctx.emit(AgentDetailsButtonEvent::Open);
|
||||
}
|
||||
}
|
||||
AgentDetailsAction::CancelTask => {
|
||||
if let Some(task_id) = self.config.cancel_task_id {
|
||||
ctx.emit(AgentDetailsButtonEvent::CancelTask { task_id });
|
||||
}
|
||||
}
|
||||
AgentDetailsAction::ForkConversation => {
|
||||
if let Some(conversation_id) = self.config.fork_conversation_id {
|
||||
ctx.emit(AgentDetailsButtonEvent::ForkConversation { conversation_id });
|
||||
}
|
||||
}
|
||||
AgentDetailsAction::ViewDetails => {
|
||||
if let Some(item_id) = &self.config.view_details_item_id {
|
||||
ctx.emit(AgentDetailsButtonEvent::ViewDetails {
|
||||
item_id: item_id.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
AgentDetailsAction::CopyLink => {
|
||||
if let Some(link) = &self.config.copy_link_url {
|
||||
ctx.emit(AgentDetailsButtonEvent::CopyLink { link: link.clone() });
|
||||
self.copy_link_button.update(ctx, |button, ctx| {
|
||||
button.set_icon(Some(Icon::Check), ctx);
|
||||
});
|
||||
let duration = COPY_FEEDBACK_DURATION;
|
||||
ctx.spawn(
|
||||
async move {
|
||||
warpui::r#async::Timer::after(duration).await;
|
||||
},
|
||||
|me, _, ctx| {
|
||||
me.copy_link_button.update(ctx, |button, ctx| {
|
||||
button.set_icon(Some(Icon::Link), ctx);
|
||||
});
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
mod agent_management_model;
|
||||
pub(crate) mod agent_type_selector;
|
||||
pub(crate) mod details_action_buttons;
|
||||
pub(crate) mod notifications;
|
||||
|
||||
pub(crate) mod cloud_setup_guide_view;
|
||||
pub(crate) mod telemetry;
|
||||
pub(crate) mod view;
|
||||
|
||||
pub(crate) use agent_management_model::{AgentManagementEvent, AgentNotificationsModel};
|
||||
|
||||
pub fn init(app: &mut warpui::AppContext) {
|
||||
view::init(app);
|
||||
agent_type_selector::init(app);
|
||||
notifications::view::NotificationMailboxView::init(app);
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
use enum_iterator::Sequence;
|
||||
use instant::Instant;
|
||||
use uuid::Uuid;
|
||||
use warpui::EntityId;
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::artifacts::Artifact;
|
||||
use crate::terminal::CLIAgent;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct NotificationId(Uuid);
|
||||
|
||||
impl NotificationId {
|
||||
fn new() -> Self {
|
||||
Self(Uuid::new_v4())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum NotificationCategory {
|
||||
/// The agent has stopped (i.e. successfully completed or was cancelled)
|
||||
Complete,
|
||||
/// The agent needs user action (i.e. blocked on some permission request or idle prompt)
|
||||
Request,
|
||||
Error,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Sequence)]
|
||||
pub enum NotificationFilter {
|
||||
All,
|
||||
Unread,
|
||||
Errors,
|
||||
}
|
||||
|
||||
impl NotificationFilter {
|
||||
pub(crate) fn label(&self) -> &'static str {
|
||||
match self {
|
||||
NotificationFilter::All => "All tabs",
|
||||
NotificationFilter::Unread => "Unread",
|
||||
NotificationFilter::Errors => "Errors",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Identifies the agent that produced a notification.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
#[allow(clippy::upper_case_acronyms)]
|
||||
pub enum NotificationSourceAgent {
|
||||
Oz,
|
||||
CLI(CLIAgent),
|
||||
}
|
||||
|
||||
/// Identifies the conversation or session a notification belongs to.
|
||||
/// Used for de-duplication (replacing stale notifications on update)
|
||||
/// and cleanup (removing notifications when the source closes).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum NotificationOrigin {
|
||||
Conversation(AIConversationId),
|
||||
/// CLI sessions are keyed by terminal view because we only track one session per pane.
|
||||
CLISession(EntityId),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NotificationItem {
|
||||
pub id: NotificationId,
|
||||
pub origin: NotificationOrigin,
|
||||
pub title: String,
|
||||
pub message: String,
|
||||
pub category: NotificationCategory,
|
||||
pub agent: NotificationSourceAgent,
|
||||
/// Whether the user has already seen this notification
|
||||
/// (either because they clicked into it or because it was emitted for a conversation
|
||||
/// that they've since navigated to).
|
||||
pub is_read: bool,
|
||||
pub created_at: Instant,
|
||||
pub terminal_view_id: EntityId,
|
||||
pub artifacts: Vec<Artifact>,
|
||||
/// The git branch associated with this notification's conversation/session.
|
||||
/// When present, the notification renders in "rich" layout with a branch header row.
|
||||
/// When absent, it falls back to the "simple" layout.
|
||||
pub branch: Option<String>,
|
||||
}
|
||||
|
||||
impl NotificationItem {
|
||||
/// Marks this notification as read. Returns true if it was previously unread.
|
||||
fn mark_as_read(&mut self) -> bool {
|
||||
if self.is_read {
|
||||
return false;
|
||||
}
|
||||
self.is_read = true;
|
||||
true
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn new(
|
||||
title: String,
|
||||
message: String,
|
||||
category: NotificationCategory,
|
||||
agent: NotificationSourceAgent,
|
||||
origin: NotificationOrigin,
|
||||
is_read: bool,
|
||||
terminal_view_id: EntityId,
|
||||
artifacts: Vec<Artifact>,
|
||||
branch: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id: NotificationId::new(),
|
||||
origin,
|
||||
title,
|
||||
message,
|
||||
category,
|
||||
agent,
|
||||
is_read,
|
||||
created_at: Instant::now(),
|
||||
terminal_view_id,
|
||||
artifacts,
|
||||
branch,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct NotificationItems {
|
||||
items: Vec<NotificationItem>,
|
||||
}
|
||||
|
||||
impl NotificationItems {
|
||||
/// Push a notification items into the mailbox list
|
||||
/// (deleting older notifications if we've exceeded the max list size).
|
||||
pub(crate) fn push(&mut self, item: NotificationItem) {
|
||||
self.remove_by_origin(item.origin);
|
||||
self.items.insert(0, item);
|
||||
self.items.truncate(100);
|
||||
}
|
||||
|
||||
pub(crate) fn remove_by_origin(&mut self, key: NotificationOrigin) -> bool {
|
||||
let before = self.items.len();
|
||||
self.items.retain(|item| item.origin != key);
|
||||
self.items.len() != before
|
||||
}
|
||||
|
||||
pub(crate) fn items_filtered(
|
||||
&self,
|
||||
filter: NotificationFilter,
|
||||
) -> impl Iterator<Item = &NotificationItem> {
|
||||
self.items.iter().filter(move |item| match filter {
|
||||
NotificationFilter::All => true,
|
||||
NotificationFilter::Unread => !item.is_read,
|
||||
NotificationFilter::Errors => item.category == NotificationCategory::Error,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn filtered_count(&self, filter: NotificationFilter) -> usize {
|
||||
self.items_filtered(filter).count()
|
||||
}
|
||||
|
||||
/// Returns the filters that should be shown as tabs.
|
||||
/// "All" is always included; other filters are included only when they have at least one item.
|
||||
pub(crate) fn visible_filters(&self) -> Vec<NotificationFilter> {
|
||||
enum_iterator::all::<NotificationFilter>()
|
||||
.filter(|f| *f == NotificationFilter::All || self.filtered_count(*f) > 0)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn get_by_id(&self, id: NotificationId) -> Option<&NotificationItem> {
|
||||
self.items.iter().find(|item| item.id == id)
|
||||
}
|
||||
|
||||
/// Marks all notifications from the given terminal view as read.
|
||||
/// Returns true if any were changed.
|
||||
pub(crate) fn mark_all_terminal_view_items_as_read(
|
||||
&mut self,
|
||||
terminal_view_id: EntityId,
|
||||
) -> bool {
|
||||
let mut any_changed = false;
|
||||
for item in &mut self.items {
|
||||
if item.terminal_view_id == terminal_view_id {
|
||||
any_changed |= item.mark_as_read();
|
||||
}
|
||||
}
|
||||
any_changed
|
||||
}
|
||||
|
||||
pub(crate) fn mark_item_read(&mut self, id: NotificationId) -> bool {
|
||||
self.items
|
||||
.iter_mut()
|
||||
.find(|item| item.id == id)
|
||||
.is_some_and(|item| item.mark_as_read())
|
||||
}
|
||||
|
||||
pub(crate) fn mark_all_items_read(&mut self) -> bool {
|
||||
let mut any_changed = false;
|
||||
for item in &mut self.items {
|
||||
any_changed |= item.mark_as_read();
|
||||
}
|
||||
any_changed
|
||||
}
|
||||
|
||||
pub(crate) fn has_unread_for_terminal_view(&self, terminal_view_id: EntityId) -> bool {
|
||||
self.items
|
||||
.iter()
|
||||
.any(|item| item.terminal_view_id == terminal_view_id && !item.is_read)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "item_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,509 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use pathfinder_color::ColorU;
|
||||
use warp_core::ui::icons::Icon;
|
||||
use warp_core::ui::theme::{Fill, WarpTheme};
|
||||
use warpui::clipboard::ClipboardContent;
|
||||
use warpui::elements::{
|
||||
ChildView, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, DispatchEventResult,
|
||||
Element, EventHandler, Flex, MainAxisAlignment, MainAxisSize, ParentElement, Radius, Rect,
|
||||
Shrinkable,
|
||||
};
|
||||
use warpui::fonts::Weight;
|
||||
use warpui::ui_components::components::{UiComponent, UiComponentStyles};
|
||||
use warpui::{View, ViewContext, ViewHandle};
|
||||
|
||||
use warp_core::ui::appearance::Appearance as CoreAppearance;
|
||||
use warp_core::ui::theme::color::internal_colors;
|
||||
|
||||
use crate::ai::agent::conversation::ConversationStatus;
|
||||
use crate::ai::agent_management::notifications::item::NotificationSourceAgent;
|
||||
use crate::ai::agent_management::notifications::{NotificationCategory, NotificationItem};
|
||||
use crate::ai::agent_management::telemetry::{AgentManagementTelemetryEvent, ArtifactType};
|
||||
use crate::ai::artifacts::{
|
||||
open_screenshot_lightbox, Artifact, ArtifactButtonsRow, ArtifactButtonsRowEvent,
|
||||
};
|
||||
use crate::appearance::Appearance;
|
||||
use crate::send_telemetry_from_ctx;
|
||||
use crate::ui_components::icon_with_status::{
|
||||
render_icon_with_status, IconWithStatusSizing, IconWithStatusVariant,
|
||||
};
|
||||
use crate::util::time_format::format_elapsed_since;
|
||||
use crate::view_components::action_button::ActionButtonTheme;
|
||||
use crate::workspace::WorkspaceAction;
|
||||
|
||||
const COLLAPSED_MAX_CHARS: usize = 100;
|
||||
const EXPANDED_MAX_CHARS: usize = 500;
|
||||
|
||||
fn truncate_text(text: &str, max_chars: usize) -> String {
|
||||
if text.chars().count() > max_chars - 3 {
|
||||
let truncated: String = text.chars().take(max_chars).collect();
|
||||
format!("{truncated}…")
|
||||
} else {
|
||||
text.to_owned()
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true when either the title or message would be truncated by `truncate_text`.
|
||||
fn content_is_truncated(title: &str, message: &str) -> bool {
|
||||
title.chars().count() > COLLAPSED_MAX_CHARS - 3
|
||||
|| message.chars().count() > COLLAPSED_MAX_CHARS - 3
|
||||
}
|
||||
|
||||
/// Determines toast-vs-mailbox rendering differences.
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum NotificationRenderContext {
|
||||
Toast,
|
||||
Mailbox,
|
||||
}
|
||||
|
||||
/// Button theme for artifact chips in notifications.
|
||||
/// Uses `outline` for the border so it's visible against `surface_2`.
|
||||
pub(crate) struct NotificationArtifactButtonTheme;
|
||||
|
||||
impl ActionButtonTheme for NotificationArtifactButtonTheme {
|
||||
fn background(&self, hovered: bool, appearance: &CoreAppearance) -> Option<Fill> {
|
||||
if hovered {
|
||||
Some(internal_colors::fg_overlay_2(appearance.theme()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn text_color(
|
||||
&self,
|
||||
_hovered: bool,
|
||||
_background: Option<Fill>,
|
||||
appearance: &CoreAppearance,
|
||||
) -> ColorU {
|
||||
appearance.theme().foreground().into_solid()
|
||||
}
|
||||
|
||||
fn border(&self, appearance: &CoreAppearance) -> Option<ColorU> {
|
||||
Some(appearance.theme().outline().into_solid())
|
||||
}
|
||||
}
|
||||
|
||||
/// Callback invoked when the user clicks the expand/collapse affordance on a clamped message.
|
||||
pub(crate) type OnExpandClick = Box<dyn Fn(&mut warpui::EventContext)>;
|
||||
|
||||
/// Renders the inner content of a notification item.
|
||||
/// Dispatches to the rich layout (with branch row) or simple layout based on `item.branch`.
|
||||
pub(crate) fn render_notification_item_content(
|
||||
item: &NotificationItem,
|
||||
artifact_buttons: Option<&ViewHandle<ArtifactButtonsRow>>,
|
||||
context: NotificationRenderContext,
|
||||
message_expanded: bool,
|
||||
on_expand_click: OnExpandClick,
|
||||
extra_content: Option<Box<dyn Element>>,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let text_column = if item.branch.is_some() {
|
||||
render_rich_text_column(
|
||||
item,
|
||||
artifact_buttons,
|
||||
context,
|
||||
message_expanded,
|
||||
on_expand_click,
|
||||
extra_content,
|
||||
appearance,
|
||||
)
|
||||
} else {
|
||||
render_simple_text_column(
|
||||
item,
|
||||
artifact_buttons,
|
||||
context,
|
||||
message_expanded,
|
||||
on_expand_click,
|
||||
extra_content,
|
||||
appearance,
|
||||
)
|
||||
};
|
||||
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Start)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(
|
||||
Container::new(render_agent_avatar(item.agent, item.category, theme))
|
||||
.with_margin_right(8.)
|
||||
.with_margin_top(2.)
|
||||
.finish(),
|
||||
)
|
||||
.with_child(Shrinkable::new(1.0, text_column).finish())
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Rich layout: branch row + clamped title + clamped message + artifact buttons.
|
||||
fn render_rich_text_column(
|
||||
item: &NotificationItem,
|
||||
artifact_buttons: Option<&ViewHandle<ArtifactButtonsRow>>,
|
||||
context: NotificationRenderContext,
|
||||
message_expanded: bool,
|
||||
on_expand_click: OnExpandClick,
|
||||
extra_content: Option<Box<dyn Element>>,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let branch = item.branch.as_deref().unwrap_or_default();
|
||||
|
||||
let branch_left = render_branch_label(branch, appearance);
|
||||
let is_truncated = content_is_truncated(&item.title, &item.message);
|
||||
|
||||
let branch_right: Box<dyn Element> = match context {
|
||||
NotificationRenderContext::Toast if is_truncated || message_expanded => {
|
||||
render_expand_chevron(message_expanded, on_expand_click, theme)
|
||||
}
|
||||
NotificationRenderContext::Toast => {
|
||||
// No chevron when content fits.
|
||||
Flex::row().finish()
|
||||
}
|
||||
NotificationRenderContext::Mailbox => render_timestamp_with_dot(item, appearance),
|
||||
};
|
||||
|
||||
let branch_row = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(branch_left)
|
||||
.with_child(branch_right)
|
||||
.finish();
|
||||
|
||||
let title = render_clamped_title(&item.title, message_expanded, appearance);
|
||||
let message = render_message_text(&item.message, message_expanded, appearance);
|
||||
|
||||
let mut content = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_child(branch_row)
|
||||
.with_child(title)
|
||||
.with_child(Container::new(message).with_margin_top(2.).finish());
|
||||
|
||||
append_trailing_content(&mut content, artifact_buttons, extra_content);
|
||||
content.finish()
|
||||
}
|
||||
|
||||
/// Simple layout: title (+ optional chevron) | timestamp row + message + artifact buttons.
|
||||
fn render_simple_text_column(
|
||||
item: &NotificationItem,
|
||||
artifact_buttons: Option<&ViewHandle<ArtifactButtonsRow>>,
|
||||
context: NotificationRenderContext,
|
||||
message_expanded: bool,
|
||||
on_expand_click: OnExpandClick,
|
||||
extra_content: Option<Box<dyn Element>>,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let is_truncated = content_is_truncated(&item.title, &item.message);
|
||||
let title_text = render_clamped_title(&item.title, message_expanded, appearance);
|
||||
|
||||
let title_row: Box<dyn Element> = if context == NotificationRenderContext::Mailbox {
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(Shrinkable::new(1.0, title_text).finish())
|
||||
.with_child(render_timestamp_with_dot(item, appearance))
|
||||
.finish()
|
||||
} else if is_truncated || message_expanded {
|
||||
let chevron = render_expand_chevron(message_expanded, on_expand_click, theme);
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Start)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(Shrinkable::new(1.0, title_text).finish())
|
||||
.with_child(Container::new(chevron).with_margin_top(2.).finish())
|
||||
.finish()
|
||||
} else {
|
||||
title_text
|
||||
};
|
||||
|
||||
let message = render_message_text(&item.message, message_expanded, appearance);
|
||||
|
||||
let mut content = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_child(title_row)
|
||||
.with_child(Container::new(message).with_margin_top(2.).finish());
|
||||
|
||||
append_trailing_content(&mut content, artifact_buttons, extra_content);
|
||||
content.finish()
|
||||
}
|
||||
|
||||
/// Appends artifact buttons and extra content to a text column.
|
||||
fn append_trailing_content(
|
||||
content: &mut Flex,
|
||||
artifact_buttons: Option<&ViewHandle<ArtifactButtonsRow>>,
|
||||
extra_content: Option<Box<dyn Element>>,
|
||||
) {
|
||||
if let Some(artifact_buttons) = artifact_buttons {
|
||||
content.add_child(
|
||||
Container::new(ChildView::new(artifact_buttons).finish())
|
||||
.with_margin_top(8.)
|
||||
.finish(),
|
||||
);
|
||||
}
|
||||
if let Some(extra) = extra_content {
|
||||
content.add_child(extra);
|
||||
}
|
||||
}
|
||||
|
||||
/// Renders a git-branch icon + branch name label.
|
||||
fn render_branch_label(branch: &str, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let color = theme.sub_text_color(theme.surface_1());
|
||||
|
||||
Shrinkable::new(
|
||||
1.0,
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_spacing(2.)
|
||||
.with_child(
|
||||
ConstrainedBox::new(Icon::GitBranch.to_warpui_icon(color).finish())
|
||||
.with_width(10.)
|
||||
.with_height(10.)
|
||||
.finish(),
|
||||
)
|
||||
.with_child(
|
||||
Shrinkable::new(
|
||||
1.0,
|
||||
appearance
|
||||
.ui_builder()
|
||||
.wrappable_text(branch.to_owned(), false)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(12.),
|
||||
font_color: Some(color.into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish(),
|
||||
)
|
||||
.finish(),
|
||||
)
|
||||
.finish(),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Renders the timestamp text + optional unread dot.
|
||||
fn render_timestamp_with_dot(item: &NotificationItem, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let mut row = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_size(MainAxisSize::Min)
|
||||
.with_child(
|
||||
appearance
|
||||
.ui_builder()
|
||||
.wrappable_text(format_elapsed_since(item.created_at), false)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(12.),
|
||||
font_color: Some(theme.disabled_text_color(theme.surface_1()).into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish(),
|
||||
);
|
||||
|
||||
if !item.is_read {
|
||||
row.add_child(
|
||||
Container::new(
|
||||
ConstrainedBox::new(
|
||||
Rect::new()
|
||||
.with_background(theme.accent())
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Percentage(50.)))
|
||||
.finish(),
|
||||
)
|
||||
.with_width(10.)
|
||||
.with_height(10.)
|
||||
.finish(),
|
||||
)
|
||||
.with_margin_left(8.)
|
||||
.finish(),
|
||||
);
|
||||
}
|
||||
|
||||
row.finish()
|
||||
}
|
||||
|
||||
/// Renders a clickable expand/collapse chevron icon.
|
||||
fn render_expand_chevron(
|
||||
expanded: bool,
|
||||
on_click: OnExpandClick,
|
||||
theme: &WarpTheme,
|
||||
) -> Box<dyn Element> {
|
||||
let icon = if expanded {
|
||||
Icon::ChevronDown
|
||||
} else {
|
||||
Icon::ChevronRight
|
||||
};
|
||||
let chevron = ConstrainedBox::new(
|
||||
icon.to_warpui_icon(theme.disabled_text_color(theme.surface_1()))
|
||||
.finish(),
|
||||
)
|
||||
.with_width(12.)
|
||||
.with_height(12.)
|
||||
.finish();
|
||||
|
||||
EventHandler::new(chevron)
|
||||
.on_left_mouse_down(move |ctx, _, _| {
|
||||
on_click(ctx);
|
||||
DispatchEventResult::StopPropagation
|
||||
})
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Renders the title text, truncated based on expanded state.
|
||||
fn render_clamped_title(title: &str, expanded: bool, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let max = if expanded {
|
||||
EXPANDED_MAX_CHARS
|
||||
} else {
|
||||
COLLAPSED_MAX_CHARS
|
||||
};
|
||||
|
||||
appearance
|
||||
.ui_builder()
|
||||
.wrappable_text(truncate_text(title, max), true)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(14.),
|
||||
font_weight: Some(Weight::Semibold),
|
||||
font_color: Some(theme.main_text_color(theme.surface_1()).into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Renders the message text, truncated based on expanded state.
|
||||
fn render_message_text(message: &str, expanded: bool, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let max = if expanded {
|
||||
EXPANDED_MAX_CHARS
|
||||
} else {
|
||||
COLLAPSED_MAX_CHARS
|
||||
};
|
||||
|
||||
appearance
|
||||
.ui_builder()
|
||||
.wrappable_text(truncate_text(message, max), true)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(14.),
|
||||
font_color: Some(theme.sub_text_color(theme.surface_1()).into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish()
|
||||
}
|
||||
|
||||
const NOTIFICATION_AVATAR_SIZING: IconWithStatusSizing = IconWithStatusSizing {
|
||||
icon_size: 16.,
|
||||
padding: 8.,
|
||||
badge_icon_size: 12.,
|
||||
badge_padding: 2.,
|
||||
overall_size_override: None,
|
||||
badge_offset: (6., 6.),
|
||||
};
|
||||
|
||||
fn render_agent_avatar(
|
||||
agent: NotificationSourceAgent,
|
||||
category: NotificationCategory,
|
||||
theme: &WarpTheme,
|
||||
) -> Box<dyn Element> {
|
||||
let status = notification_category_to_conversation_status(category);
|
||||
let variant = match agent {
|
||||
NotificationSourceAgent::Oz => IconWithStatusVariant::OzAgent {
|
||||
status: Some(status),
|
||||
is_ambient: false,
|
||||
},
|
||||
NotificationSourceAgent::CLI(cli) => IconWithStatusVariant::CLIAgent {
|
||||
agent: cli,
|
||||
status: Some(status),
|
||||
},
|
||||
};
|
||||
render_icon_with_status(
|
||||
variant,
|
||||
&NOTIFICATION_AVATAR_SIZING,
|
||||
theme,
|
||||
theme.surface_2(),
|
||||
)
|
||||
}
|
||||
|
||||
fn notification_category_to_conversation_status(
|
||||
category: NotificationCategory,
|
||||
) -> ConversationStatus {
|
||||
match category {
|
||||
NotificationCategory::Complete => ConversationStatus::Success,
|
||||
NotificationCategory::Request => ConversationStatus::Blocked {
|
||||
blocked_action: String::new(),
|
||||
},
|
||||
NotificationCategory::Error => ConversationStatus::Error,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates an `ArtifactButtonsRow` view with the notification-specific theme.
|
||||
/// The caller is responsible for subscribing to events on the returned view.
|
||||
pub(crate) fn create_notification_artifact_buttons_view(
|
||||
artifacts: &[Artifact],
|
||||
ctx: &mut ViewContext<impl View>,
|
||||
) -> Option<ViewHandle<ArtifactButtonsRow>> {
|
||||
if artifacts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let theme = Arc::new(NotificationArtifactButtonTheme);
|
||||
Some(ctx.add_typed_action_view(|ctx| ArtifactButtonsRow::with_theme(artifacts, theme, ctx)))
|
||||
}
|
||||
|
||||
/// Handles artifact button events from notification views (toasts and mailbox).
|
||||
pub(crate) fn handle_notification_artifact_buttons_event(
|
||||
event: &ArtifactButtonsRowEvent,
|
||||
ctx: &mut ViewContext<impl View>,
|
||||
) {
|
||||
match event {
|
||||
ArtifactButtonsRowEvent::OpenPlan { notebook_uid } => {
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::ArtifactClicked {
|
||||
artifact_type: ArtifactType::Plan
|
||||
},
|
||||
ctx
|
||||
);
|
||||
ctx.dispatch_typed_action(&WorkspaceAction::OpenNotebook {
|
||||
id: (*notebook_uid).into(),
|
||||
});
|
||||
}
|
||||
ArtifactButtonsRowEvent::CopyBranch { branch } => {
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::ArtifactClicked {
|
||||
artifact_type: ArtifactType::Branch
|
||||
},
|
||||
ctx
|
||||
);
|
||||
ctx.clipboard()
|
||||
.write(ClipboardContent::plain_text(branch.clone()));
|
||||
}
|
||||
ArtifactButtonsRowEvent::OpenPullRequest { url } => {
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::ArtifactClicked {
|
||||
artifact_type: ArtifactType::PullRequest
|
||||
},
|
||||
ctx
|
||||
);
|
||||
ctx.open_url(url);
|
||||
}
|
||||
ArtifactButtonsRowEvent::ViewScreenshots { artifact_uids } => {
|
||||
open_screenshot_lightbox(artifact_uids, ctx);
|
||||
}
|
||||
ArtifactButtonsRowEvent::DownloadFile { artifact_uid } => {
|
||||
send_telemetry_from_ctx!(
|
||||
AgentManagementTelemetryEvent::ArtifactClicked {
|
||||
artifact_type: ArtifactType::File
|
||||
},
|
||||
ctx
|
||||
);
|
||||
crate::ai::artifacts::download_file_artifact(artifact_uid, ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
use warpui::EntityId;
|
||||
|
||||
use super::*;
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::terminal::CLIAgent;
|
||||
|
||||
fn make_conversation_notification(
|
||||
conversation_id: AIConversationId,
|
||||
terminal_view_id: EntityId,
|
||||
) -> NotificationItem {
|
||||
NotificationItem::new(
|
||||
"test".to_owned(),
|
||||
"msg".to_owned(),
|
||||
NotificationCategory::Complete,
|
||||
NotificationSourceAgent::Oz,
|
||||
NotificationOrigin::Conversation(conversation_id),
|
||||
false,
|
||||
terminal_view_id,
|
||||
vec![],
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn make_cli_session_notification(terminal_view_id: EntityId) -> NotificationItem {
|
||||
NotificationItem::new(
|
||||
"cli test".to_owned(),
|
||||
"cli msg".to_owned(),
|
||||
NotificationCategory::Complete,
|
||||
NotificationSourceAgent::CLI(CLIAgent::Claude),
|
||||
NotificationOrigin::CLISession(terminal_view_id),
|
||||
false,
|
||||
terminal_view_id,
|
||||
vec![],
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_by_origin_cleans_up_conversation_notification() {
|
||||
let mut items = NotificationItems::default();
|
||||
let conversation_id = AIConversationId::new();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
items.push(make_conversation_notification(
|
||||
conversation_id,
|
||||
terminal_view_id,
|
||||
));
|
||||
assert_eq!(items.filtered_count(NotificationFilter::All), 1);
|
||||
|
||||
let removed = items.remove_by_origin(NotificationOrigin::Conversation(conversation_id));
|
||||
assert!(removed);
|
||||
assert_eq!(items.filtered_count(NotificationFilter::All), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_by_origin_cleans_up_cli_session_notification() {
|
||||
let mut items = NotificationItems::default();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
items.push(make_cli_session_notification(terminal_view_id));
|
||||
assert_eq!(items.filtered_count(NotificationFilter::All), 1);
|
||||
|
||||
let removed = items.remove_by_origin(NotificationOrigin::CLISession(terminal_view_id));
|
||||
assert!(removed);
|
||||
assert_eq!(items.filtered_count(NotificationFilter::All), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_by_origin_leaves_unrelated_notifications() {
|
||||
let mut items = NotificationItems::default();
|
||||
let conv_id = AIConversationId::new();
|
||||
let terminal_a = EntityId::new();
|
||||
let terminal_b = EntityId::new();
|
||||
|
||||
items.push(make_conversation_notification(conv_id, terminal_a));
|
||||
items.push(make_cli_session_notification(terminal_b));
|
||||
assert_eq!(items.filtered_count(NotificationFilter::All), 2);
|
||||
|
||||
// Remove only the conversation notification; the CLI session notification should remain.
|
||||
let removed = items.remove_by_origin(NotificationOrigin::Conversation(conv_id));
|
||||
assert!(removed);
|
||||
assert_eq!(items.filtered_count(NotificationFilter::All), 1);
|
||||
|
||||
let remaining = items
|
||||
.items_filtered(NotificationFilter::All)
|
||||
.next()
|
||||
.unwrap();
|
||||
assert_eq!(remaining.origin, NotificationOrigin::CLISession(terminal_b));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_by_origin_returns_false_when_nothing_to_remove() {
|
||||
let mut items = NotificationItems::default();
|
||||
let terminal_view_id = EntityId::new();
|
||||
|
||||
let removed = items.remove_by_origin(NotificationOrigin::CLISession(terminal_view_id));
|
||||
assert!(!removed);
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
pub(crate) mod item;
|
||||
pub(crate) mod item_rendering;
|
||||
pub(crate) mod toast_stack;
|
||||
pub(crate) mod view;
|
||||
|
||||
pub(crate) use item::{
|
||||
NotificationCategory, NotificationFilter, NotificationId, NotificationItem, NotificationItems,
|
||||
NotificationOrigin, NotificationSourceAgent,
|
||||
};
|
||||
@@ -0,0 +1,500 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use pathfinder_geometry::vector::vec2f;
|
||||
use warp_core::ui::color::blend::Blend;
|
||||
use warp_core::ui::theme::color::internal_colors;
|
||||
use warpui::elements::{
|
||||
Border, ChildAnchor, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment,
|
||||
DispatchEventResult, Element, EventHandler, Flex, Hoverable, MouseStateHandle,
|
||||
OffsetPositioning, Padding, ParentElement, PositionedElementAnchor,
|
||||
PositionedElementOffsetBounds, Radius, SavePosition, Shrinkable, Stack,
|
||||
};
|
||||
use warpui::keymap::Keystroke;
|
||||
use warpui::platform::Cursor;
|
||||
use warpui::r#async::{SpawnedFutureHandle, Timer};
|
||||
use warpui::ui_components::components::{Coords, UiComponent, UiComponentStyles};
|
||||
use warpui::ui_components::keyboard_shortcut::KeyboardShortcut;
|
||||
use warpui::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext, ViewHandle};
|
||||
|
||||
use crate::ai::agent_management::notifications::item_rendering::{
|
||||
create_notification_artifact_buttons_view, handle_notification_artifact_buttons_event,
|
||||
render_notification_item_content, NotificationRenderContext, OnExpandClick,
|
||||
};
|
||||
use crate::ai::agent_management::notifications::{NotificationId, NotificationItem};
|
||||
use crate::ai::agent_management::{AgentManagementEvent, AgentNotificationsModel};
|
||||
use crate::ai::artifacts::{Artifact, ArtifactButtonsRow, ArtifactButtonsRowEvent};
|
||||
use crate::appearance::Appearance;
|
||||
use crate::terminal::session_settings::SessionSettings;
|
||||
use crate::util::bindings::keybinding_name_to_keystroke;
|
||||
use crate::workspace::view::JUMP_TO_LATEST_TOAST_BINDING_NAME;
|
||||
use crate::workspace::WorkspaceAction;
|
||||
|
||||
const CLOSE_BUTTON_SIZE: f32 = 20.;
|
||||
|
||||
/// Tracks the state of a single visible toast in the stack
|
||||
/// (its auto-dismiss timer, hover state, and close button).
|
||||
struct NotificationToastItem {
|
||||
notification_id: NotificationId,
|
||||
abort_handle: Option<SpawnedFutureHandle>,
|
||||
mouse_state: MouseStateHandle,
|
||||
close_button_mouse_state: MouseStateHandle,
|
||||
close_button_hover_state: MouseStateHandle,
|
||||
artifact_buttons_view: Option<ViewHandle<ArtifactButtonsRow>>,
|
||||
message_expanded: bool,
|
||||
}
|
||||
|
||||
pub struct AgentNotificationToastStack {
|
||||
toasts: Vec<NotificationToastItem>,
|
||||
mailbox_is_open: bool,
|
||||
}
|
||||
|
||||
impl Entity for AgentNotificationToastStack {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AgentNotificationToastAction {
|
||||
CancelDismissalTimeout(NotificationId),
|
||||
StartDismissalTimeout(NotificationId),
|
||||
Click(NotificationId),
|
||||
Dismiss(NotificationId),
|
||||
ToggleMessageExpanded(NotificationId),
|
||||
}
|
||||
|
||||
impl AgentNotificationToastStack {
|
||||
pub fn new(ctx: &mut ViewContext<Self>) -> Self {
|
||||
let model_handle = AgentNotificationsModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&model_handle, |me, _handle, event, ctx| match event {
|
||||
AgentManagementEvent::NotificationAdded { id } => {
|
||||
me.on_notification_added(*id, ctx);
|
||||
}
|
||||
AgentManagementEvent::NotificationUpdated
|
||||
| AgentManagementEvent::AllNotificationsMarkedRead => {
|
||||
me.remove_dismissed_toasts(ctx);
|
||||
}
|
||||
AgentManagementEvent::ConversationNeedsAttention { .. } => {}
|
||||
});
|
||||
|
||||
Self {
|
||||
toasts: Vec::new(),
|
||||
mailbox_is_open: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Updates the mailbox-open state. When opening, all visible toasts are dismissed
|
||||
/// (the mailbox already shows the same notifications).
|
||||
pub fn set_mailbox_open(&mut self, open: bool, ctx: &mut ViewContext<Self>) {
|
||||
self.mailbox_is_open = open;
|
||||
if open {
|
||||
self.dismiss_all(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
/// Dismiss all visible toasts (called when the mailbox opens).
|
||||
fn dismiss_all(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
for entry in self.toasts.drain(..) {
|
||||
if let Some(handle) = entry.abort_handle {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
ctx.notify();
|
||||
}
|
||||
|
||||
fn on_notification_added(&mut self, id: NotificationId, ctx: &mut ViewContext<Self>) {
|
||||
// Don't show in-app toasts when the window is not active.
|
||||
// Native desktop notifications handle the unfocused case.
|
||||
if ctx.windows().active_window() != Some(ctx.window_id()) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Don't show toasts when the notification mailbox is already open.
|
||||
// (dismiss_all is called on open, so any new arrival would be immediately visible in the mailbox.)
|
||||
if self.mailbox_is_open {
|
||||
return;
|
||||
}
|
||||
|
||||
let notifications = AgentNotificationsModel::as_ref(ctx).notifications();
|
||||
let Some(item) = notifications.get_by_id(id) else {
|
||||
return;
|
||||
};
|
||||
|
||||
// Don't show a toast for notifications that are already read
|
||||
// (e.g. the terminal was visible when the notification was created).
|
||||
if item.is_read {
|
||||
return;
|
||||
}
|
||||
|
||||
// Clone artifacts before releasing the immutable borrow on ctx.
|
||||
let artifacts = item.artifacts.clone();
|
||||
let _ = notifications;
|
||||
|
||||
// The notification model de-dupes by origin, so a new notification for the same
|
||||
// conversation replaces the old one with a new ID.
|
||||
self.remove_dismissed_toasts(ctx);
|
||||
|
||||
let artifact_buttons_view =
|
||||
Self::create_artifact_buttons_view_from_artifacts(&artifacts, ctx);
|
||||
|
||||
self.toasts.push(NotificationToastItem {
|
||||
notification_id: id,
|
||||
abort_handle: None,
|
||||
mouse_state: MouseStateHandle::default(),
|
||||
close_button_mouse_state: MouseStateHandle::default(),
|
||||
close_button_hover_state: MouseStateHandle::default(),
|
||||
artifact_buttons_view,
|
||||
message_expanded: false,
|
||||
});
|
||||
self.start_dismissal_timeout(id, ctx);
|
||||
|
||||
// Evict the oldest toasts if we exceed the visible limit.
|
||||
while self.toasts.len() > 2 {
|
||||
let evicted = self.toasts.remove(0);
|
||||
if let Some(handle) = evicted.abort_handle {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
ctx.notify();
|
||||
}
|
||||
|
||||
/// Removes toasts for notifications that no longer exist in the model or have been read.
|
||||
fn remove_dismissed_toasts(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
let notifications = AgentNotificationsModel::as_ref(ctx).notifications();
|
||||
let before = self.toasts.len();
|
||||
self.toasts.retain(|entry| {
|
||||
let should_remove = notifications
|
||||
.get_by_id(entry.notification_id)
|
||||
.is_none_or(|item| item.is_read);
|
||||
if should_remove {
|
||||
if let Some(handle) = &entry.abort_handle {
|
||||
handle.abort();
|
||||
}
|
||||
false
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
if self.toasts.len() != before {
|
||||
ctx.notify();
|
||||
}
|
||||
}
|
||||
|
||||
fn dismiss_toast_by_id(&mut self, id: &NotificationId, ctx: &mut ViewContext<Self>) {
|
||||
if let Some(idx) = self.toasts.iter().position(|e| e.notification_id == *id) {
|
||||
let entry = self.toasts.remove(idx);
|
||||
if let Some(handle) = entry.abort_handle {
|
||||
handle.abort();
|
||||
}
|
||||
ctx.notify();
|
||||
}
|
||||
}
|
||||
|
||||
/// Pauses the auto-dismiss timer for a toast (e.g. while the user is hovering over it).
|
||||
fn cancel_dismissal_timeout(&mut self, id: &NotificationId) {
|
||||
if let Some(entry) = self.toasts.iter_mut().find(|e| e.notification_id == *id) {
|
||||
if let Some(handle) = entry.abort_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn create_artifact_buttons_view_from_artifacts(
|
||||
artifacts: &[Artifact],
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) -> Option<ViewHandle<ArtifactButtonsRow>> {
|
||||
let view = create_notification_artifact_buttons_view(artifacts, ctx)?;
|
||||
ctx.subscribe_to_view(&view, Self::handle_artifact_buttons_event);
|
||||
Some(view)
|
||||
}
|
||||
|
||||
fn handle_artifact_buttons_event(
|
||||
&mut self,
|
||||
_view: ViewHandle<ArtifactButtonsRow>,
|
||||
event: &ArtifactButtonsRowEvent,
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) {
|
||||
handle_notification_artifact_buttons_event(event, ctx);
|
||||
}
|
||||
|
||||
fn start_dismissal_timeout(&mut self, id: NotificationId, ctx: &mut ViewContext<Self>) {
|
||||
if let Some(entry) = self.toasts.iter_mut().find(|e| e.notification_id == id) {
|
||||
if let Some(handle) = entry.abort_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
let duration_secs = *SessionSettings::as_ref(ctx).notification_toast_duration_secs;
|
||||
let abort_handle = ctx.spawn_abortable(
|
||||
Timer::after(Duration::from_secs(duration_secs)),
|
||||
move |me, _, ctx| me.dismiss_toast_by_id(&id, ctx),
|
||||
|_, _| {},
|
||||
);
|
||||
entry.abort_handle = Some(abort_handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TypedActionView for AgentNotificationToastStack {
|
||||
type Action = AgentNotificationToastAction;
|
||||
|
||||
fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
|
||||
match action {
|
||||
AgentNotificationToastAction::CancelDismissalTimeout(id) => {
|
||||
self.cancel_dismissal_timeout(id);
|
||||
}
|
||||
AgentNotificationToastAction::StartDismissalTimeout(id) => {
|
||||
self.start_dismissal_timeout(*id, ctx);
|
||||
}
|
||||
AgentNotificationToastAction::Click(id) => {
|
||||
let terminal_view_id = AgentNotificationsModel::as_ref(ctx)
|
||||
.notifications()
|
||||
.get_by_id(*id)
|
||||
.map(|item| item.terminal_view_id);
|
||||
|
||||
AgentNotificationsModel::handle(ctx).update(ctx, |model, ctx| {
|
||||
model.mark_item_read(*id, ctx);
|
||||
});
|
||||
|
||||
if let Some(terminal_view_id) = terminal_view_id {
|
||||
ctx.dispatch_typed_action(&WorkspaceAction::FocusTerminalViewInWorkspace {
|
||||
terminal_view_id,
|
||||
});
|
||||
}
|
||||
|
||||
self.dismiss_toast_by_id(id, ctx);
|
||||
}
|
||||
AgentNotificationToastAction::Dismiss(id) => {
|
||||
self.dismiss_toast_by_id(id, ctx);
|
||||
}
|
||||
AgentNotificationToastAction::ToggleMessageExpanded(id) => {
|
||||
if let Some(entry) = self.toasts.iter_mut().find(|e| e.notification_id == *id) {
|
||||
entry.message_expanded = !entry.message_expanded;
|
||||
ctx.notify();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl View for AgentNotificationToastStack {
|
||||
fn ui_name() -> &'static str {
|
||||
"AgentNotificationToastStack"
|
||||
}
|
||||
|
||||
fn render(&self, app: &AppContext) -> Box<dyn Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let notifications = AgentNotificationsModel::as_ref(app).notifications();
|
||||
let keystroke = keybinding_name_to_keystroke(JUMP_TO_LATEST_TOAST_BINDING_NAME, app);
|
||||
|
||||
let mut column = Flex::column().with_cross_axis_alignment(CrossAxisAlignment::End);
|
||||
|
||||
for (i, entry) in self.toasts.iter().rev().enumerate() {
|
||||
let Some(item) = notifications.get_by_id(entry.notification_id) else {
|
||||
continue;
|
||||
};
|
||||
let is_newest = i == 0;
|
||||
let toast = render_toast(
|
||||
item,
|
||||
entry.notification_id,
|
||||
entry.mouse_state.clone(),
|
||||
entry.close_button_mouse_state.clone(),
|
||||
entry.close_button_hover_state.clone(),
|
||||
entry.artifact_buttons_view.as_ref(),
|
||||
entry.message_expanded,
|
||||
is_newest.then(|| keystroke.clone()).flatten(),
|
||||
appearance,
|
||||
);
|
||||
column.add_child(Container::new(toast).with_margin_bottom(4.).finish());
|
||||
}
|
||||
|
||||
column.finish()
|
||||
}
|
||||
}
|
||||
|
||||
fn toast_position_id(id: NotificationId) -> String {
|
||||
format!("notification_toast_{id:?}")
|
||||
}
|
||||
|
||||
fn render_close_button(
|
||||
id: NotificationId,
|
||||
close_button_mouse_state: MouseStateHandle,
|
||||
close_button_hover_state: MouseStateHandle,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
EventHandler::new(
|
||||
Hoverable::new(close_button_hover_state, |_| {
|
||||
Container::new(
|
||||
appearance
|
||||
.ui_builder()
|
||||
.close_button(CLOSE_BUTTON_SIZE, close_button_mouse_state.clone())
|
||||
.with_style(UiComponentStyles {
|
||||
font_color: Some(theme.foreground().into()),
|
||||
background: Some(theme.surface_3().into()),
|
||||
border_color: Some(theme.outline().into()),
|
||||
border_radius: Some(CornerRadius::with_all(Radius::Percentage(50.))),
|
||||
border_width: Some(1.),
|
||||
padding: Some(Coords {
|
||||
top: 2.,
|
||||
bottom: 2.,
|
||||
left: 2.,
|
||||
right: 2.,
|
||||
}),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.on_click(move |ctx, _, _| {
|
||||
ctx.dispatch_typed_action(AgentNotificationToastAction::Dismiss(id));
|
||||
})
|
||||
.finish(),
|
||||
)
|
||||
.finish()
|
||||
})
|
||||
.finish(),
|
||||
)
|
||||
.on_left_mouse_down(|_, _, _| DispatchEventResult::StopPropagation)
|
||||
.finish()
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn render_toast(
|
||||
item: &NotificationItem,
|
||||
id: NotificationId,
|
||||
mouse_state: MouseStateHandle,
|
||||
close_button_mouse_state: MouseStateHandle,
|
||||
close_button_hover_state: MouseStateHandle,
|
||||
artifact_buttons: Option<&ViewHandle<ArtifactButtonsRow>>,
|
||||
message_expanded: bool,
|
||||
keystroke: Option<Keystroke>,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let on_expand: OnExpandClick = Box::new(move |ctx: &mut warpui::EventContext| {
|
||||
ctx.dispatch_typed_action(AgentNotificationToastAction::ToggleMessageExpanded(id));
|
||||
});
|
||||
let keybinding_hint = keystroke.map(|ks| render_keybinding_hint(ks, appearance));
|
||||
|
||||
let content = render_notification_item_content(
|
||||
item,
|
||||
artifact_buttons,
|
||||
NotificationRenderContext::Toast,
|
||||
message_expanded,
|
||||
on_expand,
|
||||
keybinding_hint,
|
||||
appearance,
|
||||
);
|
||||
|
||||
let inner_column = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_child(content);
|
||||
|
||||
let position_id = toast_position_id(id);
|
||||
|
||||
EventHandler::new(
|
||||
Hoverable::new(mouse_state, move |state| {
|
||||
let bg = if state.is_hovered() {
|
||||
theme
|
||||
.surface_2()
|
||||
.blend(&internal_colors::fg_overlay_3(theme))
|
||||
} else {
|
||||
theme
|
||||
.surface_2()
|
||||
.blend(&internal_colors::fg_overlay_2(theme))
|
||||
};
|
||||
|
||||
let container = Container::new(inner_column.finish())
|
||||
.with_padding(Padding::uniform(12.))
|
||||
.with_background(bg)
|
||||
.with_border(Border::all(1.).with_border_color(theme.outline().into()))
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(8.)));
|
||||
|
||||
let sized = ConstrainedBox::new(container.finish())
|
||||
.with_width(420.)
|
||||
.finish();
|
||||
|
||||
let is_close_hovered = close_button_hover_state
|
||||
.lock()
|
||||
.is_ok_and(|s| s.is_hovered());
|
||||
|
||||
let mut stack =
|
||||
Stack::new().with_child(SavePosition::new(sized, &position_id).finish());
|
||||
|
||||
if state.is_hovered() || is_close_hovered {
|
||||
stack.add_positioned_overlay_child(
|
||||
render_close_button(
|
||||
id,
|
||||
close_button_mouse_state.clone(),
|
||||
close_button_hover_state.clone(),
|
||||
appearance,
|
||||
),
|
||||
OffsetPositioning::offset_from_save_position_element(
|
||||
&position_id,
|
||||
vec2f(-4., -4.),
|
||||
PositionedElementOffsetBounds::WindowByPosition,
|
||||
PositionedElementAnchor::TopLeft,
|
||||
ChildAnchor::TopLeft,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
stack.finish()
|
||||
})
|
||||
.with_cursor(Cursor::PointingHand)
|
||||
.on_hover(move |is_hovered, ctx, _, _| {
|
||||
if is_hovered {
|
||||
ctx.dispatch_typed_action(AgentNotificationToastAction::CancelDismissalTimeout(id));
|
||||
} else {
|
||||
ctx.dispatch_typed_action(AgentNotificationToastAction::StartDismissalTimeout(id));
|
||||
}
|
||||
})
|
||||
.finish(),
|
||||
)
|
||||
.on_left_mouse_down(move |ctx, _, _| {
|
||||
ctx.dispatch_typed_action(AgentNotificationToastAction::Click(id));
|
||||
DispatchEventResult::StopPropagation
|
||||
})
|
||||
.finish()
|
||||
}
|
||||
|
||||
fn render_keybinding_hint(keystroke: Keystroke, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let hint_text = appearance
|
||||
.ui_builder()
|
||||
.wrappable_text("Open conversation".to_string(), false)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(12.),
|
||||
font_color: Some(theme.disabled_text_color(theme.surface_2()).into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish();
|
||||
|
||||
let keybinding_style = UiComponentStyles {
|
||||
font_family_id: Some(appearance.monospace_font_family()),
|
||||
font_color: Some(theme.sub_text_color(theme.surface_2()).into()),
|
||||
font_size: Some(12.),
|
||||
background: Some(internal_colors::fg_overlay_3(theme).into()),
|
||||
border_radius: Some(CornerRadius::with_all(Radius::Pixels(3.0))),
|
||||
padding: Some(Coords {
|
||||
top: 1.0,
|
||||
bottom: 1.0,
|
||||
left: 4.0,
|
||||
right: 4.0,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let shortcut = KeyboardShortcut::new(&keystroke, keybinding_style)
|
||||
.build()
|
||||
.finish();
|
||||
|
||||
Container::new(
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_child(Shrinkable::new(1.0, hint_text).finish())
|
||||
.with_child(Container::new(shortcut).with_margin_left(8.).finish())
|
||||
.finish(),
|
||||
)
|
||||
.with_margin_top(8.)
|
||||
.finish()
|
||||
}
|
||||
@@ -0,0 +1,625 @@
|
||||
use warp_core::ui::theme::color::internal_colors;
|
||||
use warpui::elements::new_scrollable::{ScrollableAppearance, SingleAxisConfig};
|
||||
use warpui::elements::{
|
||||
Border, ChildView, ClippedScrollStateHandle, ConstrainedBox, Container, CornerRadius,
|
||||
CrossAxisAlignment, Dismiss, DispatchEventResult, Element, Empty, EventHandler,
|
||||
Fill as ElementFill, Flex, Hoverable, MainAxisAlignment, MainAxisSize, MouseStateHandle,
|
||||
NewScrollable, Padding, ParentElement, Radius, SavePosition, ScrollTarget,
|
||||
ScrollToPositionMode, ScrollbarWidth, Shrinkable,
|
||||
};
|
||||
use warpui::fonts::Weight;
|
||||
use warpui::keymap::macros::id;
|
||||
use warpui::keymap::FixedBinding;
|
||||
use warpui::platform::Cursor;
|
||||
use warpui::ui_components::components::{UiComponent, UiComponentStyles};
|
||||
use warpui::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext, ViewHandle};
|
||||
|
||||
use crate::ai::agent_management::notifications::item::NotificationFilter;
|
||||
use crate::ai::agent_management::notifications::item_rendering::{
|
||||
create_notification_artifact_buttons_view, handle_notification_artifact_buttons_event,
|
||||
render_notification_item_content, NotificationRenderContext,
|
||||
};
|
||||
use crate::ai::agent_management::notifications::{
|
||||
NotificationId, NotificationItem, NotificationItems,
|
||||
};
|
||||
use crate::ai::agent_management::{AgentManagementEvent, AgentNotificationsModel};
|
||||
use crate::ai::artifacts::{Artifact, ArtifactButtonsRow, ArtifactButtonsRowEvent};
|
||||
use crate::appearance::Appearance;
|
||||
use crate::ui_components::icons::Icon;
|
||||
use crate::view_components::action_button::{ActionButton, ButtonSize, NakedTheme};
|
||||
|
||||
const ITEM_PADDING: f32 = 12.;
|
||||
|
||||
/// Position ID prefix used with `SavePosition` so the clipped scrollable can
|
||||
/// scroll keyboard-selected items into view.
|
||||
const ITEM_POSITION_PREFIX: &str = "notification_mailbox_item_";
|
||||
|
||||
pub struct NotificationMailboxView {
|
||||
active_filter: NotificationFilter,
|
||||
scroll_state: ClippedScrollStateHandle,
|
||||
filter_button_mouse_states: Vec<MouseStateHandle>,
|
||||
close_button: ViewHandle<ActionButton>,
|
||||
mark_all_read_button: ViewHandle<ActionButton>,
|
||||
notification_mouse_states: Vec<MouseStateHandle>,
|
||||
// Cached IDs of notifications matching the active filter, in display order.
|
||||
// (Avoids re-filtering the full list on every individual item render.)
|
||||
filtered_ids: Vec<NotificationId>,
|
||||
/// Artifact button views for each filtered notification (parallel to `filtered_ids`).
|
||||
artifact_buttons_views: Vec<Option<ViewHandle<ArtifactButtonsRow>>>,
|
||||
/// Index of the currently keyboard-selected notification item, if any.
|
||||
selected_index: Option<usize>,
|
||||
}
|
||||
|
||||
impl Entity for NotificationMailboxView {
|
||||
type Event = NotificationMailboxViewEvent;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum NotificationMailboxViewEvent {
|
||||
NavigateToTerminal { terminal_view_id: warpui::EntityId },
|
||||
Dismissed,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum NotificationMailboxViewAction {
|
||||
SetFilter(NotificationFilter),
|
||||
MarkAllRead,
|
||||
ClickItem(NotificationId),
|
||||
Dismiss,
|
||||
SelectPrevious,
|
||||
SelectNext,
|
||||
CycleFilter,
|
||||
ActivateSelected,
|
||||
}
|
||||
|
||||
impl NotificationMailboxView {
|
||||
pub fn init(app: &mut AppContext) {
|
||||
app.register_fixed_bindings([
|
||||
FixedBinding::new(
|
||||
"up",
|
||||
NotificationMailboxViewAction::SelectPrevious,
|
||||
id!(NotificationMailboxView::ui_name()),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"down",
|
||||
NotificationMailboxViewAction::SelectNext,
|
||||
id!(NotificationMailboxView::ui_name()),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"shift-tab",
|
||||
NotificationMailboxViewAction::CycleFilter,
|
||||
id!(NotificationMailboxView::ui_name()),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"enter",
|
||||
NotificationMailboxViewAction::ActivateSelected,
|
||||
id!(NotificationMailboxView::ui_name()),
|
||||
),
|
||||
FixedBinding::new(
|
||||
"escape",
|
||||
NotificationMailboxViewAction::Dismiss,
|
||||
id!(NotificationMailboxView::ui_name()),
|
||||
),
|
||||
]);
|
||||
}
|
||||
|
||||
pub fn new(ctx: &mut ViewContext<Self>) -> Self {
|
||||
let model_handle = AgentNotificationsModel::handle(ctx);
|
||||
ctx.subscribe_to_model(&model_handle, |me, _handle, event, ctx| match event {
|
||||
AgentManagementEvent::NotificationAdded { .. }
|
||||
| AgentManagementEvent::NotificationUpdated
|
||||
| AgentManagementEvent::AllNotificationsMarkedRead => {
|
||||
me.rebuild_filtered_ids(ctx);
|
||||
ctx.notify();
|
||||
}
|
||||
// Legacy toast path.
|
||||
AgentManagementEvent::ConversationNeedsAttention { .. } => {}
|
||||
});
|
||||
|
||||
let close_button = ctx.add_typed_action_view(|_| {
|
||||
ActionButton::new("", NakedTheme)
|
||||
.with_icon(Icon::X)
|
||||
.with_size(ButtonSize::XSmall)
|
||||
.with_tooltip("Close")
|
||||
.with_tooltip_sublabel("Esc")
|
||||
.on_click(|ctx| {
|
||||
ctx.dispatch_typed_action(NotificationMailboxViewAction::Dismiss);
|
||||
})
|
||||
});
|
||||
|
||||
let mark_all_read_button = ctx.add_typed_action_view(|_| {
|
||||
ActionButton::new("Mark all as read", NakedTheme)
|
||||
.with_size(ButtonSize::Small)
|
||||
.on_click(|ctx| {
|
||||
ctx.dispatch_typed_action(NotificationMailboxViewAction::MarkAllRead);
|
||||
})
|
||||
});
|
||||
|
||||
Self {
|
||||
active_filter: NotificationFilter::All,
|
||||
scroll_state: Default::default(),
|
||||
filter_button_mouse_states: (0..enum_iterator::cardinality::<NotificationFilter>())
|
||||
.map(|_| MouseStateHandle::default())
|
||||
.collect(),
|
||||
close_button,
|
||||
mark_all_read_button,
|
||||
notification_mouse_states: Vec::new(),
|
||||
filtered_ids: Vec::new(),
|
||||
artifact_buttons_views: Vec::new(),
|
||||
selected_index: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resets the mailbox state when opening. Called from the workspace toggle handler.
|
||||
pub fn reset_for_open(&mut self, select_first: bool, ctx: &mut ViewContext<Self>) {
|
||||
self.rebuild_filtered_ids(ctx);
|
||||
self.selected_index = if select_first && !self.filtered_ids.is_empty() {
|
||||
Some(0)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
}
|
||||
|
||||
fn set_active_filter(&mut self, filter: NotificationFilter, ctx: &mut ViewContext<Self>) {
|
||||
self.active_filter = filter;
|
||||
self.selected_index = None;
|
||||
self.rebuild_filtered_ids(ctx);
|
||||
ctx.notify();
|
||||
}
|
||||
|
||||
fn activate_notification(&mut self, id: NotificationId, ctx: &mut ViewContext<Self>) {
|
||||
let terminal_view_id = AgentNotificationsModel::as_ref(ctx)
|
||||
.notifications()
|
||||
.get_by_id(id)
|
||||
.map(|item| item.terminal_view_id);
|
||||
|
||||
AgentNotificationsModel::handle(ctx).update(ctx, |model, ctx| {
|
||||
model.mark_item_read(id, ctx);
|
||||
});
|
||||
|
||||
if let Some(terminal_view_id) = terminal_view_id {
|
||||
ctx.emit(NotificationMailboxViewEvent::NavigateToTerminal { terminal_view_id });
|
||||
}
|
||||
}
|
||||
|
||||
/// Refreshes the cached filtered notification IDs and mouse states.
|
||||
fn rebuild_filtered_ids(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
let notifications = AgentNotificationsModel::as_ref(ctx).notifications();
|
||||
|
||||
// If the active filter's tab would be hidden (0 items), fall back to "All".
|
||||
if self.active_filter != NotificationFilter::All
|
||||
&& notifications.filtered_count(self.active_filter) == 0
|
||||
{
|
||||
self.active_filter = NotificationFilter::All;
|
||||
}
|
||||
self.filtered_ids = notifications
|
||||
.items_filtered(self.active_filter)
|
||||
.map(|item| item.id)
|
||||
.collect();
|
||||
self.notification_mouse_states
|
||||
.resize_with(self.filtered_ids.len(), MouseStateHandle::default);
|
||||
|
||||
let artifact_data: Vec<_> = notifications
|
||||
.items_filtered(self.active_filter)
|
||||
.map(|item| item.artifacts.clone())
|
||||
.collect();
|
||||
let _ = notifications;
|
||||
|
||||
self.artifact_buttons_views = artifact_data
|
||||
.iter()
|
||||
.map(|artifacts| Self::create_artifact_buttons_view_from_artifacts(artifacts, ctx))
|
||||
.collect();
|
||||
|
||||
// Clamp selection to valid range after list contents change.
|
||||
if self.filtered_ids.is_empty() {
|
||||
self.selected_index = None;
|
||||
} else if let Some(idx) = self.selected_index {
|
||||
if idx >= self.filtered_ids.len() {
|
||||
self.selected_index = Some(self.filtered_ids.len() - 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn render_item_at_index(&self, index: usize, app: &AppContext) -> Box<dyn Element> {
|
||||
let notifications = AgentNotificationsModel::as_ref(app).notifications();
|
||||
let Some(item) = self
|
||||
.filtered_ids
|
||||
.get(index)
|
||||
.and_then(|id| notifications.get_by_id(*id))
|
||||
else {
|
||||
return Empty::new().finish();
|
||||
};
|
||||
|
||||
let Some(mouse_state) = self.notification_mouse_states.get(index).cloned() else {
|
||||
log::error!("missing mouse state for notification item at index {index}");
|
||||
return Empty::new().finish();
|
||||
};
|
||||
|
||||
let artifact_buttons = self
|
||||
.artifact_buttons_views
|
||||
.get(index)
|
||||
.and_then(|v| v.as_ref());
|
||||
let is_selected = self.selected_index == Some(index);
|
||||
self.render_notification_item(
|
||||
item,
|
||||
mouse_state,
|
||||
artifact_buttons,
|
||||
is_selected,
|
||||
Appearance::as_ref(app),
|
||||
)
|
||||
}
|
||||
|
||||
fn create_artifact_buttons_view_from_artifacts(
|
||||
artifacts: &[Artifact],
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) -> Option<ViewHandle<ArtifactButtonsRow>> {
|
||||
let view = create_notification_artifact_buttons_view(artifacts, ctx)?;
|
||||
ctx.subscribe_to_view(&view, Self::handle_artifact_buttons_event);
|
||||
Some(view)
|
||||
}
|
||||
|
||||
fn handle_artifact_buttons_event(
|
||||
&mut self,
|
||||
_view: ViewHandle<ArtifactButtonsRow>,
|
||||
event: &ArtifactButtonsRowEvent,
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) {
|
||||
handle_notification_artifact_buttons_event(event, ctx);
|
||||
}
|
||||
}
|
||||
|
||||
impl TypedActionView for NotificationMailboxView {
|
||||
type Action = NotificationMailboxViewAction;
|
||||
|
||||
fn handle_action(&mut self, action: &Self::Action, ctx: &mut ViewContext<Self>) {
|
||||
match action {
|
||||
NotificationMailboxViewAction::SetFilter(filter) => {
|
||||
self.set_active_filter(*filter, ctx);
|
||||
}
|
||||
NotificationMailboxViewAction::MarkAllRead => {
|
||||
AgentNotificationsModel::handle(ctx).update(ctx, |model, ctx| {
|
||||
model.mark_all_items_read(ctx);
|
||||
});
|
||||
}
|
||||
NotificationMailboxViewAction::ClickItem(id) => {
|
||||
self.activate_notification(*id, ctx);
|
||||
}
|
||||
NotificationMailboxViewAction::Dismiss => {
|
||||
ctx.emit(NotificationMailboxViewEvent::Dismissed);
|
||||
}
|
||||
NotificationMailboxViewAction::SelectPrevious => {
|
||||
match self.selected_index {
|
||||
Some(idx) if idx > 0 => self.selected_index = Some(idx - 1),
|
||||
None if !self.filtered_ids.is_empty() => self.selected_index = Some(0),
|
||||
_ => {}
|
||||
}
|
||||
self.scroll_selected_into_view();
|
||||
ctx.notify();
|
||||
}
|
||||
NotificationMailboxViewAction::SelectNext => {
|
||||
let max = self.filtered_ids.len().saturating_sub(1);
|
||||
match self.selected_index {
|
||||
Some(idx) if idx < max => self.selected_index = Some(idx + 1),
|
||||
None if !self.filtered_ids.is_empty() => self.selected_index = Some(0),
|
||||
_ => {}
|
||||
}
|
||||
self.scroll_selected_into_view();
|
||||
ctx.notify();
|
||||
}
|
||||
NotificationMailboxViewAction::CycleFilter => {
|
||||
let notifications = AgentNotificationsModel::as_ref(ctx).notifications();
|
||||
let visible = notifications.visible_filters();
|
||||
let current_pos = visible
|
||||
.iter()
|
||||
.position(|f| *f == self.active_filter)
|
||||
.unwrap_or(0);
|
||||
let next_filter = visible[(current_pos + 1) % visible.len()];
|
||||
self.set_active_filter(next_filter, ctx);
|
||||
}
|
||||
NotificationMailboxViewAction::ActivateSelected => {
|
||||
if let Some(idx) = self.selected_index {
|
||||
if let Some(id) = self.filtered_ids.get(idx).copied() {
|
||||
self.activate_notification(id, ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl View for NotificationMailboxView {
|
||||
fn ui_name() -> &'static str {
|
||||
"NotificationMailboxView"
|
||||
}
|
||||
|
||||
fn render(&self, app: &AppContext) -> Box<dyn Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let notifications = AgentNotificationsModel::as_ref(app).notifications();
|
||||
|
||||
let mut column = Flex::column()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Stretch)
|
||||
.with_main_axis_size(MainAxisSize::Min)
|
||||
.with_child(self.render_header(appearance))
|
||||
.with_child(self.render_filter_bar(notifications, app));
|
||||
|
||||
if notifications.filtered_count(self.active_filter) == 0 {
|
||||
column.add_child(self.render_empty_state(appearance));
|
||||
} else {
|
||||
let theme = appearance.theme();
|
||||
|
||||
// Render all items directly in a column so the scrollable can
|
||||
// measure their natural height instead of always filling its
|
||||
// max constraint (which was the behaviour with the viewport-based
|
||||
// List element).
|
||||
let mut items_column =
|
||||
Flex::column().with_cross_axis_alignment(CrossAxisAlignment::Stretch);
|
||||
for index in 0..self.filtered_ids.len() {
|
||||
items_column.add_child(
|
||||
SavePosition::new(
|
||||
self.render_item_at_index(index, app),
|
||||
&format!("{ITEM_POSITION_PREFIX}{index}"),
|
||||
)
|
||||
.finish(),
|
||||
);
|
||||
}
|
||||
|
||||
let item_list = NewScrollable::vertical(
|
||||
SingleAxisConfig::Clipped {
|
||||
handle: self.scroll_state.clone(),
|
||||
child: items_column.finish(),
|
||||
},
|
||||
theme.nonactive_ui_detail().into(),
|
||||
theme.active_ui_detail().into(),
|
||||
ElementFill::None,
|
||||
)
|
||||
.with_vertical_scrollbar(ScrollableAppearance::new(ScrollbarWidth::Auto, true))
|
||||
.finish();
|
||||
column.add_child(Shrinkable::new(1.0, item_list).finish());
|
||||
}
|
||||
|
||||
let popup = Container::new(column.finish())
|
||||
.with_padding(Padding::default().with_top(4.))
|
||||
.with_background(appearance.theme().surface_2())
|
||||
.with_border(Border::all(1.).with_border_color(appearance.theme().outline().into()))
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(6.)));
|
||||
|
||||
// Wrap the popup in an EventHandler that consumes clicks on empty space
|
||||
// so they don't fall through to the Dismiss layer.
|
||||
let popup = EventHandler::new(popup.finish())
|
||||
.on_left_mouse_down(|_, _, _| DispatchEventResult::StopPropagation)
|
||||
.finish();
|
||||
|
||||
Dismiss::new(
|
||||
ConstrainedBox::new(popup)
|
||||
.with_width(420.)
|
||||
.with_max_height(500.)
|
||||
.finish(),
|
||||
)
|
||||
.on_dismiss(|ctx, _app| {
|
||||
ctx.dispatch_typed_action(NotificationMailboxViewAction::Dismiss);
|
||||
})
|
||||
.prevent_interaction_with_other_elements()
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl NotificationMailboxView {
|
||||
fn render_header(&self, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
let label = appearance
|
||||
.ui_builder()
|
||||
.wrappable_text("Notifications".to_string(), false)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(14.),
|
||||
font_color: Some(theme.main_text_color(theme.surface_2()).into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish();
|
||||
|
||||
Container::new(
|
||||
Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(label)
|
||||
.with_child(ChildView::new(&self.close_button).finish())
|
||||
.finish(),
|
||||
)
|
||||
.with_padding(
|
||||
Padding::default()
|
||||
.with_top(8.)
|
||||
.with_bottom(4.)
|
||||
.with_horizontal(12.),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
|
||||
fn render_filter_bar(
|
||||
&self,
|
||||
notifications: &NotificationItems,
|
||||
app: &AppContext,
|
||||
) -> Box<dyn Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let theme = appearance.theme();
|
||||
let has_unread = notifications.filtered_count(NotificationFilter::Unread) > 0;
|
||||
|
||||
let mut filter_buttons = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_size(MainAxisSize::Min)
|
||||
.with_spacing(2.);
|
||||
|
||||
for (i, filter) in notifications.visible_filters().into_iter().enumerate() {
|
||||
let Some(mouse_state) = self.filter_button_mouse_states.get(i).cloned() else {
|
||||
log::warn!("missing mouse state for filter button at index {i}");
|
||||
continue;
|
||||
};
|
||||
|
||||
let is_active = self.active_filter == filter;
|
||||
let count = notifications.filtered_count(filter);
|
||||
let label = if count == 0 {
|
||||
filter.label().to_string()
|
||||
} else {
|
||||
format!("{} ({count})", filter.label())
|
||||
};
|
||||
let text_color = if is_active {
|
||||
theme.main_text_color(theme.surface_2())
|
||||
} else {
|
||||
theme.sub_text_color(theme.surface_2())
|
||||
};
|
||||
let background = if is_active {
|
||||
Some(internal_colors::fg_overlay_3(theme).into())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let button = EventHandler::new(
|
||||
Hoverable::new(mouse_state, move |state| {
|
||||
let bg = if state.is_hovered() && !is_active {
|
||||
Some(internal_colors::fg_overlay_2(theme).into())
|
||||
} else {
|
||||
background
|
||||
};
|
||||
|
||||
let mut container = Container::new(
|
||||
appearance
|
||||
.ui_builder()
|
||||
.wrappable_text(label.clone(), false)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(12.),
|
||||
font_weight: Some(Weight::Semibold),
|
||||
font_color: Some(text_color.into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish(),
|
||||
)
|
||||
.with_padding(Padding::default().with_vertical(4.).with_horizontal(8.))
|
||||
.with_corner_radius(CornerRadius::with_all(Radius::Pixels(4.)));
|
||||
|
||||
if let Some(bg) = bg {
|
||||
container = container.with_background_color(bg);
|
||||
}
|
||||
|
||||
container.finish()
|
||||
})
|
||||
.with_cursor(Cursor::PointingHand)
|
||||
.finish(),
|
||||
)
|
||||
.on_left_mouse_down(move |ctx, _, _| {
|
||||
ctx.dispatch_typed_action(NotificationMailboxViewAction::SetFilter(filter));
|
||||
DispatchEventResult::StopPropagation
|
||||
})
|
||||
.finish();
|
||||
|
||||
filter_buttons.add_child(button);
|
||||
}
|
||||
|
||||
let mut row = Flex::row()
|
||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
||||
.with_main_axis_alignment(MainAxisAlignment::SpaceBetween)
|
||||
.with_main_axis_size(MainAxisSize::Max)
|
||||
.with_child(filter_buttons.finish());
|
||||
|
||||
if has_unread {
|
||||
row.add_child(ChildView::new(&self.mark_all_read_button).finish());
|
||||
}
|
||||
|
||||
Container::new(row.finish())
|
||||
.with_padding(
|
||||
Padding::default()
|
||||
.with_vertical(12.)
|
||||
.with_left(12.)
|
||||
.with_right(6.),
|
||||
)
|
||||
.with_border(Border::top(1.).with_border_color(theme.outline().into()))
|
||||
.with_border(Border::bottom(1.).with_border_color(theme.outline().into()))
|
||||
.finish()
|
||||
}
|
||||
|
||||
fn render_empty_state(&self, appearance: &Appearance) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
|
||||
Container::new(
|
||||
appearance
|
||||
.ui_builder()
|
||||
.wrappable_text("No notifications".to_string(), false)
|
||||
.with_style(UiComponentStyles {
|
||||
font_size: Some(14.),
|
||||
font_color: Some(theme.sub_text_color(theme.surface_2()).into()),
|
||||
font_family_id: Some(appearance.ui_font_family()),
|
||||
..Default::default()
|
||||
})
|
||||
.build()
|
||||
.finish(),
|
||||
)
|
||||
.with_uniform_padding(ITEM_PADDING)
|
||||
.finish()
|
||||
}
|
||||
|
||||
/// Asks the clipped scrollable to bring the currently selected item into
|
||||
/// view on the next paint pass.
|
||||
fn scroll_selected_into_view(&self) {
|
||||
if let Some(idx) = self.selected_index {
|
||||
self.scroll_state.scroll_to_position(ScrollTarget {
|
||||
position_id: format!("{ITEM_POSITION_PREFIX}{idx}"),
|
||||
mode: ScrollToPositionMode::FullyIntoView,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn render_notification_item(
|
||||
&self,
|
||||
item: &NotificationItem,
|
||||
mouse_state: MouseStateHandle,
|
||||
artifact_buttons: Option<&ViewHandle<ArtifactButtonsRow>>,
|
||||
is_selected: bool,
|
||||
appearance: &Appearance,
|
||||
) -> Box<dyn Element> {
|
||||
let theme = appearance.theme();
|
||||
let id = item.id;
|
||||
let has_branch = item.branch.is_some();
|
||||
let row = render_notification_item_content(
|
||||
item,
|
||||
artifact_buttons,
|
||||
NotificationRenderContext::Mailbox,
|
||||
false,
|
||||
Box::new(|_| {}),
|
||||
None,
|
||||
appearance,
|
||||
);
|
||||
|
||||
EventHandler::new(
|
||||
Hoverable::new(mouse_state, move |state| {
|
||||
let item_padding = if has_branch {
|
||||
Padding::uniform(ITEM_PADDING)
|
||||
} else {
|
||||
Padding::default()
|
||||
.with_vertical(ITEM_PADDING)
|
||||
.with_horizontal(16.)
|
||||
};
|
||||
let mut container = Container::new(row).with_padding(item_padding);
|
||||
|
||||
if is_selected {
|
||||
container = container
|
||||
.with_background_color(internal_colors::fg_overlay_3(theme).into());
|
||||
} else if state.is_hovered() {
|
||||
container = container
|
||||
.with_background_color(internal_colors::fg_overlay_2(theme).into());
|
||||
}
|
||||
|
||||
container.finish()
|
||||
})
|
||||
.with_cursor(Cursor::PointingHand)
|
||||
.finish(),
|
||||
)
|
||||
.on_left_mouse_down(move |ctx, _, _| {
|
||||
ctx.dispatch_typed_action(NotificationMailboxViewAction::ClickItem(id));
|
||||
DispatchEventResult::StopPropagation
|
||||
})
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,293 @@
|
||||
use serde::Serialize;
|
||||
use serde_json::json;
|
||||
use strum_macros::{EnumDiscriminants, EnumIter};
|
||||
use warp_core::telemetry::{EnablementState, TelemetryEvent, TelemetryEventDesc};
|
||||
|
||||
use crate::ai::agent_management::cloud_setup_guide_view::SetupGuideDocs;
|
||||
|
||||
/// Which setup guide workflow step the user interacted with
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SetupGuideStep {
|
||||
/// Quick start banner: Visit Oz
|
||||
VisitOz,
|
||||
/// Step 1: Create environment (slash command)
|
||||
CreateEnvironment,
|
||||
/// Step 1: Create environment (CLI command)
|
||||
CreateEnvironmentCli,
|
||||
/// Step 2: Create Slack integration
|
||||
CreateSlackIntegration,
|
||||
/// Step 2: Create Linear integration
|
||||
CreateLinearIntegration,
|
||||
}
|
||||
|
||||
/// Where the item was opened from
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum OpenedFrom {
|
||||
ManagementView,
|
||||
ConversationList,
|
||||
DetailsPanel,
|
||||
}
|
||||
|
||||
/// Type of artifact clicked
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ArtifactType {
|
||||
Plan,
|
||||
Branch,
|
||||
PullRequest,
|
||||
File,
|
||||
}
|
||||
|
||||
/// Type of filter changed
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FilterType {
|
||||
Status,
|
||||
Source,
|
||||
CreatedOn,
|
||||
Creator,
|
||||
Owner,
|
||||
Harness,
|
||||
}
|
||||
|
||||
/// Telemetry events for the agent management view
|
||||
#[derive(Serialize, Debug, EnumDiscriminants)]
|
||||
#[strum_discriminants(derive(EnumIter))]
|
||||
pub enum AgentManagementTelemetryEvent {
|
||||
/// User toggled the agent management view open or closed
|
||||
ViewToggled { is_open: bool },
|
||||
/// User opened the setup guide
|
||||
OpenSetupGuide,
|
||||
/// User dismissed the setup guide
|
||||
DismissSetupGuide,
|
||||
/// User spawned a new local agent
|
||||
SpawnNewLocalAgent,
|
||||
/// User spawned a new cloud agent
|
||||
SpawnNewCloudAgent,
|
||||
/// User opened the agent type selector modal
|
||||
AgentTypeSelectorOpened,
|
||||
/// User ran a workflow step from the setup guide
|
||||
SetupGuideStepRun { step: SetupGuideStep },
|
||||
/// User copied a workflow step from the setup guide
|
||||
SetupGuideStepCopy { step: SetupGuideStep },
|
||||
/// User clicked a URL in the setup guide
|
||||
SetupGuideDocsLink { docs: SetupGuideDocs },
|
||||
/// User opened a conversation
|
||||
ConversationOpened {
|
||||
conversation_id: String,
|
||||
opened_from: OpenedFrom,
|
||||
},
|
||||
/// User opened a cloud run
|
||||
CloudRunOpened {
|
||||
task_id: String,
|
||||
opened_from: OpenedFrom,
|
||||
},
|
||||
/// User clicked an artifact button
|
||||
ArtifactClicked { artifact_type: ArtifactType },
|
||||
/// User changed a filter
|
||||
FilterChanged { filter_type: FilterType },
|
||||
/// User clicked an item details button
|
||||
DetailsViewed {
|
||||
item_id: String,
|
||||
viewed_from: OpenedFrom,
|
||||
},
|
||||
/// User copied a conversation link
|
||||
ConversationLinkCopied {
|
||||
conversation_id: String,
|
||||
copied_from: OpenedFrom,
|
||||
},
|
||||
/// User copied a session link
|
||||
SessionLinkCopied {
|
||||
task_id: String,
|
||||
copied_from: OpenedFrom,
|
||||
},
|
||||
/// User clicked an artifact in the tombstone view
|
||||
TombstoneArtifactClicked { artifact_type: ArtifactType },
|
||||
/// User clicked "Continue locally" in the tombstone
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
TombstoneContinueLocally,
|
||||
/// User clicked "Continue locally" in the details panel
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
DetailsPanelContinueLocally,
|
||||
/// User clicked "Open in Warp" in the tombstone (wasm)
|
||||
#[cfg(target_family = "wasm")]
|
||||
TombstoneOpenInWarp,
|
||||
/// User cancelled a cloud run
|
||||
CloudRunCancelled { task_id: String },
|
||||
/// User forked a conversation
|
||||
ConversationForked { conversation_id: String },
|
||||
}
|
||||
|
||||
impl TelemetryEvent for AgentManagementTelemetryEvent {
|
||||
fn name(&self) -> &'static str {
|
||||
AgentManagementTelemetryEventDiscriminants::from(self).name()
|
||||
}
|
||||
|
||||
fn payload(&self) -> Option<serde_json::Value> {
|
||||
match self {
|
||||
AgentManagementTelemetryEvent::ViewToggled { is_open } => {
|
||||
Some(json!({ "is_open": is_open }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::OpenSetupGuide => None,
|
||||
AgentManagementTelemetryEvent::DismissSetupGuide => None,
|
||||
AgentManagementTelemetryEvent::SpawnNewLocalAgent => None,
|
||||
AgentManagementTelemetryEvent::SpawnNewCloudAgent => None,
|
||||
AgentManagementTelemetryEvent::AgentTypeSelectorOpened => None,
|
||||
AgentManagementTelemetryEvent::SetupGuideStepRun { step } => {
|
||||
Some(json!({ "step": step }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::SetupGuideStepCopy { step } => {
|
||||
Some(json!({ "step": step }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::SetupGuideDocsLink { docs } => {
|
||||
Some(json!({ "docs": docs }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::ConversationOpened {
|
||||
conversation_id,
|
||||
opened_from,
|
||||
} => Some(json!({
|
||||
"conversation_id": conversation_id,
|
||||
"opened_from": opened_from,
|
||||
})),
|
||||
AgentManagementTelemetryEvent::CloudRunOpened {
|
||||
task_id,
|
||||
opened_from,
|
||||
} => Some(json!({
|
||||
"task_id": task_id,
|
||||
"opened_from": opened_from,
|
||||
})),
|
||||
AgentManagementTelemetryEvent::ArtifactClicked { artifact_type } => {
|
||||
Some(json!({ "artifact_type": artifact_type }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::FilterChanged { filter_type } => {
|
||||
Some(json!({ "filter_type": filter_type }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::DetailsViewed {
|
||||
item_id,
|
||||
viewed_from,
|
||||
} => Some(json!({
|
||||
"item_id": item_id,
|
||||
"viewed_from": viewed_from,
|
||||
})),
|
||||
AgentManagementTelemetryEvent::ConversationLinkCopied {
|
||||
conversation_id,
|
||||
copied_from,
|
||||
} => Some(json!({
|
||||
"conversation_id": conversation_id,
|
||||
"copied_from": copied_from,
|
||||
})),
|
||||
AgentManagementTelemetryEvent::SessionLinkCopied {
|
||||
task_id,
|
||||
copied_from,
|
||||
} => Some(json!({
|
||||
"task_id": task_id,
|
||||
"copied_from": copied_from,
|
||||
})),
|
||||
AgentManagementTelemetryEvent::TombstoneArtifactClicked { artifact_type } => {
|
||||
Some(json!({ "artifact_type": artifact_type }))
|
||||
}
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
AgentManagementTelemetryEvent::TombstoneContinueLocally => None,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
AgentManagementTelemetryEvent::DetailsPanelContinueLocally => None,
|
||||
#[cfg(target_family = "wasm")]
|
||||
AgentManagementTelemetryEvent::TombstoneOpenInWarp => None,
|
||||
AgentManagementTelemetryEvent::CloudRunCancelled { task_id } => {
|
||||
Some(json!({ "task_id": task_id }))
|
||||
}
|
||||
AgentManagementTelemetryEvent::ConversationForked { conversation_id } => {
|
||||
Some(json!({ "conversation_id": conversation_id }))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn description(&self) -> &'static str {
|
||||
AgentManagementTelemetryEventDiscriminants::from(self).description()
|
||||
}
|
||||
|
||||
fn enablement_state(&self) -> EnablementState {
|
||||
AgentManagementTelemetryEventDiscriminants::from(self).enablement_state()
|
||||
}
|
||||
|
||||
fn contains_ugc(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn event_descs() -> impl Iterator<Item = Box<dyn TelemetryEventDesc>> {
|
||||
warp_core::telemetry::enum_events::<Self>()
|
||||
}
|
||||
}
|
||||
|
||||
impl TelemetryEventDesc for AgentManagementTelemetryEventDiscriminants {
|
||||
fn name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::ViewToggled => "AgentManagement.ViewToggled",
|
||||
Self::OpenSetupGuide => "AgentManagement.OpenSetupGuide",
|
||||
Self::DismissSetupGuide => "AgentManagement.DismissSetupGuide",
|
||||
Self::SpawnNewLocalAgent => "AgentManagement.SpawnNewLocalAgent",
|
||||
Self::SpawnNewCloudAgent => "AgentManagement.SpawnNewCloudAgent",
|
||||
Self::AgentTypeSelectorOpened => "AgentManagement.AgentTypeSelectorOpened",
|
||||
Self::SetupGuideStepRun => "AgentManagement.SetupGuideStepRun",
|
||||
Self::SetupGuideStepCopy => "AgentManagement.SetupGuideStepCopy",
|
||||
Self::SetupGuideDocsLink => "AgentManagement.SetupGuideDocsLink",
|
||||
Self::ConversationOpened => "AgentManagement.ConversationOpened",
|
||||
Self::CloudRunOpened => "AgentManagement.CloudRunOpened",
|
||||
Self::ArtifactClicked => "AgentManagement.ArtifactClicked",
|
||||
Self::FilterChanged => "AgentManagement.FilterChanged",
|
||||
Self::DetailsViewed => "AgentManagement.DetailsViewed",
|
||||
Self::ConversationLinkCopied => "AgentManagement.ConversationLinkCopied",
|
||||
Self::SessionLinkCopied => "AgentManagement.SessionLinkCopied",
|
||||
Self::TombstoneArtifactClicked => "AgentManagement.TombstoneArtifactClicked",
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
Self::TombstoneContinueLocally => "AgentManagement.TombstoneContinueLocally",
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
Self::DetailsPanelContinueLocally => "AgentManagement.DetailsPanelContinueLocally",
|
||||
#[cfg(target_family = "wasm")]
|
||||
Self::TombstoneOpenInWarp => "AgentManagement.TombstoneOpenInWarp",
|
||||
Self::CloudRunCancelled => "AgentManagement.CloudRunCancelled",
|
||||
Self::ConversationForked => "AgentManagement.ConversationForked",
|
||||
}
|
||||
}
|
||||
|
||||
fn description(&self) -> &'static str {
|
||||
match self {
|
||||
Self::ViewToggled => "User toggled the agent management view open or closed",
|
||||
Self::OpenSetupGuide => "User opened the ambient agent setup guide",
|
||||
Self::DismissSetupGuide => "User dismissed the ambient agent setup guide",
|
||||
Self::SpawnNewLocalAgent => "User spawned a new local agent from agent management",
|
||||
Self::SpawnNewCloudAgent => "User spawned a new cloud agent from agent management",
|
||||
Self::AgentTypeSelectorOpened => {
|
||||
"User opened the agent type selector from agent management"
|
||||
}
|
||||
Self::SetupGuideStepRun => "User ran a workflow step from the setup guide",
|
||||
Self::SetupGuideStepCopy => "User copied a workflow step from the setup guide",
|
||||
Self::SetupGuideDocsLink => "User clicked a docs URL in the setup guide",
|
||||
Self::ConversationOpened => "User opened a conversation",
|
||||
Self::CloudRunOpened => "User opened a cloud run",
|
||||
Self::ArtifactClicked => "User clicked an artifact button",
|
||||
Self::FilterChanged => "User changed a filter in the management view",
|
||||
Self::DetailsViewed => "User clicked View details",
|
||||
Self::ConversationLinkCopied => "User copied a conversation link",
|
||||
Self::SessionLinkCopied => "User copied a session link",
|
||||
Self::TombstoneArtifactClicked => "User clicked an artifact in the tombstone view",
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
Self::TombstoneContinueLocally => "User clicked Continue locally in the tombstone",
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
Self::DetailsPanelContinueLocally => {
|
||||
"User clicked Continue locally in the details panel"
|
||||
}
|
||||
#[cfg(target_family = "wasm")]
|
||||
Self::TombstoneOpenInWarp => "User clicked Open in Warp in the tombstone",
|
||||
Self::CloudRunCancelled => "User cancelled a cloud run",
|
||||
Self::ConversationForked => "User forked a conversation",
|
||||
}
|
||||
}
|
||||
|
||||
fn enablement_state(&self) -> EnablementState {
|
||||
EnablementState::Always
|
||||
}
|
||||
}
|
||||
|
||||
warp_core::register_telemetry_event!(AgentManagementTelemetryEvent);
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,236 @@
|
||||
//! General-purpose administrative commands in the Warp CLI.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::Serialize;
|
||||
use warp_cli::agent::OutputFormat;
|
||||
use warpui::{platform::TerminationMode, AppContext, SingletonEntity};
|
||||
|
||||
use crate::auth::auth_manager::{AuthManager, AuthManagerEvent};
|
||||
use crate::auth::user::PrincipalType;
|
||||
use crate::auth::AuthStateProvider;
|
||||
use crate::workspaces::user_workspaces::UserWorkspaces;
|
||||
|
||||
/// Kick off a device authorization login flow and handle auth events.
|
||||
pub fn login(ctx: &mut AppContext) -> Result<()> {
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get();
|
||||
let has_cached_credentials = auth_state.is_logged_in();
|
||||
|
||||
// If the user is already logged in, we require that the user log out before logging
|
||||
// back in to ensure their existing state isn't replaced (especially if using both the CLI
|
||||
// and the desktop app). In this case, try refreshing their credentials first. If the user
|
||||
// is trying to log in because the cached credentials are invalid, we should let them do so.
|
||||
// Track whether we've started the device auth flow. Failure events
|
||||
// that arrive before device auth has started are leftover refresh
|
||||
// errors and should be ignored rather than treated as terminal.
|
||||
let mut started_device_auth = !has_cached_credentials;
|
||||
ctx.subscribe_to_model(
|
||||
&AuthManager::handle(ctx),
|
||||
move |_, event, ctx| match event {
|
||||
AuthManagerEvent::AuthComplete => {
|
||||
if !started_device_auth {
|
||||
// Refresh succeeded - credentials are still valid.
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get();
|
||||
match (auth_state.username_for_display(), auth_state.user_email()) {
|
||||
(Some(username), Some(email)) if username != email => {
|
||||
println!("You are already logged in as {username} ({email}).")
|
||||
}
|
||||
(Some(name), _) | (None, Some(name)) => {
|
||||
println!("You are already logged in as {name}.")
|
||||
}
|
||||
(None, None) => {
|
||||
println!("You are already logged in.")
|
||||
}
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
} else {
|
||||
// Device auth succeeded.
|
||||
println!("Logged in successfully");
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
}
|
||||
AuthManagerEvent::AuthFailed(_) => {
|
||||
if !started_device_auth {
|
||||
// Refresh failed - start a fresh device auth flow.
|
||||
started_device_auth = true;
|
||||
AuthManager::handle(ctx).update(ctx, |auth_manager, ctx| {
|
||||
auth_manager.authorize_device(ctx);
|
||||
});
|
||||
} else {
|
||||
// Device auth failed.
|
||||
let err_msg = match event {
|
||||
AuthManagerEvent::AuthFailed(err) => {
|
||||
format!("Authentication failed: {err:#}")
|
||||
}
|
||||
_ => "Authentication failed".to_string(),
|
||||
};
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(err_msg))),
|
||||
);
|
||||
}
|
||||
}
|
||||
AuthManagerEvent::ReceivedDeviceAuthorizationCode {
|
||||
verification_url,
|
||||
verification_url_complete,
|
||||
user_code,
|
||||
} => {
|
||||
if let Some(url) = verification_url_complete {
|
||||
println!("To log in, open this URL in your browser:\n{url}");
|
||||
} else {
|
||||
println!(
|
||||
"To log in, visit {verification_url} and enter this code: {user_code}"
|
||||
);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
);
|
||||
|
||||
// Either refresh existing credentials or start device auth from scratch.
|
||||
AuthManager::handle(ctx).update(ctx, |auth_manager, ctx| {
|
||||
if has_cached_credentials {
|
||||
auth_manager.refresh_user(ctx);
|
||||
} else {
|
||||
auth_manager.authorize_device(ctx);
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct WhoamiOutput {
|
||||
uid: String,
|
||||
#[serde(rename = "type")]
|
||||
principal_type: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
display_name: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
email: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
team_uid: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
team_name: Option<String>,
|
||||
}
|
||||
|
||||
/// Singleton model that provides a `ModelContext` for the `whoami` command's async work.
|
||||
struct WhoamiRunner;
|
||||
|
||||
impl warpui::Entity for WhoamiRunner {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for WhoamiRunner {}
|
||||
|
||||
/// Print information about the currently authenticated principal.
|
||||
pub fn whoami(ctx: &mut AppContext, output_format: OutputFormat) -> Result<()> {
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get();
|
||||
let principal_type = auth_state.principal_type().unwrap_or_default();
|
||||
|
||||
let uid = auth_state
|
||||
.user_id()
|
||||
.map(|id| {
|
||||
let s = id.as_string();
|
||||
s.strip_prefix("serviceAccount:")
|
||||
.map(String::from)
|
||||
.unwrap_or(s)
|
||||
})
|
||||
.ok_or_else(|| anyhow::anyhow!("Could not determine user ID. Are you logged in?"))?;
|
||||
|
||||
let mut info = WhoamiOutput {
|
||||
uid,
|
||||
principal_type: match principal_type {
|
||||
PrincipalType::User => "user",
|
||||
PrincipalType::ServiceAccount => "service_account",
|
||||
},
|
||||
display_name: auth_state.display_name(),
|
||||
email: match principal_type {
|
||||
PrincipalType::User => auth_state.user_email().filter(|e| !e.is_empty()),
|
||||
PrincipalType::ServiceAccount => None,
|
||||
},
|
||||
team_uid: None,
|
||||
team_name: None,
|
||||
};
|
||||
|
||||
// Refresh workspace metadata before reading team info, so we don't print
|
||||
// stale or missing team data if the metadata hasn't been fetched yet.
|
||||
let runner = ctx.add_singleton_model(|_| WhoamiRunner);
|
||||
runner.update(ctx, move |_, ctx| {
|
||||
let refresh_future = super::common::refresh_workspace_metadata(ctx);
|
||||
ctx.spawn(refresh_future, move |_, result, ctx| {
|
||||
if let Err(err) = result {
|
||||
// Do not prevent showing user info if fetching team metadata fails.
|
||||
log::warn!("Failed to refresh team metadata for whoami: {err:#}");
|
||||
}
|
||||
|
||||
let current_team = UserWorkspaces::as_ref(ctx).current_team();
|
||||
info.team_uid = current_team.map(|t| t.uid.to_string());
|
||||
info.team_name = current_team
|
||||
.map(|t| t.name.clone())
|
||||
.filter(|n| !n.is_empty());
|
||||
|
||||
match output_format {
|
||||
OutputFormat::Json => {
|
||||
match serde_json::to_string(&info).context("whoami output should serialize") {
|
||||
Ok(json) => println!("{json}"),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
OutputFormat::Pretty => {
|
||||
match principal_type {
|
||||
PrincipalType::User => println!("User ID: {}", info.uid),
|
||||
PrincipalType::ServiceAccount => {
|
||||
println!("Service account ID: {}", info.uid)
|
||||
}
|
||||
}
|
||||
if let Some(name) = &info.display_name {
|
||||
println!("Display Name: {name}");
|
||||
}
|
||||
if let Some(email) = &info.email {
|
||||
println!("Email: {email}");
|
||||
}
|
||||
if let Some(team_uid) = &info.team_uid {
|
||||
println!("Team ID: {team_uid}");
|
||||
}
|
||||
if let Some(team_name) = &info.team_name {
|
||||
println!("Team Name: {team_name}");
|
||||
}
|
||||
}
|
||||
OutputFormat::Text => {
|
||||
println!("{}:{}", info.principal_type, info.uid);
|
||||
}
|
||||
OutputFormat::Ndjson => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"`whoami` does not support `--output-format ndjson`"
|
||||
))),
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
});
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Log out of Warp using the same logic as the app.
|
||||
pub fn logout(ctx: &mut AppContext) -> Result<()> {
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get();
|
||||
if !auth_state.is_logged_in() {
|
||||
println!("You are not logged in.");
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
crate::auth::log_out(ctx);
|
||||
println!("Logged out successfully.");
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
//! Commands to interact with available agents via the public API.
|
||||
|
||||
use crate::ai::agent_sdk::oauth_flow::poll_oauth_until_terminal;
|
||||
use crate::ai::cloud_environments::GithubRepo;
|
||||
use crate::server::server_api::ai::AgentListItem;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use warp_cli::agent::ListAgentConfigsArgs;
|
||||
use warp_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus;
|
||||
use warp_graphql::queries::user_repo_auth_status::UserRepoAuthStatusEnum;
|
||||
use warpui::{platform::TerminationMode, AppContext, ModelContext, SingletonEntity};
|
||||
|
||||
const MAX_LINE_WIDTH: usize = 90;
|
||||
const MAX_AUTH_ATTEMPTS: u32 = 8;
|
||||
|
||||
/// Singleton model that runs async work for agent CLI commands.
|
||||
struct AgentConfigRunner;
|
||||
|
||||
/// List all available agents.
|
||||
pub fn list_agents(ctx: &mut AppContext, args: ListAgentConfigsArgs) -> anyhow::Result<()> {
|
||||
let runner = ctx.add_singleton_model(|_ctx| AgentConfigRunner);
|
||||
runner.update(ctx, |runner, ctx| runner.list(args.repo.clone(), ctx))
|
||||
}
|
||||
|
||||
/// Parse a repo spec string (owner/repo or GitHub URL) into a GithubRepo.
|
||||
fn parse_repo_spec(spec: &str) -> anyhow::Result<GithubRepo> {
|
||||
let spec = spec.trim();
|
||||
|
||||
// Try URL format: https://github.com/owner/repo or https://github.com/owner/repo.git
|
||||
if spec.starts_with("https://github.com/") || spec.starts_with("http://github.com/") {
|
||||
let path = spec
|
||||
.trim_start_matches("https://github.com/")
|
||||
.trim_start_matches("http://github.com/")
|
||||
.trim_end_matches(".git")
|
||||
.trim_end_matches('/');
|
||||
|
||||
let parts: Vec<&str> = path.split('/').collect();
|
||||
if parts.len() >= 2 && !parts[0].is_empty() && !parts[1].is_empty() {
|
||||
return Ok(GithubRepo::new(parts[0].to_string(), parts[1].to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
// Try slug format: owner/repo
|
||||
let parts: Vec<&str> = spec.split('/').collect();
|
||||
if parts.len() == 2 && !parts[0].is_empty() && !parts[1].is_empty() {
|
||||
return Ok(GithubRepo::new(parts[0].to_string(), parts[1].to_string()));
|
||||
}
|
||||
|
||||
Err(anyhow::anyhow!(
|
||||
"Invalid repo format: '{}'. Expected 'owner/repo' or 'https://github.com/owner/repo'",
|
||||
spec
|
||||
))
|
||||
}
|
||||
|
||||
impl AgentConfigRunner {
|
||||
fn list(&self, repo: Option<String>, ctx: &mut ModelContext<Self>) -> anyhow::Result<()> {
|
||||
// If a repo is specified, check auth first
|
||||
if let Some(ref repo_spec) = repo {
|
||||
let github_repo = parse_repo_spec(repo_spec)?;
|
||||
self.auth_then_list(vec![github_repo], 1, repo, ctx);
|
||||
} else {
|
||||
// No repo specified - just list from environments
|
||||
self.fetch_and_display_agents(repo, ctx);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check GitHub auth for repos, then list agents.
|
||||
fn auth_then_list(
|
||||
&self,
|
||||
repos: Vec<GithubRepo>,
|
||||
attempt: u32,
|
||||
repo_spec: Option<String>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if attempt > MAX_AUTH_ATTEMPTS {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"Exceeded maximum number of authorization attempts ({}). Please try again later.",
|
||||
MAX_AUTH_ATTEMPTS
|
||||
))),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let integrations_client = ServerApiProvider::handle(ctx)
|
||||
.as_ref(ctx)
|
||||
.get_integrations_client();
|
||||
|
||||
let repo_tuples: Vec<(String, String)> = repos
|
||||
.iter()
|
||||
.map(|repo| (repo.owner.clone(), repo.repo.clone()))
|
||||
.collect();
|
||||
|
||||
let auth_check_future = async move {
|
||||
integrations_client
|
||||
.check_user_repo_auth_status(repo_tuples)
|
||||
.await
|
||||
};
|
||||
|
||||
ctx.spawn(auth_check_future, move |runner, auth_result, ctx| {
|
||||
match auth_result {
|
||||
Ok(response) => {
|
||||
let mut has_blocking_private_issues = false;
|
||||
|
||||
for status in &response.statuses {
|
||||
match status.status {
|
||||
UserRepoAuthStatusEnum::Success => {}
|
||||
UserRepoAuthStatusEnum::NoInstallationOrAccessForRepo => {
|
||||
if !status.is_public {
|
||||
eprintln!(
|
||||
"Cannot access private repo {}/{}",
|
||||
status.owner, status.repo,
|
||||
);
|
||||
has_blocking_private_issues = true;
|
||||
}
|
||||
// Public repos without auth are fine - no warning needed
|
||||
}
|
||||
UserRepoAuthStatusEnum::UserNotConnectedToGithub => {
|
||||
eprintln!("User not connected to GitHub");
|
||||
has_blocking_private_issues = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !has_blocking_private_issues {
|
||||
// No blocking issues - proceed with listing
|
||||
runner.fetch_and_display_agents(repo_spec, ctx);
|
||||
return;
|
||||
}
|
||||
|
||||
// Handle OAuth flow if server provides auth_url + tx_id
|
||||
match (response.auth_url, response.tx_id) {
|
||||
(Some(auth_url), Some(tx_id)) => {
|
||||
println!("\nAuthorization required for private repository access.");
|
||||
println!("Opening browser for GitHub authorization: {auth_url}\n");
|
||||
ctx.open_url(&auth_url);
|
||||
|
||||
let integrations_client = ServerApiProvider::handle(ctx)
|
||||
.as_ref(ctx)
|
||||
.get_integrations_client();
|
||||
let tx_id = tx_id.into_inner();
|
||||
let poll_future = poll_oauth_until_terminal(integrations_client, tx_id);
|
||||
|
||||
let next_attempt = attempt + 1;
|
||||
|
||||
ctx.spawn(poll_future, move |runner, poll_result, ctx| {
|
||||
match poll_result {
|
||||
Ok(OauthConnectTxStatus::Completed) => {
|
||||
// OAuth completed, retry
|
||||
runner.auth_then_list(repos, next_attempt, repo_spec, ctx);
|
||||
}
|
||||
Ok(OauthConnectTxStatus::Failed) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"GitHub authorization failed. Please try again."
|
||||
))),
|
||||
);
|
||||
}
|
||||
Ok(OauthConnectTxStatus::Expired) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"GitHub authorization expired. Please try again."
|
||||
))),
|
||||
);
|
||||
}
|
||||
Ok(_) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"Unexpected OAuth status"
|
||||
))),
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"Error polling OAuth status: {err}"
|
||||
))),
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
(Some(auth_url), None) => {
|
||||
println!("\nAuthorize access here: {auth_url}\n");
|
||||
println!("After authorizing, please re-run this command.");
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
_ => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"Cannot list agents: authorization required but no auth flow provided"
|
||||
))),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(e.context("Failed to check GitHub auth status"))),
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn fetch_and_display_agents(&self, repo: Option<String>, ctx: &mut ModelContext<Self>) {
|
||||
let ai_client = ServerApiProvider::handle(ctx).as_ref(ctx).get_ai_client();
|
||||
|
||||
if repo.is_some() {
|
||||
println!("Fetching agent skills from the specified repository...");
|
||||
} else {
|
||||
println!("Fetching agent skills from your Warp environments...");
|
||||
}
|
||||
|
||||
let list_future = async move { ai_client.list_agents(repo).await };
|
||||
|
||||
ctx.spawn(list_future, |_, result, ctx| match result {
|
||||
Ok(agents) => {
|
||||
Self::print_agents_table(&agents);
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(err, ctx);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Print a list of agents in a card-style format.
|
||||
fn print_agents_table(agents: &[AgentListItem]) {
|
||||
if agents.is_empty() {
|
||||
println!("No agents found.");
|
||||
return;
|
||||
}
|
||||
|
||||
if agents.len() == 1 {
|
||||
println!("\nAgent:");
|
||||
} else {
|
||||
println!("\nAgents ({}):", agents.len());
|
||||
}
|
||||
|
||||
for agent in agents {
|
||||
println!("\n{}", agent.name);
|
||||
|
||||
for variant in &agent.variants {
|
||||
let mut table = super::output::standard_table();
|
||||
|
||||
// ID
|
||||
table.add_row(vec![format!("ID: {}", variant.id)]);
|
||||
|
||||
// Description
|
||||
if !variant.description.is_empty() {
|
||||
let description_cell = super::text_layout::render_labeled_wrapped_field(
|
||||
"Description",
|
||||
&variant.description,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![description_cell]);
|
||||
}
|
||||
|
||||
// Base prompt (truncated)
|
||||
if !variant.base_prompt.is_empty() {
|
||||
let mut chars = variant.base_prompt.chars();
|
||||
let truncated: String = chars.by_ref().take(100).collect();
|
||||
let truncated_prompt = if chars.next().is_some() {
|
||||
format!("{truncated}...")
|
||||
} else {
|
||||
truncated
|
||||
};
|
||||
let prompt_cell = super::text_layout::render_labeled_wrapped_field(
|
||||
"Base Prompt",
|
||||
&truncated_prompt,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![prompt_cell]);
|
||||
}
|
||||
|
||||
// Source
|
||||
table.add_row(vec![format!(
|
||||
"Source: {}/{}",
|
||||
variant.source.owner, variant.source.name
|
||||
)]);
|
||||
|
||||
// Environments
|
||||
if !variant.environments.is_empty() {
|
||||
let env_entries: Vec<_> = variant
|
||||
.environments
|
||||
.iter()
|
||||
.map(|e| format!("{} ({})", e.name, e.uid))
|
||||
.collect();
|
||||
table.add_row(vec![format!("Environments: {}", env_entries.join(", "))]);
|
||||
}
|
||||
|
||||
println!("{table}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl warpui::Entity for AgentConfigRunner {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for AgentConfigRunner {}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,169 @@
|
||||
//! Unit tests for `filter_from_args`. Verifies the clap enums are faithfully translated into
|
||||
//! `TaskListFilter` without dropping any fields.
|
||||
|
||||
use chrono::{TimeZone, Utc};
|
||||
|
||||
use warp_cli::json_filter::JsonOutput;
|
||||
use warp_cli::task::{
|
||||
ArtifactTypeArg, ExecutionLocationArg, ListTasksArgs, RunSortByArg, RunSortOrderArg,
|
||||
RunSourceArg, RunStateArg,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use crate::server::server_api::ai::{ArtifactType, ExecutionLocation, RunSortBy, RunSortOrder};
|
||||
|
||||
/// A `ListTasksArgs` whose fields are all at their defaults.
|
||||
fn empty_args() -> ListTasksArgs {
|
||||
ListTasksArgs {
|
||||
limit: 10,
|
||||
state: vec![],
|
||||
source: None,
|
||||
execution_location: None,
|
||||
creator: None,
|
||||
environment: None,
|
||||
skill: None,
|
||||
schedule: None,
|
||||
ancestor_run: None,
|
||||
name: None,
|
||||
model: None,
|
||||
artifact_type: None,
|
||||
created_after: None,
|
||||
created_before: None,
|
||||
updated_after: None,
|
||||
query: None,
|
||||
sort_by: None,
|
||||
sort_order: None,
|
||||
cursor: None,
|
||||
json_output: JsonOutput::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_args_yields_default_filter() {
|
||||
let filter = filter_from_args(&empty_args());
|
||||
assert!(filter.creator_uid.is_none());
|
||||
assert!(filter.updated_after.is_none());
|
||||
assert!(filter.created_after.is_none());
|
||||
assert!(filter.created_before.is_none());
|
||||
assert!(filter.states.is_none());
|
||||
assert!(filter.source.is_none());
|
||||
assert!(filter.execution_location.is_none());
|
||||
assert!(filter.environment_id.is_none());
|
||||
assert!(filter.skill_spec.is_none());
|
||||
assert!(filter.schedule_id.is_none());
|
||||
assert!(filter.ancestor_run_id.is_none());
|
||||
assert!(filter.config_name.is_none());
|
||||
assert!(filter.model_id.is_none());
|
||||
assert!(filter.artifact_type.is_none());
|
||||
assert!(filter.search_query.is_none());
|
||||
assert!(filter.sort_by.is_none());
|
||||
assert!(filter.sort_order.is_none());
|
||||
assert!(filter.cursor.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_flags_map_to_filter() {
|
||||
let args = ListTasksArgs {
|
||||
state: vec![
|
||||
RunStateArg::Failed,
|
||||
RunStateArg::Error,
|
||||
RunStateArg::Cancelled,
|
||||
],
|
||||
..empty_args()
|
||||
};
|
||||
let filter = filter_from_args(&args);
|
||||
assert_eq!(
|
||||
filter.states.as_deref(),
|
||||
Some(
|
||||
[
|
||||
AmbientAgentTaskState::Failed,
|
||||
AmbientAgentTaskState::Error,
|
||||
AmbientAgentTaskState::Cancelled,
|
||||
]
|
||||
.as_slice()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn source_cli_maps_to_cli() {
|
||||
let args = ListTasksArgs {
|
||||
source: Some(RunSourceArg::Cli),
|
||||
..empty_args()
|
||||
};
|
||||
let filter = filter_from_args(&args);
|
||||
assert_eq!(filter.source, Some(AgentSource::Cli));
|
||||
// Sanity-check the wire value: `--source CLI` must send `source=CLI`.
|
||||
assert_eq!(filter.source.as_ref().map(AgentSource::as_str), Some("CLI"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn source_interactive_maps_to_local() {
|
||||
// The public API uses `LOCAL` as the source value for local interactive
|
||||
// tasks. The CLI exposes this as `--source INTERACTIVE` for readability,
|
||||
// but the request sent to the server must use `LOCAL`.
|
||||
let args = ListTasksArgs {
|
||||
source: Some(RunSourceArg::Interactive),
|
||||
..empty_args()
|
||||
};
|
||||
let filter = filter_from_args(&args);
|
||||
assert_eq!(filter.source, Some(AgentSource::Interactive));
|
||||
assert_eq!(
|
||||
filter.source.as_ref().map(AgentSource::as_str),
|
||||
Some("LOCAL")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_field_maps_through() {
|
||||
let created_after = Utc.with_ymd_and_hms(2026, 4, 1, 0, 0, 0).unwrap();
|
||||
let created_before = Utc.with_ymd_and_hms(2026, 4, 2, 0, 0, 0).unwrap();
|
||||
let updated_after = Utc.with_ymd_and_hms(2026, 4, 3, 12, 30, 0).unwrap();
|
||||
|
||||
let args = ListTasksArgs {
|
||||
limit: 20,
|
||||
state: vec![RunStateArg::InProgress],
|
||||
source: Some(RunSourceArg::Api),
|
||||
execution_location: Some(ExecutionLocationArg::Remote),
|
||||
creator: Some("user-uid".to_string()),
|
||||
environment: Some("env-123".to_string()),
|
||||
skill: Some("owner/repo:SKILL.md".to_string()),
|
||||
schedule: Some("sched-1".to_string()),
|
||||
ancestor_run: Some("run-parent".to_string()),
|
||||
name: Some("nightly".to_string()),
|
||||
model: Some("claude-4-5".to_string()),
|
||||
artifact_type: Some(ArtifactTypeArg::PullRequest),
|
||||
created_after: Some(created_after),
|
||||
created_before: Some(created_before),
|
||||
updated_after: Some(updated_after),
|
||||
query: Some("oz run".to_string()),
|
||||
sort_by: Some(RunSortByArg::CreatedAt),
|
||||
sort_order: Some(RunSortOrderArg::Asc),
|
||||
cursor: Some("abcd==".to_string()),
|
||||
json_output: JsonOutput::default(),
|
||||
};
|
||||
|
||||
let filter = filter_from_args(&args);
|
||||
|
||||
assert_eq!(filter.creator_uid.as_deref(), Some("user-uid"));
|
||||
assert_eq!(filter.updated_after, Some(updated_after));
|
||||
assert_eq!(filter.created_after, Some(created_after));
|
||||
assert_eq!(filter.created_before, Some(created_before));
|
||||
assert_eq!(
|
||||
filter.states.as_deref(),
|
||||
Some([AmbientAgentTaskState::InProgress].as_slice())
|
||||
);
|
||||
assert_eq!(filter.source, Some(AgentSource::AgentWebhook));
|
||||
assert_eq!(filter.execution_location, Some(ExecutionLocation::Remote));
|
||||
assert_eq!(filter.environment_id.as_deref(), Some("env-123"));
|
||||
assert_eq!(filter.skill_spec.as_deref(), Some("owner/repo:SKILL.md"));
|
||||
assert_eq!(filter.schedule_id.as_deref(), Some("sched-1"));
|
||||
assert_eq!(filter.ancestor_run_id.as_deref(), Some("run-parent"));
|
||||
assert_eq!(filter.config_name.as_deref(), Some("nightly"));
|
||||
assert_eq!(filter.model_id.as_deref(), Some("claude-4-5"));
|
||||
assert_eq!(filter.artifact_type, Some(ArtifactType::PullRequest));
|
||||
assert_eq!(filter.search_query.as_deref(), Some("oz run"));
|
||||
assert_eq!(filter.sort_by, Some(RunSortBy::CreatedAt));
|
||||
assert_eq!(filter.sort_order, Some(RunSortOrder::Asc));
|
||||
assert_eq!(filter.cursor.as_deref(), Some("abcd=="));
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
use std::{path::PathBuf, sync::Arc};
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::Serialize;
|
||||
use warp_cli::agent::OutputFormat;
|
||||
use warp_cli::artifact::{
|
||||
ArtifactCommand, DownloadArtifactArgs, GetArtifactArgs, UploadArtifactArgs,
|
||||
};
|
||||
use warp_cli::GlobalOptions;
|
||||
use warpui::{platform::TerminationMode, AppContext, ModelContext, SingletonEntity};
|
||||
|
||||
use crate::ai::artifact_download::{download_artifact_bytes, download_destination};
|
||||
#[cfg(test)]
|
||||
use crate::server::server_api::ai::FileArtifactRecord;
|
||||
use crate::server::server_api::ai::{AIClient, ArtifactDownloadResponse};
|
||||
use crate::server::server_api::{ServerApi, ServerApiProvider};
|
||||
|
||||
use super::artifact_upload::{
|
||||
CompletedFileArtifactUpload, FileArtifactUploadRequest, FileArtifactUploader,
|
||||
};
|
||||
|
||||
/// Run artifact-related commands.
|
||||
pub fn run(
|
||||
ctx: &mut AppContext,
|
||||
global_options: GlobalOptions,
|
||||
command: ArtifactCommand,
|
||||
) -> Result<()> {
|
||||
let runner = ctx.add_singleton_model(|_| ArtifactCommandRunner);
|
||||
match command {
|
||||
ArtifactCommand::Upload(args) => {
|
||||
runner.update(ctx, |runner, ctx| {
|
||||
runner.upload(args, global_options.output_format, ctx);
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
ArtifactCommand::Get(args) => {
|
||||
runner.update(ctx, |runner, ctx| {
|
||||
runner.get(args, global_options.output_format, ctx);
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
ArtifactCommand::Download(args) => {
|
||||
runner.update(ctx, |runner, ctx| {
|
||||
runner.download(args, global_options.output_format, ctx);
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ArtifactCommandRunner;
|
||||
|
||||
impl ArtifactCommandRunner {
|
||||
fn get(
|
||||
&self,
|
||||
args: GetArtifactArgs,
|
||||
output_format: OutputFormat,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
|
||||
|
||||
ctx.spawn(
|
||||
async move { get_artifact(ai_client, &args.artifact_uid).await },
|
||||
move |_, result, ctx| match result {
|
||||
Ok(artifact) => {
|
||||
if let Err(err) = write_get_output(&artifact, output_format) {
|
||||
super::report_fatal_error(err, ctx);
|
||||
return;
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => super::report_fatal_error(err, ctx),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn download(
|
||||
&self,
|
||||
args: DownloadArtifactArgs,
|
||||
output_format: OutputFormat,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
|
||||
let server_api = ServerApiProvider::as_ref(ctx).get();
|
||||
|
||||
ctx.spawn(
|
||||
async move { download_artifact(ai_client, server_api, args).await },
|
||||
move |_, result, ctx| match result {
|
||||
Ok(output) => {
|
||||
if let Err(err) = write_download_output(&output, output_format) {
|
||||
super::report_fatal_error(err, ctx);
|
||||
return;
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => super::report_fatal_error(err, ctx),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn upload(
|
||||
&self,
|
||||
args: UploadArtifactArgs,
|
||||
output_format: OutputFormat,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let server_api = ServerApiProvider::as_ref(ctx).get();
|
||||
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
|
||||
let uploader = FileArtifactUploader::new(ai_client, server_api.clone());
|
||||
|
||||
ctx.spawn(
|
||||
async move {
|
||||
let request = FileArtifactUploadRequest::try_from(args)?;
|
||||
let association = uploader.resolve_upload_association(&request).await?;
|
||||
server_api.set_ambient_agent_task_id(Some(association.ambient_task_id));
|
||||
uploader.upload_with_association(request, association).await
|
||||
},
|
||||
move |_, result, ctx| match result {
|
||||
Ok(artifact) => {
|
||||
if let Err(err) = write_upload_output(&artifact, output_format) {
|
||||
super::report_fatal_error(err, ctx);
|
||||
return;
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => super::report_fatal_error(err, ctx),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl warpui::Entity for ArtifactCommandRunner {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for ArtifactCommandRunner {}
|
||||
|
||||
async fn get_artifact(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
artifact_uid: &str,
|
||||
) -> Result<ArtifactDownloadResponse> {
|
||||
ai_client
|
||||
.get_artifact_download(artifact_uid)
|
||||
.await
|
||||
.with_context(|| format!("Failed to get artifact '{artifact_uid}'"))
|
||||
}
|
||||
|
||||
async fn download_artifact(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
server_api: Arc<ServerApi>,
|
||||
args: DownloadArtifactArgs,
|
||||
) -> Result<DownloadArtifactOutput> {
|
||||
let artifact = get_artifact(ai_client, &args.artifact_uid).await?;
|
||||
let path = download_destination(&artifact, args.out);
|
||||
download_artifact_bytes(server_api.http_client(), &artifact, &path).await?;
|
||||
let path = std::path::absolute(&path).unwrap_or(path);
|
||||
Ok(DownloadArtifactOutput::new(&artifact, path))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ArtifactMetadataOutput {
|
||||
artifact_uid: String,
|
||||
artifact_type: String,
|
||||
created_at: String,
|
||||
download_url: String,
|
||||
expires_at: String,
|
||||
content_type: String,
|
||||
filepath: Option<String>,
|
||||
filename: Option<String>,
|
||||
description: Option<String>,
|
||||
size_bytes: Option<i64>,
|
||||
}
|
||||
|
||||
impl ArtifactMetadataOutput {
|
||||
fn new(artifact: &ArtifactDownloadResponse) -> Self {
|
||||
Self {
|
||||
artifact_uid: artifact.artifact_uid().to_string(),
|
||||
artifact_type: artifact.artifact_type().to_string(),
|
||||
created_at: artifact.created_at().to_rfc3339(),
|
||||
download_url: artifact.download_url().to_string(),
|
||||
expires_at: artifact.expires_at().to_rfc3339(),
|
||||
content_type: artifact.content_type().to_string(),
|
||||
filepath: artifact.filepath().map(ToString::to_string),
|
||||
filename: artifact.filename().map(ToString::to_string),
|
||||
description: artifact.description().map(ToString::to_string),
|
||||
size_bytes: artifact.size_bytes(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct DownloadArtifactOutput {
|
||||
artifact_uid: String,
|
||||
artifact_type: String,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl DownloadArtifactOutput {
|
||||
fn new(artifact: &ArtifactDownloadResponse, path: PathBuf) -> Self {
|
||||
Self {
|
||||
artifact_uid: artifact.artifact_uid().to_string(),
|
||||
artifact_type: artifact.artifact_type().to_string(),
|
||||
path,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UploadArtifactOutput {
|
||||
artifact_uid: String,
|
||||
filepath: String,
|
||||
description: Option<String>,
|
||||
mime_type: String,
|
||||
size_bytes: Option<i64>,
|
||||
}
|
||||
|
||||
fn write_get_output(
|
||||
artifact: &ArtifactDownloadResponse,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let mut stdout = std::io::stdout();
|
||||
write_get_output_to(&mut stdout, artifact, output_format)
|
||||
}
|
||||
|
||||
fn write_get_output_to<W: std::io::Write>(
|
||||
output: &mut W,
|
||||
artifact: &ArtifactDownloadResponse,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let output_record = ArtifactMetadataOutput::new(artifact);
|
||||
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
serde_json::to_writer(&mut *output, &output_record)
|
||||
.context("unable to write JSON output")?;
|
||||
writeln!(&mut *output)?;
|
||||
}
|
||||
OutputFormat::Pretty => {
|
||||
writeln!(&mut *output, "Artifact UID: {}", output_record.artifact_uid)?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"Artifact type: {}",
|
||||
output_record.artifact_type
|
||||
)?;
|
||||
writeln!(&mut *output, "Created at: {}", output_record.created_at)?;
|
||||
writeln!(&mut *output, "Download URL: {}", output_record.download_url)?;
|
||||
writeln!(&mut *output, "Expires at: {}", output_record.expires_at)?;
|
||||
writeln!(&mut *output, "Content type: {}", output_record.content_type)?;
|
||||
if let Some(filepath) = output_record.filepath {
|
||||
writeln!(&mut *output, "Filepath: {filepath}")?;
|
||||
}
|
||||
if let Some(filename) = output_record.filename {
|
||||
writeln!(&mut *output, "Filename: {filename}")?;
|
||||
}
|
||||
if let Some(description) = output_record.description {
|
||||
writeln!(&mut *output, "Description: {description}")?;
|
||||
}
|
||||
if let Some(size_bytes) = output_record.size_bytes {
|
||||
writeln!(&mut *output, "Size bytes: {size_bytes}")?;
|
||||
}
|
||||
}
|
||||
OutputFormat::Text => {
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"Artifact UID\tArtifact type\tCreated at\tDownload URL\tExpires at\tContent type\tFilepath\tFilename\tDescription\tSize bytes"
|
||||
)?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}",
|
||||
output_record.artifact_uid,
|
||||
output_record.artifact_type,
|
||||
output_record.created_at,
|
||||
output_record.download_url,
|
||||
output_record.expires_at,
|
||||
output_record.content_type,
|
||||
output_record.filepath.unwrap_or_default(),
|
||||
output_record.filename.unwrap_or_default(),
|
||||
output_record.description.unwrap_or_default(),
|
||||
output_record
|
||||
.size_bytes
|
||||
.map(|size| size.to_string())
|
||||
.unwrap_or_default()
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_download_output(
|
||||
output_record: &DownloadArtifactOutput,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let mut stdout = std::io::stdout();
|
||||
write_download_output_to(&mut stdout, output_record, output_format)
|
||||
}
|
||||
|
||||
fn write_download_output_to<W: std::io::Write>(
|
||||
output: &mut W,
|
||||
output_record: &DownloadArtifactOutput,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
serde_json::to_writer(&mut *output, output_record)
|
||||
.context("unable to write JSON output")?;
|
||||
writeln!(&mut *output)?;
|
||||
}
|
||||
OutputFormat::Pretty => {
|
||||
writeln!(&mut *output, "Artifact downloaded")?;
|
||||
writeln!(&mut *output, "Artifact UID: {}", output_record.artifact_uid)?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"Artifact type: {}",
|
||||
output_record.artifact_type
|
||||
)?;
|
||||
writeln!(&mut *output, "Path: {}", output_record.path.display())?;
|
||||
}
|
||||
OutputFormat::Text => {
|
||||
writeln!(&mut *output, "Artifact UID\tArtifact type\tPath")?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"{}\t{}\t{}",
|
||||
output_record.artifact_uid,
|
||||
output_record.artifact_type,
|
||||
output_record.path.display()
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_upload_output(
|
||||
artifact: &CompletedFileArtifactUpload,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let mut stdout = std::io::stdout();
|
||||
write_upload_output_to(&mut stdout, artifact, output_format)
|
||||
}
|
||||
|
||||
fn write_upload_output_to<W: std::io::Write>(
|
||||
output: &mut W,
|
||||
artifact: &CompletedFileArtifactUpload,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
let output_record = UploadArtifactOutput {
|
||||
artifact_uid: artifact.artifact.artifact_uid.clone(),
|
||||
filepath: artifact.artifact.filepath.clone(),
|
||||
description: artifact.artifact.description.clone(),
|
||||
mime_type: artifact.artifact.mime_type.clone(),
|
||||
size_bytes: Some(artifact.size_bytes),
|
||||
};
|
||||
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
serde_json::to_writer(&mut *output, &output_record)
|
||||
.context("unable to write JSON output")?;
|
||||
writeln!(&mut *output)?;
|
||||
}
|
||||
OutputFormat::Pretty => {
|
||||
writeln!(&mut *output, "Artifact uploaded")?;
|
||||
writeln!(&mut *output, "Artifact UID: {}", output_record.artifact_uid)?;
|
||||
writeln!(&mut *output, "Filepath: {}", output_record.filepath)?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"Description: {}",
|
||||
output_record.description.as_deref().unwrap_or("")
|
||||
)?;
|
||||
writeln!(&mut *output, "MIME type: {}", output_record.mime_type)?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"Size bytes: {}",
|
||||
output_record
|
||||
.size_bytes
|
||||
.map(|size| size.to_string())
|
||||
.unwrap_or_default()
|
||||
)?;
|
||||
}
|
||||
OutputFormat::Text => {
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"Artifact UID\tFilepath\tDescription\tMIME type\tSize bytes"
|
||||
)?;
|
||||
writeln!(
|
||||
&mut *output,
|
||||
"{}\t{}\t{}\t{}\t{}",
|
||||
output_record.artifact_uid,
|
||||
output_record.filepath,
|
||||
output_record.description.unwrap_or_default(),
|
||||
output_record.mime_type,
|
||||
output_record
|
||||
.size_bytes
|
||||
.map(|size| size.to_string())
|
||||
.unwrap_or_default()
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "artifact_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,255 @@
|
||||
use std::path::PathBuf;
|
||||
use warp_cli::agent::OutputFormat;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn sample_completed_upload() -> CompletedFileArtifactUpload {
|
||||
CompletedFileArtifactUpload {
|
||||
artifact: sample_artifact_record(),
|
||||
size_bytes: 42,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_artifact_record() -> FileArtifactRecord {
|
||||
FileArtifactRecord {
|
||||
artifact_uid: "artifact-123".to_string(),
|
||||
filepath: "outputs/report.txt".to_string(),
|
||||
description: Some("daily summary".to_string()),
|
||||
mime_type: "text/plain".to_string(),
|
||||
size_bytes: Some(42),
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_file_download_response() -> ArtifactDownloadResponse {
|
||||
serde_json::from_str(
|
||||
r#"{
|
||||
"artifact_uid": "artifact-123",
|
||||
"artifact_type": "FILE",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"data": {
|
||||
"download_url": "https://storage.example.com/report.txt",
|
||||
"expires_at": "2024-01-15T11:30:00Z",
|
||||
"content_type": "text/plain",
|
||||
"filepath": "outputs/report.txt",
|
||||
"filename": "report.txt",
|
||||
"description": "daily summary",
|
||||
"size_bytes": 42
|
||||
}
|
||||
}"#,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn sample_screenshot_download_response() -> ArtifactDownloadResponse {
|
||||
serde_json::from_str(
|
||||
r#"{
|
||||
"artifact_uid": "screenshot-123",
|
||||
"artifact_type": "SCREENSHOT",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"data": {
|
||||
"download_url": "https://storage.example.com/screenshot.png",
|
||||
"expires_at": "2024-01-15T11:30:00Z",
|
||||
"content_type": "image/png",
|
||||
"description": "dashboard screenshot"
|
||||
}
|
||||
}"#,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_get_output_to_writes_json_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_get_output_to(
|
||||
&mut output,
|
||||
&sample_file_download_response(),
|
||||
OutputFormat::Json,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"{\"artifact_uid\":\"artifact-123\",\"artifact_type\":\"FILE\",\"created_at\":\"2024-01-15T10:30:00+00:00\",\"download_url\":\"https://storage.example.com/report.txt\",\"expires_at\":\"2024-01-15T11:30:00+00:00\",\"content_type\":\"text/plain\",\"filepath\":\"outputs/report.txt\",\"filename\":\"report.txt\",\"description\":\"daily summary\",\"size_bytes\":42}\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_get_output_to_writes_ndjson_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_get_output_to(
|
||||
&mut output,
|
||||
&sample_file_download_response(),
|
||||
OutputFormat::Ndjson,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"{\"artifact_uid\":\"artifact-123\",\"artifact_type\":\"FILE\",\"created_at\":\"2024-01-15T10:30:00+00:00\",\"download_url\":\"https://storage.example.com/report.txt\",\"expires_at\":\"2024-01-15T11:30:00+00:00\",\"content_type\":\"text/plain\",\"filepath\":\"outputs/report.txt\",\"filename\":\"report.txt\",\"description\":\"daily summary\",\"size_bytes\":42}\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_get_output_to_writes_pretty_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_get_output_to(
|
||||
&mut output,
|
||||
&sample_file_download_response(),
|
||||
OutputFormat::Pretty,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"Artifact UID: artifact-123\nArtifact type: FILE\nCreated at: 2024-01-15T10:30:00+00:00\nDownload URL: https://storage.example.com/report.txt\nExpires at: 2024-01-15T11:30:00+00:00\nContent type: text/plain\nFilepath: outputs/report.txt\nFilename: report.txt\nDescription: daily summary\nSize bytes: 42\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_get_output_to_writes_text_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_get_output_to(
|
||||
&mut output,
|
||||
&sample_file_download_response(),
|
||||
OutputFormat::Text,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"Artifact UID\tArtifact type\tCreated at\tDownload URL\tExpires at\tContent type\tFilepath\tFilename\tDescription\tSize bytes\nartifact-123\tFILE\t2024-01-15T10:30:00+00:00\thttps://storage.example.com/report.txt\t2024-01-15T11:30:00+00:00\ttext/plain\toutputs/report.txt\treport.txt\tdaily summary\t42\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_download_output_to_writes_pretty_output() {
|
||||
let artifact = sample_file_download_response();
|
||||
let path = std::path::absolute("report.txt").unwrap();
|
||||
let output_record = DownloadArtifactOutput::new(&artifact, path.clone());
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_download_output_to(&mut output, &output_record, OutputFormat::Pretty).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
format!(
|
||||
"Artifact downloaded\nArtifact UID: artifact-123\nArtifact type: FILE\nPath: {}\n",
|
||||
path.display()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn download_destination_uses_explicit_path() {
|
||||
assert_eq!(
|
||||
download_destination(
|
||||
&sample_file_download_response(),
|
||||
Some(PathBuf::from("downloads/report.txt"))
|
||||
),
|
||||
PathBuf::from("downloads/report.txt")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn download_destination_defaults_to_file_artifact_filename() {
|
||||
assert_eq!(
|
||||
download_destination(&sample_file_download_response(), None),
|
||||
PathBuf::from("report.txt")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn download_destination_defaults_screenshot_to_artifact_uid_with_extension() {
|
||||
assert_eq!(
|
||||
download_destination(&sample_screenshot_download_response(), None),
|
||||
PathBuf::from("artifact-screenshot-123.png")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn download_destination_defaults_pdf_to_artifact_uid_with_extension() {
|
||||
let artifact: ArtifactDownloadResponse = serde_json::from_str(
|
||||
r#"{
|
||||
"artifact_uid": "artifact-pdf-123",
|
||||
"artifact_type": "FILE",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"data": {
|
||||
"download_url": "https://storage.example.com/report.pdf",
|
||||
"expires_at": "2024-01-15T11:30:00Z",
|
||||
"content_type": "application/pdf",
|
||||
"filepath": "outputs/report.pdf",
|
||||
"filename": "",
|
||||
"description": "pdf report",
|
||||
"size_bytes": 42
|
||||
}
|
||||
}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
download_destination(&artifact, None),
|
||||
PathBuf::from("artifact-artifact-pdf-123.pdf")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_upload_output_to_writes_json_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_upload_output_to(&mut output, &sample_completed_upload(), OutputFormat::Json).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"{\"artifact_uid\":\"artifact-123\",\"filepath\":\"outputs/report.txt\",\"description\":\"daily summary\",\"mime_type\":\"text/plain\",\"size_bytes\":42}\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_upload_output_to_writes_ndjson_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_upload_output_to(
|
||||
&mut output,
|
||||
&sample_completed_upload(),
|
||||
OutputFormat::Ndjson,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"{\"artifact_uid\":\"artifact-123\",\"filepath\":\"outputs/report.txt\",\"description\":\"daily summary\",\"mime_type\":\"text/plain\",\"size_bytes\":42}\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_upload_output_to_writes_pretty_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_upload_output_to(
|
||||
&mut output,
|
||||
&sample_completed_upload(),
|
||||
OutputFormat::Pretty,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"Artifact uploaded\nArtifact UID: artifact-123\nFilepath: outputs/report.txt\nDescription: daily summary\nMIME type: text/plain\nSize bytes: 42\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_upload_output_to_writes_text_output() {
|
||||
let mut output = Vec::new();
|
||||
|
||||
write_upload_output_to(&mut output, &sample_completed_upload(), OutputFormat::Text).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
"Artifact UID\tFilepath\tDescription\tMIME type\tSize bytes\nartifact-123\toutputs/report.txt\tdaily summary\ttext/plain\t42\n"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,386 @@
|
||||
use std::env;
|
||||
use std::fs::File;
|
||||
use std::io::Read as _;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, bail, Context, Result};
|
||||
use blocking::unblock;
|
||||
use mime_guess::from_path;
|
||||
use warp_cli::artifact::UploadArtifactArgs;
|
||||
|
||||
use super::common::parse_ambient_task_id;
|
||||
use crate::ai::agent::api::ServerConversationToken;
|
||||
use crate::ai::agent::conversation::ServerAIConversationMetadata;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::server::server_api::ai::{
|
||||
AIClient, CreateFileArtifactUploadRequest, CreateFileArtifactUploadResponse,
|
||||
FileArtifactRecord, FileArtifactUploadTargetInfo,
|
||||
};
|
||||
use crate::server::server_api::presigned_upload::upload_file_to_target;
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
const MIME_SNIFF_BYTES: usize = 8 * 1024;
|
||||
const OZ_RUN_ID_ENV_VAR: &str = "OZ_RUN_ID";
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub(crate) struct FileArtifactUploadRequest {
|
||||
pub(crate) path: PathBuf,
|
||||
pub(crate) run_id: Option<AmbientAgentTaskId>,
|
||||
pub(crate) conversation_id: Option<ServerConversationToken>,
|
||||
pub(crate) description: Option<String>,
|
||||
}
|
||||
|
||||
impl TryFrom<UploadArtifactArgs> for FileArtifactUploadRequest {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: UploadArtifactArgs) -> Result<Self> {
|
||||
let run_id = match value.run_id {
|
||||
Some(run_id) => Some(parse_run_id(&run_id, "Invalid run ID")?),
|
||||
None => None,
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
path: value.path,
|
||||
run_id,
|
||||
conversation_id: value.conversation_id.map(ServerConversationToken::new),
|
||||
description: value.description,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct CompletedFileArtifactUpload {
|
||||
pub(crate) artifact: FileArtifactRecord,
|
||||
pub(crate) size_bytes: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub(crate) struct ResolvedUploadAssociation {
|
||||
conversation_id: Option<ServerConversationToken>,
|
||||
run_id: Option<AmbientAgentTaskId>,
|
||||
pub(crate) ambient_task_id: AmbientAgentTaskId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct PreparedUploadArtifact {
|
||||
path: PathBuf,
|
||||
filepath: String,
|
||||
mime_type: String,
|
||||
file_size: u64,
|
||||
}
|
||||
|
||||
impl PreparedUploadArtifact {
|
||||
fn from_path(path: PathBuf) -> Result<Self> {
|
||||
// `infer` only needs leading signature bytes, so avoid buffering the whole artifact
|
||||
// before we stream the file body to the upload target.
|
||||
let (file_size, mime_sniff_bytes) = file_size_and_prefix_for_path(&path, MIME_SNIFF_BYTES)?;
|
||||
|
||||
Ok(Self {
|
||||
filepath: normalize_artifact_filepath(&path),
|
||||
mime_type: infer_mime_type(&path, &mime_sniff_bytes),
|
||||
file_size,
|
||||
path,
|
||||
})
|
||||
}
|
||||
|
||||
fn graphql_size_bytes(&self) -> Option<i32> {
|
||||
checked_graphql_size_bytes_for_upload(&self.path, self.file_size)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct FileArtifactUploader {
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
server_api: Arc<ServerApi>,
|
||||
}
|
||||
|
||||
impl FileArtifactUploader {
|
||||
pub(crate) fn new(ai_client: Arc<dyn AIClient>, server_api: Arc<ServerApi>) -> Self {
|
||||
Self {
|
||||
ai_client,
|
||||
server_api,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn upload_with_association(
|
||||
&self,
|
||||
request: FileArtifactUploadRequest,
|
||||
association: ResolvedUploadAssociation,
|
||||
) -> Result<CompletedFileArtifactUpload> {
|
||||
let FileArtifactUploadRequest {
|
||||
path, description, ..
|
||||
} = request;
|
||||
|
||||
let artifact = self.prepare_upload_artifact(path).await?;
|
||||
let create_response = self
|
||||
.create_upload_target(association, description, &artifact)
|
||||
.await?;
|
||||
|
||||
let checksum = self
|
||||
.upload_artifact_bytes(&create_response.upload_target, &artifact)
|
||||
.await?;
|
||||
let uploaded_artifact = self
|
||||
.confirm_upload(create_response.artifact.artifact_uid, checksum)
|
||||
.await?;
|
||||
let size_bytes = i64::try_from(artifact.file_size)
|
||||
.context("Artifact file size exceeds supported range")?;
|
||||
|
||||
Ok(CompletedFileArtifactUpload {
|
||||
artifact: uploaded_artifact,
|
||||
size_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
async fn prepare_upload_artifact(&self, path: PathBuf) -> Result<PreparedUploadArtifact> {
|
||||
unblock(move || PreparedUploadArtifact::from_path(path)).await
|
||||
}
|
||||
|
||||
async fn create_upload_target(
|
||||
&self,
|
||||
association: ResolvedUploadAssociation,
|
||||
description: Option<String>,
|
||||
artifact: &PreparedUploadArtifact,
|
||||
) -> Result<CreateFileArtifactUploadResponse> {
|
||||
self.ai_client
|
||||
.create_file_artifact_upload_target(CreateFileArtifactUploadRequest {
|
||||
conversation_id: association
|
||||
.conversation_id
|
||||
.as_ref()
|
||||
.map(|token| token.as_str().to_string()),
|
||||
run_id: association.run_id.as_ref().map(ToString::to_string),
|
||||
filepath: artifact.filepath.clone(),
|
||||
description,
|
||||
mime_type: Some(artifact.mime_type.clone()),
|
||||
size_bytes: artifact.graphql_size_bytes(),
|
||||
})
|
||||
.await
|
||||
.context("Failed to create file artifact upload target")
|
||||
}
|
||||
|
||||
async fn upload_artifact_bytes(
|
||||
&self,
|
||||
target: &FileArtifactUploadTargetInfo,
|
||||
artifact: &PreparedUploadArtifact,
|
||||
) -> Result<String> {
|
||||
upload_file_to_target(
|
||||
self.server_api.http_client(),
|
||||
target,
|
||||
&artifact.path,
|
||||
artifact.file_size,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn confirm_upload(
|
||||
&self,
|
||||
artifact_uid: String,
|
||||
checksum: String,
|
||||
) -> Result<FileArtifactRecord> {
|
||||
self.ai_client
|
||||
.confirm_file_artifact_upload(artifact_uid, checksum)
|
||||
.await
|
||||
.context("Failed to confirm file artifact upload")
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_upload_association(
|
||||
&self,
|
||||
request: &FileArtifactUploadRequest,
|
||||
) -> Result<ResolvedUploadAssociation> {
|
||||
let conversation_task_id = match (request.run_id.as_ref(), request.conversation_id.as_ref())
|
||||
{
|
||||
// we were given a conversation id, so we need to resolve the task id from the conversation via the api
|
||||
(None, Some(conversation_id)) => {
|
||||
Some(self.resolve_conversation_task_id(conversation_id).await)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
resolve_upload_association_from_sources(
|
||||
request.run_id,
|
||||
request.conversation_id.clone(),
|
||||
conversation_task_id,
|
||||
load_env_run_id()?,
|
||||
)
|
||||
}
|
||||
|
||||
async fn resolve_conversation_task_id(
|
||||
&self,
|
||||
conversation_id: &ServerConversationToken,
|
||||
) -> Result<AmbientAgentTaskId> {
|
||||
let metadata = self
|
||||
.ai_client
|
||||
.list_ai_conversation_metadata(Some(vec![conversation_id.as_str().to_string()]))
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"Failed to load conversation '{}' to resolve artifact upload headers",
|
||||
conversation_id.as_str()
|
||||
)
|
||||
})?;
|
||||
|
||||
let metadata = single_conversation_metadata(conversation_id.as_str(), metadata)
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"Failed to load conversation '{}' to resolve artifact upload headers",
|
||||
conversation_id.as_str()
|
||||
)
|
||||
})?;
|
||||
|
||||
ambient_task_id_from_conversation_metadata(conversation_id.as_str(), metadata)
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_artifact_filepath(path: &Path) -> String {
|
||||
path.to_string_lossy().replace('\\', "/")
|
||||
}
|
||||
|
||||
fn infer_mime_type(path: &Path, file_bytes: &[u8]) -> String {
|
||||
infer::get(file_bytes)
|
||||
.map(|kind| kind.mime_type().to_string())
|
||||
.unwrap_or_else(|| from_path(path).first_or_octet_stream().to_string())
|
||||
}
|
||||
|
||||
fn file_size_and_prefix_for_path(path: &Path, max_bytes: usize) -> Result<(u64, Vec<u8>)> {
|
||||
let mut file = File::open(path)
|
||||
.with_context(|| format!("Failed to open artifact file '{}'", path.display()))?;
|
||||
let file_size = file
|
||||
.metadata()
|
||||
.with_context(|| format!("Failed to stat artifact file '{}'", path.display()))?
|
||||
.len();
|
||||
let mut bytes = vec![0; max_bytes];
|
||||
let bytes_read = file
|
||||
.read(&mut bytes)
|
||||
.with_context(|| format!("Failed to read artifact file '{}'", path.display()))?;
|
||||
bytes.truncate(bytes_read);
|
||||
Ok((file_size, bytes))
|
||||
}
|
||||
|
||||
fn checked_graphql_size_bytes_for_upload(path: &Path, size_bytes: u64) -> Option<i32> {
|
||||
let graphql_size_bytes = i32::try_from(size_bytes).ok();
|
||||
if graphql_size_bytes.is_none() {
|
||||
// The backing upload can handle large files, but the GraphQL field is still `Int`.
|
||||
// Dropping `size_bytes` preserves the upload request instead of failing on conversion.
|
||||
log::warn!(
|
||||
"Artifact file '{}' is {} bytes, which exceeds the GraphQL size_bytes limit of {} bytes; omitting size_bytes from the upload target request",
|
||||
path.display(),
|
||||
size_bytes,
|
||||
i32::MAX,
|
||||
);
|
||||
}
|
||||
|
||||
graphql_size_bytes
|
||||
}
|
||||
|
||||
fn single_conversation_metadata(
|
||||
conversation_id: &str,
|
||||
mut metadata: Vec<ServerAIConversationMetadata>,
|
||||
) -> Result<ServerAIConversationMetadata> {
|
||||
match metadata.len() {
|
||||
0 => bail!("Conversation not found"),
|
||||
1 => Ok(metadata.pop().expect("metadata length checked")),
|
||||
_ => bail!("Multiple conversations found for '{conversation_id}'"),
|
||||
}
|
||||
}
|
||||
|
||||
fn ambient_task_id_from_conversation_metadata(
|
||||
conversation_id: &str,
|
||||
metadata: ServerAIConversationMetadata,
|
||||
) -> Result<AmbientAgentTaskId> {
|
||||
metadata.ambient_agent_task_id.ok_or_else(|| {
|
||||
anyhow!("Conversation '{conversation_id}' is not backed by a cloud agent task")
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_run_id(run_id: &str, error_prefix: &str) -> Result<AmbientAgentTaskId> {
|
||||
parse_ambient_task_id(run_id, error_prefix)
|
||||
}
|
||||
|
||||
fn load_env_run_id() -> Result<Option<String>> {
|
||||
match env::var(OZ_RUN_ID_ENV_VAR) {
|
||||
Ok(run_id) => Ok(Some(run_id)),
|
||||
Err(env::VarError::NotPresent) => Ok(None),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(anyhow!(
|
||||
"{OZ_RUN_ID_ENV_VAR} is set but is not valid Unicode"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_env_run_id(env_run_id: Option<String>) -> Result<AmbientAgentTaskId> {
|
||||
let Some(run_id) = env_run_id else {
|
||||
bail!("{OZ_RUN_ID_ENV_VAR} is not set");
|
||||
};
|
||||
|
||||
parse_run_id(&run_id, "Invalid OZ_RUN_ID")
|
||||
}
|
||||
|
||||
fn resolve_upload_association_from_sources(
|
||||
explicit_run_id: Option<AmbientAgentTaskId>,
|
||||
explicit_conversation_id: Option<ServerConversationToken>,
|
||||
conversation_task_id: Option<Result<AmbientAgentTaskId>>,
|
||||
env_run_id: Option<String>,
|
||||
) -> Result<ResolvedUploadAssociation> {
|
||||
// Precedence is deliberate:
|
||||
// 1. An explicit run ID is authoritative and must not silently fall back.
|
||||
// 2. A conversation ID stays attached to the artifact even if we have to borrow the ambient
|
||||
// task ID from `OZ_RUN_ID` because the conversation lacks cloud-task metadata.
|
||||
// 3. `OZ_RUN_ID` becomes the sole source of truth only when the caller supplied nothing else.
|
||||
if let Some(run_id) = explicit_run_id {
|
||||
let ambient_task_id = run_id;
|
||||
return Ok(ResolvedUploadAssociation {
|
||||
conversation_id: None,
|
||||
run_id: Some(run_id),
|
||||
ambient_task_id,
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(conversation_id) = explicit_conversation_id {
|
||||
match conversation_task_id
|
||||
.ok_or_else(|| anyhow!("conversation resolution should be provided"))?
|
||||
{
|
||||
Ok(ambient_task_id) => {
|
||||
return Ok(ResolvedUploadAssociation {
|
||||
conversation_id: Some(conversation_id),
|
||||
run_id: None,
|
||||
ambient_task_id,
|
||||
});
|
||||
}
|
||||
Err(conversation_err) => {
|
||||
let env_err = match resolve_env_run_id(env_run_id) {
|
||||
Ok(ambient_task_id) => {
|
||||
log::warn!(
|
||||
"Conversation '{}' task resolution failed ({conversation_err}); falling back to {OZ_RUN_ID_ENV_VAR} for ambient task context",
|
||||
conversation_id.as_str()
|
||||
);
|
||||
return Ok(ResolvedUploadAssociation {
|
||||
conversation_id: Some(conversation_id),
|
||||
run_id: None,
|
||||
ambient_task_id,
|
||||
});
|
||||
}
|
||||
Err(env_err) => env_err,
|
||||
};
|
||||
|
||||
return Err(anyhow!(
|
||||
"Failed to resolve artifact upload association for conversation '{}': {conversation_err}; also failed to use {OZ_RUN_ID_ENV_VAR}: {env_err}",
|
||||
conversation_id.as_str()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ambient_task_id = resolve_env_run_id(env_run_id).map_err(|env_err| {
|
||||
anyhow!(
|
||||
"Failed to resolve artifact upload association: no usable --run-id or --conversation-id was provided, and {OZ_RUN_ID_ENV_VAR}: {env_err}"
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(ResolvedUploadAssociation {
|
||||
conversation_id: None,
|
||||
run_id: Some(ambient_task_id),
|
||||
ambient_task_id,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "artifact_upload_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,269 @@
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use chrono::Utc;
|
||||
use tempfile::tempdir;
|
||||
use warp_cli::artifact::UploadArtifactArgs;
|
||||
|
||||
use super::*;
|
||||
use crate::ai::agent::api::ServerConversationToken;
|
||||
use crate::ai::agent::conversation::{AIAgentHarness, ServerAIConversationMetadata};
|
||||
use crate::cloud_object::{Revision, ServerMetadata, ServerPermissions};
|
||||
use crate::persistence::model::ConversationUsageMetadata;
|
||||
use crate::server::ids::ServerId;
|
||||
|
||||
fn create_mock_server_metadata() -> ServerMetadata {
|
||||
ServerMetadata {
|
||||
uid: ServerId::default(),
|
||||
revision: Revision::now(),
|
||||
metadata_last_updated_ts: Utc::now().into(),
|
||||
trashed_ts: None,
|
||||
folder_id: None,
|
||||
is_welcome_object: false,
|
||||
creator_uid: None,
|
||||
last_editor_uid: None,
|
||||
current_editor_uid: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn create_conversation_metadata(
|
||||
conversation_id: &str,
|
||||
ambient_task_id: Option<&str>,
|
||||
) -> ServerAIConversationMetadata {
|
||||
ServerAIConversationMetadata {
|
||||
title: "Artifact upload".to_string(),
|
||||
working_directory: None,
|
||||
harness: AIAgentHarness::Oz,
|
||||
usage: ConversationUsageMetadata {
|
||||
was_summarized: false,
|
||||
context_window_usage: 0.0,
|
||||
credits_spent: 0.0,
|
||||
credits_spent_for_last_block: None,
|
||||
token_usage: vec![],
|
||||
tool_usage_metadata: Default::default(),
|
||||
},
|
||||
metadata: create_mock_server_metadata(),
|
||||
permissions: ServerPermissions::mock_personal(),
|
||||
ambient_agent_task_id: ambient_task_id.map(|task_id| task_id.parse().unwrap()),
|
||||
server_conversation_token: ServerConversationToken::new(conversation_id.to_string()),
|
||||
artifacts: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_artifact_filepath_preserves_shape_and_normalizes_separators() {
|
||||
let path = PathBuf::from(r"outputs\reports/final.txt");
|
||||
assert_eq!(
|
||||
normalize_artifact_filepath(&path),
|
||||
"outputs/reports/final.txt"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn checked_graphql_size_bytes_for_upload_returns_none_for_overflow() {
|
||||
let path = PathBuf::from("outputs/large-artifact.bin");
|
||||
|
||||
assert_eq!(
|
||||
checked_graphql_size_bytes_for_upload(&path, i32::MAX as u64),
|
||||
Some(i32::MAX)
|
||||
);
|
||||
assert_eq!(
|
||||
checked_graphql_size_bytes_for_upload(&path, i32::MAX as u64 + 1),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_size_and_prefix_for_path_returns_truncated_prefix() {
|
||||
let tempdir = tempdir().unwrap();
|
||||
let path = tempdir.path().join("artifact.bin");
|
||||
fs::write(&path, b"0123456789").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
file_size_and_prefix_for_path(&path, 4).unwrap(),
|
||||
(10, b"0123".to_vec())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_size_and_prefix_for_path_returns_full_contents_when_prefix_exceeds_file() {
|
||||
let tempdir = tempdir().unwrap();
|
||||
let path = tempdir.path().join("artifact.bin");
|
||||
fs::write(&path, b"0123456789").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
file_size_and_prefix_for_path(&path, 32).unwrap(),
|
||||
(10, b"0123456789".to_vec())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_conversation_metadata_returns_the_only_metadata_record() {
|
||||
let metadata = single_conversation_metadata(
|
||||
"conversation-123",
|
||||
vec![create_conversation_metadata(
|
||||
"conversation-123",
|
||||
Some("550e8400-e29b-41d4-a716-446655440000"),
|
||||
)],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let task_id = ambient_task_id_from_conversation_metadata("conversation-123", metadata).unwrap();
|
||||
assert_eq!(
|
||||
task_id,
|
||||
"550e8400-e29b-41d4-a716-446655440000".parse().unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_conversation_metadata_errors_when_no_metadata_is_returned() {
|
||||
let err = single_conversation_metadata("conversation-123", Vec::new()).unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("Conversation not found"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ambient_task_id_from_conversation_metadata_requires_cloud_task_metadata() {
|
||||
let err = ambient_task_id_from_conversation_metadata(
|
||||
"conversation-123",
|
||||
create_conversation_metadata("conversation-123", None),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err
|
||||
.to_string()
|
||||
.contains("Conversation 'conversation-123' is not backed by a cloud agent task"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_run_id_wins_over_env_fallback() {
|
||||
let resolved = resolve_upload_association_from_sources(
|
||||
Some("550e8400-e29b-41d4-a716-446655440000".parse().unwrap()),
|
||||
None,
|
||||
None,
|
||||
Some("11111111-1111-1111-1111-111111111111".to_string()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved,
|
||||
ResolvedUploadAssociation {
|
||||
conversation_id: None,
|
||||
run_id: Some("550e8400-e29b-41d4-a716-446655440000".parse().unwrap()),
|
||||
ambient_task_id: "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_explicit_run_id_errors_even_if_env_fallback_exists() {
|
||||
let err = FileArtifactUploadRequest::try_from(UploadArtifactArgs {
|
||||
path: PathBuf::from("outputs/report.txt"),
|
||||
run_id: Some("not-a-run-id".to_string()),
|
||||
conversation_id: None,
|
||||
description: None,
|
||||
})
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("Invalid run ID 'not-a-run-id'"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_conversation_resolution_ignores_env_fallback() {
|
||||
let resolved = resolve_upload_association_from_sources(
|
||||
None,
|
||||
Some(ServerConversationToken::new("conversation-123".to_string())),
|
||||
Some(Ok("550e8400-e29b-41d4-a716-446655440000".parse().unwrap())),
|
||||
Some("11111111-1111-1111-1111-111111111111".to_string()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved,
|
||||
ResolvedUploadAssociation {
|
||||
conversation_id: Some(ServerConversationToken::new("conversation-123".to_string())),
|
||||
run_id: None,
|
||||
ambient_task_id: "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_conversation_resolution_falls_back_to_env_run_id() {
|
||||
let resolved = resolve_upload_association_from_sources(
|
||||
None,
|
||||
Some(ServerConversationToken::new("conversation-123".to_string())),
|
||||
Some(Err(anyhow!(
|
||||
"Conversation 'conversation-123' is not backed by a cloud agent task"
|
||||
))),
|
||||
Some("550e8400-e29b-41d4-a716-446655440000".to_string()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved,
|
||||
ResolvedUploadAssociation {
|
||||
conversation_id: Some(ServerConversationToken::new("conversation-123".to_string())),
|
||||
run_id: None,
|
||||
ambient_task_id: "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_args_fall_back_to_env_run_id_for_request_association() {
|
||||
let resolved = resolve_upload_association_from_sources(
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("550e8400-e29b-41d4-a716-446655440000".to_string()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved,
|
||||
ResolvedUploadAssociation {
|
||||
conversation_id: None,
|
||||
run_id: Some("550e8400-e29b-41d4-a716-446655440000".parse().unwrap()),
|
||||
ambient_task_id: "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_args_and_missing_env_return_clear_error() {
|
||||
let err = resolve_upload_association_from_sources(None, None, None, None).unwrap_err();
|
||||
|
||||
assert!(err
|
||||
.to_string()
|
||||
.contains("no usable --run-id or --conversation-id was provided"));
|
||||
assert!(err.to_string().contains("OZ_RUN_ID"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_env_run_id_returns_clear_error() {
|
||||
let err =
|
||||
resolve_upload_association_from_sources(None, None, None, Some("not-a-run-id".to_string()))
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("Invalid OZ_RUN_ID 'not-a-run-id'"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_env_run_id_reads_variable() {
|
||||
let previous = env::var_os(OZ_RUN_ID_ENV_VAR);
|
||||
env::set_var(OZ_RUN_ID_ENV_VAR, "550e8400-e29b-41d4-a716-446655440000");
|
||||
|
||||
let loaded = load_env_run_id().unwrap();
|
||||
|
||||
match previous {
|
||||
Some(value) => env::set_var(OZ_RUN_ID_ENV_VAR, value),
|
||||
None => env::remove_var(OZ_RUN_ID_ENV_VAR),
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
loaded.as_deref(),
|
||||
Some("550e8400-e29b-41d4-a716-446655440000")
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,350 @@
|
||||
//! Common utilities for agent SDK commands.
|
||||
|
||||
use std::fmt;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::TryFutureExt;
|
||||
use inquire::{InquireError, Select};
|
||||
use warp_cli::agent::Harness;
|
||||
use warp_cli::environment::{EnvironmentCreateArgs, EnvironmentUpdateArgs};
|
||||
use warpui::r#async::FutureExt;
|
||||
use warpui::{AppContext, GetSingletonModelHandle, SingletonEntity as _, UpdateModel};
|
||||
|
||||
use crate::ai::agent::conversation::ServerAIConversationMetadata;
|
||||
use crate::ai::agent_sdk::driver::{AgentDriverError, WARP_DRIVE_SYNC_TIMEOUT};
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::ai::cloud_environments::CloudAmbientAgentEnvironment;
|
||||
use crate::ai::llms::{LLMId, LLMPreferences};
|
||||
use crate::auth::auth_state::AuthStateProvider;
|
||||
use crate::cloud_object::{CloudObject, Owner};
|
||||
use crate::server::cloud_objects::update_manager::UpdateManager;
|
||||
use crate::server::ids::{ServerId, SyncId};
|
||||
use crate::server::server_api::ai::AIClient;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use crate::workspaces::update_manager::TeamUpdateManager;
|
||||
use crate::workspaces::user_workspaces::UserWorkspaces;
|
||||
|
||||
/// How long to wait for workspace metadata to refresh.
|
||||
pub const WORKSPACE_METADATA_REFRESH_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
pub fn validate_agent_mode_base_model_id(
|
||||
model_id: &str,
|
||||
ctx: &AppContext,
|
||||
) -> anyhow::Result<LLMId> {
|
||||
let llm_prefs = LLMPreferences::as_ref(ctx);
|
||||
|
||||
let llm_id: LLMId = model_id.into();
|
||||
let valid_ids = llm_prefs
|
||||
.get_base_llm_choices_for_agent_mode()
|
||||
.map(|info| info.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if valid_ids.contains(&llm_id) {
|
||||
Ok(llm_id)
|
||||
} else {
|
||||
let suggestions = valid_ids
|
||||
.into_iter()
|
||||
.map(|id| id.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
Err(anyhow::anyhow!(
|
||||
"Unknown model id '{model_id}'. Try one of: {suggestions}"
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn parse_ambient_task_id(
|
||||
run_id: &str,
|
||||
error_prefix: &str,
|
||||
) -> anyhow::Result<AmbientAgentTaskId> {
|
||||
run_id
|
||||
.parse()
|
||||
.map_err(|err| anyhow::anyhow!("{error_prefix} '{run_id}': {err}"))
|
||||
}
|
||||
|
||||
pub(super) fn set_ambient_task_context_from_run_id(
|
||||
ctx: &AppContext,
|
||||
run_id: &str,
|
||||
) -> anyhow::Result<AmbientAgentTaskId> {
|
||||
let task_id = parse_ambient_task_id(run_id, "Invalid run ID")?;
|
||||
ServerApiProvider::handle(ctx)
|
||||
.as_ref(ctx)
|
||||
.get()
|
||||
.set_ambient_agent_task_id(Some(task_id));
|
||||
Ok(task_id)
|
||||
}
|
||||
|
||||
/// Resolve the owner of a new cloud object. This resolution is based on the CLI `--team` and `--personal` flags.
|
||||
///
|
||||
/// If `team_flag` is true, attempts to get the current team UID (errors if not on a team).
|
||||
/// If `user_flag` is true, gets the current user's UID.
|
||||
/// Otherwise, defaults to team if available, falling back to user.
|
||||
pub fn resolve_owner(team_flag: bool, user_flag: bool, ctx: &AppContext) -> anyhow::Result<Owner> {
|
||||
if team_flag {
|
||||
let team_id = UserWorkspaces::as_ref(ctx)
|
||||
.current_team_uid()
|
||||
.ok_or_else(|| anyhow::anyhow!("User is not on a team"))?;
|
||||
return Ok(Owner::Team { team_uid: team_id });
|
||||
}
|
||||
|
||||
if user_flag {
|
||||
let user_id = AuthStateProvider::as_ref(ctx)
|
||||
.get()
|
||||
.user_id()
|
||||
.ok_or_else(|| anyhow::anyhow!("User should be logged in"))?;
|
||||
return Ok(Owner::User { user_uid: user_id });
|
||||
}
|
||||
|
||||
// Default: try team first, fall back to user
|
||||
if let Some(team_uid) = UserWorkspaces::as_ref(ctx).current_team_uid() {
|
||||
return Ok(Owner::Team { team_uid });
|
||||
}
|
||||
|
||||
log::warn!("Tried to default to creating team object, team could not be found.");
|
||||
let user_id = AuthStateProvider::as_ref(ctx)
|
||||
.get()
|
||||
.user_id()
|
||||
.ok_or_else(|| anyhow::anyhow!("User should be logged in"))?;
|
||||
Ok(Owner::User { user_uid: user_id })
|
||||
}
|
||||
|
||||
/// Refresh workspace metadata before executing an operation.
|
||||
///
|
||||
/// This ensures that team state is up-to-date before creating cloud objects or performing
|
||||
/// other operations that depend on team membership.
|
||||
pub fn refresh_workspace_metadata<C>(
|
||||
ctx: &mut C,
|
||||
) -> impl Future<Output = anyhow::Result<()>> + Send + 'static
|
||||
where
|
||||
C: GetSingletonModelHandle + UpdateModel,
|
||||
{
|
||||
let refresh_future = TeamUpdateManager::handle(ctx).update(ctx, |manager, ctx| {
|
||||
manager
|
||||
.refresh_workspace_metadata(ctx)
|
||||
.with_timeout(WORKSPACE_METADATA_REFRESH_TIMEOUT)
|
||||
});
|
||||
|
||||
async move {
|
||||
let _ = refresh_future
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Timed out refreshing team metadata"))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Refresh Warp Drive before executing an operation.
|
||||
pub fn refresh_warp_drive(
|
||||
ctx: &AppContext,
|
||||
) -> impl Future<Output = anyhow::Result<()>> + Send + 'static {
|
||||
UpdateManager::as_ref(ctx)
|
||||
.initial_load_complete()
|
||||
.with_timeout(WARP_DRIVE_SYNC_TIMEOUT)
|
||||
.map_err(|_| anyhow::anyhow!("Timed out waiting for Warp Drive to sync"))
|
||||
}
|
||||
|
||||
/// Fetch the conversation's server metadata and validate that its harness matches the caller's
|
||||
/// `--harness` choice. Returns the metadata on success so the caller can reuse it (e.g. for the
|
||||
/// server conversation token).
|
||||
///
|
||||
/// Called up-front before any task/config-build logic consumes `args.harness`, so a mismatch
|
||||
/// error surfaces before side effects like task creation. We deliberately do NOT auto-upgrade
|
||||
/// the harness: `Harness::Oz` default with a Claude conversation id is treated as a mismatch
|
||||
/// and errors out.
|
||||
pub(super) async fn fetch_and_validate_conversation_harness(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
conversation_id: &str,
|
||||
args_harness: Harness,
|
||||
) -> Result<ServerAIConversationMetadata, AgentDriverError> {
|
||||
let metadata = ai_client
|
||||
.list_ai_conversation_metadata(Some(vec![conversation_id.to_string()]))
|
||||
.await
|
||||
.map_err(|e| AgentDriverError::ConversationLoadFailed(format!("{e:#}")))?
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| {
|
||||
AgentDriverError::ConversationLoadFailed(format!(
|
||||
"conversation {conversation_id} not found or not accessible"
|
||||
))
|
||||
})?;
|
||||
|
||||
if metadata.harness != args_harness {
|
||||
return Err(AgentDriverError::ConversationHarnessMismatch {
|
||||
conversation_id: conversation_id.to_string(),
|
||||
expected: Harness::from(metadata.harness).to_string(),
|
||||
got: args_harness.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(metadata)
|
||||
}
|
||||
|
||||
/// Format an object owner for display in the CLI.
|
||||
pub fn format_owner(owner: &Owner) -> &'static str {
|
||||
// TODO: For potentially-shared objects, consider looking up the particular user/team name.
|
||||
match owner {
|
||||
Owner::User { .. } => "Personal",
|
||||
Owner::Team { .. } => "Team",
|
||||
}
|
||||
}
|
||||
|
||||
/// An error resolving an agent option, which we may have prompted the user for.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ResolveConfigurationError {
|
||||
/// The user canceled the operation, and we should exit.
|
||||
#[error("Operation canceled")]
|
||||
Canceled,
|
||||
#[error("{id} is not a valid {kind} identifier")]
|
||||
InvalidId { id: String, kind: &'static str },
|
||||
#[error("{kind} {id} not found")]
|
||||
ObjectNotFound { id: String, kind: &'static str },
|
||||
#[error(transparent)]
|
||||
Other(anyhow::Error),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum EnvironmentChoice {
|
||||
/// The user explicitly chose not to use an environment.
|
||||
None,
|
||||
/// The user chose a specific environment.
|
||||
Environment { id: String, name: String },
|
||||
}
|
||||
|
||||
impl EnvironmentChoice {
|
||||
/// Resolve the environment to use when creating an agent integration.
|
||||
/// Warp Drive *must* have been synced first.
|
||||
pub fn resolve_for_create(
|
||||
args: EnvironmentCreateArgs,
|
||||
ctx: &AppContext,
|
||||
) -> Result<Self, ResolveConfigurationError> {
|
||||
if args.no_environment {
|
||||
Ok(EnvironmentChoice::None)
|
||||
} else if let Some(id) = args.environment {
|
||||
Self::get_by_id(id, ctx)
|
||||
} else {
|
||||
let all_environments = CloudAmbientAgentEnvironment::get_all(ctx);
|
||||
let mut synced_environments: Vec<(ServerId, &CloudAmbientAgentEnvironment)> =
|
||||
all_environments
|
||||
.iter()
|
||||
.filter_map(|env| {
|
||||
if let SyncId::ServerId(server_id) = env.sync_id() {
|
||||
Some((server_id, env))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
synced_environments
|
||||
.sort_by_key(|(_, env)| env.model().string_model.name.to_lowercase());
|
||||
|
||||
let environments: Vec<EnvironmentChoice> = synced_environments
|
||||
.into_iter()
|
||||
.map(|(server_id, env)| EnvironmentChoice::Environment {
|
||||
id: server_id.to_string(),
|
||||
name: env.model().string_model.name.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut options = vec![EnvironmentChoice::None];
|
||||
options.extend(environments);
|
||||
|
||||
// If there are no synced environments, require the user to create one or use --no-environment.
|
||||
if options.len() == 1 {
|
||||
let cli_name = warp_cli::binary_name().unwrap_or_else(|| "warp".to_string());
|
||||
return Err(ResolveConfigurationError::Other(anyhow::anyhow!(
|
||||
"No environments are configured for this account.\n\
|
||||
You can create an environment with `{cli_name} environment create`.\n\
|
||||
Or, re-run this command with `--no-environment` to not use an environment.\n\
|
||||
Without an environment, the agent will not be able to access private repositories or create pull requests.",
|
||||
)));
|
||||
}
|
||||
|
||||
let prompt = "Select an environment to run the agent in (or 'No environment'):";
|
||||
|
||||
let choice = Select::new(prompt, options).prompt();
|
||||
|
||||
match choice {
|
||||
Ok(choice) => Ok(choice),
|
||||
Err(InquireError::OperationCanceled | InquireError::OperationInterrupted) => {
|
||||
Err(ResolveConfigurationError::Canceled)
|
||||
}
|
||||
Err(err) => Err(ResolveConfigurationError::Other(anyhow::anyhow!(
|
||||
"Error selecting environment: {err}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the environment to use when updating an agent integration. If the user did not
|
||||
/// request any changes to the environment, this returns `Ok(None)`.
|
||||
/// Warp Drive *must* have been synced first.
|
||||
pub fn resolve_for_update(
|
||||
args: EnvironmentUpdateArgs,
|
||||
ctx: &AppContext,
|
||||
) -> Result<Option<Self>, ResolveConfigurationError> {
|
||||
if args.remove_environment {
|
||||
Ok(Some(EnvironmentChoice::None))
|
||||
} else if let Some(id) = args.environment {
|
||||
Self::get_by_id(id, ctx).map(Some)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn get_by_id(id: String, ctx: &AppContext) -> Result<Self, ResolveConfigurationError> {
|
||||
let sync_id = SyncId::ServerId(ServerId::try_from(id.as_str()).map_err(|_| {
|
||||
ResolveConfigurationError::InvalidId {
|
||||
id: id.clone(),
|
||||
kind: "environment",
|
||||
}
|
||||
})?);
|
||||
|
||||
let environment =
|
||||
CloudAmbientAgentEnvironment::get_by_id(&sync_id, ctx).ok_or_else(|| {
|
||||
ResolveConfigurationError::ObjectNotFound {
|
||||
id: id.clone(),
|
||||
kind: "environment",
|
||||
}
|
||||
})?;
|
||||
|
||||
Ok(EnvironmentChoice::Environment {
|
||||
id,
|
||||
name: environment.model().string_model.name.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for EnvironmentChoice {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
EnvironmentChoice::None => write!(
|
||||
f,
|
||||
"No environment (agent will not be able to access private repositories or create pull requests)",
|
||||
),
|
||||
EnvironmentChoice::Environment { id, name } => write!(f, "{name} ({id})"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::parse_ambient_task_id;
|
||||
|
||||
#[test]
|
||||
fn parse_ambient_task_id_accepts_valid_ids() {
|
||||
let task_id =
|
||||
parse_ambient_task_id("550e8400-e29b-41d4-a716-446655440000", "Invalid run ID")
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(task_id.to_string(), "550e8400-e29b-41d4-a716-446655440000");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_ambient_task_id_preserves_error_prefix() {
|
||||
let err = parse_ambient_task_id("not-a-run-id", "Invalid run ID").unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("Invalid run ID 'not-a-run-id'"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
use std::path::Path;
|
||||
|
||||
use anyhow::Context as _;
|
||||
use serde_json::{Map, Value};
|
||||
use warp_cli::mcp::MCPSpec;
|
||||
|
||||
use crate::ai::ambient_agents::AgentConfigSnapshot;
|
||||
|
||||
/// A strict, file-based representation of `AgentConfigSnapshot`.
|
||||
///
|
||||
/// Notes:
|
||||
/// - Keys are snake_case and unknown keys are rejected.
|
||||
/// - MCP configuration must be provided only under the `mcp_servers` key and must be the
|
||||
/// unwrapped server map `{ <server_name>: <server_config>, ... }`.
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct AgentConfigSnapshotFile {
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub environment_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub model_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub base_prompt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub mcp_servers: Option<Map<String, Value>>,
|
||||
#[serde(default)]
|
||||
pub host: Option<String>,
|
||||
#[serde(default)]
|
||||
pub computer_use_enabled: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoadedAgentConfigSnapshotFile {
|
||||
pub file: AgentConfigSnapshotFile,
|
||||
}
|
||||
|
||||
/// Load an `AgentConfigSnapshotFile` from disk.
|
||||
///
|
||||
/// Parsing rules:
|
||||
/// - `.json` => JSON
|
||||
/// - `.yml` / `.yaml` => YAML
|
||||
/// - otherwise: try JSON, then YAML
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub fn load_config_file(path: &Path) -> anyhow::Result<LoadedAgentConfigSnapshotFile> {
|
||||
let contents = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("Failed to read config file '{}'", path.display()))?;
|
||||
|
||||
let ext = path
|
||||
.extension()
|
||||
.and_then(|s| s.to_str())
|
||||
.map(|s| s.to_ascii_lowercase());
|
||||
|
||||
let file = match ext.as_deref() {
|
||||
Some("json") => parse_json(&contents)
|
||||
.with_context(|| format!("Invalid JSON in config file '{}'", path.display()))?,
|
||||
Some("yml") | Some("yaml") => parse_yaml(&contents)
|
||||
.with_context(|| format!("Invalid YAML in config file '{}'", path.display()))?,
|
||||
_ => parse_json(&contents)
|
||||
.or_else(|_| parse_yaml(&contents))
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"Failed to parse config file '{}' as JSON or YAML",
|
||||
path.display()
|
||||
)
|
||||
})?,
|
||||
};
|
||||
|
||||
if let Some(mcp_servers) = &file.mcp_servers {
|
||||
super::mcp_config::validate_mcp_servers(mcp_servers)
|
||||
.with_context(|| format!("Invalid mcp_servers in '{}'", path.display()))?;
|
||||
}
|
||||
|
||||
Ok(LoadedAgentConfigSnapshotFile { file })
|
||||
}
|
||||
|
||||
/// WASM builds don't use CLI command execution / local file access.
|
||||
#[cfg(target_family = "wasm")]
|
||||
pub fn load_config_file(_path: &Path) -> anyhow::Result<LoadedAgentConfigSnapshotFile> {
|
||||
Err(anyhow::anyhow!(
|
||||
"Config files are not supported in WASM builds"
|
||||
))
|
||||
}
|
||||
|
||||
fn parse_json(input: &str) -> anyhow::Result<AgentConfigSnapshotFile> {
|
||||
serde_json::from_str::<AgentConfigSnapshotFile>(input).with_context(supported_keys_context)
|
||||
}
|
||||
|
||||
fn parse_yaml(input: &str) -> anyhow::Result<AgentConfigSnapshotFile> {
|
||||
// `serde_yaml` can deserialize into `serde_json::Value` directly.
|
||||
serde_yaml::from_str::<AgentConfigSnapshotFile>(input).with_context(supported_keys_context)
|
||||
}
|
||||
|
||||
fn supported_keys_context() -> String {
|
||||
"Supported keys: name, environment_id, model_id, base_prompt, mcp_servers, host, computer_use_enabled".to_string()
|
||||
}
|
||||
|
||||
/// Convert an unwrapped `mcp_servers` map into runtime MCP specs for AgentDriver.
|
||||
///
|
||||
/// Behavior:
|
||||
/// - Entries with `warp_id` become `MCPSpec::Uuid`.
|
||||
/// - Entries with `command`/`url` remain as inline JSON (`MCPSpec::Json`) containing the unwrapped server map.
|
||||
pub fn mcp_specs_from_mcp_servers(
|
||||
mcp_servers: &Map<String, Value>,
|
||||
) -> anyhow::Result<Vec<MCPSpec>> {
|
||||
let mut uuids: Vec<uuid::Uuid> = Vec::new();
|
||||
let mut json_map: Map<String, Value> = Map::new();
|
||||
|
||||
for (name, config) in mcp_servers {
|
||||
let obj = config
|
||||
.as_object()
|
||||
.ok_or_else(|| anyhow::anyhow!("MCP server '{name}' config must be a JSON object"))?;
|
||||
|
||||
if let Some(warp_id) = obj.get("warp_id").and_then(Value::as_str) {
|
||||
let uuid = uuid::Uuid::parse_str(warp_id).map_err(|_| {
|
||||
anyhow::anyhow!("MCP server '{name}' field 'warp_id' must be a UUID")
|
||||
})?;
|
||||
uuids.push(uuid);
|
||||
} else {
|
||||
json_map.insert(name.clone(), config.clone());
|
||||
}
|
||||
}
|
||||
|
||||
uuids.sort();
|
||||
uuids.dedup();
|
||||
|
||||
let mut specs: Vec<MCPSpec> = uuids.into_iter().map(MCPSpec::Uuid).collect();
|
||||
|
||||
if !json_map.is_empty() {
|
||||
let json =
|
||||
serde_json::to_string(&json_map).context("Failed to serialize MCP server map")?;
|
||||
specs.push(MCPSpec::Json(json));
|
||||
}
|
||||
|
||||
Ok(specs)
|
||||
}
|
||||
|
||||
/// Merge config file settings with CLI-provided overrides.
|
||||
///
|
||||
/// Precedence: CLI > file > default.
|
||||
pub fn merge_with_precedence(
|
||||
file: Option<&LoadedAgentConfigSnapshotFile>,
|
||||
cli: AgentConfigSnapshot,
|
||||
) -> AgentConfigSnapshot {
|
||||
let default_file = AgentConfigSnapshotFile::default();
|
||||
let file = file.map(|loaded| &loaded.file).unwrap_or(&default_file);
|
||||
|
||||
let name = cli.name.or_else(|| file.name.clone());
|
||||
let environment_id = cli.environment_id.or_else(|| file.environment_id.clone());
|
||||
let model_id = cli.model_id.or_else(|| file.model_id.clone());
|
||||
let base_prompt = cli.base_prompt.or_else(|| file.base_prompt.clone());
|
||||
|
||||
let mcp_servers = merge_mcp_servers(file.mcp_servers.clone(), cli.mcp_servers);
|
||||
let worker_host = cli.worker_host.or_else(|| file.host.clone());
|
||||
let computer_use_enabled = cli.computer_use_enabled.or(file.computer_use_enabled);
|
||||
|
||||
AgentConfigSnapshot {
|
||||
name,
|
||||
environment_id,
|
||||
model_id,
|
||||
base_prompt,
|
||||
mcp_servers,
|
||||
profile_id: None,
|
||||
worker_host,
|
||||
skill_spec: cli.skill_spec,
|
||||
computer_use_enabled,
|
||||
harness: cli.harness,
|
||||
harness_auth_secrets: cli.harness_auth_secrets,
|
||||
}
|
||||
}
|
||||
|
||||
/// Merge MCP servers from two sources.
|
||||
///
|
||||
/// Returns the merged map, or None if both inputs are None/empty.
|
||||
pub fn merge_mcp_servers(
|
||||
file_mcp: Option<Map<String, Value>>,
|
||||
cli_mcp: Option<Map<String, Value>>,
|
||||
) -> Option<Map<String, Value>> {
|
||||
match (file_mcp, cli_mcp) {
|
||||
(None, None) => None,
|
||||
(Some(map), None) => {
|
||||
if map.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(map)
|
||||
}
|
||||
}
|
||||
(None, Some(map)) => {
|
||||
if map.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(map)
|
||||
}
|
||||
}
|
||||
(Some(mut file_map), Some(cli_map)) => {
|
||||
for (k, v) in cli_map {
|
||||
file_map.insert(k, v);
|
||||
}
|
||||
if file_map.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(file_map)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "config_file_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,214 @@
|
||||
#![cfg(not(target_family = "wasm"))]
|
||||
|
||||
use std::io::Write as _;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai::ambient_agents::AgentConfigSnapshot;
|
||||
use warp_cli::mcp::MCPSpec;
|
||||
|
||||
fn write_temp(suffix: &str, contents: &str) -> tempfile::NamedTempFile {
|
||||
let mut file = tempfile::Builder::new().suffix(suffix).tempfile().unwrap();
|
||||
file.write_all(contents.as_bytes()).unwrap();
|
||||
file
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loads_json_and_validates_mcp_servers() {
|
||||
let contents = json!({
|
||||
"model_id": "gpt-4o",
|
||||
"environment_id": "env-123",
|
||||
"base_prompt": "be helpful",
|
||||
"mcp_servers": {
|
||||
"s": { "command": "npx", "args": [] }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
assert_eq!(loaded.file.model_id.as_deref(), Some("gpt-4o"));
|
||||
assert_eq!(loaded.file.environment_id.as_deref(), Some("env-123"));
|
||||
assert_eq!(loaded.file.base_prompt.as_deref(), Some("be helpful"));
|
||||
assert!(loaded.file.mcp_servers.is_some());
|
||||
assert!(loaded.file.mcp_servers.as_ref().unwrap().contains_key("s"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loads_yaml() {
|
||||
let contents = r#"
|
||||
model_id: gpt-4o
|
||||
mcp_servers:
|
||||
s:
|
||||
command: npx
|
||||
args: []
|
||||
"#;
|
||||
|
||||
let file = write_temp(".yaml", contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
assert_eq!(loaded.file.model_id.as_deref(), Some("gpt-4o"));
|
||||
assert!(loaded.file.mcp_servers.as_ref().unwrap().contains_key("s"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_keys_are_rejected() {
|
||||
let contents = json!({
|
||||
"model_id": "gpt-4o",
|
||||
"typo_model": "oops"
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let err = super::load_config_file(file.path()).unwrap_err();
|
||||
let err_str = format!("{err:#}");
|
||||
assert!(err_str.contains("Supported keys"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_must_be_under_mcp_servers_key() {
|
||||
let contents = json!({
|
||||
"model_id": "gpt-4o",
|
||||
"mcpServers": { "s": { "command": "npx", "args": [] } }
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let err = super::load_config_file(file.path()).unwrap_err();
|
||||
let err_str = format!("{err:#}");
|
||||
assert!(err_str.contains("Supported keys"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_precedence_cli_over_file_and_merges_mcp() {
|
||||
let contents = json!({
|
||||
"model_id": "file-model",
|
||||
"mcp_servers": {
|
||||
"a": { "url": "https://example.com/mcp" }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
let cli = AgentConfigSnapshot {
|
||||
name: Some("cli-name".to_string()),
|
||||
environment_id: None,
|
||||
model_id: Some("cli-model".to_string()),
|
||||
base_prompt: None,
|
||||
mcp_servers: Some(serde_json::Map::from_iter([(
|
||||
"a".to_string(),
|
||||
json!({"command": "npx", "args": []}),
|
||||
)])),
|
||||
profile_id: None,
|
||||
worker_host: None,
|
||||
skill_spec: None,
|
||||
computer_use_enabled: None,
|
||||
harness: None,
|
||||
harness_auth_secrets: None,
|
||||
};
|
||||
|
||||
let merged = super::merge_with_precedence(Some(&loaded), cli);
|
||||
|
||||
assert_eq!(merged.name.as_deref(), Some("cli-name"));
|
||||
assert_eq!(merged.model_id.as_deref(), Some("cli-model"));
|
||||
|
||||
let a = merged.mcp_servers.as_ref().unwrap().get("a").unwrap();
|
||||
assert_eq!(a.get("command").and_then(|v| v.as_str()), Some("npx"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_empty_mcp_servers_is_loaded_as_empty_map() {
|
||||
let contents = json!({
|
||||
"mcp_servers": {}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
assert!(loaded.file.mcp_servers.is_some());
|
||||
assert!(loaded.file.mcp_servers.as_ref().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_servers_map_converts_to_runtime_specs() {
|
||||
let contents = json!({
|
||||
"mcp_servers": {
|
||||
"existing": { "warp_id": "550e8400-e29b-41d4-a716-446655440000" },
|
||||
"ephemeral": { "command": "npx", "args": [] }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
let map = loaded.file.mcp_servers.as_ref().unwrap();
|
||||
let specs = super::mcp_specs_from_mcp_servers(map).unwrap();
|
||||
|
||||
assert!(specs.iter().any(|s| matches!(s, MCPSpec::Uuid(_))));
|
||||
assert!(specs.iter().any(|s| matches!(s, MCPSpec::Json(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loads_computer_use_enabled_from_json() {
|
||||
let contents = json!({
|
||||
"computer_use_enabled": true
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
assert_eq!(loaded.file.computer_use_enabled, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loads_computer_use_enabled_from_yaml() {
|
||||
let contents = "computer_use_enabled: false\n";
|
||||
|
||||
let file = write_temp(".yaml", contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
assert_eq!(loaded.file.computer_use_enabled, Some(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_precedence_cli_computer_use_enabled_over_file() {
|
||||
let contents = json!({
|
||||
"computer_use_enabled": false
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
let cli = AgentConfigSnapshot {
|
||||
computer_use_enabled: Some(true),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let merged = super::merge_with_precedence(Some(&loaded), cli);
|
||||
|
||||
assert_eq!(merged.computer_use_enabled, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_precedence_file_computer_use_enabled_when_cli_none() {
|
||||
let contents = json!({
|
||||
"computer_use_enabled": true
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let file = write_temp(".json", &contents);
|
||||
let loaded = super::load_config_file(file.path()).unwrap();
|
||||
|
||||
let cli = AgentConfigSnapshot::default();
|
||||
|
||||
let merged = super::merge_with_precedence(Some(&loaded), cli);
|
||||
|
||||
assert_eq!(merged.computer_use_enabled, Some(true));
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,313 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Context;
|
||||
use base64::{engine::general_purpose, Engine};
|
||||
use futures::future::join_all;
|
||||
use futures::TryStreamExt as _;
|
||||
use mime_guess::from_path;
|
||||
use tokio::fs;
|
||||
use tokio_util::io::StreamReader;
|
||||
use warp_core::features::FeatureFlag;
|
||||
|
||||
use crate::ai::agent_sdk::retry::with_bounded_retry;
|
||||
use crate::ai::ambient_agents::task::{AttachmentInput, TaskAttachment};
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::ai::attachment_utils::MAX_ATTACHMENT_SIZE_BYTES;
|
||||
use crate::server::server_api::ai::AIClient;
|
||||
use crate::server::server_api::presigned_upload::HttpStatusError;
|
||||
use crate::server::server_api::ServerApi;
|
||||
use crate::util::image::MIN_IMAGE_HEADER_SIZE;
|
||||
|
||||
/// Maximum number of file attachments for a cloud agent task.
|
||||
pub const MAX_ATTACHMENT_COUNT_FOR_CLOUD_QUERY: usize = 25;
|
||||
|
||||
/// Fetches task attachments via GraphQL and downloads them to the filesystem.
|
||||
/// Returns the attachments directory path if any attachments were downloaded,
|
||||
/// so the caller can pass it to the server via `StartFromAmbientRunPrompt`.
|
||||
///
|
||||
/// `attachments_dir` is the per-session directory where files should be downloaded.
|
||||
///
|
||||
/// Makes a best-effort attempt to download all attachments.
|
||||
/// Individual download failures are logged but don't cause the entire function to fail.
|
||||
pub(crate) async fn fetch_and_download_attachments(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
http_client: Arc<ServerApi>,
|
||||
task_id: String,
|
||||
attachments_dir: PathBuf,
|
||||
) -> anyhow::Result<Option<String>> {
|
||||
if !FeatureFlag::AmbientAgentsImageUpload.is_enabled() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let attachments = ai_client
|
||||
.get_task_attachments(task_id.clone())
|
||||
.await
|
||||
.context("Failed to fetch task attachments")?;
|
||||
|
||||
log::info!("Fetched {} task attachments", attachments.len());
|
||||
|
||||
if attachments.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
download_and_write_attachments(attachments, &attachments_dir, &http_client).await?;
|
||||
|
||||
Ok(Some(attachments_dir.to_string_lossy().into_owned()))
|
||||
}
|
||||
|
||||
/// Fetches handoff snapshot attachments for the active execution and downloads
|
||||
/// them into `{attachments_dir}/handoff/{attachment_uuid}` so the runtime's
|
||||
/// rehydration prompt references always point at a file that exists on disk.
|
||||
///
|
||||
/// Returns `Some(attachments_dir)` when at least one attachment wrote to disk, mirroring
|
||||
/// the contract of the sibling [`fetch_and_download_attachments`]. Partial failures are
|
||||
/// logged at WARN level inside this function; per-file errors are not surfaced to callers.
|
||||
///
|
||||
/// Fatal failures (listing the attachments, creating the handoff dir) return `Err`.
|
||||
pub(crate) async fn fetch_and_download_handoff_snapshot_attachments(
|
||||
ai_client: Arc<dyn AIClient>,
|
||||
http_client: &http_client::Client,
|
||||
task_id: AmbientAgentTaskId,
|
||||
attachments_dir: PathBuf,
|
||||
) -> anyhow::Result<Option<String>> {
|
||||
if !FeatureFlag::OzHandoff.is_enabled() {
|
||||
log::error!(
|
||||
"fetch_and_download_handoff_snapshot_attachments called with OzHandoff disabled; \
|
||||
call sites should gate on the flag before invoking"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let attachments = ai_client
|
||||
.get_handoff_snapshot_attachments(&task_id)
|
||||
.await
|
||||
.context("Failed to fetch handoff snapshot attachments")?;
|
||||
|
||||
if attachments.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let handoff_dir = attachments_dir.join("handoff");
|
||||
fs::create_dir_all(&handoff_dir)
|
||||
.await
|
||||
.context("Failed to create handoff attachments directory")?;
|
||||
|
||||
let attempts = attachments.len();
|
||||
let download_futures = attachments.into_iter().map(|attachment| {
|
||||
let file_path = handoff_dir.join(&attachment.file_id);
|
||||
download_handoff_entry(attachment, file_path, http_client)
|
||||
});
|
||||
let results = join_all(download_futures).await;
|
||||
|
||||
let mut succeeded: usize = 0;
|
||||
let mut failures: Vec<(String, String)> = Vec::new();
|
||||
for result in results {
|
||||
match result {
|
||||
Ok(()) => succeeded += 1,
|
||||
Err((filename, err)) => failures.push((filename, err)),
|
||||
}
|
||||
}
|
||||
|
||||
if failures.is_empty() {
|
||||
log::info!("Handoff snapshot attachments: {succeeded}/{attempts} downloaded");
|
||||
} else {
|
||||
let detail = failures
|
||||
.iter()
|
||||
.map(|(filename, err)| format!("{filename}: {err}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("; ");
|
||||
log::warn!(
|
||||
"Handoff snapshot attachments: {succeeded}/{attempts} downloaded; {} failed ({detail})",
|
||||
failures.len()
|
||||
);
|
||||
}
|
||||
|
||||
// Only surface the attachments dir if at least one file made it to disk. Passing a dir
|
||||
// with zero usable entries downstream would make the rehydration prompt reference a
|
||||
// phantom path.
|
||||
if succeeded == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(attachments_dir.to_string_lossy().into_owned()))
|
||||
}
|
||||
|
||||
/// Downloads task attachments from presigned URLs and writes them to the filesystem.
|
||||
/// Downloads are performed concurrently using `join_all`.
|
||||
/// Makes a best-effort attempt to download all attachments, logging warnings for failures.
|
||||
/// The filename is already formatted by the server with UUID prefix (e.g., "uuid_filename.png").
|
||||
async fn download_and_write_attachments(
|
||||
attachments: Vec<TaskAttachment>,
|
||||
attachment_dir: &Path,
|
||||
http_client: &ServerApi,
|
||||
) -> anyhow::Result<()> {
|
||||
fs::create_dir_all(attachment_dir)
|
||||
.await
|
||||
.context("Failed to create attachments directory")?;
|
||||
log::info!(
|
||||
"Created attachments directory at: {}",
|
||||
attachment_dir.display()
|
||||
);
|
||||
|
||||
let http = http_client.http_client();
|
||||
let download_futures = attachments
|
||||
.into_iter()
|
||||
.map(|attachment| download_task_attachment(attachment, attachment_dir, http));
|
||||
let results = join_all(download_futures).await;
|
||||
|
||||
let mut successful = 0;
|
||||
let mut failed = 0;
|
||||
for result in results {
|
||||
match result {
|
||||
Ok(()) => successful += 1,
|
||||
Err(_) => failed += 1,
|
||||
}
|
||||
}
|
||||
|
||||
log::info!("Attachment download complete: {successful} successful, {failed} failed");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Download a single task attachment into `attachment_dir/<sanitized filename>`.
|
||||
///
|
||||
/// Delegates to [`download_attachment`] so transient failures retry on the shared schedule.
|
||||
async fn download_task_attachment(
|
||||
attachment: TaskAttachment,
|
||||
attachment_dir: &Path,
|
||||
http_client: &http_client::Client,
|
||||
) -> anyhow::Result<()> {
|
||||
let safe_filename = Path::new(&attachment.filename)
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("Invalid filename for file_id={}", attachment.file_id))?
|
||||
.to_string();
|
||||
|
||||
let file_path = attachment_dir.join(&safe_filename);
|
||||
log::info!(
|
||||
"Downloading attachment: {} -> {}",
|
||||
attachment.filename,
|
||||
file_path.display()
|
||||
);
|
||||
|
||||
download_attachment(http_client, &attachment.download_url, &file_path).await?;
|
||||
|
||||
log::info!("Successfully wrote attachment to: {}", file_path.display());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Download a single handoff attachment into `file_path`, mapping failure to
|
||||
/// `(filename, error_message)` so the aggregator in
|
||||
/// [`fetch_and_download_handoff_snapshot_attachments`] can log and count per-file outcomes.
|
||||
async fn download_handoff_entry(
|
||||
attachment: TaskAttachment,
|
||||
file_path: PathBuf,
|
||||
http_client: &http_client::Client,
|
||||
) -> Result<(), (String, String)> {
|
||||
// Factor `file_id` and `download_url` out before the retry closure so `attachment` is fully
|
||||
// consumed up-front. The closure borrows the two fields it needs as references.
|
||||
let TaskAttachment {
|
||||
file_id,
|
||||
download_url,
|
||||
..
|
||||
} = attachment;
|
||||
download_attachment(http_client, &download_url, &file_path)
|
||||
.await
|
||||
.map_err(|e| (file_id, format!("{e:#}")))
|
||||
}
|
||||
|
||||
/// Shared download primitive: GET `download_url`, write the body to `file_path`, and retry
|
||||
/// transient HTTP failures on the shared bounded-backoff schedule. Non-2xx responses surface
|
||||
/// an [`HttpStatusError`] so the retry classifier can decide whether to retry.
|
||||
async fn download_attachment(
|
||||
http_client: &http_client::Client,
|
||||
download_url: &str,
|
||||
file_path: &Path,
|
||||
) -> anyhow::Result<()> {
|
||||
let operation = format!("download attachment '{}'", file_path.display());
|
||||
with_bounded_retry(&operation, || async {
|
||||
let response = http_client
|
||||
.get(download_url)
|
||||
.send()
|
||||
.await
|
||||
.context("Failed to send download request")?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(anyhow::Error::new(HttpStatusError {
|
||||
status: status.as_u16(),
|
||||
body: body.clone(),
|
||||
})
|
||||
.context(format!("Download failed with status {status}: {body}")));
|
||||
}
|
||||
|
||||
// Stream the response body directly to disk instead of buffering the full payload
|
||||
// in memory.
|
||||
let mut file = fs::File::create(file_path)
|
||||
.await
|
||||
.context("Failed to create file")?;
|
||||
let mut response_stream =
|
||||
StreamReader::new(response.bytes_stream().map_err(std::io::Error::other));
|
||||
tokio::io::copy(&mut response_stream, &mut file)
|
||||
.await
|
||||
.context("Failed to write file")?;
|
||||
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Process a file attachment for ambient agent upload.
|
||||
/// Returns AttachmentInput with base64-encoded data.
|
||||
/// All file types share the same 10MB size limit.
|
||||
pub fn process_attachment(
|
||||
attachment_path: &PathBuf,
|
||||
index: usize,
|
||||
) -> anyhow::Result<AttachmentInput> {
|
||||
let file_bytes = std::fs::read(attachment_path).map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to read attachment file '{}': {e}",
|
||||
attachment_path.display()
|
||||
)
|
||||
})?;
|
||||
|
||||
// Detect MIME type from file data using infer crate, fall back to file extension
|
||||
let mime_type = if file_bytes.len() >= MIN_IMAGE_HEADER_SIZE {
|
||||
infer::get(&file_bytes).map(|kind| kind.mime_type().to_string())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// If infer couldn't detect, fall back to file extension
|
||||
let mime_type = mime_type.unwrap_or_else(|| {
|
||||
from_path(attachment_path)
|
||||
.first_or_octet_stream()
|
||||
.to_string()
|
||||
});
|
||||
|
||||
if file_bytes.len() > MAX_ATTACHMENT_SIZE_BYTES {
|
||||
return Err(anyhow::anyhow!(
|
||||
"File is too large ({}MB). Maximum size is 10MB.",
|
||||
file_bytes.len() / (1024 * 1024)
|
||||
));
|
||||
}
|
||||
|
||||
let base64_data = general_purpose::STANDARD.encode(&file_bytes);
|
||||
|
||||
let file_name = attachment_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| format!("task_attachment_{index}"));
|
||||
|
||||
Ok(AttachmentInput {
|
||||
file_name,
|
||||
mime_type,
|
||||
data: base64_data,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "attachments_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,389 @@
|
||||
use std::fs;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
|
||||
use mockito::{Matcher, Server};
|
||||
use tempfile::{Builder as TempDirBuilder, NamedTempFile, TempDir};
|
||||
|
||||
use super::*;
|
||||
use crate::ai::agent_sdk::test_support::build_test_http_client;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::server::server_api::ai::MockAIClient;
|
||||
|
||||
#[test]
|
||||
fn process_attachment_text_file() {
|
||||
let mut f = NamedTempFile::with_suffix(".txt").unwrap();
|
||||
write!(f, "hello world").unwrap();
|
||||
|
||||
let result = process_attachment(&f.path().to_path_buf(), 0).unwrap();
|
||||
assert_eq!(
|
||||
result.file_name,
|
||||
f.path().file_name().unwrap().to_str().unwrap()
|
||||
);
|
||||
assert_eq!(result.mime_type, "text/plain");
|
||||
assert_eq!(
|
||||
general_purpose::STANDARD.decode(&result.data).unwrap(),
|
||||
b"hello world"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn process_attachment_too_large() {
|
||||
let mut f = NamedTempFile::with_suffix(".bin").unwrap();
|
||||
let data = vec![0u8; MAX_ATTACHMENT_SIZE_BYTES + 1];
|
||||
f.write_all(&data).unwrap();
|
||||
|
||||
let err = process_attachment(&f.path().to_path_buf(), 0).unwrap_err();
|
||||
assert!(err.to_string().contains("too large"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn process_attachment_nonexistent_file() {
|
||||
let path = std::path::PathBuf::from("/tmp/nonexistent_attachment_test_file.xyz");
|
||||
let err = process_attachment(&path, 0).unwrap_err();
|
||||
assert!(err.to_string().contains("Failed to read"));
|
||||
}
|
||||
|
||||
// End-to-end handoff snapshot download tests. Each test drives the real
|
||||
// `fetch_and_download_handoff_snapshot_attachments` pipeline (including the shared
|
||||
// `with_bounded_retry` helper) against a `mockito::Server` + a real `http_client::Client`,
|
||||
// with `MockAIClient` stubbing only the listing call.
|
||||
|
||||
fn handoff_tempdir() -> TempDir {
|
||||
TempDirBuilder::new()
|
||||
.prefix("handoff-test")
|
||||
.tempdir()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Construct a `TaskAttachment` pointing at `{server_base_url}/download/{file_id}` so each test
|
||||
/// can register matching mocks without fishing for URLs.
|
||||
fn make_attachment(server_base_url: &str, file_id: &str, filename: &str) -> TaskAttachment {
|
||||
TaskAttachment {
|
||||
file_id: file_id.to_string(),
|
||||
filename: filename.to_string(),
|
||||
download_url: format!("{server_base_url}/download/{file_id}"),
|
||||
mime_type: "application/octet-stream".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Regex path matcher for `/download/<file_id>`.
|
||||
fn download_path(file_id: &str) -> Matcher {
|
||||
Matcher::Regex(format!("^/download/{file_id}$"))
|
||||
}
|
||||
|
||||
/// Pre-parsed task id for tests that exercise the outer function. Any valid UUID works; the
|
||||
/// mocked `AIClient` consumes this opaquely.
|
||||
fn fake_task_id() -> AmbientAgentTaskId {
|
||||
"550e8400-e29b-41d4-a716-446655440000".parse().unwrap()
|
||||
}
|
||||
|
||||
/// Build a `MockAIClient` whose `get_handoff_snapshot_attachments` returns `attachments`.
|
||||
fn mock_client_returning(attachments: Vec<TaskAttachment>) -> Arc<MockAIClient> {
|
||||
let mut mock = MockAIClient::new();
|
||||
mock.expect_get_handoff_snapshot_attachments()
|
||||
.times(1)
|
||||
.returning(move |_task_id| Ok(attachments.clone()));
|
||||
Arc::new(mock)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_happy_path_downloads_all_and_writes_to_disk() {
|
||||
// Two attachments, both served 200 with distinct payloads. The pipeline must write each
|
||||
// byte stream to `{attachments_dir}/handoff/{file_id}` and report the dir back.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
let mut server = Server::new_async().await;
|
||||
|
||||
let first_mock = server
|
||||
.mock("GET", download_path("alpha-uuid"))
|
||||
.with_status(200)
|
||||
.with_body("alpha-body")
|
||||
.expect(1)
|
||||
.create_async()
|
||||
.await;
|
||||
let second_mock = server
|
||||
.mock("GET", download_path("beta-uuid"))
|
||||
.with_status(200)
|
||||
.with_body("beta-body")
|
||||
.expect(1)
|
||||
.create_async()
|
||||
.await;
|
||||
|
||||
let attachments = vec![
|
||||
make_attachment(&server.url(), "alpha-uuid", "alpha.patch"),
|
||||
make_attachment(&server.url(), "beta-uuid", "beta.patch"),
|
||||
];
|
||||
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
mock_client_returning(attachments),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("should not be fatal");
|
||||
|
||||
assert_eq!(result.as_deref(), Some(&*attachments_dir.to_string_lossy()));
|
||||
assert_eq!(
|
||||
fs::read(attachments_dir.join("handoff").join("alpha-uuid")).unwrap(),
|
||||
b"alpha-body"
|
||||
);
|
||||
assert_eq!(
|
||||
fs::read(attachments_dir.join("handoff").join("beta-uuid")).unwrap(),
|
||||
b"beta-body"
|
||||
);
|
||||
first_mock.assert();
|
||||
second_mock.assert();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_transient_5xx_retried_then_succeeds() {
|
||||
// Declare the failing 503 mock first, then the success mock. Mockito serves them in
|
||||
// registration order, so attempt #1 hits 503 and attempt #2 hits 200.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
let mut server = Server::new_async().await;
|
||||
|
||||
let flaky_fail = server
|
||||
.mock("GET", download_path("flaky-uuid"))
|
||||
.with_status(503)
|
||||
.with_body("temporarily unavailable")
|
||||
.expect(1)
|
||||
.create_async()
|
||||
.await;
|
||||
let flaky_ok = server
|
||||
.mock("GET", download_path("flaky-uuid"))
|
||||
.with_status(200)
|
||||
.with_body("finally-here")
|
||||
.expect(1)
|
||||
.create_async()
|
||||
.await;
|
||||
|
||||
let attachments = vec![make_attachment(&server.url(), "flaky-uuid", "flaky.patch")];
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
mock_client_returning(attachments),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.as_deref(), Some(&*attachments_dir.to_string_lossy()));
|
||||
assert_eq!(
|
||||
fs::read(attachments_dir.join("handoff").join("flaky-uuid")).unwrap(),
|
||||
b"finally-here"
|
||||
);
|
||||
flaky_fail.assert_async().await;
|
||||
flaky_ok.assert_async().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_permanent_4xx_fails_fast_without_retries() {
|
||||
// 404 is a permanent error; the retry loop must NOT retry. Exactly one GET is expected.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
let mut server = Server::new_async().await;
|
||||
|
||||
let missing = server
|
||||
.mock("GET", download_path("missing-uuid"))
|
||||
.with_status(404)
|
||||
.with_body("not found")
|
||||
.expect(1) // exactly one attempt
|
||||
.create_async()
|
||||
.await;
|
||||
|
||||
let attachments = vec![make_attachment(&server.url(), "missing-uuid", "gone.patch")];
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
mock_client_returning(attachments),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_none(), "no file landed, dir should be None");
|
||||
assert!(
|
||||
!attachments_dir
|
||||
.join("handoff")
|
||||
.join("missing-uuid")
|
||||
.exists(),
|
||||
"no file should be written on permanent failure"
|
||||
);
|
||||
missing.assert_async().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_retry_exhaustion_marks_failed() {
|
||||
// Three persistent 5xxs: the retry loop bails out after MAX_ATTEMPTS (3); no file lands.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
let mut server = Server::new_async().await;
|
||||
|
||||
let persistent = server
|
||||
.mock("GET", download_path("dead-uuid"))
|
||||
.with_status(502)
|
||||
.with_body("bad gateway")
|
||||
.expect(3) // exactly MAX_ATTEMPTS
|
||||
.create_async()
|
||||
.await;
|
||||
|
||||
let attachments = vec![make_attachment(&server.url(), "dead-uuid", "dead.patch")];
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
mock_client_returning(attachments),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_none());
|
||||
assert!(!attachments_dir.join("handoff").join("dead-uuid").exists());
|
||||
persistent.assert_async().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_partial_success_returns_dir_with_downloaded_subset() {
|
||||
// One attachment succeeds, one fails permanently. The dir is returned so the caller can
|
||||
// still see `{handoff_dir}/ok-uuid` downstream; the failed sibling's file is absent.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
let mut server = Server::new_async().await;
|
||||
|
||||
let ok_mock = server
|
||||
.mock("GET", download_path("ok-uuid"))
|
||||
.with_status(200)
|
||||
.with_body("present")
|
||||
.expect(1)
|
||||
.create_async()
|
||||
.await;
|
||||
let bad_mock = server
|
||||
.mock("GET", download_path("bad-uuid"))
|
||||
.with_status(403)
|
||||
.with_body("forbidden")
|
||||
.expect(1)
|
||||
.create_async()
|
||||
.await;
|
||||
|
||||
let attachments = vec![
|
||||
make_attachment(&server.url(), "ok-uuid", "ok.patch"),
|
||||
make_attachment(&server.url(), "bad-uuid", "bad.patch"),
|
||||
];
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
mock_client_returning(attachments),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.as_deref(), Some(&*attachments_dir.to_string_lossy()));
|
||||
assert_eq!(
|
||||
fs::read(attachments_dir.join("handoff").join("ok-uuid")).unwrap(),
|
||||
b"present"
|
||||
);
|
||||
assert!(!attachments_dir.join("handoff").join("bad-uuid").exists());
|
||||
ok_mock.assert_async().await;
|
||||
bad_mock.assert_async().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_empty_attachment_list_returns_none_without_creating_dir() {
|
||||
// With zero attachments listed, the function returns None early and does NOT create the
|
||||
// handoff dir — nothing to land there.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
mock_client_returning(Vec::new()),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("empty list should not be a fatal error");
|
||||
|
||||
assert!(result.is_none());
|
||||
assert!(
|
||||
!attachments_dir.join("handoff").exists(),
|
||||
"handoff dir should not be created when there are no attachments"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_get_handoff_snapshot_attachments_failure_is_fatal() {
|
||||
// When the listing call errors, the function must return Err wrapping the underlying
|
||||
// message with a context describing where it happened.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(true);
|
||||
let tempdir = handoff_tempdir();
|
||||
let http = build_test_http_client();
|
||||
|
||||
let mut mock = MockAIClient::new();
|
||||
mock.expect_get_handoff_snapshot_attachments()
|
||||
.times(1)
|
||||
.returning(|_task_id| Err(anyhow::anyhow!("simulated listing failure")));
|
||||
|
||||
let err = fetch_and_download_handoff_snapshot_attachments(
|
||||
Arc::new(mock),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
tempdir.path().to_path_buf(),
|
||||
)
|
||||
.await
|
||||
.expect_err("listing failure must be fatal");
|
||||
|
||||
let chain: Vec<String> = err.chain().map(|c| c.to_string()).collect();
|
||||
assert!(
|
||||
chain
|
||||
.iter()
|
||||
.any(|s| s.contains("Failed to fetch handoff snapshot attachments")),
|
||||
"expected context-wrapped error in chain: {chain:?}"
|
||||
);
|
||||
assert!(
|
||||
chain
|
||||
.iter()
|
||||
.any(|s| s.contains("simulated listing failure")),
|
||||
"expected underlying error in chain: {chain:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn e2e_returns_none_when_oz_handoff_flag_is_disabled() {
|
||||
// With the feature flag off, the function short-circuits to None without calling the
|
||||
// AIClient. Any call site that forgot to gate on the flag would log an error; here we
|
||||
// just verify the return value.
|
||||
let _guard = FeatureFlag::OzHandoff.override_enabled(false);
|
||||
let tempdir = handoff_tempdir();
|
||||
let attachments_dir = tempdir.path().to_path_buf();
|
||||
let http = build_test_http_client();
|
||||
|
||||
// No expect_get_handoff_snapshot_attachments: if the function calls it, the mock panics.
|
||||
let mock = MockAIClient::new();
|
||||
let result = fetch_and_download_handoff_snapshot_attachments(
|
||||
Arc::new(mock),
|
||||
&http,
|
||||
fake_task_id(),
|
||||
attachments_dir.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("flag-disabled path should not be fatal");
|
||||
|
||||
assert!(result.is_none());
|
||||
assert!(!attachments_dir.join("handoff").exists());
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
use std::{collections::HashMap, ffi::OsString, future::Future, pin::Pin};
|
||||
|
||||
use anyhow::Error;
|
||||
use warpui::ModelSpawner;
|
||||
|
||||
use super::terminal::TerminalDriver;
|
||||
use crate::ai::cloud_environments::ProvidersConfig;
|
||||
|
||||
mod aws;
|
||||
mod gcp;
|
||||
|
||||
pub(crate) type Result<T> = std::result::Result<T, CloudProviderSetupError>;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("{provider_name} setup failed")]
|
||||
pub(crate) struct CloudProviderSetupError {
|
||||
provider_name: &'static str,
|
||||
#[source]
|
||||
source: Error,
|
||||
}
|
||||
|
||||
impl CloudProviderSetupError {
|
||||
pub(crate) fn new(provider_name: &'static str, source: impl Into<Error>) -> Self {
|
||||
Self {
|
||||
provider_name,
|
||||
source: source.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A cloud provider that we configure automatic Oz access to.
|
||||
pub(crate) trait CloudProvider: Send {
|
||||
/// Return environment variables that should be injected into the terminal
|
||||
/// session.
|
||||
fn env_vars(&self) -> Result<HashMap<OsString, OsString>>;
|
||||
|
||||
/// Perform any async setup that requires the terminal session to be running.
|
||||
fn setup(
|
||||
&mut self,
|
||||
_spawner: ModelSpawner<TerminalDriver>,
|
||||
) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
/// Best-effort cleanup of any resources created during setup.
|
||||
///
|
||||
/// The default implementation is a no-op.
|
||||
fn cleanup(self: Box<Self>) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the set of cloud providers from an environment's provider configuration.
|
||||
pub(crate) fn load_providers(
|
||||
providers: &ProvidersConfig,
|
||||
run_id: &str,
|
||||
) -> Result<Vec<Box<dyn CloudProvider>>> {
|
||||
let mut result: Vec<Box<dyn CloudProvider>> = Vec::new();
|
||||
|
||||
if let Some(aws) = &providers.aws {
|
||||
result.push(Box::new(aws::AwsCloudProvider::new(aws, run_id)?));
|
||||
}
|
||||
|
||||
if let Some(gcp) = &providers.gcp {
|
||||
result.push(Box::new(gcp::GcpCloudProvider::new(gcp, run_id)?));
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Collect all environment variables from a list of providers.
|
||||
pub(crate) fn collect_env_vars(
|
||||
providers: &[Box<dyn CloudProvider>],
|
||||
vars: &mut HashMap<OsString, OsString>,
|
||||
) -> Result<()> {
|
||||
for provider in providers {
|
||||
vars.extend(provider.env_vars()?);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "cloud_provider_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,130 @@
|
||||
use std::{collections::HashMap, ffi::OsString, future::Future, pin::Pin, time::Duration};
|
||||
|
||||
use anyhow::Context;
|
||||
use tempfile::{Builder, NamedTempFile};
|
||||
use vec1::Vec1;
|
||||
use warp_core::safe_info;
|
||||
use warp_managed_secrets::ManagedSecretManager;
|
||||
use warpui::{ModelSpawner, SingletonEntity};
|
||||
|
||||
use crate::ai::aws_credentials::aws_role_session_name;
|
||||
use crate::ai::cloud_environments::AwsProviderConfig;
|
||||
|
||||
use super::super::terminal::TerminalDriver;
|
||||
use super::{CloudProvider, CloudProviderSetupError, Result};
|
||||
|
||||
/// Default duration for OIDC identity tokens issued for cloud provider auth.
|
||||
/// The AWS CLI doesn't offer a mechanism for refreshing web identity tokens, so we
|
||||
/// set this to the current maximum task duration.
|
||||
const IDENTITY_TOKEN_DURATION: Duration = Duration::from_hours(3);
|
||||
|
||||
/// AWS STS audience for Warp Oz OIDC federation.
|
||||
const AWS_AUDIENCE: &str = "sts.amazonaws.com";
|
||||
|
||||
/// Provides AWS Web Identity credentials for the agent session.
|
||||
pub(crate) struct AwsCloudProvider {
|
||||
/// ARN of the role to assume.
|
||||
role_arn: String,
|
||||
session_name: String,
|
||||
/// File containing the OIDC token that the AWS CLI will use to assume the role.
|
||||
token_file: NamedTempFile,
|
||||
}
|
||||
|
||||
impl AwsCloudProvider {
|
||||
const PROVIDER_NAME: &'static str = "aws";
|
||||
|
||||
pub fn new(config: &AwsProviderConfig, run_id: &str) -> Result<Self> {
|
||||
// The `tempfile` crate defaults to creating temporary files with user-only permissions.
|
||||
let token_file = Builder::new()
|
||||
.prefix(&format!("oz_aws_oidc_{run_id}_"))
|
||||
.suffix(".token")
|
||||
.tempfile()
|
||||
.context("Failed to create temporary AWS OIDC token file")
|
||||
.map_err(|error| CloudProviderSetupError::new(Self::PROVIDER_NAME, error))?;
|
||||
|
||||
Ok(Self {
|
||||
role_arn: config.role_arn.clone(),
|
||||
session_name: aws_role_session_name(run_id),
|
||||
token_file,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl CloudProvider for AwsCloudProvider {
|
||||
fn env_vars(&self) -> Result<HashMap<OsString, OsString>> {
|
||||
// Set variables that the AWS CLI and SDKs check for assuming a role with web identity:
|
||||
// https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-role.html#cli-configure-role-oidc
|
||||
let mut vars = HashMap::new();
|
||||
vars.insert(
|
||||
OsString::from("AWS_ROLE_ARN"),
|
||||
OsString::from(&self.role_arn),
|
||||
);
|
||||
vars.insert(
|
||||
OsString::from("AWS_ROLE_SESSION_NAME"),
|
||||
OsString::from(&self.session_name),
|
||||
);
|
||||
vars.insert(
|
||||
OsString::from("AWS_WEB_IDENTITY_TOKEN_FILE"),
|
||||
self.token_file.path().as_os_str().to_owned(),
|
||||
);
|
||||
Ok(vars)
|
||||
}
|
||||
|
||||
fn setup(
|
||||
&mut self,
|
||||
spawner: ModelSpawner<TerminalDriver>,
|
||||
) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
|
||||
Box::pin(async move {
|
||||
let token_file_path = self.token_file.path();
|
||||
safe_info!(
|
||||
safe: ("Setting up AWS cloud provider credentials"),
|
||||
full: ("Setting up AWS cloud provider: role_arn={}, session={}", self.role_arn, self.session_name)
|
||||
);
|
||||
|
||||
// 1. Issue an OIDC identity token.
|
||||
let audience = AWS_AUDIENCE.to_string();
|
||||
let duration = IDENTITY_TOKEN_DURATION;
|
||||
|
||||
// Use the scoped principal as the subject, since AWS can't match directly
|
||||
// on the team claim.
|
||||
let subject_template = Vec1::new("scoped_principal".into());
|
||||
let token = spawner
|
||||
.spawn(move |_, ctx| {
|
||||
ManagedSecretManager::handle(ctx)
|
||||
.as_ref(ctx)
|
||||
.issue_task_identity_token(
|
||||
warp_managed_secrets::client::IdentityTokenOptions {
|
||||
audience,
|
||||
requested_duration: duration,
|
||||
subject_template,
|
||||
},
|
||||
)
|
||||
})
|
||||
.await
|
||||
.map_err(|err| CloudProviderSetupError::new(Self::PROVIDER_NAME, err))?
|
||||
.await
|
||||
.map_err(|err| CloudProviderSetupError::new(Self::PROVIDER_NAME, err))?;
|
||||
|
||||
// 2. Write the token to the pre-created temporary file.
|
||||
async_fs::write(&token_file_path, token.token.as_bytes())
|
||||
.await
|
||||
.map_err(|err| CloudProviderSetupError::new(Self::PROVIDER_NAME, err))?;
|
||||
|
||||
safe_info!(
|
||||
safe: ("AWS cloud provider setup complete"),
|
||||
full: ("AWS cloud provider setup complete: token_file={}", token_file_path.display())
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn cleanup(self: Box<Self>) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
|
||||
Box::pin(async move {
|
||||
let Self { token_file, .. } = *self;
|
||||
token_file
|
||||
.close()
|
||||
.context("Failed to remove AWS OIDC token file")
|
||||
.map_err(|err| CloudProviderSetupError::new(Self::PROVIDER_NAME, err))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
use std::{collections::HashMap, ffi::OsString, future::Future, pin::Pin, time::Duration};
|
||||
|
||||
use anyhow::Context as _;
|
||||
use warp_managed_secrets::{GcpCredentials, GcpFederationConfig};
|
||||
|
||||
use crate::ai::cloud_environments::GcpProviderConfig;
|
||||
|
||||
use super::{CloudProvider, CloudProviderSetupError, Result};
|
||||
|
||||
/// Token lifetime for GCP executable-sourced credentials. The GCP client
|
||||
/// libraries handle refreshing automatically, so we keep this short.
|
||||
const TOKEN_LIFETIME: Duration = Duration::from_secs(30 * 60);
|
||||
|
||||
/// Provides GCP Workload Identity Federation credentials for the agent session.
|
||||
///
|
||||
/// The credential config file is written eagerly during construction. GCP SDKs
|
||||
/// discover it via `GOOGLE_APPLICATION_CREDENTIALS` and invoke the embedded
|
||||
/// executable to obtain tokens on demand.
|
||||
pub(crate) struct GcpCloudProvider {
|
||||
credentials: GcpCredentials,
|
||||
}
|
||||
|
||||
impl GcpCloudProvider {
|
||||
const PROVIDER_NAME: &'static str = "gcp";
|
||||
|
||||
pub fn new(config: &GcpProviderConfig, run_id: &str) -> Result<Self> {
|
||||
let federation_config = GcpFederationConfig {
|
||||
project_number: config.project_number.clone(),
|
||||
pool_id: config.workload_identity_federation_pool_id.clone(),
|
||||
provider_id: config.workload_identity_federation_provider_id.clone(),
|
||||
service_account_email: config.service_account_email.clone(),
|
||||
token_lifetime: Some(TOKEN_LIFETIME),
|
||||
};
|
||||
|
||||
let credentials = GcpCredentials::federated(run_id, &federation_config)
|
||||
.context("Failed to prepare GCP federation credentials")
|
||||
.map_err(|error| CloudProviderSetupError::new(Self::PROVIDER_NAME, error))?;
|
||||
|
||||
Ok(Self { credentials })
|
||||
}
|
||||
}
|
||||
|
||||
impl CloudProvider for GcpCloudProvider {
|
||||
fn env_vars(&self) -> Result<HashMap<OsString, OsString>> {
|
||||
Ok(self.credentials.env_vars())
|
||||
}
|
||||
|
||||
fn cleanup(self: Box<Self>) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
|
||||
Box::pin(async move {
|
||||
self.credentials
|
||||
.cleanup()
|
||||
.context("Failed to remove GCP credential files")
|
||||
.map_err(|err| CloudProviderSetupError::new(Self::PROVIDER_NAME, err))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
use std::{collections::HashMap, ffi::OsString, path::PathBuf};
|
||||
|
||||
use crate::ai::cloud_environments::{AwsProviderConfig, GcpProviderConfig, ProvidersConfig};
|
||||
|
||||
use super::{
|
||||
aws::AwsCloudProvider, collect_env_vars, gcp::GcpCloudProvider, load_providers, CloudProvider,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn aws_provider_env_vars_before_setup() {
|
||||
let config = AwsProviderConfig {
|
||||
role_arn: "arn:aws:iam::123456789012:role/MyRole".to_string(),
|
||||
};
|
||||
let provider = AwsCloudProvider::new(&config, "abc-123").unwrap();
|
||||
let vars = provider.env_vars().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
vars.get(&OsString::from("AWS_ROLE_ARN")),
|
||||
Some(&OsString::from("arn:aws:iam::123456789012:role/MyRole"))
|
||||
);
|
||||
assert_eq!(
|
||||
vars.get(&OsString::from("AWS_ROLE_SESSION_NAME")),
|
||||
Some(&OsString::from("Oz_Run_abc-123"))
|
||||
);
|
||||
let token_file = PathBuf::from(
|
||||
vars.get(&OsString::from("AWS_WEB_IDENTITY_TOKEN_FILE"))
|
||||
.unwrap(),
|
||||
);
|
||||
assert!(token_file
|
||||
.extension()
|
||||
.is_some_and(|extension| extension == "token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_cloud_providers_empty_when_no_providers() {
|
||||
let config = ProvidersConfig {
|
||||
gcp: None,
|
||||
aws: None,
|
||||
};
|
||||
let providers = load_providers(&config, "run-1").unwrap();
|
||||
assert!(providers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_cloud_providers_creates_aws_provider() {
|
||||
let config = ProvidersConfig {
|
||||
gcp: None,
|
||||
aws: Some(AwsProviderConfig {
|
||||
role_arn: "arn:aws:iam::111111111111:role/TestRole".to_string(),
|
||||
}),
|
||||
};
|
||||
let providers = load_providers(&config, "run-42").unwrap();
|
||||
assert_eq!(providers.len(), 1);
|
||||
|
||||
let vars = providers[0].env_vars().unwrap();
|
||||
assert_eq!(
|
||||
vars.get(&OsString::from("AWS_ROLE_ARN")),
|
||||
Some(&OsString::from("arn:aws:iam::111111111111:role/TestRole"))
|
||||
);
|
||||
assert_eq!(
|
||||
vars.get(&OsString::from("AWS_ROLE_SESSION_NAME")),
|
||||
Some(&OsString::from("Oz_Run_run-42"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gcp_provider_env_vars() {
|
||||
let config = GcpProviderConfig {
|
||||
project_number: "123456789".to_string(),
|
||||
workload_identity_federation_pool_id: "my-pool".to_string(),
|
||||
workload_identity_federation_provider_id: "my-provider".to_string(),
|
||||
service_account_email: None,
|
||||
};
|
||||
let provider = GcpCloudProvider::new(&config, "run-99").unwrap();
|
||||
let vars = provider.env_vars().unwrap();
|
||||
|
||||
assert!(vars.contains_key(&OsString::from("GOOGLE_APPLICATION_CREDENTIALS")));
|
||||
assert_eq!(
|
||||
vars.get(&OsString::from("GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES")),
|
||||
Some(&OsString::from("1"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_cloud_providers_creates_gcp_provider() {
|
||||
let config = ProvidersConfig {
|
||||
gcp: Some(GcpProviderConfig {
|
||||
project_number: "111".to_string(),
|
||||
workload_identity_federation_pool_id: "pool".to_string(),
|
||||
workload_identity_federation_provider_id: "prov".to_string(),
|
||||
service_account_email: None,
|
||||
}),
|
||||
aws: None,
|
||||
};
|
||||
let providers = load_providers(&config, "run-1").unwrap();
|
||||
assert_eq!(providers.len(), 1);
|
||||
|
||||
let vars = providers[0].env_vars().unwrap();
|
||||
assert!(vars.contains_key(&OsString::from("GOOGLE_APPLICATION_CREDENTIALS")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_provider_env_vars_merges_all_providers() {
|
||||
let config = ProvidersConfig {
|
||||
gcp: Some(GcpProviderConfig {
|
||||
project_number: "222".to_string(),
|
||||
workload_identity_federation_pool_id: "pool".to_string(),
|
||||
workload_identity_federation_provider_id: "prov".to_string(),
|
||||
service_account_email: None,
|
||||
}),
|
||||
aws: Some(AwsProviderConfig {
|
||||
role_arn: "arn:aws:iam::999:role/R".to_string(),
|
||||
}),
|
||||
};
|
||||
let providers = load_providers(&config, "id-7").unwrap();
|
||||
assert_eq!(providers.len(), 2);
|
||||
|
||||
let mut vars = HashMap::new();
|
||||
collect_env_vars(&providers, &mut vars).unwrap();
|
||||
|
||||
// AWS variables.
|
||||
assert!(vars.contains_key(&OsString::from("AWS_ROLE_ARN")));
|
||||
assert!(vars.contains_key(&OsString::from("AWS_ROLE_SESSION_NAME")));
|
||||
// GCP variables.
|
||||
assert!(vars.contains_key(&OsString::from("GOOGLE_APPLICATION_CREDENTIALS")));
|
||||
}
|
||||
@@ -0,0 +1,537 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
future::Future,
|
||||
path::{Path, PathBuf},
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use crate::ai::cloud_environments::{AmbientAgentEnvironment, GithubRepo};
|
||||
use crate::terminal::model::session::command_executor::shell_escape_single_quotes;
|
||||
use crate::terminal::shell::ShellType;
|
||||
use ai::index::full_source_code_embedding::manager::{
|
||||
CodebaseIndexManager, CodebaseIndexManagerEvent,
|
||||
};
|
||||
use futures::{channel::oneshot, future::join_all};
|
||||
use repo_metadata::repositories::{DetectedRepositories, RepoDetectionSource};
|
||||
use warp_completer::completer::CommandExitStatus;
|
||||
use warp_core::{command::ExitCode, safe_info, safe_warn};
|
||||
use warpui::{r#async::FutureExt, ModelContext, ModelSpawner, SingletonEntity};
|
||||
|
||||
use super::{terminal::TerminalDriver, AgentDriverError};
|
||||
use warp_cli::agent::Harness;
|
||||
|
||||
const CODEBASE_INDEX_SYNC_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum PrepareEnvironmentError {
|
||||
#[error("Invalid runtime state - please file a bug report.")]
|
||||
InvalidRuntimeState,
|
||||
#[error("Failed to clone {repo_name}")]
|
||||
CloneRepo { repo_name: String },
|
||||
#[error("Failed to run setup command: {command}")]
|
||||
SetupCommand { command: String },
|
||||
#[error("Failed to change directory into {repo_name}")]
|
||||
ChangeDirectory { repo_name: String },
|
||||
#[error("Terminal driver error while preparing environment: {source}")]
|
||||
TerminalDriver { source: AgentDriverError },
|
||||
}
|
||||
|
||||
/// Prepare a cloud agent environment within a terminal session. This will:
|
||||
/// 1. Clone all repositories, skipping any that are already cloned.
|
||||
/// 2. Begin codebase indexing for all repositories (Oz harness only).
|
||||
/// 3. Run any setup commands.
|
||||
/// 4. If there is only one repository, navigate into it.
|
||||
///
|
||||
/// `is_sandbox` tells the preparer that `working_dir` only exists inside a
|
||||
/// Docker sandbox container and therefore the host filesystem can't be used
|
||||
/// for repo detection or indexing. This is an explicit signal from the
|
||||
/// caller rather than a path-prefix inference, so non-sandbox callers that
|
||||
/// happen to pass a path like `/home/agent/...` don't silently flip into
|
||||
/// sandbox-only mode.
|
||||
pub fn prepare_environment(
|
||||
environment: AmbientAgentEnvironment,
|
||||
working_dir: PathBuf,
|
||||
is_sandbox: bool,
|
||||
harness: Harness,
|
||||
ctx: &mut ModelContext<TerminalDriver>,
|
||||
) -> impl Future<Output = Result<(), PrepareEnvironmentError>> {
|
||||
let spawner = ctx.spawner();
|
||||
async move {
|
||||
let AmbientAgentEnvironment {
|
||||
github_repos,
|
||||
setup_commands,
|
||||
..
|
||||
} = environment;
|
||||
|
||||
// Only index the codebase for the Oz harness; third-party harnesses (e.g. Claude)
|
||||
// have their own methods for navigating a codebase.
|
||||
let should_index_codebase = harness == Harness::Oz;
|
||||
let should_subscribe_to_index_updates = should_index_codebase && !github_repos.is_empty();
|
||||
let repo_channels = Arc::new(Mutex::new(HashMap::<PathBuf, oneshot::Sender<()>>::new()));
|
||||
|
||||
if should_subscribe_to_index_updates {
|
||||
subscribe_to_codebase_index_events(&spawner, Arc::clone(&repo_channels)).await?;
|
||||
}
|
||||
|
||||
let result = prepare_environment_impl(
|
||||
&spawner,
|
||||
working_dir.as_path(),
|
||||
is_sandbox,
|
||||
&github_repos,
|
||||
setup_commands,
|
||||
should_index_codebase,
|
||||
Arc::clone(&repo_channels),
|
||||
)
|
||||
.await;
|
||||
|
||||
if should_subscribe_to_index_updates {
|
||||
let _ = spawner
|
||||
.spawn(|_, ctx| {
|
||||
ctx.unsubscribe_from_model(&CodebaseIndexManager::handle(ctx));
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
async fn prepare_environment_impl(
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
working_dir: &Path,
|
||||
is_sandbox: bool,
|
||||
github_repos: &[GithubRepo],
|
||||
setup_commands: Vec<String>,
|
||||
should_index_codebase: bool,
|
||||
repo_channels: Arc<Mutex<HashMap<PathBuf, oneshot::Sender<()>>>>,
|
||||
) -> Result<(), PrepareEnvironmentError> {
|
||||
let working_dir_string = working_dir.to_string_lossy().to_string();
|
||||
|
||||
// Position the session in `working_dir` before running any probes / clones.
|
||||
// Routed through the silent executor so we don't add a user-visible `cd`
|
||||
// block to the blocklist — in the common case (cloud agents) the session
|
||||
// is already cd'd here by its startup dir, so this is a no-op re-cd and
|
||||
// shouldn't appear in the user's terminal history.
|
||||
if !cd_in_terminal_silent(working_dir_string.clone(), spawner).await? {
|
||||
return Err(PrepareEnvironmentError::ChangeDirectory {
|
||||
repo_name: working_dir_string,
|
||||
});
|
||||
}
|
||||
let mut codebase_context_receivers = Vec::new();
|
||||
|
||||
for repo in github_repos {
|
||||
let repo_name = format!("{}/{}", repo.owner, repo.repo);
|
||||
let repo_url = format!("https://github.com/{repo_name}.git");
|
||||
// We do a partial clone here to speed up environment setup time.
|
||||
let command = format!("git clone --filter=tree:0 {repo_url}");
|
||||
|
||||
let repo_dir = working_dir.join(&repo.repo);
|
||||
// Always ask the session whether the repo dir already exists, rather
|
||||
// than stat'ing from the host. The session knows about sandbox-only
|
||||
// paths, and this goes through the silent executor so `test -d` is
|
||||
// not added to the user-visible blocklist. Pass the absolute path
|
||||
// explicitly so the probe doesn't rely on the session's CWD.
|
||||
let dir_exists = terminal_directory_exists(&repo_dir.to_string_lossy(), spawner).await?;
|
||||
|
||||
if dir_exists {
|
||||
safe_warn!(
|
||||
safe: ("We already have a directory with the same repository name in the terminal working directory, skipping clone..."),
|
||||
full: (
|
||||
"We already have a directory with the name {} in the terminal working directory, skipping clone...",
|
||||
repo.repo)
|
||||
);
|
||||
} else {
|
||||
safe_info!(
|
||||
safe: ("Cloning repository via terminal"),
|
||||
full: ("Cloning repository via terminal: {repo_name}")
|
||||
);
|
||||
|
||||
let exit_code = execute_command(command, spawner).await?;
|
||||
if exit_code != 0.into() {
|
||||
return Err(PrepareEnvironmentError::CloneRepo {
|
||||
repo_name: repo_name.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
safe_info!(
|
||||
safe: ("Successfully cloned repository"),
|
||||
full: ("Successfully cloned: {repo_name}")
|
||||
);
|
||||
}
|
||||
|
||||
// Register the repo with DetectedRepositories so that the skill watcher
|
||||
// and other repo-aware subsystems can discover it before the first query.
|
||||
//
|
||||
// TODO(advait): When the remote code server lands for Docker sandboxes,
|
||||
// sandbox-only working directories will be reachable from the host and
|
||||
// we should register + index them here too (likely via a remote-aware
|
||||
// path instead of `detect_possible_git_repo`/`index_directory`, which
|
||||
// both assume a local filesystem). For now, skip so we don't try to
|
||||
// stat paths that only exist inside the sandbox.
|
||||
if is_sandbox {
|
||||
safe_info!(
|
||||
safe: ("Skipping local repo detection for sandbox-only working directory"),
|
||||
full: (
|
||||
"Skipping local repo detection and indexing for sandbox-only working directory {}",
|
||||
working_dir.display()
|
||||
)
|
||||
);
|
||||
} else {
|
||||
let repo_dir_str = repo_dir.to_string_lossy().to_string();
|
||||
let detect_future = spawner
|
||||
.spawn(move |_, ctx| {
|
||||
DetectedRepositories::handle(ctx).update(ctx, |repos, ctx| {
|
||||
repos.detect_possible_git_repo(
|
||||
&repo_dir_str,
|
||||
RepoDetectionSource::CloudEnvironmentPrep,
|
||||
ctx,
|
||||
)
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)?;
|
||||
// Await detection so the repo is registered in DirectoryWatcher
|
||||
// before the agent's first query.
|
||||
if detect_future.await.is_none() {
|
||||
safe_warn!(
|
||||
safe: ("Repository detection returned no path"),
|
||||
full: ("Repository detection returned no path for {}", repo_dir.display())
|
||||
);
|
||||
}
|
||||
|
||||
if should_index_codebase {
|
||||
let receiver = index_repo_codebase(
|
||||
&repo.repo,
|
||||
working_dir,
|
||||
Arc::clone(&repo_channels),
|
||||
spawner,
|
||||
)
|
||||
.await?;
|
||||
if let Some(receiver) = receiver {
|
||||
codebase_context_receivers.push(receiver);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let has_setup_commands = !setup_commands.is_empty();
|
||||
if has_setup_commands {
|
||||
// Set CI=true so setup commands run in a CI-like environment. This should help us run
|
||||
// non-interactive versions of setup commands, as many command line tools recognize the CI
|
||||
// environment variable.
|
||||
execute_command("export CI=true".to_string(), spawner).await?;
|
||||
}
|
||||
|
||||
for command in setup_commands {
|
||||
let command_for_error = command.clone();
|
||||
safe_info!(
|
||||
safe: ("Running setup command"),
|
||||
full: ("Running setup command: {command}")
|
||||
);
|
||||
|
||||
let exit_code = execute_command(command, spawner).await?;
|
||||
if exit_code != 0.into() {
|
||||
return Err(PrepareEnvironmentError::SetupCommand {
|
||||
command: command_for_error,
|
||||
});
|
||||
}
|
||||
|
||||
let working_dir_string = working_dir.to_string_lossy().to_string();
|
||||
if let Err(error) = cd_in_terminal(working_dir_string, spawner).await {
|
||||
log::warn!("Failed to reset working directory after setup command: {error}");
|
||||
}
|
||||
|
||||
safe_info!(
|
||||
safe: ("Successfully completed setup command"),
|
||||
full: ("Successfully completed setup command: {command_for_error}")
|
||||
);
|
||||
}
|
||||
|
||||
if has_setup_commands {
|
||||
// Unset CI after setup commands complete so the agent session
|
||||
// does not run with CI=true.
|
||||
execute_command("unset CI".to_string(), spawner).await?;
|
||||
}
|
||||
|
||||
if !github_repos.is_empty() {
|
||||
// Wait for codebase indexing for all repositories after running setup commands.
|
||||
// We skip this if running in Docker sandboxes since they don't have a cache volume.
|
||||
// We also skip this in Namespace to reduce startup time.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
let should_wait_for_indexing = !matches!(
|
||||
warp_isolation_platform::detect(),
|
||||
Some(
|
||||
warp_isolation_platform::IsolationPlatformType::DockerSandbox
|
||||
| warp_isolation_platform::IsolationPlatformType::Namespace
|
||||
)
|
||||
);
|
||||
#[cfg(target_family = "wasm")]
|
||||
let should_wait_for_indexing = true;
|
||||
|
||||
if should_wait_for_indexing {
|
||||
let repos_indexed = join_all(codebase_context_receivers);
|
||||
if repos_indexed
|
||||
.with_timeout(CODEBASE_INDEX_SYNC_TIMEOUT)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
log::warn!(
|
||||
"Timed out waiting for codebase index sync; continuing without guaranteed codebase context",
|
||||
);
|
||||
}
|
||||
} else {
|
||||
drop(codebase_context_receivers);
|
||||
log::info!("Not waiting for codebase index sync");
|
||||
}
|
||||
}
|
||||
|
||||
// If there's only one repo in the environment, start the agent in that repo.
|
||||
// This way, it doesn't have to locate the correct repo to work on.
|
||||
if let Some(repo_name) = single_repo_name(github_repos) {
|
||||
safe_info!(
|
||||
safe: ("Changing directory into single repository"),
|
||||
full: ("Changing directory into single repository: {repo_name}")
|
||||
);
|
||||
let exit_code = cd_in_terminal(repo_name.clone(), spawner).await?;
|
||||
if exit_code != 0.into() {
|
||||
return Err(PrepareEnvironmentError::ChangeDirectory { repo_name });
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn subscribe_to_codebase_index_events(
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
repo_channels: Arc<Mutex<HashMap<PathBuf, oneshot::Sender<()>>>>,
|
||||
) -> Result<(), PrepareEnvironmentError> {
|
||||
spawner
|
||||
.spawn(move |_, ctx| {
|
||||
let repo_channels = Arc::clone(&repo_channels);
|
||||
ctx.subscribe_to_model(
|
||||
&CodebaseIndexManager::handle(ctx),
|
||||
move |_, event, ctx| {
|
||||
if !matches!(event, CodebaseIndexManagerEvent::SyncStateUpdated) {
|
||||
return;
|
||||
}
|
||||
|
||||
let manager = CodebaseIndexManager::as_ref(ctx);
|
||||
let mut repos_to_notify = Vec::new();
|
||||
let mut channels = repo_channels
|
||||
.lock()
|
||||
.expect("repo channel map lock should not be poisoned");
|
||||
|
||||
for repo in channels.keys() {
|
||||
let Some(status) =
|
||||
manager.get_codebase_index_status_for_path(repo, ctx)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
if status.has_synced_version() {
|
||||
repos_to_notify.push(repo.clone());
|
||||
continue;
|
||||
}
|
||||
|
||||
if !status.has_pending() && status.last_sync_successful() == Some(false) {
|
||||
safe_warn!(
|
||||
safe: ("Codebase index sync failed for a repo; unblocking environment setup"),
|
||||
full: ("Codebase index sync failed for {repo:?}; unblocking environment setup")
|
||||
);
|
||||
repos_to_notify.push(repo.clone());
|
||||
}
|
||||
}
|
||||
|
||||
for repo in repos_to_notify {
|
||||
if let Some(tx) = channels.remove(&repo) {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
})
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)
|
||||
}
|
||||
|
||||
async fn index_repo_codebase(
|
||||
repo_name: &str,
|
||||
working_dir: &Path,
|
||||
repo_channels: Arc<Mutex<HashMap<PathBuf, oneshot::Sender<()>>>>,
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
) -> Result<Option<oneshot::Receiver<()>>, PrepareEnvironmentError> {
|
||||
let repo_path = working_dir.join(repo_name);
|
||||
|
||||
safe_info!(
|
||||
safe: ("Trying to index repository for codebase context"),
|
||||
full: ("Trying to index {:?} for codebase context", repo_path)
|
||||
);
|
||||
|
||||
let repo_path_for_spawn = repo_path.clone();
|
||||
spawner
|
||||
.spawn(move |_, ctx| {
|
||||
CodebaseIndexManager::handle(ctx).update(ctx, |manager, ctx| {
|
||||
manager.index_directory(repo_path_for_spawn.clone(), ctx);
|
||||
});
|
||||
|
||||
let status = CodebaseIndexManager::as_ref(ctx)
|
||||
.get_codebase_index_status_for_path(&repo_path_for_spawn, ctx);
|
||||
|
||||
match status {
|
||||
Some(status) if status.has_synced_version() => {
|
||||
safe_info!(
|
||||
safe: ("Not waiting on codebase index for repository; we have one already"),
|
||||
full: ("Not waiting on codebase index for {:?}, we have one already", repo_path_for_spawn)
|
||||
);
|
||||
None
|
||||
}
|
||||
_ => {
|
||||
safe_info!(
|
||||
safe: ("Waiting on codebase index for repository"),
|
||||
full: ("Waiting on codebase index for {:?}", repo_path_for_spawn)
|
||||
);
|
||||
let (tx, rx) = oneshot::channel::<()>();
|
||||
repo_channels
|
||||
.lock()
|
||||
.expect("repo channel map lock should not be poisoned")
|
||||
.insert(repo_path_for_spawn, tx);
|
||||
Some(rx)
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)
|
||||
}
|
||||
|
||||
/// Execute a command in the context of a terminal session.
|
||||
async fn execute_command(
|
||||
command: String,
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
) -> Result<ExitCode, PrepareEnvironmentError> {
|
||||
spawner
|
||||
.spawn(move |terminal_driver, ctx| terminal_driver.execute_command(&command, ctx))
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)?
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})?
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})?
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})
|
||||
}
|
||||
|
||||
/// Change the current directory in the context of a terminal session (using `cd {dir}`).
|
||||
async fn cd_in_terminal(
|
||||
target: String,
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
) -> Result<ExitCode, PrepareEnvironmentError> {
|
||||
spawner
|
||||
.spawn(move |terminal_driver, ctx| terminal_driver.cd(&target, ctx))
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)?
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})?
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})?
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})
|
||||
}
|
||||
|
||||
fn single_repo_name(repos: &[GithubRepo]) -> Option<String> {
|
||||
if repos.len() != 1 {
|
||||
return None;
|
||||
}
|
||||
Some(repos[0].repo.clone())
|
||||
}
|
||||
|
||||
/// Change the active terminal session's working directory via `cd <target>`,
|
||||
/// silently.
|
||||
///
|
||||
/// Thin wrapper around [`TerminalDriver::cd_silent`] so the call stays
|
||||
/// consistent with the other `*_in_terminal` / `terminal_*` helpers in this
|
||||
/// module. Uses the same [`ShellFamily::shell_escape`] logic as the visible
|
||||
/// [`TerminalDriver::cd`] path, so it's safe across bash/zsh/fish/pwsh host
|
||||
/// shells.
|
||||
///
|
||||
/// Returns `true` if the `cd` exited successfully.
|
||||
async fn cd_in_terminal_silent(
|
||||
target: String,
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
) -> Result<bool, PrepareEnvironmentError> {
|
||||
let output = spawner
|
||||
.spawn(move |driver, ctx| driver.cd_silent(&target, ctx))
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)?
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})?;
|
||||
Ok(output.status == CommandExitStatus::Success)
|
||||
}
|
||||
|
||||
/// Returns whether the given path resolves to an existing directory from the
|
||||
/// perspective of the active terminal session.
|
||||
///
|
||||
/// Runs `test -d <path>` through the session's in-band command executor, so
|
||||
/// the check is invisible in the user-facing blocklist and works for paths
|
||||
/// that only exist inside a remote/sandbox filesystem. The path is escaped
|
||||
/// using the *session's* actual shell type (bash/zsh use the `'"'"'` trick,
|
||||
/// fish uses a backslash, PowerShell doubles the quote) rather than assuming
|
||||
/// bash.
|
||||
///
|
||||
/// Prefer passing an absolute path: relative paths resolve against the
|
||||
/// session's current working directory, which couples the caller to
|
||||
/// whatever `cd` state the session happens to be in.
|
||||
///
|
||||
/// TODO(advait): `test -d ...` itself is POSIX-only. When we support
|
||||
/// environment prep on Windows host shells (PowerShell / cmd.exe), also
|
||||
/// branch on `ShellType` to emit the appropriate probe (e.g.
|
||||
/// `Test-Path -PathType Container <path>` for PowerShell).
|
||||
async fn terminal_directory_exists(
|
||||
path: &str,
|
||||
spawner: &ModelSpawner<TerminalDriver>,
|
||||
) -> Result<bool, PrepareEnvironmentError> {
|
||||
let path = path.to_owned();
|
||||
let output = spawner
|
||||
.spawn(move |driver, ctx| {
|
||||
// Fall back to Bash if the session's shell type isn't known yet
|
||||
// (e.g. pre-bootstrap). Bash-style escaping is a safe default for
|
||||
// every POSIX shell we currently support.
|
||||
let shell_type = driver
|
||||
.active_session_shell_type(ctx)
|
||||
.unwrap_or(ShellType::Bash);
|
||||
let escaped = shell_escape_single_quotes(&path, shell_type);
|
||||
let command = format!("test -d '{escaped}'");
|
||||
driver.execute_silent_command(command, ctx)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)?
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
AgentDriverError::InvalidRuntimeState => PrepareEnvironmentError::InvalidRuntimeState,
|
||||
source => PrepareEnvironmentError::TerminalDriver { source },
|
||||
})?;
|
||||
Ok(output.status == CommandExitStatus::Success)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "environment_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,24 @@
|
||||
use super::single_repo_name;
|
||||
use crate::ai::cloud_environments::GithubRepo;
|
||||
|
||||
#[test]
|
||||
fn single_repo_name_returns_repo_when_exactly_one_repo() {
|
||||
let repos = vec![GithubRepo::new(
|
||||
"warpdotdev".to_string(),
|
||||
"warp-internal".to_string(),
|
||||
)];
|
||||
let selected_repo = single_repo_name(&repos);
|
||||
assert_eq!(selected_repo, Some("warp-internal".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_repo_name_returns_none_for_zero_or_many_repos() {
|
||||
let no_repos = Vec::<GithubRepo>::new();
|
||||
assert_eq!(single_repo_name(&no_repos), None);
|
||||
|
||||
let two_repos = vec![
|
||||
GithubRepo::new("warpdotdev".to_string(), "warp-internal".to_string()),
|
||||
GithubRepo::new("warpdotdev".to_string(), "warp-server".to_string()),
|
||||
];
|
||||
assert_eq!(single_repo_name(&two_repos), None);
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
use crate::ai::blocklist::task_status_sync_model::classify_renderable_error;
|
||||
use crate::server::server_api::ai::TaskStatusUpdate;
|
||||
use warp_graphql::ai::{AgentTaskState, PlatformErrorCode};
|
||||
|
||||
use super::terminal::ShareSessionError;
|
||||
use super::AgentDriverError;
|
||||
|
||||
/// Classify an `AgentDriverError` into a task state and a `TaskStatusUpdate`
|
||||
/// suitable for reporting via `update_agent_task`.
|
||||
pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskStatusUpdate) {
|
||||
match error {
|
||||
// --- Warp-side errors (task → ERROR) ---
|
||||
AgentDriverError::TerminalUnavailable | AgentDriverError::InvalidRuntimeState => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
"An internal error occurred. Please try running your task again. If the issue persists, contact support.",
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::BootstrapFailed => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
"Terminal session failed to start. Please try running your task again.",
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::ShareSessionFailed { error: share_err } => {
|
||||
let message = match share_err {
|
||||
ShareSessionError::Internal(_) => {
|
||||
"Failed to share agent session due to an internal error. Please try running your task again.".to_string()
|
||||
}
|
||||
ShareSessionError::Failed(reason) => {
|
||||
// The reason string comes from the session-sharing layer and is aimed at
|
||||
// interactive users (e.g. "try sharing again"). Provide a cloud-agent-
|
||||
// appropriate message instead of wrapping it, which would produce
|
||||
// repetitive "try again" text.
|
||||
format!("Failed to share agent session: {reason}")
|
||||
}
|
||||
ShareSessionError::Disabled => {
|
||||
"Session sharing is not enabled for your account. This is likely because \
|
||||
an administrator has disabled session sharing for your team. Please \
|
||||
verify that session sharing is enabled in your team settings, or try \
|
||||
running without the --share flag."
|
||||
.to_string()
|
||||
}
|
||||
ShareSessionError::Timeout => {
|
||||
"Failed to share agent session: timed out waiting for the session sharing \
|
||||
server to respond. Please check your network connection and try again."
|
||||
.to_string()
|
||||
}
|
||||
ShareSessionError::Interrupted => {
|
||||
"Session sharing was interrupted before it could complete. Please try running your task again.".to_string()
|
||||
}
|
||||
};
|
||||
(
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
message,
|
||||
match share_err {
|
||||
ShareSessionError::Disabled => PlatformErrorCode::FeatureNotAvailable,
|
||||
_ => PlatformErrorCode::InternalError,
|
||||
},
|
||||
),
|
||||
)
|
||||
}
|
||||
AgentDriverError::WarpDriveSyncFailed => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
"Warp Drive failed to sync. Please check your network connection and try again.",
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::NotLoggedIn => {
|
||||
let bin = warp_cli::binary_name().unwrap_or_else(|| "warp".to_string());
|
||||
(
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Authentication required. Log in via '{bin} login', provide an API key via '--api-key', or set the WARP_API_KEY environment variable."
|
||||
),
|
||||
PlatformErrorCode::AuthenticationRequired,
|
||||
),
|
||||
)
|
||||
}
|
||||
AgentDriverError::CloudProviderSetupFailed(err) => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Error configuring cloud access: {err:#}"),
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
|
||||
// --- User-side errors (task → FAILED) ---
|
||||
AgentDriverError::MCPServerNotFound(uuid) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"MCP server {uuid} was not found. Verify the server exists in your Warp Drive and the UUID is correct."
|
||||
),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::MCPStartupFailed => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
"One or more MCP servers failed to start. Check that your MCP server configuration is valid and the server process is runnable.",
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::MCPJsonParseError(msg) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Failed to parse MCP server JSON configuration: {msg}"),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::MCPMissingVariables => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
"MCP server configuration is missing required variables. Provide all required environment variables or template values.",
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::ProfileError(name) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Agent profile \"{name}\" not found. Check the profile ID and ensure it exists in your team's Warp Drive."
|
||||
),
|
||||
PlatformErrorCode::ResourceNotFound,
|
||||
),
|
||||
),
|
||||
AgentDriverError::AIWorkflowNotFound(id) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Saved prompt not found for ID {id}. Verify the prompt exists in your Warp Drive."
|
||||
),
|
||||
PlatformErrorCode::ResourceNotFound,
|
||||
),
|
||||
),
|
||||
AgentDriverError::EnvironmentNotFound(id) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Environment '{id}' not found. Verify the environment ID and ensure it exists in your team settings."
|
||||
),
|
||||
PlatformErrorCode::ResourceNotFound,
|
||||
),
|
||||
),
|
||||
AgentDriverError::EnvironmentSetupFailed(msg) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Environment setup failed: {msg}. Check your repository URLs and setup commands."
|
||||
),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::InvalidWorkingDirectory { path, .. } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Working directory '{}' does not exist or is not a directory. Verify the path in your environment configuration.",
|
||||
path.display()
|
||||
),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
|
||||
// --- Conversation errors ---
|
||||
// Delegate to classify_renderable_error for proper ERROR vs FAILED
|
||||
// distinction and PlatformErrorCode. This is a belt-and-suspenders
|
||||
// fallback — TaskStatusSyncModel handles most conversation errors,
|
||||
// but the driver catches them too if the conversation ends with an error.
|
||||
AgentDriverError::ConversationError { error } => {
|
||||
let (state, update) = classify_renderable_error(error);
|
||||
(
|
||||
state,
|
||||
update.unwrap_or_else(|| {
|
||||
TaskStatusUpdate::with_error_code(
|
||||
error.to_string(),
|
||||
PlatformErrorCode::InternalError,
|
||||
)
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
// --- Cancellation / Blocked (no error code) ---
|
||||
AgentDriverError::ConversationCancelled { .. } => (
|
||||
AgentTaskState::Cancelled,
|
||||
TaskStatusUpdate::message("Task cancelled."),
|
||||
),
|
||||
AgentDriverError::ConversationBlocked { blocked_action } => (
|
||||
AgentTaskState::Blocked,
|
||||
TaskStatusUpdate::message(format!(
|
||||
"The agent got stuck waiting for user confirmation on the action: {blocked_action}"
|
||||
)),
|
||||
),
|
||||
|
||||
// --- Setup errors ---
|
||||
AgentDriverError::TeamMetadataRefreshTimeout => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
"Timed out refreshing team metadata. Please check your network connection and try again.",
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::SkillResolutionFailed(msg) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Skill resolution failed: {msg}"),
|
||||
PlatformErrorCode::ResourceNotFound,
|
||||
),
|
||||
),
|
||||
AgentDriverError::ConfigBuildFailed(err) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Failed to build agent configuration: {err}"),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::PromptResolutionFailed(err) => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Failed to resolve prompt for the run: {err}"),
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::SecretsFetchFailed(err) => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Failed to fetch task secrets: {err}"),
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::AwsBedrockCredentialsFailed(msg) => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Failed to initialize AWS Bedrock credentials: {msg}"),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::ConversationLoadFailed(msg) => (
|
||||
AgentTaskState::Error,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Failed to load conversation: {msg}"),
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::ConversationHarnessMismatch { conversation_id, expected, got } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Conversation {conversation_id} was produced by the {expected} harness, but --harness {got} was requested. \
|
||||
Re-run with --harness {expected} (or omit --harness) to continue this conversation."
|
||||
),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::TaskHarnessMismatch { task_id, expected, got } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Task {task_id} was created with the {expected} harness, but --harness {got} was requested. \
|
||||
Re-run with --harness {expected} (or omit --harness) to continue this task."
|
||||
),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::ConversationResumeStateMissing { harness, conversation_id } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!(
|
||||
"Conversation {conversation_id} has no stored transcript for the {harness} harness. \
|
||||
The prior run may have crashed before saving any state."
|
||||
),
|
||||
PlatformErrorCode::ResourceNotFound,
|
||||
),
|
||||
),
|
||||
AgentDriverError::HarnessCommandFailed { exit_code } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Harness command exited with code {exit_code}"),
|
||||
PlatformErrorCode::InternalError,
|
||||
),
|
||||
),
|
||||
AgentDriverError::HarnessSetupFailed { harness, reason } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Harness '{harness}' validation failed: {reason}"),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
AgentDriverError::HarnessConfigSetupFailed { harness, error } => (
|
||||
AgentTaskState::Failed,
|
||||
TaskStatusUpdate::with_error_code(
|
||||
format!("Harness '{harness}' config setup failed: {error}"),
|
||||
PlatformErrorCode::EnvironmentSetupFailed,
|
||||
),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "error_classification_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,155 @@
|
||||
use warp_graphql::ai::{AgentTaskState, PlatformErrorCode};
|
||||
|
||||
use super::classify_driver_error;
|
||||
use crate::ai::agent_sdk::driver::terminal::ShareSessionError;
|
||||
use crate::ai::agent_sdk::driver::AgentDriverError;
|
||||
|
||||
fn assert_state_and_code(
|
||||
error: AgentDriverError,
|
||||
expected_state: AgentTaskState,
|
||||
expected_code: Option<PlatformErrorCode>,
|
||||
) {
|
||||
let (state, update) = classify_driver_error(&error);
|
||||
assert_eq!(state, expected_state, "unexpected state for {error}");
|
||||
assert_eq!(
|
||||
update.error_code, expected_code,
|
||||
"unexpected error_code for {error}"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Infrastructure errors → ERROR ---
|
||||
|
||||
#[test]
|
||||
fn bootstrap_failed_is_error_with_internal() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::BootstrapFailed,
|
||||
AgentTaskState::Error,
|
||||
Some(PlatformErrorCode::InternalError),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_unavailable_is_error_with_internal() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::TerminalUnavailable,
|
||||
AgentTaskState::Error,
|
||||
Some(PlatformErrorCode::InternalError),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn not_logged_in_is_error_with_auth_required() {
|
||||
let (state, update) = classify_driver_error(&AgentDriverError::NotLoggedIn);
|
||||
assert_eq!(state, AgentTaskState::Error);
|
||||
assert_eq!(
|
||||
update.error_code,
|
||||
Some(PlatformErrorCode::AuthenticationRequired)
|
||||
);
|
||||
assert!(
|
||||
update.message.contains("WARP_API_KEY"),
|
||||
"message should mention WARP_API_KEY: {:?}",
|
||||
update.message
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warp_drive_sync_failed_is_error() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::WarpDriveSyncFailed,
|
||||
AgentTaskState::Error,
|
||||
Some(PlatformErrorCode::InternalError),
|
||||
);
|
||||
}
|
||||
|
||||
// --- Config/user errors → FAILED ---
|
||||
|
||||
#[test]
|
||||
fn mcp_server_not_found_is_failed_with_env_setup() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::MCPServerNotFound(uuid::Uuid::nil()),
|
||||
AgentTaskState::Failed,
|
||||
Some(PlatformErrorCode::EnvironmentSetupFailed),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn environment_setup_failed_is_failed() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::EnvironmentSetupFailed("bad repo".into()),
|
||||
AgentTaskState::Failed,
|
||||
Some(PlatformErrorCode::EnvironmentSetupFailed),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn profile_error_is_failed_with_resource_not_found() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::ProfileError("my-profile".into()),
|
||||
AgentTaskState::Failed,
|
||||
Some(PlatformErrorCode::ResourceNotFound),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn environment_not_found_is_failed_with_resource_not_found() {
|
||||
assert_state_and_code(
|
||||
AgentDriverError::EnvironmentNotFound("env-123".into()),
|
||||
AgentTaskState::Failed,
|
||||
Some(PlatformErrorCode::ResourceNotFound),
|
||||
);
|
||||
}
|
||||
|
||||
// --- ShareSessionFailed variants ---
|
||||
|
||||
#[test]
|
||||
fn share_session_disabled_gets_feature_not_available() {
|
||||
let (state, update) = classify_driver_error(&AgentDriverError::ShareSessionFailed {
|
||||
error: ShareSessionError::Disabled,
|
||||
});
|
||||
assert_eq!(state, AgentTaskState::Error);
|
||||
assert_eq!(
|
||||
update.error_code,
|
||||
Some(PlatformErrorCode::FeatureNotAvailable)
|
||||
);
|
||||
assert!(update.message.contains("not enabled"));
|
||||
assert!(update.message.contains("--share flag"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn share_session_timeout_gets_internal_error() {
|
||||
let (state, update) = classify_driver_error(&AgentDriverError::ShareSessionFailed {
|
||||
error: ShareSessionError::Timeout,
|
||||
});
|
||||
assert_eq!(state, AgentTaskState::Error);
|
||||
assert_eq!(update.error_code, Some(PlatformErrorCode::InternalError));
|
||||
assert!(update.message.contains("timed out"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn share_session_failed_includes_reason() {
|
||||
let (state, update) = classify_driver_error(&AgentDriverError::ShareSessionFailed {
|
||||
error: ShareSessionError::Failed("server rejected".into()),
|
||||
});
|
||||
assert_eq!(state, AgentTaskState::Error);
|
||||
assert!(update.message.contains("server rejected"));
|
||||
}
|
||||
|
||||
// --- Conversation-level outcomes ---
|
||||
|
||||
#[test]
|
||||
fn conversation_cancelled_is_cancelled() {
|
||||
let (state, update) = classify_driver_error(&AgentDriverError::ConversationCancelled {
|
||||
reason: crate::ai::agent::CancellationReason::ManuallyCancelled,
|
||||
});
|
||||
assert_eq!(state, AgentTaskState::Cancelled);
|
||||
assert!(update.error_code.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_blocked_is_blocked() {
|
||||
let (state, update) = classify_driver_error(&AgentDriverError::ConversationBlocked {
|
||||
blocked_action: "rm -rf /".into(),
|
||||
});
|
||||
assert_eq!(state, AgentTaskState::Blocked);
|
||||
assert!(update.message.contains("rm -rf /"));
|
||||
}
|
||||
@@ -0,0 +1,664 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Write as _;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
use tempfile::NamedTempFile;
|
||||
use uuid::Uuid;
|
||||
use warp_cli::agent::Harness;
|
||||
use warpui::{ModelHandle, ModelSpawner};
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::server::server_api::harness_support::{upload_to_target, HarnessSupportClient};
|
||||
use crate::server::server_api::ServerApi;
|
||||
use crate::terminal::model::block::BlockId;
|
||||
use crate::terminal::model::session::ExecuteCommandOptions;
|
||||
use crate::terminal::CLIAgent;
|
||||
|
||||
use super::super::terminal::{CommandHandle, TerminalDriver};
|
||||
use super::super::{AgentDriver, AgentDriverError};
|
||||
use super::claude_transcript::{
|
||||
claude_config_dir, read_envelope, write_envelope, write_session_index_entry, ClaudeResumeInfo,
|
||||
ClaudeTranscriptEnvelope,
|
||||
};
|
||||
use super::json_utils::{read_json_file_or_default, write_json_file};
|
||||
use super::{
|
||||
write_temp_file, HarnessRunner, ManagedSecretValue, ResumePayload, SavePoint, ThirdPartyHarness,
|
||||
};
|
||||
mod parent_bridge;
|
||||
|
||||
#[cfg(test)]
|
||||
use super::super::OZ_MESSAGE_LISTENER_STATE_ROOT_ENV;
|
||||
use parent_bridge::MessageBridge;
|
||||
#[cfg(test)]
|
||||
use parent_bridge::{
|
||||
acknowledge_parent_bridge_hook_output, ensure_parent_bridge_state_dir,
|
||||
parent_bridge_char_count, parent_bridge_hook_output_ack_file, parent_bridge_hook_output_file,
|
||||
parent_bridge_root, parent_bridge_staged_message_path, parent_bridge_surfaced_message_path,
|
||||
prepare_parent_bridge_hook_output, render_parent_bridge_message_block,
|
||||
stage_parent_bridge_message, MessageBridgeHookOutput, MessageBridgeMessageRecord,
|
||||
MESSAGE_BRIDGE_CONTEXT_PREAMBLE,
|
||||
};
|
||||
|
||||
pub(crate) struct ClaudeHarness;
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl ThirdPartyHarness for ClaudeHarness {
|
||||
fn harness(&self) -> Harness {
|
||||
Harness::Claude
|
||||
}
|
||||
|
||||
fn cli_agent(&self) -> CLIAgent {
|
||||
CLIAgent::Claude
|
||||
}
|
||||
|
||||
fn install_docs_url(&self) -> Option<&'static str> {
|
||||
Some("https://code.claude.com/docs/en/quickstart")
|
||||
}
|
||||
|
||||
fn prepare_environment_config(
|
||||
&self,
|
||||
working_dir: &Path,
|
||||
_system_prompt: Option<&str>,
|
||||
secrets: &HashMap<String, ManagedSecretValue>,
|
||||
) -> Result<(), AgentDriverError> {
|
||||
prepare_claude_environment_config(working_dir, secrets).map_err(|error| {
|
||||
AgentDriverError::HarnessConfigSetupFailed {
|
||||
harness: self.cli_agent().command_prefix().to_owned(),
|
||||
error,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch the Claude Code transcript for the current task's conversation and wrap it
|
||||
/// into a [`ResumePayload::Claude`]. Maps a server 404 to
|
||||
/// [`AgentDriverError::ConversationResumeStateMissing`] tagged as the `claude` harness
|
||||
/// so the user sees a resume-specific error rather than a generic load failure.
|
||||
async fn fetch_resume_payload(
|
||||
&self,
|
||||
conversation_id: &AIConversationId,
|
||||
harness_support_client: Arc<dyn HarnessSupportClient>,
|
||||
) -> Result<Option<ResumePayload>, AgentDriverError> {
|
||||
let conversation_id_str = conversation_id.to_string();
|
||||
let bytes = harness_support_client
|
||||
.fetch_transcript()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
// A 404 from the server maps to "no stored transcript" so the CLI can tell
|
||||
// the user the prior run never saved state.
|
||||
let message = format!("{err:#}").to_lowercase();
|
||||
if message.contains("status 404") {
|
||||
AgentDriverError::ConversationResumeStateMissing {
|
||||
harness: "claude".to_string(),
|
||||
conversation_id: conversation_id_str.clone(),
|
||||
}
|
||||
} else {
|
||||
AgentDriverError::ConversationLoadFailed(format!("{err:#}"))
|
||||
}
|
||||
})?;
|
||||
let envelope: ClaudeTranscriptEnvelope = serde_json::from_slice(&bytes).map_err(|err| {
|
||||
AgentDriverError::ConversationLoadFailed(format!(
|
||||
"Failed to deserialize Claude transcript for {conversation_id_str}: {err:#}"
|
||||
))
|
||||
})?;
|
||||
let session_id = envelope.uuid;
|
||||
Ok(Some(ResumePayload::Claude(ClaudeResumeInfo {
|
||||
conversation_id: *conversation_id,
|
||||
session_id,
|
||||
envelope,
|
||||
})))
|
||||
}
|
||||
|
||||
fn build_runner(
|
||||
&self,
|
||||
prompt: &str,
|
||||
system_prompt: Option<&str>,
|
||||
resumption_prompt: Option<&str>,
|
||||
working_dir: &Path,
|
||||
task_id: Option<AmbientAgentTaskId>,
|
||||
server_api: Arc<ServerApi>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
resume: Option<ResumePayload>,
|
||||
) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
|
||||
// Extract the Claude variant; any other variant is ignored since it belongs to a
|
||||
// different harness. Today there are no other variants, but this keeps the shape
|
||||
// ready for future CLI-specific payloads.
|
||||
let claude_resume = resume.map(|payload| match payload {
|
||||
ResumePayload::Claude(info) => info,
|
||||
});
|
||||
// Claude treats the user-turn message as immediate intent, so the resumption preamble
|
||||
// is most reliable when prepended directly to the prompt that gets piped into the CLI.
|
||||
let owned_prompt = match resumption_prompt {
|
||||
Some(preamble) if !preamble.is_empty() => format!("{preamble}\n\n{prompt}"),
|
||||
_ => prompt.to_string(),
|
||||
};
|
||||
Ok(Box::new(ClaudeHarnessRunner::new(
|
||||
self.cli_agent().command_prefix(),
|
||||
&owned_prompt,
|
||||
system_prompt,
|
||||
working_dir,
|
||||
task_id,
|
||||
server_api,
|
||||
terminal_driver,
|
||||
claude_resume,
|
||||
)?))
|
||||
}
|
||||
}
|
||||
|
||||
/// Format slug sent to the server when creating a Claude Code conversation.
|
||||
const CLAUDE_CODE_FORMAT: &str = "claude_code_cli";
|
||||
/// Command used to exit claude.
|
||||
const CLAUDE_EXIT_COMMAND: &str = "/exit";
|
||||
|
||||
/// Build the shell command that launches the Claude CLI for a given session and
|
||||
/// prompt file.
|
||||
///
|
||||
/// When `resuming` is true we pass `--resume <uuid>` so Claude picks up the
|
||||
/// existing on-disk session; otherwise we pass `--session-id <uuid>` to pin a
|
||||
/// fresh session to that id. If `system_prompt_path` is provided, the CLI is
|
||||
/// told to append its contents to the base system prompt.
|
||||
fn claude_command(
|
||||
cli_name: &str,
|
||||
session_id: &Uuid,
|
||||
prompt_path: &str,
|
||||
system_prompt_path: Option<&str>,
|
||||
resuming: bool,
|
||||
) -> String {
|
||||
let flag = if resuming { "--resume" } else { "--session-id" };
|
||||
let mut cmd = format!("{cli_name} {flag} {session_id} --dangerously-skip-permissions");
|
||||
if let Some(sp_path) = system_prompt_path {
|
||||
let _ = write!(cmd, " --append-system-prompt-file '{sp_path}'");
|
||||
}
|
||||
format!("{cmd} < '{prompt_path}'")
|
||||
}
|
||||
|
||||
/// Runtime state of a [`ClaudeHarnessRunner`].
|
||||
enum ClaudeRunnerState {
|
||||
/// Runner is built but [`HarnessRunner::start`] has not been called yet.
|
||||
Preexec,
|
||||
/// The harness command is running (or has finished).
|
||||
Running {
|
||||
conversation_id: AIConversationId,
|
||||
block_id: BlockId,
|
||||
},
|
||||
}
|
||||
|
||||
struct ClaudeHarnessRunner {
|
||||
command: String,
|
||||
/// The CLI name used to invoke Claude Code.
|
||||
cli_name: String,
|
||||
/// Held so the temp file is cleaned up when the runner is dropped.
|
||||
_temp_prompt_file: NamedTempFile,
|
||||
/// Held so the system prompt temp file is cleaned up when the runner is dropped.
|
||||
_temp_system_prompt_file: Option<NamedTempFile>,
|
||||
client: Arc<dyn HarnessSupportClient>,
|
||||
server_api: Arc<ServerApi>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
state: Mutex<ClaudeRunnerState>,
|
||||
session_id: Uuid,
|
||||
working_dir: PathBuf,
|
||||
parent_bridge: Option<MessageBridge>,
|
||||
/// Lazily cached output of `claude --version`.
|
||||
claude_version: Mutex<Option<String>>,
|
||||
/// When resuming an existing conversation, we pin the runner's server conversation id
|
||||
/// up front instead of calling `create_external_conversation` in [`HarnessRunner::start`].
|
||||
/// Subsequent saves overwrite the same GCS objects keyed by this id.
|
||||
preexisting_conversation_id: Option<AIConversationId>,
|
||||
}
|
||||
|
||||
impl ClaudeHarnessRunner {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn new(
|
||||
cli_command: &str,
|
||||
prompt: &str,
|
||||
system_prompt: Option<&str>,
|
||||
working_dir: &Path,
|
||||
task_id: Option<AmbientAgentTaskId>,
|
||||
server_api: Arc<ServerApi>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
resume: Option<ClaudeResumeInfo>,
|
||||
) -> Result<Self, AgentDriverError> {
|
||||
// Write the prompt to a temp file so we can feed it via stdin redirect,
|
||||
// avoiding shell-quoting issues with complex content (e.g. skill instructions).
|
||||
let temp_file = write_temp_file("oz_prompt_", prompt)?;
|
||||
let prompt_path = temp_file.path().display().to_string();
|
||||
|
||||
let (session_id, preexisting_conversation_id, resuming) = match resume {
|
||||
Some(ClaudeResumeInfo {
|
||||
conversation_id,
|
||||
session_id,
|
||||
mut envelope,
|
||||
}) => {
|
||||
// Rehydrate the stored envelope under the current working directory so
|
||||
// `claude --resume <uuid>` finds the jsonl under ~/.claude/projects/<encoded_cwd>/.
|
||||
// The original envelope's cwd usually points at the cloud sandbox path, which
|
||||
// doesn't exist locally.
|
||||
envelope.cwd = working_dir.to_path_buf();
|
||||
let config_root = claude_config_dir().map_err(|e| {
|
||||
AgentDriverError::ConfigBuildFailed(
|
||||
e.context("Failed to resolve Claude config dir"),
|
||||
)
|
||||
})?;
|
||||
write_envelope(&envelope, &config_root).map_err(|e| {
|
||||
AgentDriverError::ConfigBuildFailed(
|
||||
e.context("Failed to rehydrate Claude transcript"),
|
||||
)
|
||||
})?;
|
||||
// Index write is best-effort: upstream Claude versions vary in how they use
|
||||
// `sessions-index.json`, so losing the index entry shouldn't abort the run.
|
||||
if let Err(e) = write_session_index_entry(session_id, working_dir, &config_root) {
|
||||
log::warn!("Failed to update Claude sessions-index.json: {e:#}");
|
||||
}
|
||||
(session_id, Some(conversation_id), true)
|
||||
}
|
||||
None => (Uuid::new_v4(), None, false),
|
||||
};
|
||||
|
||||
let temp_system_prompt_file = system_prompt
|
||||
.map(|sp| write_temp_file("oz_system_prompt_", sp))
|
||||
.transpose()?;
|
||||
let system_prompt_path = temp_system_prompt_file
|
||||
.as_ref()
|
||||
.map(|f| f.path().display().to_string());
|
||||
let parent_bridge = task_id
|
||||
.map(|task_id| MessageBridge::new(task_id.to_string(), session_id))
|
||||
.transpose()
|
||||
.map_err(AgentDriverError::ConfigBuildFailed)?;
|
||||
let client: Arc<dyn HarnessSupportClient> = server_api.clone();
|
||||
|
||||
Ok(Self {
|
||||
command: claude_command(
|
||||
cli_command,
|
||||
&session_id,
|
||||
&prompt_path,
|
||||
system_prompt_path.as_deref(),
|
||||
resuming,
|
||||
),
|
||||
cli_name: cli_command.to_string(),
|
||||
_temp_prompt_file: temp_file,
|
||||
_temp_system_prompt_file: temp_system_prompt_file,
|
||||
client,
|
||||
server_api,
|
||||
terminal_driver,
|
||||
state: Mutex::new(ClaudeRunnerState::Preexec),
|
||||
session_id,
|
||||
working_dir: working_dir.to_path_buf(),
|
||||
parent_bridge,
|
||||
claude_version: Mutex::new(None),
|
||||
preexisting_conversation_id,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ClaudeHarnessRunner {
|
||||
async fn handle_parent_bridge_session_update(&self) -> Result<()> {
|
||||
let Some(parent_bridge) = self.parent_bridge.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
parent_bridge
|
||||
.handle_session_update(self.server_api.clone())
|
||||
.await
|
||||
}
|
||||
|
||||
async fn flush_parent_bridge_acks(&self) -> Result<()> {
|
||||
let Some(parent_bridge) = self.parent_bridge.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
parent_bridge.flush_acks(self.server_api.clone()).await
|
||||
}
|
||||
/// Return the cached Claude Code version, or resolve it by running
|
||||
/// `<cli_name> --version`.
|
||||
async fn resolve_claude_version(
|
||||
&self,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Option<String> {
|
||||
if let Some(cached) = self.claude_version.lock().clone() {
|
||||
return Some(cached);
|
||||
}
|
||||
|
||||
let terminal_driver = self.terminal_driver.clone();
|
||||
let session = foreground
|
||||
.spawn(move |_, ctx| {
|
||||
let tv = terminal_driver.as_ref(ctx).terminal_view().as_ref(ctx);
|
||||
tv.active_session().as_ref(ctx).session(ctx)
|
||||
})
|
||||
.await
|
||||
.ok()?;
|
||||
let session = session?;
|
||||
|
||||
let cli_name = &self.cli_name;
|
||||
let output = session
|
||||
.execute_command(
|
||||
&format!("{cli_name} --version"),
|
||||
None,
|
||||
None,
|
||||
ExecuteCommandOptions::default(),
|
||||
)
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
let version = output.to_string().ok()?.trim().to_string();
|
||||
if version.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
*self.claude_version.lock() = Some(version.clone());
|
||||
Some(version)
|
||||
}
|
||||
|
||||
async fn start_parent_bridge(&self, foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
let Some(parent_bridge) = self.parent_bridge.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
parent_bridge
|
||||
.start(foreground, self.server_api.clone())
|
||||
.await
|
||||
}
|
||||
|
||||
fn cleanup_parent_bridge(&self) -> Result<()> {
|
||||
if let Some(parent_bridge) = self.parent_bridge.as_ref() {
|
||||
parent_bridge.cleanup()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl HarnessRunner for ClaudeHarnessRunner {
|
||||
async fn start(
|
||||
&self,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Result<CommandHandle, AgentDriverError> {
|
||||
// When resuming, we already have a server conversation id from the prior run.
|
||||
// Otherwise create a fresh external conversation record for this run.
|
||||
// TODO(REMOTE-1149): `create_external_conversation` currently won't work for local CLI
|
||||
// runs. We should either support it or have a fallback.
|
||||
let conversation_id = match self.preexisting_conversation_id {
|
||||
Some(id) => {
|
||||
log::info!("Resuming external conversation {id}");
|
||||
id
|
||||
}
|
||||
None => {
|
||||
let id = self
|
||||
.client
|
||||
.create_external_conversation(CLAUDE_CODE_FORMAT)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
log::error!("Failed to create external conversation: {e}");
|
||||
AgentDriverError::ConfigBuildFailed(e)
|
||||
})?;
|
||||
log::info!("Created external conversation {id}");
|
||||
id
|
||||
}
|
||||
};
|
||||
self.start_parent_bridge(foreground)
|
||||
.await
|
||||
.map_err(AgentDriverError::ConfigBuildFailed)?;
|
||||
|
||||
let command = self.command.clone();
|
||||
let terminal_driver = self.terminal_driver.clone();
|
||||
let command_handle = match foreground
|
||||
.spawn(move |_, ctx| {
|
||||
terminal_driver.update(ctx, |driver, ctx| driver.execute_command(&command, ctx))
|
||||
})
|
||||
.await??
|
||||
.await
|
||||
{
|
||||
Ok(command_handle) => command_handle,
|
||||
Err(err) => {
|
||||
self.cleanup_parent_bridge()
|
||||
.map_err(AgentDriverError::ConfigBuildFailed)?;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
// Only store conversation info once the CLI command has started.
|
||||
*self.state.lock() = ClaudeRunnerState::Running {
|
||||
conversation_id,
|
||||
block_id: command_handle.block_id().clone(),
|
||||
};
|
||||
|
||||
Ok(command_handle)
|
||||
}
|
||||
|
||||
async fn exit(&self, foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
log::info!("Sending /exit to Claude Code CLI");
|
||||
let terminal_driver = self.terminal_driver.clone();
|
||||
foreground
|
||||
.spawn(move |_, ctx| {
|
||||
terminal_driver.update(ctx, |driver, ctx| {
|
||||
driver.send_text_to_cli(CLAUDE_EXIT_COMMAND.to_string(), ctx);
|
||||
});
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Agent driver dropped while sending /exit"))
|
||||
}
|
||||
|
||||
async fn handle_session_update(&self, _foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
self.handle_parent_bridge_session_update().await
|
||||
}
|
||||
|
||||
async fn save_conversation(
|
||||
&self,
|
||||
save_point: SavePoint,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Result<()> {
|
||||
if matches!(save_point, SavePoint::Periodic)
|
||||
&& !super::has_running_cli_agent(&self.terminal_driver, foreground).await
|
||||
{
|
||||
log::debug!("Will not save conversation, Claude Code not in progress");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let (conversation_id, block_id) = match &*self.state.lock() {
|
||||
ClaudeRunnerState::Preexec => {
|
||||
log::warn!("save_conversation called before start");
|
||||
return Ok(());
|
||||
}
|
||||
ClaudeRunnerState::Running {
|
||||
conversation_id,
|
||||
block_id,
|
||||
} => (*conversation_id, block_id.clone()),
|
||||
};
|
||||
|
||||
let claude_version = self.resolve_claude_version(foreground).await;
|
||||
|
||||
let client = self.client.as_ref();
|
||||
let session_id = self.session_id;
|
||||
let working_dir = &self.working_dir;
|
||||
|
||||
futures::try_join!(
|
||||
super::upload_current_block_snapshot(
|
||||
foreground,
|
||||
&self.terminal_driver,
|
||||
client,
|
||||
conversation_id,
|
||||
block_id,
|
||||
),
|
||||
upload_transcript(
|
||||
client,
|
||||
conversation_id,
|
||||
session_id,
|
||||
working_dir,
|
||||
claude_version
|
||||
),
|
||||
)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
async fn cleanup(&self, _foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
self.flush_parent_bridge_acks().await?;
|
||||
self.cleanup_parent_bridge()
|
||||
}
|
||||
}
|
||||
|
||||
/// Upload the Claude Code session transcript to the server.
|
||||
async fn upload_transcript(
|
||||
client: &dyn HarnessSupportClient,
|
||||
conversation_id: AIConversationId,
|
||||
session_id: Uuid,
|
||||
working_dir: &Path,
|
||||
claude_version: Option<String>,
|
||||
) -> Result<()> {
|
||||
log::info!("Uploading Claude Code transcript to conversation {conversation_id}");
|
||||
|
||||
let config_dir = claude_config_dir().context("Failed to resolve Claude config dir")?;
|
||||
let working_dir = working_dir.to_path_buf();
|
||||
let body = tokio::task::spawn_blocking(move || {
|
||||
let mut envelope = read_envelope(session_id, &working_dir, &config_dir)
|
||||
.with_context(|| format!("Failed to read transcript for session {session_id}"))?;
|
||||
envelope.claude_version = claude_version;
|
||||
serde_json::to_vec(&envelope).context("Failed to serialize transcript envelope")
|
||||
})
|
||||
.await
|
||||
.context("read_envelope task panicked")??;
|
||||
let target = client
|
||||
.get_transcript_upload_target(&conversation_id)
|
||||
.await
|
||||
.with_context(|| format!("Failed to get transcript upload target for {conversation_id}"))?;
|
||||
upload_to_target(client.http_client(), &target, body).await
|
||||
}
|
||||
|
||||
fn prepare_claude_environment_config(
|
||||
working_dir: &Path,
|
||||
secrets: &HashMap<String, ManagedSecretValue>,
|
||||
) -> Result<()> {
|
||||
let home_dir =
|
||||
dirs::home_dir().ok_or_else(|| anyhow::anyhow!("could not determine home directory"))?;
|
||||
let claude_json_path = home_dir.join(CLAUDE_JSON_FILE_NAME);
|
||||
let claude_settings_path = claude_config_dir()?.join(CLAUDE_SETTINGS_FILE_NAME);
|
||||
let api_key_suffix = resolve_anthropic_api_key_suffix(secrets);
|
||||
prepare_claude_config(&claude_json_path, working_dir, api_key_suffix.as_deref())?;
|
||||
prepare_claude_settings(&claude_settings_path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prepare_claude_config(
|
||||
claude_json_path: &Path,
|
||||
working_dir: &Path,
|
||||
api_key_suffix: Option<&str>,
|
||||
) -> Result<()> {
|
||||
let mut claude_config: ClaudeConfig = read_json_file_or_default(claude_json_path)?;
|
||||
claude_config.has_completed_onboarding = true;
|
||||
claude_config.lsp_recommendation_disabled = true;
|
||||
claude_config
|
||||
.projects
|
||||
.entry(working_dir.to_string_lossy().into_owned())
|
||||
.or_default()
|
||||
.has_trust_dialog_accepted = true;
|
||||
if let Some(suffix) = api_key_suffix {
|
||||
let responses = claude_config
|
||||
.custom_api_key_responses
|
||||
.get_or_insert_with(CustomApiKeyResponses::default);
|
||||
if !responses.approved.iter().any(|s| s == suffix) {
|
||||
responses.approved.push(suffix.to_owned());
|
||||
}
|
||||
}
|
||||
write_json_file(
|
||||
claude_json_path,
|
||||
&claude_config,
|
||||
"Failed to serialize Claude config",
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prepare_claude_settings(claude_settings_path: &Path) -> Result<()> {
|
||||
let mut settings: ClaudeSettings = read_json_file_or_default(claude_settings_path)?;
|
||||
settings.skip_dangerous_mode_permission_prompt = true;
|
||||
write_json_file(
|
||||
claude_settings_path,
|
||||
&settings,
|
||||
"Failed to serialize Claude settings",
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
|
||||
const CLAUDE_JSON_FILE_NAME: &str = ".claude.json";
|
||||
const CLAUDE_SETTINGS_FILE_NAME: &str = "settings.json";
|
||||
const ANTHROPIC_API_KEY_SUFFIX_LEN: usize = 20;
|
||||
|
||||
#[derive(Default, Deserialize, Serialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ClaudeConfig {
|
||||
#[serde(default)]
|
||||
has_completed_onboarding: bool,
|
||||
#[serde(default)]
|
||||
lsp_recommendation_disabled: bool,
|
||||
#[serde(default)]
|
||||
projects: HashMap<String, ClaudeProjectConfig>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
custom_api_key_responses: Option<CustomApiKeyResponses>,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct CustomApiKeyResponses {
|
||||
#[serde(default)]
|
||||
approved: Vec<String>,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ClaudeProjectConfig {
|
||||
#[serde(default)]
|
||||
has_trust_dialog_accepted: bool,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ClaudeSettings {
|
||||
#[serde(default)]
|
||||
skip_dangerous_mode_permission_prompt: bool,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
/// Try to get the last 20 chars of the ANTHROPIC_API_KEY from the secrets map,
|
||||
/// where 20 chars is the suffix length that Claude Code truncates keys to.
|
||||
/// Falls back to the environment variable.
|
||||
fn resolve_anthropic_api_key_suffix(
|
||||
secrets: &HashMap<String, ManagedSecretValue>,
|
||||
) -> Option<String> {
|
||||
// First, check for an AnthropicApiKey variant anywhere in the secrets map,
|
||||
// since the secret name doesn't necessarily match the env var.
|
||||
for secret in secrets.values() {
|
||||
if let ManagedSecretValue::AnthropicApiKey { api_key } = secret {
|
||||
return suffix_of(api_key).map(str::to_owned);
|
||||
}
|
||||
}
|
||||
// Then check for a RawValue stored under the env var name.
|
||||
if let Some(ManagedSecretValue::RawValue { value }) = secrets.get(ANTHROPIC_API_KEY_ENV) {
|
||||
return suffix_of(value).map(str::to_owned);
|
||||
}
|
||||
// Fall back to the environment variable, which a user may have set separately in the env.
|
||||
std::env::var(ANTHROPIC_API_KEY_ENV)
|
||||
.ok()
|
||||
.and_then(|k| suffix_of(&k).map(str::to_owned))
|
||||
}
|
||||
|
||||
fn suffix_of(key: &str) -> Option<&str> {
|
||||
if key.len() >= ANTHROPIC_API_KEY_SUFFIX_LEN {
|
||||
key.get(key.len() - ANTHROPIC_API_KEY_SUFFIX_LEN..)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "claude_code_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,621 @@
|
||||
//! Bridges durable lead-agent messages into Claude Code's hook-driven
|
||||
//! next-turn context.
|
||||
//!
|
||||
//! The bridge uses an on-disk three-stage state machine inside the per-session
|
||||
//! state directory:
|
||||
//! - `staged/` holds newly observed message IDs from the event stream.
|
||||
//! - `surfaced/` holds the fully hydrated records currently exposed to Claude.
|
||||
//! - `pending-hook-output.json` plus `pending-hook-output.ack` coordinates the
|
||||
//! handoff between Warp's driver and the Claude hook process.
|
||||
use std::fmt::Write as _;
|
||||
use std::fs;
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tempfile::NamedTempFile;
|
||||
use tokio::sync::Mutex as AsyncMutex;
|
||||
use uuid::Uuid;
|
||||
use warpui::r#async::SpawnedFutureHandle;
|
||||
use warpui::ModelSpawner;
|
||||
|
||||
use crate::ai::agent_events::{
|
||||
run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow,
|
||||
AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource,
|
||||
};
|
||||
use crate::ai::agent_sdk::driver::{AgentDriver, OZ_MESSAGE_LISTENER_STATE_ROOT_ENV};
|
||||
use crate::server::server_api::ai::AgentRunEvent;
|
||||
use crate::server::server_api::ServerApi;
|
||||
|
||||
const LEGACY_MESSAGE_LISTENER_STATE_ROOT_ENV: &str = "OZ_PARENT_STATE_ROOT";
|
||||
const PARENT_BRIDGE_DEFAULT_STATE_ROOT: &str = ".claude-code/oz-parent-bridge";
|
||||
const PARENT_BRIDGE_SURFACED_DIR_NAME: &str = "surfaced";
|
||||
const PARENT_BRIDGE_HOOK_OUTPUT_FILE_NAME: &str = "pending-hook-output.json";
|
||||
const PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME: &str = "pending-hook-output.ack";
|
||||
const PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV: &str = "OZ_PARENT_MAX_CONTEXT_CHARS";
|
||||
const PARENT_BRIDGE_DEFAULT_MAX_CONTEXT_CHARS: usize = 6000;
|
||||
pub(super) const MESSAGE_BRIDGE_CONTEXT_PREAMBLE: &str = "Lead-agent updates arrived from Oz. Treat the latest lead-agent instructions below as authoritative.\n";
|
||||
const PARENT_BRIDGE_REMAINING_MESSAGES_NOTE: &str =
|
||||
"\n\nMore lead-agent messages are still staged and will be surfaced on a later turn.";
|
||||
|
||||
pub(super) struct MessageBridge {
|
||||
run_id: String,
|
||||
state_dir: PathBuf,
|
||||
runtime: Mutex<Option<MessageBridgeRuntime>>,
|
||||
state_lock: AsyncMutex<()>,
|
||||
}
|
||||
struct MessageBridgeRuntime {
|
||||
task: SpawnedFutureHandle,
|
||||
}
|
||||
|
||||
struct MessageBridgeEventConsumer {
|
||||
run_id: String,
|
||||
state_dir: PathBuf,
|
||||
}
|
||||
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
impl AgentEventConsumer for MessageBridgeEventConsumer {
|
||||
async fn on_event(
|
||||
&mut self,
|
||||
event: AgentRunEvent,
|
||||
) -> anyhow::Result<AgentEventConsumerControlFlow> {
|
||||
if event.event_type != "new_message" || event.run_id != self.run_id {
|
||||
return Ok(AgentEventConsumerControlFlow::Continue);
|
||||
}
|
||||
|
||||
let Some(message_id) = event.ref_id else {
|
||||
return Ok(AgentEventConsumerControlFlow::Continue);
|
||||
};
|
||||
|
||||
if let Err(err) = stage_parent_bridge_message(
|
||||
&self.state_dir,
|
||||
&MessageBridgeMessageRecord {
|
||||
sequence: event.sequence,
|
||||
message_id: message_id.clone(),
|
||||
sender_run_id: String::new(),
|
||||
subject: String::new(),
|
||||
body: String::new(),
|
||||
occurred_at: event.occurred_at,
|
||||
},
|
||||
) {
|
||||
log::warn!(
|
||||
"Failed to stage Claude lead-agent message {message_id} at sequence {}: {err:#}",
|
||||
event.sequence
|
||||
);
|
||||
}
|
||||
|
||||
Ok(AgentEventConsumerControlFlow::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub(super) struct MessageBridgeMessageRecord {
|
||||
pub sequence: i64,
|
||||
pub message_id: String,
|
||||
#[serde(default)]
|
||||
pub sender_run_id: String,
|
||||
#[serde(default)]
|
||||
pub subject: String,
|
||||
#[serde(default)]
|
||||
pub body: String,
|
||||
pub occurred_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub(super) struct MessageBridgeHookOutput {
|
||||
pub additional_context: String,
|
||||
pub remaining_staged_count: usize,
|
||||
pub surfaced_count: usize,
|
||||
}
|
||||
|
||||
struct RenderedMessageBridgeMessage {
|
||||
block: String,
|
||||
block_chars: usize,
|
||||
}
|
||||
struct SelectedMessageBridgeMessage {
|
||||
path: PathBuf,
|
||||
record: MessageBridgeMessageRecord,
|
||||
rendered: RenderedMessageBridgeMessage,
|
||||
}
|
||||
|
||||
struct SelectedMessageBridgeMessages {
|
||||
messages: Vec<SelectedMessageBridgeMessage>,
|
||||
context_chars: usize,
|
||||
total_available_count: usize,
|
||||
}
|
||||
|
||||
impl MessageBridge {
|
||||
pub(super) fn new(run_id: String, session_id: Uuid) -> Result<Self> {
|
||||
Ok(Self {
|
||||
run_id,
|
||||
state_dir: parent_bridge_root()?.join(session_id.to_string()),
|
||||
runtime: Mutex::new(None),
|
||||
state_lock: AsyncMutex::new(()),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn start(
|
||||
&self,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
server_api: Arc<ServerApi>,
|
||||
) -> Result<()> {
|
||||
if self.runtime.lock().is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
ensure_parent_bridge_state_dir(&self.state_dir)?;
|
||||
let run_id = self.run_id.clone();
|
||||
let state_dir = self.state_dir.clone();
|
||||
let task = foreground
|
||||
.spawn(move |_, ctx| {
|
||||
ctx.spawn(
|
||||
async move {
|
||||
if let Err(err) =
|
||||
run_parent_bridge_forever(server_api, run_id, state_dir.clone()).await
|
||||
{
|
||||
log::warn!(
|
||||
"Claude message bridge stopped for {}: {err:#}",
|
||||
state_dir.display()
|
||||
);
|
||||
}
|
||||
},
|
||||
|_, _, _| {},
|
||||
)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow!("Agent driver dropped while starting Claude message bridge"))?;
|
||||
*self.runtime.lock() = Some(MessageBridgeRuntime { task });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn handle_session_update(&self, server_api: Arc<ServerApi>) -> Result<()> {
|
||||
if !self.state_dir.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let hydrator = MessageHydrator::new(server_api);
|
||||
let _guard = self.state_lock.lock().await;
|
||||
acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await?;
|
||||
prepare_parent_bridge_hook_output(
|
||||
&hydrator,
|
||||
&self.state_dir,
|
||||
parent_bridge_max_context_chars(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(super) async fn flush_acks(&self, server_api: Arc<ServerApi>) -> Result<()> {
|
||||
if !self.state_dir.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let hydrator = MessageHydrator::new(server_api);
|
||||
let _guard = self.state_lock.lock().await;
|
||||
acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await
|
||||
}
|
||||
|
||||
pub(super) fn cleanup(&self) -> Result<()> {
|
||||
if let Some(runtime) = self.runtime.lock().take() {
|
||||
runtime.task.abort();
|
||||
}
|
||||
match fs::remove_dir_all(&self.state_dir) {
|
||||
Ok(()) => {}
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
|
||||
Err(err) => {
|
||||
return Err(anyhow::Error::from(err).context(format!(
|
||||
"Failed to remove Claude message bridge state dir {}",
|
||||
self.state_dir.display()
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn parent_bridge_root() -> Result<PathBuf> {
|
||||
for env_name in [
|
||||
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
LEGACY_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
] {
|
||||
if let Ok(dir) = std::env::var(env_name) {
|
||||
if !dir.is_empty() {
|
||||
return Ok(PathBuf::from(dir));
|
||||
}
|
||||
}
|
||||
}
|
||||
dirs::home_dir()
|
||||
.map(|home| home.join(PARENT_BRIDGE_DEFAULT_STATE_ROOT))
|
||||
.ok_or_else(|| anyhow!("could not determine home directory"))
|
||||
}
|
||||
|
||||
fn parent_bridge_staged_dir(state_dir: &Path) -> PathBuf {
|
||||
state_dir.join("staged")
|
||||
}
|
||||
|
||||
fn parent_bridge_surfaced_dir(state_dir: &Path) -> PathBuf {
|
||||
state_dir.join(PARENT_BRIDGE_SURFACED_DIR_NAME)
|
||||
}
|
||||
|
||||
pub(super) fn parent_bridge_hook_output_file(state_dir: &Path) -> PathBuf {
|
||||
state_dir.join(PARENT_BRIDGE_HOOK_OUTPUT_FILE_NAME)
|
||||
}
|
||||
|
||||
pub(super) fn parent_bridge_hook_output_ack_file(state_dir: &Path) -> PathBuf {
|
||||
state_dir.join(PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME)
|
||||
}
|
||||
|
||||
fn parent_bridge_message_path(dir: &Path, sequence: i64, message_id: &str) -> PathBuf {
|
||||
dir.join(format!("{sequence:020}-{message_id}.json"))
|
||||
}
|
||||
|
||||
pub(super) fn parent_bridge_staged_message_path(
|
||||
state_dir: &Path,
|
||||
sequence: i64,
|
||||
message_id: &str,
|
||||
) -> PathBuf {
|
||||
parent_bridge_message_path(&parent_bridge_staged_dir(state_dir), sequence, message_id)
|
||||
}
|
||||
|
||||
pub(super) fn parent_bridge_surfaced_message_path(
|
||||
state_dir: &Path,
|
||||
sequence: i64,
|
||||
message_id: &str,
|
||||
) -> PathBuf {
|
||||
parent_bridge_message_path(&parent_bridge_surfaced_dir(state_dir), sequence, message_id)
|
||||
}
|
||||
|
||||
pub(super) fn ensure_parent_bridge_state_dir(state_dir: &Path) -> Result<()> {
|
||||
fs::create_dir_all(parent_bridge_staged_dir(state_dir))
|
||||
.with_context(|| format!("Failed to create {}", state_dir.display()))?;
|
||||
fs::create_dir_all(parent_bridge_surfaced_dir(state_dir))
|
||||
.with_context(|| format!("Failed to create {}", state_dir.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn stage_parent_bridge_message(
|
||||
state_dir: &Path,
|
||||
record: &MessageBridgeMessageRecord,
|
||||
) -> Result<()> {
|
||||
let target = parent_bridge_staged_message_path(state_dir, record.sequence, &record.message_id);
|
||||
if !target.exists() {
|
||||
write_parent_bridge_json_atomically(&target, record)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parent_bridge_max_context_chars() -> usize {
|
||||
std::env::var(PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(PARENT_BRIDGE_DEFAULT_MAX_CONTEXT_CHARS)
|
||||
}
|
||||
|
||||
fn parent_bridge_sorted_message_paths(dir: &Path) -> Result<Vec<PathBuf>> {
|
||||
if !dir.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut paths = fs::read_dir(dir)
|
||||
.with_context(|| format!("Failed to read {}", dir.display()))?
|
||||
.filter_map(|entry| entry.ok().map(|entry| entry.path()))
|
||||
.filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("json"))
|
||||
.collect::<Vec<_>>();
|
||||
paths.sort();
|
||||
Ok(paths)
|
||||
}
|
||||
|
||||
fn parent_bridge_message_records(dir: &Path) -> Result<Vec<(PathBuf, MessageBridgeMessageRecord)>> {
|
||||
parent_bridge_sorted_message_paths(dir)?
|
||||
.into_iter()
|
||||
.map(|path| {
|
||||
let record = serde_json::from_slice::<MessageBridgeMessageRecord>(
|
||||
&fs::read(&path).with_context(|| format!("Failed to read {}", path.display()))?,
|
||||
)
|
||||
.with_context(|| format!("Failed to parse {}", path.display()))?;
|
||||
Ok((path, record))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn parent_bridge_char_count(text: &str) -> usize {
|
||||
text.chars().count()
|
||||
}
|
||||
|
||||
fn parent_bridge_truncate_chars(text: &str, limit: usize) -> String {
|
||||
text.chars().take(limit).collect()
|
||||
}
|
||||
|
||||
pub(super) fn render_parent_bridge_message_block(record: &MessageBridgeMessageRecord) -> String {
|
||||
let subject = if record.subject.is_empty() {
|
||||
"(no subject)"
|
||||
} else {
|
||||
record.subject.as_str()
|
||||
};
|
||||
|
||||
let mut block = String::from("---\nLead-agent message");
|
||||
if record.sequence != 0 {
|
||||
let _ = write!(block, " #{}", record.sequence);
|
||||
}
|
||||
if !record.sender_run_id.is_empty() {
|
||||
let _ = write!(block, " from {}", record.sender_run_id);
|
||||
}
|
||||
let _ = write!(block, "\nSubject: {subject}\n\n{}", record.body);
|
||||
block
|
||||
}
|
||||
|
||||
fn render_parent_bridge_message(
|
||||
record: &MessageBridgeMessageRecord,
|
||||
) -> RenderedMessageBridgeMessage {
|
||||
let block = render_parent_bridge_message_block(record);
|
||||
let block_chars = parent_bridge_char_count(&block);
|
||||
RenderedMessageBridgeMessage { block, block_chars }
|
||||
}
|
||||
|
||||
fn truncate_parent_bridge_message(rendered: &mut RenderedMessageBridgeMessage, max_chars: usize) {
|
||||
if rendered.block_chars <= max_chars || max_chars <= 3 {
|
||||
return;
|
||||
}
|
||||
|
||||
rendered.block = parent_bridge_truncate_chars(&rendered.block, max_chars - 3);
|
||||
rendered.block.push_str("...");
|
||||
rendered.block_chars = parent_bridge_char_count(&rendered.block);
|
||||
}
|
||||
|
||||
fn build_parent_bridge_hook_output(
|
||||
selected: &SelectedMessageBridgeMessages,
|
||||
max_context_chars: usize,
|
||||
) -> Option<MessageBridgeHookOutput> {
|
||||
if selected.messages.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut additional_context = String::from(MESSAGE_BRIDGE_CONTEXT_PREAMBLE);
|
||||
for (index, message) in selected.messages.iter().enumerate() {
|
||||
if index > 0 {
|
||||
additional_context.push_str("\n\n");
|
||||
}
|
||||
additional_context.push_str(&message.rendered.block);
|
||||
}
|
||||
|
||||
let remaining_staged_count = selected
|
||||
.total_available_count
|
||||
.saturating_sub(selected.messages.len());
|
||||
let remaining_note_chars = parent_bridge_char_count(PARENT_BRIDGE_REMAINING_MESSAGES_NOTE);
|
||||
if remaining_staged_count > 0
|
||||
&& selected.context_chars + remaining_note_chars <= max_context_chars
|
||||
{
|
||||
additional_context.push_str(PARENT_BRIDGE_REMAINING_MESSAGES_NOTE);
|
||||
}
|
||||
|
||||
Some(MessageBridgeHookOutput {
|
||||
additional_context,
|
||||
remaining_staged_count,
|
||||
surfaced_count: selected.messages.len(),
|
||||
})
|
||||
}
|
||||
|
||||
fn write_parent_bridge_hook_output(
|
||||
state_dir: &Path,
|
||||
output: &MessageBridgeHookOutput,
|
||||
) -> Result<()> {
|
||||
let path = parent_bridge_hook_output_file(state_dir);
|
||||
write_parent_bridge_json_atomically(&path, output)
|
||||
}
|
||||
|
||||
fn remove_file_if_exists(path: &Path) -> Result<()> {
|
||||
match fs::remove_file(path) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||
Err(err) => {
|
||||
Err(anyhow::Error::from(err).context(format!("Failed to remove {}", path.display())))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn hydrate_parent_bridge_message_record(
|
||||
hydrator: &MessageHydrator,
|
||||
record: &MessageBridgeMessageRecord,
|
||||
) -> Result<MessageBridgeMessageRecord> {
|
||||
if !record.sender_run_id.is_empty() {
|
||||
return Ok(record.clone());
|
||||
}
|
||||
|
||||
let message = hydrator
|
||||
.read_message_with_timeout(&record.message_id)
|
||||
.await
|
||||
.with_context(|| format!("Failed to read lead-agent message {}", record.message_id))?;
|
||||
Ok(MessageBridgeMessageRecord {
|
||||
sequence: record.sequence,
|
||||
message_id: message.message_id,
|
||||
sender_run_id: message.sender_run_id,
|
||||
subject: message.subject,
|
||||
body: message.body,
|
||||
occurred_at: record.occurred_at.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn select_parent_bridge_messages_for_hook_output(
|
||||
hydrator: &MessageHydrator,
|
||||
records: Vec<(PathBuf, MessageBridgeMessageRecord)>,
|
||||
max_context_chars: usize,
|
||||
) -> Result<SelectedMessageBridgeMessages> {
|
||||
let total_available_count = records.len();
|
||||
let mut messages = Vec::new();
|
||||
let mut context_chars = parent_bridge_char_count(MESSAGE_BRIDGE_CONTEXT_PREAMBLE);
|
||||
|
||||
for (path, record) in records {
|
||||
let separator_chars = if messages.is_empty() { 0 } else { 2 };
|
||||
let remaining = max_context_chars.saturating_sub(context_chars + separator_chars);
|
||||
if remaining <= 3 && !messages.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
let record = hydrate_parent_bridge_message_record(hydrator, &record).await?;
|
||||
let mut rendered = render_parent_bridge_message(&record);
|
||||
if context_chars + separator_chars + rendered.block_chars > max_context_chars {
|
||||
if remaining > 3 && rendered.block_chars > remaining {
|
||||
truncate_parent_bridge_message(&mut rendered, remaining);
|
||||
} else if !messages.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
context_chars += separator_chars + rendered.block_chars;
|
||||
messages.push(SelectedMessageBridgeMessage {
|
||||
path,
|
||||
record,
|
||||
rendered,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(SelectedMessageBridgeMessages {
|
||||
messages,
|
||||
context_chars,
|
||||
total_available_count,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn prepare_parent_bridge_hook_output(
|
||||
hydrator: &MessageHydrator,
|
||||
state_dir: &Path,
|
||||
max_context_chars: usize,
|
||||
) -> Result<()> {
|
||||
let hook_output_path = parent_bridge_hook_output_file(state_dir);
|
||||
if hook_output_path.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let surfaced_records = parent_bridge_message_records(&parent_bridge_surfaced_dir(state_dir))?;
|
||||
if !surfaced_records.is_empty() {
|
||||
let selected = select_parent_bridge_messages_for_hook_output(
|
||||
hydrator,
|
||||
surfaced_records,
|
||||
max_context_chars,
|
||||
)
|
||||
.await?;
|
||||
for message in &selected.messages {
|
||||
write_parent_bridge_json_atomically(&message.path, &message.record)?;
|
||||
}
|
||||
if let Some(output) = build_parent_bridge_hook_output(&selected, max_context_chars) {
|
||||
write_parent_bridge_hook_output(state_dir, &output)?;
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let staged_records = parent_bridge_message_records(&parent_bridge_staged_dir(state_dir))?;
|
||||
if staged_records.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let selected =
|
||||
select_parent_bridge_messages_for_hook_output(hydrator, staged_records, max_context_chars)
|
||||
.await?;
|
||||
let Some(output) = build_parent_bridge_hook_output(&selected, max_context_chars) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
for message in &selected.messages {
|
||||
let target = parent_bridge_surfaced_message_path(
|
||||
state_dir,
|
||||
message.record.sequence,
|
||||
&message.record.message_id,
|
||||
);
|
||||
fs::rename(&message.path, &target).with_context(|| {
|
||||
format!(
|
||||
"Failed to move message bridge record {} to {}",
|
||||
message.path.display(),
|
||||
target.display()
|
||||
)
|
||||
})?;
|
||||
write_parent_bridge_json_atomically(&target, &message.record)?;
|
||||
}
|
||||
|
||||
remove_file_if_exists(&parent_bridge_hook_output_ack_file(state_dir))?;
|
||||
write_parent_bridge_hook_output(state_dir, &output)
|
||||
}
|
||||
|
||||
pub(super) async fn acknowledge_parent_bridge_hook_output(
|
||||
hydrator: &MessageHydrator,
|
||||
state_dir: &Path,
|
||||
) -> Result<()> {
|
||||
let ack_path = parent_bridge_hook_output_ack_file(state_dir);
|
||||
if !ack_path.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Remove the hook output first so an acknowledged block cannot be re-emitted
|
||||
// if the harness restarts while delivery cleanup is still in progress.
|
||||
remove_file_if_exists(&parent_bridge_hook_output_file(state_dir))?;
|
||||
|
||||
let surfaced_records = parent_bridge_message_records(&parent_bridge_surfaced_dir(state_dir))?;
|
||||
let message_ids = surfaced_records
|
||||
.iter()
|
||||
.map(|(_, record)| record.message_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let delivery_failures = hydrator
|
||||
.mark_messages_delivered_best_effort(message_ids.iter().map(String::as_str))
|
||||
.await;
|
||||
for (message_id, err) in delivery_failures {
|
||||
log::warn!(
|
||||
"Failed to mark Claude message bridge message {message_id} as delivered: {err:#}"
|
||||
);
|
||||
}
|
||||
|
||||
for (path, _) in surfaced_records {
|
||||
remove_file_if_exists(&path)?;
|
||||
}
|
||||
remove_file_if_exists(&ack_path)
|
||||
}
|
||||
|
||||
async fn run_parent_bridge_forever(
|
||||
server_api: Arc<ServerApi>,
|
||||
run_id: String,
|
||||
state_dir: PathBuf,
|
||||
) -> Result<()> {
|
||||
ensure_parent_bridge_state_dir(&state_dir)?;
|
||||
// The shared driver keeps `since_sequence` in memory across its own retry
|
||||
// loop, which is all this per-session bridge needs because the state dir is
|
||||
// not reused across sessions.
|
||||
let config = AgentEventDriverConfig::retry_forever(vec![run_id.clone()], 0);
|
||||
let source = ServerApiAgentEventSource::new(server_api);
|
||||
let mut consumer = MessageBridgeEventConsumer { run_id, state_dir };
|
||||
run_agent_event_driver(source, config, &mut consumer).await
|
||||
}
|
||||
|
||||
fn write_parent_bridge_json_atomically<T: Serialize>(path: &Path, value: &T) -> Result<()> {
|
||||
write_parent_bridge_bytes_atomically(path, &serde_json::to_vec(value)?)
|
||||
}
|
||||
|
||||
fn write_parent_bridge_bytes_atomically(path: &Path, bytes: &[u8]) -> Result<()> {
|
||||
let Some(parent) = path.parent() else {
|
||||
return Err(anyhow!("{} has no parent directory", path.display()));
|
||||
};
|
||||
fs::create_dir_all(parent).with_context(|| format!("Failed to create {}", parent.display()))?;
|
||||
|
||||
let prefix = path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.unwrap_or("parent-bridge");
|
||||
let mut temp_file = NamedTempFile::new_in(parent)
|
||||
.with_context(|| format!("Failed to create temp file for {}", path.display()))?;
|
||||
temp_file
|
||||
.write_all(bytes)
|
||||
.with_context(|| format!("Failed to write temp file for {}", path.display()))?;
|
||||
temp_file
|
||||
.flush()
|
||||
.with_context(|| format!("Failed to flush temp file for {}", path.display()))?;
|
||||
temp_file
|
||||
.persist(path)
|
||||
.map(|_| ())
|
||||
.map_err(|err| {
|
||||
anyhow::Error::from(err.error).context(format!("Failed to write {}", path.display()))
|
||||
})
|
||||
.with_context(|| format!("Failed to persist temporary {prefix} file"))?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,590 @@
|
||||
use mockall::predicate::eq;
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use tempfile::TempDir;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::*;
|
||||
use crate::ai::agent_events::MessageHydrator;
|
||||
use crate::server::server_api::ai::{MockAIClient, ReadAgentMessageResponse};
|
||||
|
||||
fn sample_parent_bridge_message(
|
||||
sequence: i64,
|
||||
message_id: &str,
|
||||
subject: &str,
|
||||
body: &str,
|
||||
) -> MessageBridgeMessageRecord {
|
||||
MessageBridgeMessageRecord {
|
||||
sequence,
|
||||
message_id: message_id.to_string(),
|
||||
sender_run_id: "parent-run-456".to_string(),
|
||||
subject: subject.to_string(),
|
||||
body: body.to_string(),
|
||||
occurred_at: "2026-04-17T15:46:00Z".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_staged_parent_bridge_message(
|
||||
sequence: i64,
|
||||
message_id: &str,
|
||||
) -> MessageBridgeMessageRecord {
|
||||
MessageBridgeMessageRecord {
|
||||
sequence,
|
||||
message_id: message_id.to_string(),
|
||||
sender_run_id: String::new(),
|
||||
subject: String::new(),
|
||||
body: String::new(),
|
||||
occurred_at: "2026-04-17T15:46:00Z".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn write_surfaced_parent_bridge_message(state_dir: &Path, record: &MessageBridgeMessageRecord) {
|
||||
fs::write(
|
||||
parent_bridge_surfaced_message_path(state_dir, record.sequence, &record.message_id),
|
||||
serde_json::to_vec(record).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_command_uses_session_id_when_not_resuming() {
|
||||
let uuid = Uuid::new_v4();
|
||||
let cmd = claude_command("claude", &uuid, "/tmp/prompt.txt", None, false);
|
||||
assert!(
|
||||
cmd.contains(&format!("--session-id {uuid}")),
|
||||
"expected --session-id flag in non-resume command, got: {cmd}"
|
||||
);
|
||||
assert!(
|
||||
!cmd.contains("--resume"),
|
||||
"non-resume command should not contain --resume, got: {cmd}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_command_uses_resume_flag_when_resuming() {
|
||||
let uuid = Uuid::new_v4();
|
||||
let cmd = claude_command("claude", &uuid, "/tmp/prompt.txt", None, true);
|
||||
assert!(
|
||||
cmd.contains(&format!("--resume {uuid}")),
|
||||
"expected --resume flag in resume command, got: {cmd}"
|
||||
);
|
||||
assert!(
|
||||
!cmd.contains("--session-id"),
|
||||
"resume command should not contain --session-id, got: {cmd}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_command_pipes_prompt_path() {
|
||||
let uuid = Uuid::new_v4();
|
||||
let cmd = claude_command("claude", &uuid, "/tmp/prompt with spaces.txt", None, true);
|
||||
assert!(
|
||||
cmd.contains("< '/tmp/prompt with spaces.txt'"),
|
||||
"expected single-quoted stdin redirect of the prompt path, got: {cmd}"
|
||||
);
|
||||
assert!(
|
||||
cmd.contains("--dangerously-skip-permissions"),
|
||||
"expected --dangerously-skip-permissions, got: {cmd}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn parent_bridge_root_prefers_environment_override() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path());
|
||||
let root = parent_bridge_root().unwrap();
|
||||
std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV);
|
||||
|
||||
assert_eq!(root, tmp.path());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stage_parent_bridge_message_writes_message_record() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let state_dir = tmp.path().join("session-123");
|
||||
ensure_parent_bridge_state_dir(&state_dir).unwrap();
|
||||
let record = sample_staged_parent_bridge_message(42, "msg-123");
|
||||
|
||||
stage_parent_bridge_message(&state_dir, &record).unwrap();
|
||||
|
||||
let staged_path = parent_bridge_staged_message_path(&state_dir, 42, "msg-123");
|
||||
let staged_record: MessageBridgeMessageRecord =
|
||||
serde_json::from_slice(&fs::read(&staged_path).unwrap()).unwrap();
|
||||
assert_eq!(staged_record.sequence, 42);
|
||||
assert_eq!(staged_record.message_id, "msg-123");
|
||||
assert!(staged_record.sender_run_id.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_parent_bridge_hook_output_moves_selected_messages_to_surfaced_dir() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let state_dir = tmp.path().join("session-123");
|
||||
ensure_parent_bridge_state_dir(&state_dir).unwrap();
|
||||
|
||||
let first = sample_parent_bridge_message(
|
||||
42,
|
||||
"msg-123",
|
||||
"Please pivot",
|
||||
"Inspect the failing tests first.",
|
||||
);
|
||||
stage_parent_bridge_message(
|
||||
&state_dir,
|
||||
&sample_staged_parent_bridge_message(42, "msg-123"),
|
||||
)
|
||||
.unwrap();
|
||||
stage_parent_bridge_message(
|
||||
&state_dir,
|
||||
&sample_staged_parent_bridge_message(43, "msg-456"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let mut ai_client = MockAIClient::new();
|
||||
let expected_first = first.clone();
|
||||
ai_client
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.times(1)
|
||||
.returning(move |_| {
|
||||
Ok(ReadAgentMessageResponse {
|
||||
message_id: expected_first.message_id.clone(),
|
||||
sender_run_id: expected_first.sender_run_id.clone(),
|
||||
subject: expected_first.subject.clone(),
|
||||
body: expected_first.body.clone(),
|
||||
sent_at: "2026-04-17T15:46:00Z".to_string(),
|
||||
delivered_at: None,
|
||||
read_at: Some("2026-04-17T15:46:02Z".to_string()),
|
||||
})
|
||||
});
|
||||
let hydrator = MessageHydrator::new(
|
||||
Arc::new(ai_client) as Arc<dyn crate::server::server_api::ai::AIClient>
|
||||
);
|
||||
|
||||
let max_context_chars = parent_bridge_char_count(MESSAGE_BRIDGE_CONTEXT_PREAMBLE)
|
||||
+ parent_bridge_char_count(&render_parent_bridge_message_block(&first));
|
||||
prepare_parent_bridge_hook_output(&hydrator, &state_dir, max_context_chars)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let hook_output: MessageBridgeHookOutput =
|
||||
serde_json::from_slice(&fs::read(parent_bridge_hook_output_file(&state_dir)).unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(hook_output.surfaced_count, 1);
|
||||
assert_eq!(hook_output.remaining_staged_count, 1);
|
||||
assert!(hook_output.additional_context.contains("Please pivot"));
|
||||
assert!(!hook_output.additional_context.contains("Second update"));
|
||||
let surfaced_path = parent_bridge_surfaced_message_path(&state_dir, 42, "msg-123");
|
||||
assert!(surfaced_path.exists());
|
||||
assert!(parent_bridge_staged_message_path(&state_dir, 43, "msg-456").exists());
|
||||
assert!(!parent_bridge_staged_message_path(&state_dir, 42, "msg-123").exists());
|
||||
let surfaced_record: MessageBridgeMessageRecord =
|
||||
serde_json::from_slice(&fs::read(&surfaced_path).unwrap()).unwrap();
|
||||
assert_eq!(surfaced_record.subject, first.subject);
|
||||
assert_eq!(surfaced_record.body, first.body);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn acknowledge_parent_bridge_hook_output_marks_messages_delivered_and_clears_state() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let state_dir = tmp.path().join("session-123");
|
||||
ensure_parent_bridge_state_dir(&state_dir).unwrap();
|
||||
|
||||
let record = sample_parent_bridge_message(
|
||||
42,
|
||||
"msg-123",
|
||||
"Please pivot",
|
||||
"Inspect the failing tests first.",
|
||||
);
|
||||
write_surfaced_parent_bridge_message(&state_dir, &record);
|
||||
fs::write(
|
||||
parent_bridge_hook_output_file(&state_dir),
|
||||
serde_json::to_vec(&MessageBridgeHookOutput {
|
||||
additional_context: "context".to_string(),
|
||||
remaining_staged_count: 0,
|
||||
surfaced_count: 1,
|
||||
})
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
fs::write(parent_bridge_hook_output_ack_file(&state_dir), "").unwrap();
|
||||
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.expect_mark_message_delivered()
|
||||
.with(eq("msg-123"))
|
||||
.times(1)
|
||||
.returning(|_| Ok(()));
|
||||
let hydrator = MessageHydrator::new(
|
||||
Arc::new(ai_client) as Arc<dyn crate::server::server_api::ai::AIClient>
|
||||
);
|
||||
|
||||
acknowledge_parent_bridge_hook_output(&hydrator, &state_dir)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!parent_bridge_surfaced_message_path(&state_dir, 42, "msg-123").exists());
|
||||
assert!(!parent_bridge_hook_output_file(&state_dir).exists());
|
||||
assert!(!parent_bridge_hook_output_ack_file(&state_dir).exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_parent_bridge_hook_output_reuses_surfaced_records_without_rehydrating() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let state_dir = tmp.path().join("session-123");
|
||||
ensure_parent_bridge_state_dir(&state_dir).unwrap();
|
||||
|
||||
let record = sample_parent_bridge_message(
|
||||
42,
|
||||
"msg-123",
|
||||
"Please pivot",
|
||||
"Inspect the failing tests first.",
|
||||
);
|
||||
write_surfaced_parent_bridge_message(&state_dir, &record);
|
||||
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client.expect_read_agent_message().times(0);
|
||||
let hydrator = MessageHydrator::new(
|
||||
Arc::new(ai_client) as Arc<dyn crate::server::server_api::ai::AIClient>
|
||||
);
|
||||
|
||||
let max_context_chars = parent_bridge_char_count(MESSAGE_BRIDGE_CONTEXT_PREAMBLE)
|
||||
+ parent_bridge_char_count(&render_parent_bridge_message_block(&record));
|
||||
prepare_parent_bridge_hook_output(&hydrator, &state_dir, max_context_chars)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let hook_output: MessageBridgeHookOutput =
|
||||
serde_json::from_slice(&fs::read(parent_bridge_hook_output_file(&state_dir)).unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(hook_output.surfaced_count, 1);
|
||||
assert_eq!(hook_output.remaining_staged_count, 0);
|
||||
assert!(hook_output.additional_context.contains(&record.subject));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_parent_bridge_hook_output_truncates_single_large_message() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let state_dir = tmp.path().join("session-123");
|
||||
ensure_parent_bridge_state_dir(&state_dir).unwrap();
|
||||
stage_parent_bridge_message(
|
||||
&state_dir,
|
||||
&sample_staged_parent_bridge_message(42, "msg-123"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let long_body = "x".repeat(200);
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client
|
||||
.expect_read_agent_message()
|
||||
.with(eq("msg-123"))
|
||||
.times(1)
|
||||
.returning(move |_| {
|
||||
Ok(ReadAgentMessageResponse {
|
||||
message_id: "msg-123".to_string(),
|
||||
sender_run_id: "parent-run-456".to_string(),
|
||||
subject: "Please pivot".to_string(),
|
||||
body: long_body.clone(),
|
||||
sent_at: "2026-04-17T15:46:00Z".to_string(),
|
||||
delivered_at: None,
|
||||
read_at: Some("2026-04-17T15:46:02Z".to_string()),
|
||||
})
|
||||
});
|
||||
let hydrator = MessageHydrator::new(
|
||||
Arc::new(ai_client) as Arc<dyn crate::server::server_api::ai::AIClient>
|
||||
);
|
||||
|
||||
let max_context_chars = parent_bridge_char_count(MESSAGE_BRIDGE_CONTEXT_PREAMBLE) + 48;
|
||||
prepare_parent_bridge_hook_output(&hydrator, &state_dir, max_context_chars)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let hook_output: MessageBridgeHookOutput =
|
||||
serde_json::from_slice(&fs::read(parent_bridge_hook_output_file(&state_dir)).unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(hook_output.surfaced_count, 1);
|
||||
assert!(
|
||||
hook_output.additional_context.ends_with("..."),
|
||||
"expected truncated context, got: {}",
|
||||
hook_output.additional_context
|
||||
);
|
||||
assert!(parent_bridge_char_count(&hook_output.additional_context) <= max_context_chars);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn acknowledge_parent_bridge_hook_output_ignores_missing_ack_marker() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let state_dir = tmp.path().join("session-123");
|
||||
ensure_parent_bridge_state_dir(&state_dir).unwrap();
|
||||
|
||||
let record = sample_parent_bridge_message(
|
||||
42,
|
||||
"msg-123",
|
||||
"Please pivot",
|
||||
"Inspect the failing tests first.",
|
||||
);
|
||||
write_surfaced_parent_bridge_message(&state_dir, &record);
|
||||
|
||||
let mut ai_client = MockAIClient::new();
|
||||
ai_client.expect_mark_message_delivered().times(0);
|
||||
let hydrator = MessageHydrator::new(
|
||||
Arc::new(ai_client) as Arc<dyn crate::server::server_api::ai::AIClient>
|
||||
);
|
||||
|
||||
acknowledge_parent_bridge_hook_output(&hydrator, &state_dir)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(parent_bridge_surfaced_message_path(&state_dir, 42, "msg-123").exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_config_creates_config_file_without_api_suffix() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_json_path = tmp.path().join(".claude.json");
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
|
||||
prepare_claude_config(&claude_json_path, &working_dir, None).unwrap();
|
||||
|
||||
let claude_config: Value =
|
||||
serde_json::from_slice(&fs::read(claude_json_path).unwrap()).unwrap();
|
||||
assert_eq!(claude_config["hasCompletedOnboarding"], Value::Bool(true));
|
||||
assert_eq!(
|
||||
claude_config["lspRecommendationDisabled"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
let working_dir_key = working_dir.to_string_lossy().to_string();
|
||||
assert_eq!(
|
||||
claude_config["projects"][working_dir_key]["hasTrustDialogAccepted"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
assert_eq!(claude_config.get("customApiKeyResponses"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_config_creates_config_file_with_api_suffix() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_json_path = tmp.path().join(".claude.json");
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
|
||||
prepare_claude_config(
|
||||
&claude_json_path,
|
||||
&working_dir,
|
||||
Some("QLWn-dUnuwQ-hIhDiAAA"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let claude_config: Value =
|
||||
serde_json::from_slice(&fs::read(claude_json_path).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
claude_config["customApiKeyResponses"]["approved"],
|
||||
serde_json::json!(["QLWn-dUnuwQ-hIhDiAAA"]),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_config_merges_existing_config() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_json_path = tmp.path().join(".claude.json");
|
||||
fs::write(
|
||||
&claude_json_path,
|
||||
r#"{"theme":"dark","projects":{"/existing/project":{"allowedTools":["Bash"],"nested":{"value":2}}},"customApiKeyResponses":{"approved":["existing-suffix-12345"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
prepare_claude_config(
|
||||
&claude_json_path,
|
||||
&working_dir,
|
||||
Some("new-suffix-1234567890"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let claude_config: Value =
|
||||
serde_json::from_slice(&fs::read(claude_json_path).unwrap()).unwrap();
|
||||
assert_eq!(claude_config["theme"], "dark");
|
||||
assert_eq!(
|
||||
claude_config["lspRecommendationDisabled"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
assert_eq!(
|
||||
claude_config["projects"]["/existing/project"]["allowedTools"],
|
||||
serde_json::json!(["Bash"])
|
||||
);
|
||||
assert_eq!(
|
||||
claude_config["projects"]["/existing/project"]["nested"]["value"],
|
||||
2
|
||||
);
|
||||
// Both existing and new suffixes should be present.
|
||||
assert_eq!(
|
||||
claude_config["customApiKeyResponses"]["approved"],
|
||||
serde_json::json!(["existing-suffix-12345", "new-suffix-1234567890"]),
|
||||
);
|
||||
let working_dir_key = working_dir.to_string_lossy().to_string();
|
||||
assert_eq!(
|
||||
claude_config["projects"][working_dir_key]["hasTrustDialogAccepted"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_config_no_duplicate_suffix() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_json_path = tmp.path().join(".claude.json");
|
||||
fs::write(
|
||||
&claude_json_path,
|
||||
r#"{"customApiKeyResponses":{"approved":["QLWn-dUnuwQ-hIhDiAAA"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
prepare_claude_config(
|
||||
&claude_json_path,
|
||||
&working_dir,
|
||||
Some("QLWn-dUnuwQ-hIhDiAAA"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let claude_config: Value =
|
||||
serde_json::from_slice(&fs::read(claude_json_path).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
claude_config["customApiKeyResponses"]["approved"],
|
||||
serde_json::json!(["QLWn-dUnuwQ-hIhDiAAA"]),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_config_none_suffix_preserves_existing_responses() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_json_path = tmp.path().join(".claude.json");
|
||||
fs::write(
|
||||
&claude_json_path,
|
||||
r#"{"customApiKeyResponses":{"approved":["existing-suffix-12345"],"rejected":["bad-key"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
prepare_claude_config(&claude_json_path, &working_dir, None).unwrap();
|
||||
|
||||
let claude_config: Value =
|
||||
serde_json::from_slice(&fs::read(claude_json_path).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
claude_config["customApiKeyResponses"]["approved"],
|
||||
serde_json::json!(["existing-suffix-12345"]),
|
||||
);
|
||||
assert_eq!(
|
||||
claude_config["customApiKeyResponses"]["rejected"],
|
||||
serde_json::json!(["bad-key"]),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_suffix_from_raw_value_secret() {
|
||||
let key = "sk-ant-api03-abcdefghij1234567890ABCDEFGHIJ1234567890abcdefghij1234567890QLWn-dUnuwQ-hIhDiAAA";
|
||||
let secrets = HashMap::from([(
|
||||
"ANTHROPIC_API_KEY".to_string(),
|
||||
ManagedSecretValue::raw_value(key),
|
||||
)]);
|
||||
let suffix = resolve_anthropic_api_key_suffix(&secrets);
|
||||
assert_eq!(suffix.as_deref(), Some("QLWn-dUnuwQ-hIhDiAAA"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_suffix_from_anthropic_api_key_secret() {
|
||||
let key = "sk-ant-api03-abcdefghij1234567890ABCDEFGHIJ1234567890abcdefghij1234567890QLWn-dUnuwQ-hIhDiAAA";
|
||||
let secrets = HashMap::from([(
|
||||
"ANTHROPIC_API_KEY".to_string(),
|
||||
ManagedSecretValue::anthropic_api_key(key),
|
||||
)]);
|
||||
let suffix = resolve_anthropic_api_key_suffix(&secrets);
|
||||
assert_eq!(suffix.as_deref(), Some("QLWn-dUnuwQ-hIhDiAAA"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_suffix_from_anthropic_api_key_with_different_secret_name() {
|
||||
let key = "sk-ant-api03-abcdefghij1234567890ABCDEFGHIJ1234567890abcdefghij1234567890QLWn-dUnuwQ-hIhDiAAA";
|
||||
// Secret name doesn't match the env var, but the AnthropicApiKey variant
|
||||
// should still be found by iterating all secrets.
|
||||
let secrets = HashMap::from([(
|
||||
"my-anthropic-key".to_string(),
|
||||
ManagedSecretValue::anthropic_api_key(key),
|
||||
)]);
|
||||
let suffix = resolve_anthropic_api_key_suffix(&secrets);
|
||||
assert_eq!(suffix.as_deref(), Some("QLWn-dUnuwQ-hIhDiAAA"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_suffix_prefers_anthropic_api_key_variant_over_raw_value() {
|
||||
let anthropic_key = "sk-ant-api03-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA-anthropic-suffix";
|
||||
let raw_key = "sk-ant-api03-BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB-raw-suffix";
|
||||
let secrets = HashMap::from([
|
||||
(
|
||||
"my-anthropic-key".to_string(),
|
||||
ManagedSecretValue::anthropic_api_key(anthropic_key),
|
||||
),
|
||||
(
|
||||
"ANTHROPIC_API_KEY".to_string(),
|
||||
ManagedSecretValue::raw_value(raw_key),
|
||||
),
|
||||
]);
|
||||
let suffix = resolve_anthropic_api_key_suffix(&secrets);
|
||||
// AnthropicApiKey variant should be preferred.
|
||||
assert_eq!(suffix.as_deref(), Some("AAA-anthropic-suffix"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_suffix_returns_none_for_short_key() {
|
||||
let secrets = HashMap::from([(
|
||||
"ANTHROPIC_API_KEY".to_string(),
|
||||
ManagedSecretValue::raw_value("short"),
|
||||
)]);
|
||||
assert_eq!(resolve_anthropic_api_key_suffix(&secrets), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_suffix_returns_none_for_short_anthropic_api_key() {
|
||||
let secrets = HashMap::from([(
|
||||
"ANTHROPIC_API_KEY".to_string(),
|
||||
ManagedSecretValue::anthropic_api_key("short"),
|
||||
)]);
|
||||
assert_eq!(resolve_anthropic_api_key_suffix(&secrets), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_settings_creates_settings_file() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_settings_path = tmp.path().join(".claude/settings.json");
|
||||
|
||||
prepare_claude_settings(&claude_settings_path).unwrap();
|
||||
|
||||
let claude_settings: Value =
|
||||
serde_json::from_slice(&fs::read(claude_settings_path).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
claude_settings["skipDangerousModePermissionPrompt"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_claude_settings_merges_existing_settings() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let claude_settings_path = tmp.path().join("settings.json");
|
||||
fs::write(
|
||||
&claude_settings_path,
|
||||
r#"{"editor":"vim","nested":{"value":1}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
prepare_claude_settings(&claude_settings_path).unwrap();
|
||||
|
||||
let claude_settings: Value =
|
||||
serde_json::from_slice(&fs::read(claude_settings_path).unwrap()).unwrap();
|
||||
assert_eq!(claude_settings["editor"], "vim");
|
||||
assert_eq!(claude_settings["nested"]["value"], 1);
|
||||
assert_eq!(
|
||||
claude_settings["skipDangerousModePermissionPrompt"],
|
||||
Value::Bool(true)
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,343 @@
|
||||
//! Claude Code transcript layout + rehydration helpers.
|
||||
//!
|
||||
//! This module owns:
|
||||
//! - [`ClaudeTranscriptEnvelope`] — the on-wire/on-GCS shape of a saved Claude session
|
||||
//! (main jsonl entries + subagent jsonl files + per-agent todo JSONs), plus reader/writer
|
||||
//! functions that interoperate with Claude's own `~/.claude` layout.
|
||||
//! - [`ClaudeResumeInfo`] — everything the harness runner needs to resume an existing
|
||||
//! Claude conversation: the Warp server conversation id to reuse, the Claude session uuid
|
||||
//! to pass to `claude --resume`, and the decoded envelope to rehydrate onto disk.
|
||||
//! - [`write_session_index_entry`] — best-effort update of `~/.claude/sessions-index.json`
|
||||
//! so Claude's `--resume <uuid>` lookup can find the freshly-rehydrated jsonl. Upstream
|
||||
//! versions vary in how they use this index (claude-code#33912, #39667, #5768); we write
|
||||
//! a conservative entry and log on failure.
|
||||
//!
|
||||
//! Split out from `claude_code.rs` so the `AIClient` transcript-fetch impl can deserialize
|
||||
//! envelopes without pulling in the rest of the harness runner.
|
||||
use std::collections::HashMap;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use uuid::Uuid;
|
||||
use warp_core::safe_warn;
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
|
||||
/// JSON envelope sent to the server representing a complete Claude Code session.
|
||||
///
|
||||
/// Bundles the main session transcript, any subagent transcripts, and
|
||||
/// per-agent TODO lists assembled from the Claude state directory.
|
||||
#[derive(Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub(crate) struct ClaudeTranscriptEnvelope {
|
||||
/// The directory that the Claude Code session started in.
|
||||
pub(crate) cwd: PathBuf,
|
||||
/// Unique session identifier.
|
||||
pub(crate) uuid: Uuid,
|
||||
/// Claude Code version, if available.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) claude_version: Option<String>,
|
||||
/// List of messages in the main agent conversation.
|
||||
pub(crate) entries: Vec<Value>,
|
||||
/// Messages in each subagent conversation, keyed by the agent filename (e.g. `"agent-aac0b7f3db6bccfaf"`).
|
||||
pub(crate) subagents: HashMap<String, Vec<Value>>,
|
||||
/// TODO lists for each agent, keyed on the session and agent (e.g. `"<session_uuid>-agent-<agent_id>"`).
|
||||
pub(crate) todos: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
/// Everything needed to resume an existing Claude conversation.
|
||||
///
|
||||
/// Populated from a `--conversation` id after the client fetches the stored envelope from
|
||||
/// the server. Passed into `ClaudeHarnessRunner::new` so the runner reuses the existing
|
||||
/// session and server conversation ids instead of minting fresh ones.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct ClaudeResumeInfo {
|
||||
/// The Warp server-side conversation id. The runner stores this instead of calling
|
||||
/// `create_external_conversation` so subsequent transcript/block-snapshot uploads overwrite
|
||||
/// the same GCS objects.
|
||||
pub(crate) conversation_id: AIConversationId,
|
||||
/// The Claude session uuid to pass to `claude --resume`. Matches `envelope.uuid`.
|
||||
pub(crate) session_id: Uuid,
|
||||
/// Envelope from the server. Its `cwd` field is rewritten to the current run's working
|
||||
/// directory before being written to disk, so `claude --resume <uuid>` finds the jsonl under
|
||||
/// `~/.claude/projects/<encoded(new_cwd)>/`.
|
||||
pub(crate) envelope: ClaudeTranscriptEnvelope,
|
||||
}
|
||||
|
||||
/// Encode a filesystem path as a Claude config directory name, matching the
|
||||
/// Claude CLI convention of replacing every `/` with `-`.
|
||||
///
|
||||
/// Example: `/Users/ben/src/foo` → `-Users-ben-src-foo`
|
||||
pub(crate) fn encode_cwd(cwd: &Path) -> String {
|
||||
cwd.to_string_lossy().replace(['/', '.'], "-")
|
||||
}
|
||||
|
||||
/// Resolve the Claude config directory.
|
||||
///
|
||||
/// Reads `$CLAUDE_CONFIG_DIR` if set, otherwise falls back to `~/.claude`.
|
||||
//
|
||||
/// TODO(REMOTE-1209): Use the transcript path reported by our hook.
|
||||
pub(crate) fn claude_config_dir() -> Result<PathBuf> {
|
||||
if let Ok(dir) = std::env::var("CLAUDE_CONFIG_DIR") {
|
||||
return Ok(PathBuf::from(dir));
|
||||
}
|
||||
dirs::home_dir()
|
||||
.map(|h| h.join(".claude"))
|
||||
.ok_or_else(|| anyhow::anyhow!("could not determine home directory"))
|
||||
}
|
||||
|
||||
/// Assemble a [`ClaudeTranscriptEnvelope`] from the Claude config directory.
|
||||
///
|
||||
/// Reads:
|
||||
/// - `<config_root>/projects/<encoded_cwd>/<session_uuid>.jsonl` - main transcript
|
||||
/// - `<config_root>/projects/<encoded_cwd>/<session_uuid>/subagents/*.jsonl` - subagents
|
||||
/// - `<config_root>/todos/<session_uuid>-agent-*.json` - per-agent todo lists
|
||||
///
|
||||
/// If the main JSONL does not exist yet (e.g. during an early periodic save)
|
||||
/// the envelope is returned with an empty `entries` list rather than an error.
|
||||
pub(crate) fn read_envelope(
|
||||
session_uuid: Uuid,
|
||||
cwd: &Path,
|
||||
config_root: &Path,
|
||||
) -> Result<ClaudeTranscriptEnvelope> {
|
||||
let encoded = encode_cwd(cwd);
|
||||
let projects_dir = config_root.join("projects").join(&encoded);
|
||||
|
||||
// Main session transcript.
|
||||
let session_file = projects_dir.join(format!("{session_uuid}.jsonl"));
|
||||
let entries = read_jsonl(&session_file)?;
|
||||
|
||||
// Subagents are stored in a directory named after the session UUID.
|
||||
let mut subagents: HashMap<String, Vec<Value>> = HashMap::new();
|
||||
let subagents_dir = projects_dir
|
||||
.join(session_uuid.to_string())
|
||||
.join("subagents");
|
||||
if subagents_dir.is_dir() {
|
||||
for entry in std::fs::read_dir(&subagents_dir)
|
||||
.with_context(|| format!("Failed to read subagents dir {}", subagents_dir.display()))?
|
||||
{
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("jsonl") {
|
||||
continue;
|
||||
}
|
||||
let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
subagents.insert(stem.to_owned(), read_jsonl(&path)?);
|
||||
}
|
||||
}
|
||||
|
||||
// Per-agent todo lists.
|
||||
let mut todos: HashMap<String, Value> = HashMap::new();
|
||||
let todos_dir = config_root.join("todos");
|
||||
let todos_prefix = format!("{session_uuid}-agent-");
|
||||
if todos_dir.is_dir() {
|
||||
for entry in std::fs::read_dir(&todos_dir)
|
||||
.with_context(|| format!("Failed to read todos dir {}", todos_dir.display()))?
|
||||
{
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("json") {
|
||||
continue;
|
||||
}
|
||||
let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
if !stem.starts_with(&todos_prefix) {
|
||||
continue;
|
||||
}
|
||||
match std::fs::read_to_string(&path) {
|
||||
Ok(content) => match serde_json::from_str(&content) {
|
||||
Ok(value) => {
|
||||
todos.insert(stem.to_owned(), value);
|
||||
}
|
||||
Err(e) => log::warn!("Failed to parse todos file {}: {e}", path.display()),
|
||||
},
|
||||
Err(e) => log::warn!("Failed to read todos file {}: {e}", path.display()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ClaudeTranscriptEnvelope {
|
||||
cwd: cwd.to_path_buf(),
|
||||
uuid: session_uuid,
|
||||
claude_version: None,
|
||||
entries,
|
||||
subagents,
|
||||
todos,
|
||||
})
|
||||
}
|
||||
|
||||
/// Write a [`ClaudeTranscriptEnvelope`] back to disk using the same layout
|
||||
/// that Claude Code uses.
|
||||
///
|
||||
/// Creates:
|
||||
/// - `<config_root>/projects/<encoded_cwd>/<uuid>.jsonl` - main transcript
|
||||
/// - `<config_root>/projects/<encoded_cwd>/<uuid>/subagents/<stem>.jsonl` - subagents
|
||||
/// - `<config_root>/todos/<stem>.json` - per-agent todo lists
|
||||
pub(crate) fn write_envelope(
|
||||
envelope: &ClaudeTranscriptEnvelope,
|
||||
config_root: &Path,
|
||||
) -> Result<()> {
|
||||
let encoded = encode_cwd(&envelope.cwd);
|
||||
let projects_dir = config_root.join("projects").join(&encoded);
|
||||
std::fs::create_dir_all(&projects_dir)
|
||||
.with_context(|| format!("Failed to create {}", projects_dir.display()))?;
|
||||
|
||||
// Main session JSONL.
|
||||
let session_file = projects_dir.join(format!("{}.jsonl", envelope.uuid));
|
||||
std::fs::write(&session_file, entries_to_jsonl(&envelope.entries)?)
|
||||
.with_context(|| format!("Failed to write {}", session_file.display()))?;
|
||||
|
||||
// Subagent JSONLs.
|
||||
if !envelope.subagents.is_empty() {
|
||||
let subagents_dir = projects_dir
|
||||
.join(envelope.uuid.to_string())
|
||||
.join("subagents");
|
||||
std::fs::create_dir_all(&subagents_dir)
|
||||
.with_context(|| format!("Failed to create {}", subagents_dir.display()))?;
|
||||
for (stem, entries) in &envelope.subagents {
|
||||
let path = subagents_dir.join(format!("{stem}.jsonl"));
|
||||
std::fs::write(&path, entries_to_jsonl(entries)?)
|
||||
.with_context(|| format!("Failed to write {}", path.display()))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Per-agent todo lists.
|
||||
if !envelope.todos.is_empty() {
|
||||
let todos_dir = config_root.join("todos");
|
||||
std::fs::create_dir_all(&todos_dir)
|
||||
.with_context(|| format!("Failed to create {}", todos_dir.display()))?;
|
||||
for (stem, value) in &envelope.todos {
|
||||
let path = todos_dir.join(format!("{stem}.json"));
|
||||
std::fs::write(&path, serde_json::to_vec(value)?)
|
||||
.with_context(|| format!("Failed to write {}", path.display()))?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Filename of Claude's global session index.
|
||||
const SESSIONS_INDEX_FILENAME: &str = "sessions-index.json";
|
||||
|
||||
/// Upsert an entry for `session_uuid` into `<config_root>/sessions-index.json` so Claude's
|
||||
/// `claude --resume <uuid>` lookup can find the rehydrated jsonl.
|
||||
///
|
||||
/// Upstream Claude versions vary in how the index is keyed and what fields they read; this
|
||||
/// writer uses a conservative session-uuid-keyed schema (session id, cwd, jsonl path) that
|
||||
/// mirrors the fragments documented in claude-code#33912 / #39667 / #5768. Unknown fields are
|
||||
/// preserved on existing entries, and we never remove other entries.
|
||||
///
|
||||
/// Best-effort: callers should log a warning on failure rather than aborting the run — if the
|
||||
/// index is missing or wrong, `--resume` simply falls back to "No conversation found" and the
|
||||
/// resumed run surfaces the expected resume-failure error.
|
||||
pub(crate) fn write_session_index_entry(
|
||||
session_uuid: Uuid,
|
||||
cwd: &Path,
|
||||
config_root: &Path,
|
||||
) -> Result<()> {
|
||||
let index_path = config_root.join(SESSIONS_INDEX_FILENAME);
|
||||
|
||||
// Read the existing index if present. Missing or malformed files are treated as empty —
|
||||
// we'd rather clobber an unparseable file than fail the whole resume.
|
||||
let mut index: serde_json::Map<String, Value> = match std::fs::read_to_string(&index_path) {
|
||||
Ok(content) => match serde_json::from_str::<Value>(&content) {
|
||||
Ok(Value::Object(map)) => map,
|
||||
Ok(_) => {
|
||||
safe_warn!(
|
||||
safe: ("sessions-index.json is not a JSON object; overwriting"),
|
||||
full: ("sessions-index.json at {} is not a JSON object; overwriting", index_path.display())
|
||||
);
|
||||
serde_json::Map::new()
|
||||
}
|
||||
Err(e) => {
|
||||
safe_warn!(
|
||||
safe: ("Failed to parse sessions-index.json; overwriting"),
|
||||
full: ("Failed to parse sessions-index.json at {}: {e}; overwriting", index_path.display())
|
||||
);
|
||||
serde_json::Map::new()
|
||||
}
|
||||
},
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => serde_json::Map::new(),
|
||||
Err(e) => {
|
||||
return Err(
|
||||
anyhow::Error::from(e).context(format!("Failed to read {}", index_path.display()))
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let encoded = encode_cwd(cwd);
|
||||
let transcript_path = format!("projects/{encoded}/{session_uuid}.jsonl");
|
||||
let entry = serde_json::json!({
|
||||
"sessionId": session_uuid.to_string(),
|
||||
"cwd": cwd.to_string_lossy(),
|
||||
"projectPath": encoded,
|
||||
"transcriptPath": transcript_path,
|
||||
});
|
||||
index.insert(session_uuid.to_string(), entry);
|
||||
|
||||
if let Some(parent) = index_path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.with_context(|| format!("Failed to create {}", parent.display()))?;
|
||||
}
|
||||
std::fs::write(
|
||||
&index_path,
|
||||
serde_json::to_vec_pretty(&Value::Object(index))
|
||||
.context("Failed to serialize sessions-index.json")?,
|
||||
)
|
||||
.with_context(|| format!("Failed to write {}", index_path.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Serialize a slice of JSON values as a JSONL byte string (one value per line).
|
||||
fn entries_to_jsonl(entries: &[Value]) -> Result<Vec<u8>> {
|
||||
let mut buf = Vec::new();
|
||||
for entry in entries {
|
||||
serde_json::to_writer(&mut buf, entry)?;
|
||||
buf.push(b'\n');
|
||||
}
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
/// Read a JSONL file, returning one parsed [`Value`] per non-blank line.
|
||||
///
|
||||
/// Lines that fail to parse as JSON are skipped with a warning rather than
|
||||
/// causing the entire read to fail. A missing file returns an empty [`Vec`].
|
||||
pub(crate) fn read_jsonl(path: &Path) -> Result<Vec<Value>> {
|
||||
let file = match std::fs::File::open(path) {
|
||||
Ok(f) => f,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
|
||||
Err(e) => {
|
||||
return Err(
|
||||
anyhow::Error::from(e).context(format!("Failed to open {}", path.display()))
|
||||
);
|
||||
}
|
||||
};
|
||||
let reader = BufReader::new(file);
|
||||
let mut entries = Vec::new();
|
||||
for line in reader.lines() {
|
||||
let line = line.with_context(|| format!("Failed to read line from {}", path.display()))?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
match serde_json::from_str(trimmed) {
|
||||
Ok(value) => entries.push(value),
|
||||
Err(e) => {
|
||||
safe_warn!(
|
||||
safe: ("Skipping malformed JSONL entry"),
|
||||
full: ("Skipping malformed JSONL entry in {}: {e}", path.display())
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "claude_transcript_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,246 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use tempfile::TempDir;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn write_file(dir: &Path, name: &str, content: &str) {
|
||||
fs::write(dir.join(name), content).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encode_cwd_replaces_slashes_and_dots() {
|
||||
assert_eq!(
|
||||
encode_cwd(Path::new("/Users/ben/src/foo")),
|
||||
"-Users-ben-src-foo"
|
||||
);
|
||||
assert_eq!(encode_cwd(Path::new("/")), "-");
|
||||
assert_eq!(encode_cwd(Path::new("/a/b")), "-a-b");
|
||||
assert_eq!(
|
||||
encode_cwd(Path::new("/Users/ben/.config/foo")),
|
||||
"-Users-ben--config-foo"
|
||||
);
|
||||
assert_eq!(encode_cwd(Path::new("/a.b/c.d")), "-a-b-c-d");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_envelope_main_only() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = Path::new("/my/project");
|
||||
let uuid = Uuid::new_v4();
|
||||
|
||||
let encoded = encode_cwd(cwd);
|
||||
let projects_dir = tmp.path().join("projects").join(&encoded);
|
||||
fs::create_dir_all(&projects_dir).unwrap();
|
||||
write_file(
|
||||
&projects_dir,
|
||||
&format!("{uuid}.jsonl"),
|
||||
"{\"type\":\"user\"}\n{\"type\":\"assistant\"}\n",
|
||||
);
|
||||
|
||||
let envelope = read_envelope(uuid, cwd, tmp.path()).unwrap();
|
||||
assert_eq!(
|
||||
envelope.entries,
|
||||
vec![
|
||||
serde_json::json!({"type": "user"}),
|
||||
serde_json::json!({"type": "assistant"}),
|
||||
]
|
||||
);
|
||||
assert_eq!(envelope.uuid, uuid);
|
||||
assert_eq!(envelope.cwd, cwd);
|
||||
assert!(envelope.subagents.is_empty());
|
||||
assert!(envelope.todos.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_envelope_with_subagents() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = Path::new("/my/project");
|
||||
let uuid = Uuid::new_v4();
|
||||
|
||||
let encoded = encode_cwd(cwd);
|
||||
let projects_dir = tmp.path().join("projects").join(&encoded);
|
||||
fs::create_dir_all(&projects_dir).unwrap();
|
||||
write_file(&projects_dir, &format!("{uuid}.jsonl"), "");
|
||||
|
||||
let subagents_dir = projects_dir.join(uuid.to_string()).join("subagents");
|
||||
fs::create_dir_all(&subagents_dir).unwrap();
|
||||
write_file(
|
||||
&subagents_dir,
|
||||
"agent-abc123def456.jsonl",
|
||||
"{\"type\":\"user\"}\n",
|
||||
);
|
||||
|
||||
let envelope = read_envelope(uuid, cwd, tmp.path()).unwrap();
|
||||
assert_eq!(
|
||||
envelope.subagents["agent-abc123def456"],
|
||||
vec![serde_json::json!({"type": "user"})]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_envelope_missing_session_file() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = Path::new("/my/project");
|
||||
let uuid = Uuid::new_v4();
|
||||
|
||||
// No files created - should return Ok with empty entries rather than an error.
|
||||
let envelope = read_envelope(uuid, cwd, tmp.path()).unwrap();
|
||||
assert!(envelope.entries.is_empty());
|
||||
assert!(envelope.subagents.is_empty());
|
||||
assert!(envelope.todos.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_envelope_creates_files() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = Path::new("/my/project");
|
||||
let uuid = Uuid::new_v4();
|
||||
let todo_stem = format!("{uuid}-agent-{uuid}");
|
||||
|
||||
let envelope = ClaudeTranscriptEnvelope {
|
||||
cwd: cwd.to_path_buf(),
|
||||
uuid,
|
||||
claude_version: None,
|
||||
entries: vec![serde_json::json!({"type": "user"})],
|
||||
subagents: HashMap::from([(
|
||||
"agent-abc".to_string(),
|
||||
vec![serde_json::json!({"type": "assistant"})],
|
||||
)]),
|
||||
todos: HashMap::from([(
|
||||
todo_stem.clone(),
|
||||
serde_json::json!([{"id": "1", "title": "Do it"}]),
|
||||
)]),
|
||||
};
|
||||
|
||||
write_envelope(&envelope, tmp.path()).unwrap();
|
||||
|
||||
let encoded = encode_cwd(cwd);
|
||||
let projects_dir = tmp.path().join("projects").join(&encoded);
|
||||
|
||||
// Main session JSONL.
|
||||
let session_file = projects_dir.join(format!("{uuid}.jsonl"));
|
||||
assert!(session_file.exists(), "session JSONL missing");
|
||||
assert_eq!(read_jsonl(&session_file).unwrap(), envelope.entries);
|
||||
|
||||
// Subagent JSONL.
|
||||
let subagent_file = projects_dir
|
||||
.join(uuid.to_string())
|
||||
.join("subagents")
|
||||
.join("agent-abc.jsonl");
|
||||
assert!(subagent_file.exists(), "subagent JSONL missing");
|
||||
assert_eq!(
|
||||
read_jsonl(&subagent_file).unwrap(),
|
||||
envelope.subagents["agent-abc"]
|
||||
);
|
||||
|
||||
// Todo JSON.
|
||||
let todo_file = tmp.path().join("todos").join(format!("{todo_stem}.json"));
|
||||
assert!(todo_file.exists(), "todo JSON missing");
|
||||
let todo_on_disk: serde_json::Value =
|
||||
serde_json::from_slice(&fs::read(&todo_file).unwrap()).unwrap();
|
||||
assert_eq!(todo_on_disk, envelope.todos[&todo_stem]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_envelope_round_trip() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let cwd = Path::new("/test/dir");
|
||||
let uuid = Uuid::new_v4();
|
||||
|
||||
let original = ClaudeTranscriptEnvelope {
|
||||
cwd: cwd.to_path_buf(),
|
||||
uuid,
|
||||
claude_version: None,
|
||||
entries: vec![serde_json::json!({"type": "user"})],
|
||||
subagents: HashMap::new(),
|
||||
todos: HashMap::new(),
|
||||
};
|
||||
|
||||
write_envelope(&original, tmp.path()).unwrap();
|
||||
|
||||
let decoded = read_envelope(uuid, cwd, tmp.path()).unwrap();
|
||||
assert_eq!(decoded, original);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_session_index_entry_creates_missing_file() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let uuid = Uuid::new_v4();
|
||||
let cwd = Path::new("/my/project");
|
||||
|
||||
write_session_index_entry(uuid, cwd, tmp.path()).unwrap();
|
||||
|
||||
let index: serde_json::Value =
|
||||
serde_json::from_slice(&fs::read(tmp.path().join("sessions-index.json")).unwrap()).unwrap();
|
||||
let entry = &index[uuid.to_string()];
|
||||
assert_eq!(entry["sessionId"], uuid.to_string());
|
||||
assert_eq!(entry["cwd"], "/my/project");
|
||||
assert_eq!(entry["projectPath"], "-my-project");
|
||||
assert_eq!(
|
||||
entry["transcriptPath"],
|
||||
format!("projects/-my-project/{uuid}.jsonl")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_session_index_entry_preserves_other_entries() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let other_uuid = Uuid::new_v4();
|
||||
fs::write(
|
||||
tmp.path().join("sessions-index.json"),
|
||||
serde_json::to_vec_pretty(&serde_json::json!({
|
||||
other_uuid.to_string(): {"sessionId": other_uuid.to_string(), "custom_field": 42},
|
||||
}))
|
||||
.unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let new_uuid = Uuid::new_v4();
|
||||
write_session_index_entry(new_uuid, Path::new("/my/project"), tmp.path()).unwrap();
|
||||
|
||||
let index: serde_json::Value =
|
||||
serde_json::from_slice(&fs::read(tmp.path().join("sessions-index.json")).unwrap()).unwrap();
|
||||
// New entry landed.
|
||||
assert_eq!(
|
||||
index[new_uuid.to_string()]["sessionId"],
|
||||
new_uuid.to_string()
|
||||
);
|
||||
// Old entry preserved verbatim, including unknown fields.
|
||||
assert_eq!(
|
||||
index[other_uuid.to_string()]["sessionId"],
|
||||
other_uuid.to_string()
|
||||
);
|
||||
assert_eq!(index[other_uuid.to_string()]["custom_field"], 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_session_index_entry_overwrites_same_session() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let uuid = Uuid::new_v4();
|
||||
|
||||
write_session_index_entry(uuid, Path::new("/old/cwd"), tmp.path()).unwrap();
|
||||
write_session_index_entry(uuid, Path::new("/new/cwd"), tmp.path()).unwrap();
|
||||
|
||||
let index: serde_json::Value =
|
||||
serde_json::from_slice(&fs::read(tmp.path().join("sessions-index.json")).unwrap()).unwrap();
|
||||
// Only one entry for this session id, with the newer cwd.
|
||||
assert_eq!(index.as_object().unwrap().len(), 1);
|
||||
assert_eq!(index[uuid.to_string()]["cwd"], "/new/cwd");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_session_index_entry_overwrites_malformed_file() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
fs::write(tmp.path().join("sessions-index.json"), b"not json").unwrap();
|
||||
|
||||
let uuid = Uuid::new_v4();
|
||||
write_session_index_entry(uuid, Path::new("/my/project"), tmp.path()).unwrap();
|
||||
|
||||
let index: serde_json::Value =
|
||||
serde_json::from_slice(&fs::read(tmp.path().join("sessions-index.json")).unwrap()).unwrap();
|
||||
assert_eq!(index[uuid.to_string()]["sessionId"], uuid.to_string());
|
||||
}
|
||||
@@ -0,0 +1,335 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
use tempfile::NamedTempFile;
|
||||
use warp_cli::agent::Harness;
|
||||
use warp_managed_secrets::ManagedSecretValue;
|
||||
use warpui::{ModelHandle, ModelSpawner};
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::server::server_api::harness_support::HarnessSupportClient;
|
||||
use crate::server::server_api::ServerApi;
|
||||
use crate::terminal::model::block::BlockId;
|
||||
use crate::terminal::CLIAgent;
|
||||
|
||||
use super::super::terminal::{CommandHandle, TerminalDriver};
|
||||
use super::super::{AgentDriver, AgentDriverError};
|
||||
use super::json_utils::{read_json_file_or_default, write_json_file};
|
||||
use super::{write_temp_file, HarnessRunner, ResumePayload, SavePoint, ThirdPartyHarness};
|
||||
|
||||
pub(crate) struct GeminiHarness;
|
||||
|
||||
/// Format slug sent to the server when creating a Gemini conversation.
|
||||
const GEMINI_CLI_FORMAT: &str = "gemini_cli";
|
||||
/// Slash command Gemini's TUI recognises as a graceful shutdown.
|
||||
const GEMINI_EXIT_COMMAND: &str = "/quit";
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl ThirdPartyHarness for GeminiHarness {
|
||||
fn harness(&self) -> Harness {
|
||||
Harness::Gemini
|
||||
}
|
||||
|
||||
fn cli_agent(&self) -> CLIAgent {
|
||||
CLIAgent::Gemini
|
||||
}
|
||||
|
||||
fn install_docs_url(&self) -> Option<&'static str> {
|
||||
Some("https://geminicli.com/")
|
||||
}
|
||||
|
||||
fn prepare_environment_config(
|
||||
&self,
|
||||
working_dir: &Path,
|
||||
system_prompt: Option<&str>,
|
||||
_secrets: &HashMap<String, ManagedSecretValue>,
|
||||
) -> Result<(), AgentDriverError> {
|
||||
prepare_gemini_environment_config(working_dir, system_prompt).map_err(|error| {
|
||||
AgentDriverError::HarnessConfigSetupFailed {
|
||||
harness: self.cli_agent().command_prefix().to_owned(),
|
||||
error,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_runner(
|
||||
&self,
|
||||
prompt: &str,
|
||||
system_prompt: Option<&str>,
|
||||
_resumption_prompt: Option<&str>,
|
||||
working_dir: &Path,
|
||||
_task_id: Option<AmbientAgentTaskId>,
|
||||
server_api: Arc<ServerApi>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
_resume: Option<ResumePayload>,
|
||||
) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
|
||||
// Gemini does not support conversation resume yet. When it does, it will add its
|
||||
// own `ResumePayload::Gemini(..)` variant and override `fetch_resume_payload`,
|
||||
// and decide how to surface the user-turn resumption preamble.
|
||||
let client: Arc<dyn HarnessSupportClient> = server_api;
|
||||
Ok(Box::new(GeminiHarnessRunner::new(
|
||||
self.cli_agent().command_prefix(),
|
||||
prompt,
|
||||
system_prompt,
|
||||
working_dir,
|
||||
client,
|
||||
terminal_driver,
|
||||
)?))
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the shell command that launches the Gemini TUI.
|
||||
///
|
||||
/// `--yolo` auto-approves tool call. `-i` seeds the initial prompt and
|
||||
/// continues in interactive TUI mode.
|
||||
fn gemini_command(cli_name: &str, prompt_path: &str) -> String {
|
||||
format!("{cli_name} --yolo -i \"$(cat '{prompt_path}')\"")
|
||||
}
|
||||
|
||||
enum GeminiRunnerState {
|
||||
Preexec,
|
||||
Running {
|
||||
conversation_id: AIConversationId,
|
||||
block_id: BlockId,
|
||||
},
|
||||
}
|
||||
|
||||
struct GeminiHarnessRunner {
|
||||
command: String,
|
||||
/// Held so the temp file is cleaned up when the runner is dropped.
|
||||
_temp_prompt_file: NamedTempFile,
|
||||
client: Arc<dyn HarnessSupportClient>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
state: Mutex<GeminiRunnerState>,
|
||||
}
|
||||
|
||||
impl GeminiHarnessRunner {
|
||||
fn new(
|
||||
cli_command: &str,
|
||||
prompt: &str,
|
||||
_system_prompt: Option<&str>,
|
||||
_working_dir: &Path,
|
||||
client: Arc<dyn HarnessSupportClient>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
) -> Result<Self, AgentDriverError> {
|
||||
let temp_file = write_temp_file("oz_prompt_", prompt)?;
|
||||
let prompt_path = temp_file.path().display().to_string();
|
||||
|
||||
Ok(Self {
|
||||
command: gemini_command(cli_command, &prompt_path),
|
||||
_temp_prompt_file: temp_file,
|
||||
client,
|
||||
terminal_driver,
|
||||
state: Mutex::new(GeminiRunnerState::Preexec),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl HarnessRunner for GeminiHarnessRunner {
|
||||
async fn start(
|
||||
&self,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Result<CommandHandle, AgentDriverError> {
|
||||
// Create the external conversation record on the server.
|
||||
let conversation_id = self
|
||||
.client
|
||||
.create_external_conversation(GEMINI_CLI_FORMAT)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
log::error!("Failed to create external conversation: {e}");
|
||||
AgentDriverError::ConfigBuildFailed(e)
|
||||
})?;
|
||||
log::info!("Created external conversation {conversation_id}");
|
||||
|
||||
let command = self.command.clone();
|
||||
let terminal_driver = self.terminal_driver.clone();
|
||||
let command_handle = foreground
|
||||
.spawn(move |_, ctx| {
|
||||
terminal_driver.update(ctx, |driver, ctx| driver.execute_command(&command, ctx))
|
||||
})
|
||||
.await??
|
||||
.await?;
|
||||
|
||||
// Only store conversation info once the CLI command has started.
|
||||
*self.state.lock() = GeminiRunnerState::Running {
|
||||
conversation_id,
|
||||
block_id: command_handle.block_id().clone(),
|
||||
};
|
||||
|
||||
Ok(command_handle)
|
||||
}
|
||||
|
||||
async fn exit(&self, foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
log::info!("Sending /quit to Gemini CLI");
|
||||
let terminal_driver = self.terminal_driver.clone();
|
||||
foreground
|
||||
.spawn(move |_, ctx| {
|
||||
terminal_driver.update(ctx, |driver, ctx| {
|
||||
driver.send_text_to_cli(GEMINI_EXIT_COMMAND.to_string(), ctx);
|
||||
});
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Agent driver dropped while sending /quit"))
|
||||
}
|
||||
|
||||
async fn save_conversation(
|
||||
&self,
|
||||
save_point: SavePoint,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Result<()> {
|
||||
if matches!(save_point, SavePoint::Periodic)
|
||||
&& !super::has_running_cli_agent(&self.terminal_driver, foreground).await
|
||||
{
|
||||
log::debug!("Will not save conversation, Gemini not in progress");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let (conversation_id, block_id) = match &*self.state.lock() {
|
||||
GeminiRunnerState::Preexec => {
|
||||
log::warn!("save_conversation called before start");
|
||||
return Ok(());
|
||||
}
|
||||
GeminiRunnerState::Running {
|
||||
conversation_id,
|
||||
block_id,
|
||||
} => (*conversation_id, block_id.clone()),
|
||||
};
|
||||
|
||||
// TODO(REMOTE-1408) Also save the conversation transcript.
|
||||
super::upload_current_block_snapshot(
|
||||
foreground,
|
||||
&self.terminal_driver,
|
||||
self.client.as_ref(),
|
||||
conversation_id,
|
||||
block_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare_gemini_environment_config(
|
||||
working_dir: &Path,
|
||||
system_prompt: Option<&str>,
|
||||
) -> Result<()> {
|
||||
let home_dir =
|
||||
dirs::home_dir().ok_or_else(|| anyhow::anyhow!("could not determine home directory"))?;
|
||||
let gemini_dir = home_dir.join(GEMINI_CONFIG_DIR);
|
||||
prepare_gemini_settings(
|
||||
&gemini_dir.join(GEMINI_SETTINGS_FILE_NAME),
|
||||
system_prompt.is_some(),
|
||||
)?;
|
||||
prepare_gemini_trusted_folders(
|
||||
&gemini_dir.join(GEMINI_TRUSTED_FOLDERS_FILE_NAME),
|
||||
working_dir,
|
||||
)?;
|
||||
if let Some(prompt) = system_prompt {
|
||||
let prompt_path = gemini_dir.join(GEMINI_SYSTEM_PROMPT_FILE_NAME);
|
||||
std::fs::write(&prompt_path, prompt).with_context(|| {
|
||||
format!(
|
||||
"Failed to write Gemini system prompt to {}",
|
||||
prompt_path.display()
|
||||
)
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prepare_gemini_settings(settings_path: &Path, has_system_prompt: bool) -> Result<()> {
|
||||
let mut settings: GeminiSettings = read_json_file_or_default(settings_path)?;
|
||||
settings
|
||||
.security
|
||||
.get_or_insert_with(GeminiSecurity::default)
|
||||
.auth
|
||||
.get_or_insert_with(GeminiAuth::default)
|
||||
.selected_type = Some(GEMINI_API_KEY_AUTH_TYPE.to_owned());
|
||||
|
||||
if has_system_prompt {
|
||||
let context = settings.context.get_or_insert_with(GeminiContext::default);
|
||||
let file_name = GEMINI_SYSTEM_PROMPT_FILE_NAME.to_owned();
|
||||
if !context.file_name.contains(&file_name) {
|
||||
context.file_name.push(file_name);
|
||||
}
|
||||
}
|
||||
|
||||
write_json_file(
|
||||
settings_path,
|
||||
&settings,
|
||||
"Failed to serialize Gemini settings",
|
||||
)
|
||||
}
|
||||
|
||||
fn prepare_gemini_trusted_folders(trusted_path: &Path, working_dir: &Path) -> Result<()> {
|
||||
let mut trusted: HashMap<String, String> = read_json_file_or_default(trusted_path)?;
|
||||
trusted.insert(
|
||||
working_dir.to_string_lossy().into_owned(),
|
||||
GEMINI_TRUST_LEVEL_FOLDER.to_owned(),
|
||||
);
|
||||
write_json_file(
|
||||
trusted_path,
|
||||
&trusted,
|
||||
"Failed to serialize Gemini trusted folders",
|
||||
)
|
||||
}
|
||||
|
||||
const GEMINI_CONFIG_DIR: &str = ".gemini";
|
||||
const GEMINI_SETTINGS_FILE_NAME: &str = "settings.json";
|
||||
const GEMINI_TRUSTED_FOLDERS_FILE_NAME: &str = "trustedFolders.json";
|
||||
const GEMINI_SYSTEM_PROMPT_FILE_NAME: &str = "OZ_SYSTEM_PROMPT.md";
|
||||
/// Auth-type discriminant for API-key auth — matches `AuthType.USE_GEMINI` in
|
||||
/// Gemini's `packages/core/src/core/contentGenerator.ts`.
|
||||
const GEMINI_API_KEY_AUTH_TYPE: &str = "gemini-api-key";
|
||||
/// Trust level discriminant that grants full trust to a single folder — matches
|
||||
/// Gemini's `TrustLevel.TRUST_FOLDER` in `packages/cli/src/config/trustedFolders.ts`.
|
||||
const GEMINI_TRUST_LEVEL_FOLDER: &str = "TRUST_FOLDER";
|
||||
|
||||
#[derive(Default, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GeminiSettings {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
security: Option<GeminiSecurity>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
context: Option<GeminiContext>,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GeminiSecurity {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
auth: Option<GeminiAuth>,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GeminiAuth {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
selected_type: Option<String>,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GeminiContext {
|
||||
#[serde(default)]
|
||||
file_name: Vec<String>,
|
||||
#[serde(flatten)]
|
||||
extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "gemini_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,166 @@
|
||||
use std::fs;
|
||||
|
||||
use serde_json::Value;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_creates_file_with_api_key_auth() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
|
||||
prepare_gemini_settings(&settings_path, false).unwrap();
|
||||
|
||||
let settings: Value = serde_json::from_slice(&fs::read(settings_path).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
settings["security"]["auth"]["selectedType"],
|
||||
Value::String("gemini-api-key".to_owned()),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_preserves_unrelated_keys() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
fs::write(
|
||||
&settings_path,
|
||||
r#"{
|
||||
"ui": {"theme": "dark"},
|
||||
"security": {
|
||||
"folderTrust": {"enabled": true},
|
||||
"auth": {"enforcedType": "vertex-ai"}
|
||||
}
|
||||
}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
prepare_gemini_settings(&settings_path, false).unwrap();
|
||||
|
||||
let settings: Value = serde_json::from_slice(&fs::read(settings_path).unwrap()).unwrap();
|
||||
assert_eq!(settings["ui"]["theme"], "dark");
|
||||
assert_eq!(settings["security"]["folderTrust"]["enabled"], true);
|
||||
// Sibling auth fields survive, and selectedType is set alongside them.
|
||||
assert_eq!(settings["security"]["auth"]["enforcedType"], "vertex-ai");
|
||||
assert_eq!(
|
||||
settings["security"]["auth"]["selectedType"],
|
||||
"gemini-api-key",
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_surfaces_malformed_json_as_error() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
// `security` typed as a string instead of an object is a parse error;
|
||||
// we prefer surfacing that to silently rewriting user-owned state.
|
||||
fs::write(
|
||||
&settings_path,
|
||||
r#"{"ui":{"theme":"dark"},"security":"broken"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(prepare_gemini_settings(&settings_path, false).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_trusted_folders_creates_file_with_working_dir() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let trusted_path = tmp.path().join("trustedFolders.json");
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
|
||||
prepare_gemini_trusted_folders(&trusted_path, &working_dir).unwrap();
|
||||
|
||||
let trusted: Value = serde_json::from_slice(&fs::read(trusted_path).unwrap()).unwrap();
|
||||
let working_dir_key = working_dir.to_string_lossy().to_string();
|
||||
assert_eq!(trusted[working_dir_key], "TRUST_FOLDER");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_trusted_folders_preserves_existing_entries() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let trusted_path = tmp.path().join("trustedFolders.json");
|
||||
fs::write(
|
||||
&trusted_path,
|
||||
r#"{"/other/project":"TRUST_PARENT","/do/not/trust":"DO_NOT_TRUST"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let working_dir = tmp.path().join("workspace/project");
|
||||
|
||||
prepare_gemini_trusted_folders(&trusted_path, &working_dir).unwrap();
|
||||
|
||||
let trusted: Value = serde_json::from_slice(&fs::read(trusted_path).unwrap()).unwrap();
|
||||
assert_eq!(trusted["/other/project"], "TRUST_PARENT");
|
||||
assert_eq!(trusted["/do/not/trust"], "DO_NOT_TRUST");
|
||||
let working_dir_key = working_dir.to_string_lossy().to_string();
|
||||
assert_eq!(trusted[working_dir_key], "TRUST_FOLDER");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_adds_context_file_name_when_system_prompt_present() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
|
||||
prepare_gemini_settings(&settings_path, true).unwrap();
|
||||
|
||||
let settings: Value = serde_json::from_slice(&fs::read(settings_path).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
settings["context"]["fileName"],
|
||||
Value::Array(vec![Value::String("OZ_SYSTEM_PROMPT.md".to_owned())]),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_appends_to_existing_context_file_name() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
fs::write(
|
||||
&settings_path,
|
||||
r#"{"context":{"fileName":["existing.md"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
prepare_gemini_settings(&settings_path, true).unwrap();
|
||||
|
||||
let settings: Value = serde_json::from_slice(&fs::read(settings_path).unwrap()).unwrap();
|
||||
let file_names: Vec<String> = settings["context"]["fileName"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|v| v.as_str().unwrap().to_owned())
|
||||
.collect();
|
||||
assert_eq!(file_names, vec!["existing.md", "OZ_SYSTEM_PROMPT.md"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_does_not_duplicate_system_prompt_file_name() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
fs::write(
|
||||
&settings_path,
|
||||
r#"{"context":{"fileName":["OZ_SYSTEM_PROMPT.md"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
prepare_gemini_settings(&settings_path, true).unwrap();
|
||||
|
||||
let settings: Value = serde_json::from_slice(&fs::read(settings_path).unwrap()).unwrap();
|
||||
let file_names: Vec<String> = settings["context"]["fileName"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|v| v.as_str().unwrap().to_owned())
|
||||
.collect();
|
||||
assert_eq!(file_names, vec!["OZ_SYSTEM_PROMPT.md"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_gemini_settings_omits_context_when_no_system_prompt() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let settings_path = tmp.path().join("settings.json");
|
||||
|
||||
prepare_gemini_settings(&settings_path, false).unwrap();
|
||||
|
||||
let settings: Value = serde_json::from_slice(&fs::read(settings_path).unwrap()).unwrap();
|
||||
assert!(settings.get("context").is_none());
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
//! Shared JSON read/merge/write helpers for third-party harness config prep.
|
||||
//!
|
||||
//! Third-party CLIs like Claude Code and Gemini CLI persist onboarding, trust,
|
||||
//! and auth state in JSON files. The harness preparation step
|
||||
//! needs to set a few keys on those files without clobbering
|
||||
//! user-owned state. These helpers allow us to read and merge with existing
|
||||
//! JSON file state easily.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Read a JSON file as `T`, or return `T::default()` if the file does not exist.
|
||||
///
|
||||
/// Returns an error if the file exists but cannot be read or parsed.
|
||||
pub(super) fn read_json_file_or_default<T>(path: &Path) -> Result<T>
|
||||
where
|
||||
T: Default + for<'de> Deserialize<'de>,
|
||||
{
|
||||
let content = match std::fs::read_to_string(path) {
|
||||
Ok(content) => content,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
return Ok(T::default());
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(
|
||||
anyhow::Error::from(e).context(format!("Failed to read {}", path.display()))
|
||||
);
|
||||
}
|
||||
};
|
||||
serde_json::from_str(&content).with_context(|| format!("Failed to parse {}", path.display()))
|
||||
}
|
||||
|
||||
/// Serialize `value` as pretty JSON and write it to `path`, creating parent
|
||||
/// directories as needed. `serialize_error` is used as the context for the
|
||||
/// serialization step so the caller-facing error is specific to the config
|
||||
/// file being written.
|
||||
pub(super) fn write_json_file<T>(
|
||||
path: &Path,
|
||||
value: &T,
|
||||
serialize_error: &'static str,
|
||||
) -> Result<()>
|
||||
where
|
||||
T: Serialize,
|
||||
{
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.with_context(|| format!("Failed to create {}", parent.display()))?;
|
||||
}
|
||||
std::fs::write(
|
||||
path,
|
||||
serde_json::to_vec_pretty(value).context(serialize_error)?,
|
||||
)
|
||||
.with_context(|| format!("Failed to write {}", path.display()))
|
||||
}
|
||||
@@ -0,0 +1,461 @@
|
||||
use std::borrow::Cow;
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsString;
|
||||
use std::fmt;
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use tempfile::NamedTempFile;
|
||||
use warp_cli::agent::Harness;
|
||||
use warp_managed_secrets::ManagedSecretValue;
|
||||
use warpui::{ModelHandle, ModelSpawner, SingletonEntity};
|
||||
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::server::server_api::harness_support::{upload_to_target, HarnessSupportClient};
|
||||
use crate::server::server_api::ServerApi;
|
||||
use crate::terminal::cli_agent_sessions::{CLIAgentSessionStatus, CLIAgentSessionsModel};
|
||||
use crate::terminal::model::block::{BlockId, SerializedBlock};
|
||||
use crate::terminal::CLIAgent;
|
||||
use crate::util::path::resolve_executable;
|
||||
use warp_cli::{
|
||||
OZ_CLI_ENV, OZ_HARNESS_ENV, OZ_PARENT_RUN_ID_ENV, OZ_RUN_ID_ENV, SERVER_ROOT_URL_OVERRIDE_ENV,
|
||||
SESSION_SHARING_SERVER_URL_OVERRIDE_ENV, WS_SERVER_URL_OVERRIDE_ENV,
|
||||
};
|
||||
use warp_core::channel::ChannelState;
|
||||
|
||||
use super::terminal::{CommandHandle, TerminalDriver};
|
||||
use super::{
|
||||
AgentDriver, AgentDriverError, LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV,
|
||||
LEGACY_OZ_PARENT_STATE_ROOT_ENV, OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV,
|
||||
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
};
|
||||
|
||||
mod claude_code;
|
||||
pub(crate) mod claude_transcript;
|
||||
mod gemini;
|
||||
mod json_utils;
|
||||
|
||||
pub(crate) use claude_code::ClaudeHarness;
|
||||
use claude_transcript::ClaudeResumeInfo;
|
||||
use gemini::GeminiHarness;
|
||||
|
||||
/// Harness-agnostic payload describing how to resume an existing conversation.
|
||||
///
|
||||
/// Each variant carries the data a specific harness needs to rehydrate state before its CLI
|
||||
/// launches. Harnesses match on the variant they produce and ignore others; new CLIs that
|
||||
/// want resume support add a new variant and override [`ThirdPartyHarness::fetch_resume_payload`].
|
||||
pub(crate) enum ResumePayload {
|
||||
/// Claude Code session state fetched from the server's transcript endpoint.
|
||||
Claude(ClaudeResumeInfo),
|
||||
}
|
||||
|
||||
/// Trait for third-party agent harnesses that execute prompts via their own CLIs.
|
||||
///
|
||||
/// Each new external harness (e.g. Claude, Codex) implements this to be used with cloud agents.
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub(crate) trait ThirdPartyHarness: Send + Sync {
|
||||
/// Returns the [`Harness`] variant this implementation corresponds to.
|
||||
fn harness(&self) -> Harness;
|
||||
|
||||
/// Returns the CLIAgent type associated with this harness.
|
||||
fn cli_agent(&self) -> CLIAgent;
|
||||
|
||||
/// URL to install instructions for this harness's CLI, surfaced in the
|
||||
/// default [`validate`] impl when the CLI is not on `PATH`.
|
||||
fn install_docs_url(&self) -> Option<&'static str> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Validate that the harness is ready to run. Default impl checks that the
|
||||
/// CLI is installed on `PATH`; override for additional checks.
|
||||
fn validate(&self) -> Result<(), AgentDriverError> {
|
||||
validate_cli_installed(self.cli_agent().command_prefix(), self.install_docs_url())
|
||||
}
|
||||
|
||||
/// Prepare CLI-specific config files before launching the harness command.
|
||||
fn prepare_environment_config(
|
||||
&self,
|
||||
_working_dir: &Path,
|
||||
_system_prompt: Option<&str>,
|
||||
_secrets: &HashMap<String, ManagedSecretValue>,
|
||||
) -> Result<(), AgentDriverError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch the harness-specific resume payload for an existing conversation.
|
||||
///
|
||||
/// The driver calls this when the user passes `--conversation <id>` and the harness
|
||||
/// matches the stored conversation's harness. Harnesses that don't support resume
|
||||
/// use the default impl, which returns `Ok(None)` and causes the run to start fresh.
|
||||
///
|
||||
/// Implementations download the raw transcript via [`HarnessSupportClient::fetch_transcript`]
|
||||
/// (which derives the conversation from the current task's `agent_conversation_id`) and
|
||||
/// own all harness-specific deserialization and error mapping (e.g. a 404 maps to
|
||||
/// [`AgentDriverError::ConversationResumeStateMissing`] tagged with the harness label).
|
||||
async fn fetch_resume_payload(
|
||||
&self,
|
||||
_conversation_id: &AIConversationId,
|
||||
_harness_support_client: Arc<dyn HarnessSupportClient>,
|
||||
) -> Result<Option<ResumePayload>, AgentDriverError> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Build a runner for executing this harness with the given prompt.
|
||||
///
|
||||
/// If `resume` is `Some`, the harness matches on its own [`ResumePayload`] variant and
|
||||
/// reuses the stored session/conversation ids instead of minting fresh ones. Variants
|
||||
/// belonging to other harnesses are ignored.
|
||||
///
|
||||
/// `resumption_prompt`, when non-empty, is a short user-turn preamble the server emits
|
||||
/// during a resumed session. Each harness decides exactly how to surface it (e.g. Claude
|
||||
/// prepends it to the user-turn prompt that gets piped into the CLI). Harnesses that
|
||||
/// don't yet support resumption can ignore it.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn build_runner(
|
||||
&self,
|
||||
prompt: &str,
|
||||
system_prompt: Option<&str>,
|
||||
resumption_prompt: Option<&str>,
|
||||
working_dir: &Path,
|
||||
task_id: Option<AmbientAgentTaskId>,
|
||||
server_api: Arc<ServerApi>,
|
||||
terminal_driver: ModelHandle<TerminalDriver>,
|
||||
resume: Option<ResumePayload>,
|
||||
) -> Result<Box<dyn HarnessRunner>, AgentDriverError>;
|
||||
}
|
||||
|
||||
/// Harness type for driver dispatch.
|
||||
pub(crate) enum HarnessKind {
|
||||
Oz,
|
||||
/// Third-party CLI-backed harness (e.g. Claude, Gemini).
|
||||
ThirdParty(Box<dyn ThirdPartyHarness>),
|
||||
/// Harnesses that exist in the shared CLI enum but are not supported by the
|
||||
/// standalone agent driver.
|
||||
Unsupported(Harness),
|
||||
}
|
||||
|
||||
impl HarnessKind {
|
||||
/// Corresponding [`Harness`] enum value.
|
||||
pub(crate) fn harness(&self) -> Harness {
|
||||
match self {
|
||||
HarnessKind::Oz => Harness::Oz,
|
||||
HarnessKind::ThirdParty(h) => h.harness(),
|
||||
HarnessKind::Unsupported(harness) => *harness,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for HarnessKind {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
// Use the `Display` method on the [`Harness`] enum.
|
||||
write!(f, "{}", self.harness())
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a [`HarnessKind`] for the given [`Harness`].
|
||||
///
|
||||
/// We shouldn't ever get a `--harness unknown` here because clap should handle
|
||||
/// it.
|
||||
pub(crate) fn harness_kind(harness: Harness) -> Result<HarnessKind, AgentDriverError> {
|
||||
match harness {
|
||||
Harness::Oz => Ok(HarnessKind::Oz),
|
||||
Harness::Claude => Ok(HarnessKind::ThirdParty(Box::new(ClaudeHarness))),
|
||||
Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)),
|
||||
Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))),
|
||||
Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check that `cli` is installed and on PATH, returning a `HarnessSetupFailed`
|
||||
/// error with an optional install-docs link when it isn't.
|
||||
pub(crate) fn validate_cli_installed(
|
||||
cli: &str,
|
||||
install_docs_url: Option<&str>,
|
||||
) -> Result<(), AgentDriverError> {
|
||||
if resolve_executable(cli).is_none() {
|
||||
let mut reason = format!("'{cli}' CLI not found on your machine.");
|
||||
if let Some(url) = install_docs_url {
|
||||
reason.push_str(&format!(" Install it first: {url}"));
|
||||
}
|
||||
return Err(AgentDriverError::HarnessSetupFailed {
|
||||
harness: cli.into(),
|
||||
reason,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn insert_non_empty_task_env_var(
|
||||
env_vars: &mut HashMap<OsString, OsString>,
|
||||
key: &'static str,
|
||||
value: String,
|
||||
) {
|
||||
if value.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
env_vars.insert(OsString::from(key), OsString::from(value));
|
||||
}
|
||||
|
||||
fn insert_task_env_var_aliases(
|
||||
env_vars: &mut HashMap<OsString, OsString>,
|
||||
keys: &[&'static str],
|
||||
value: &str,
|
||||
) {
|
||||
for key in keys {
|
||||
env_vars.insert(OsString::from(key), OsString::from(value));
|
||||
}
|
||||
}
|
||||
|
||||
fn message_listener_state_root() -> Option<String> {
|
||||
[
|
||||
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
LEGACY_OZ_PARENT_STATE_ROOT_ENV,
|
||||
]
|
||||
.into_iter()
|
||||
.find_map(|key| std::env::var(key).ok().filter(|value| !value.is_empty()))
|
||||
}
|
||||
|
||||
fn task_env_vars_for_harness_name(
|
||||
task_id: Option<&AmbientAgentTaskId>,
|
||||
parent_run_id: Option<&str>,
|
||||
selected_harness: Harness,
|
||||
) -> HashMap<OsString, OsString> {
|
||||
let mut env_vars = HashMap::with_capacity(7);
|
||||
|
||||
if let Some(id) = task_id {
|
||||
env_vars.insert(
|
||||
OsString::from(OZ_RUN_ID_ENV),
|
||||
OsString::from(id.to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(parent_run_id) = parent_run_id.filter(|id| !id.is_empty()) {
|
||||
env_vars.insert(
|
||||
OsString::from(OZ_PARENT_RUN_ID_ENV),
|
||||
OsString::from(parent_run_id),
|
||||
);
|
||||
}
|
||||
|
||||
env_vars.insert(
|
||||
OsString::from(OZ_CLI_ENV),
|
||||
OsString::from(
|
||||
std::env::current_exe()
|
||||
.unwrap_or_else(|_| ChannelState::channel().cli_command_name().into()),
|
||||
),
|
||||
);
|
||||
// `OZ_HARNESS` is only consumed by child orchestration telemetry when the child
|
||||
// CLI emits `run message *` events.
|
||||
env_vars.insert(
|
||||
OsString::from(OZ_HARNESS_ENV),
|
||||
OsString::from(selected_harness.to_string()),
|
||||
);
|
||||
if selected_harness == Harness::Claude && task_id.is_some() {
|
||||
insert_task_env_var_aliases(
|
||||
&mut env_vars,
|
||||
&[
|
||||
OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV,
|
||||
LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV,
|
||||
],
|
||||
"1",
|
||||
);
|
||||
if let Some(state_root) = message_listener_state_root() {
|
||||
insert_task_env_var_aliases(
|
||||
&mut env_vars,
|
||||
&[
|
||||
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
LEGACY_OZ_PARENT_STATE_ROOT_ENV,
|
||||
],
|
||||
&state_root,
|
||||
);
|
||||
}
|
||||
}
|
||||
// Server URL overrides are disabled on release channels, so there's no
|
||||
// override to propagate to child processes there.
|
||||
if ChannelState::channel().allows_server_url_overrides() {
|
||||
insert_non_empty_task_env_var(
|
||||
&mut env_vars,
|
||||
SERVER_ROOT_URL_OVERRIDE_ENV,
|
||||
ChannelState::server_root_url().into_owned(),
|
||||
);
|
||||
insert_non_empty_task_env_var(
|
||||
&mut env_vars,
|
||||
WS_SERVER_URL_OVERRIDE_ENV,
|
||||
ChannelState::ws_server_url().into_owned(),
|
||||
);
|
||||
if let Some(url) = ChannelState::session_sharing_server_url()
|
||||
.map(Cow::into_owned)
|
||||
.filter(|url| !url.is_empty())
|
||||
{
|
||||
env_vars.insert(
|
||||
OsString::from(SESSION_SHARING_SERVER_URL_OVERRIDE_ENV),
|
||||
OsString::from(url),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
env_vars
|
||||
}
|
||||
|
||||
pub(crate) fn task_env_vars(
|
||||
task_id: Option<&AmbientAgentTaskId>,
|
||||
parent_run_id: Option<&str>,
|
||||
selected_harness: Harness,
|
||||
) -> HashMap<OsString, OsString> {
|
||||
task_env_vars_for_harness_name(task_id, parent_run_id, selected_harness)
|
||||
}
|
||||
|
||||
/// Indicates when the harness conversation is being saved.
|
||||
/// Implementations may use this to customize the saved data, such as
|
||||
/// recording additional metadata on completion.
|
||||
pub(crate) enum SavePoint {
|
||||
/// A periodic auto-save to minimize data loss.
|
||||
Periodic,
|
||||
/// The final save of conversation state, after the harness has completed.
|
||||
Final,
|
||||
/// A save after the harness reports it finished an agent turn.
|
||||
PostTurn,
|
||||
}
|
||||
|
||||
/// Stateful per-run representation of an external harness produced
|
||||
/// by [`ThirdPartyHarness::build_runner`].
|
||||
///
|
||||
/// All `HarnessRunner` methods take `&self` as a parameter, but may mutate internal
|
||||
/// state. There are no `&mut self` methods, as this would require that the `AgentDriver`
|
||||
/// store the runner in a mutex and lock it across `await` points.
|
||||
///
|
||||
/// The driver uses this to manage the lifecycle of a particular third-party harness.
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub(crate) trait HarnessRunner: Send + Sync {
|
||||
/// Create the external conversation on the server and start the harness
|
||||
/// command in the terminal.
|
||||
///
|
||||
/// Returns a [`CommandHandle`] that resolves to the exit code. The runner
|
||||
/// stores the conversation ID and block ID internally for use in
|
||||
/// [`save_conversation`].
|
||||
async fn start(
|
||||
&self,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Result<CommandHandle, AgentDriverError>;
|
||||
|
||||
/// Save the current conversation state (transcript upload, etc.).
|
||||
async fn save_conversation(
|
||||
&self,
|
||||
save_point: SavePoint,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> Result<()>;
|
||||
|
||||
/// Gracefully ask the harness to exit.
|
||||
async fn exit(&self, foreground: &ModelSpawner<AgentDriver>) -> Result<()>;
|
||||
/// Handle a CLI session update such as a prompt submit or completed tool use.
|
||||
async fn handle_session_update(&self, _foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clean up any harness-owned background state after the harness exits.
|
||||
async fn cleanup(&self, _foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns `true` if the terminal tracked by `terminal_driver` has a CLI agent session
|
||||
/// that is currently in progress.
|
||||
pub(crate) async fn has_running_cli_agent(
|
||||
terminal_driver: &ModelHandle<TerminalDriver>,
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
) -> bool {
|
||||
let driver = terminal_driver.clone();
|
||||
let Ok(running) = foreground
|
||||
.spawn(move |_, ctx| {
|
||||
let terminal_view_id = driver.as_ref(ctx).terminal_view().id();
|
||||
CLIAgentSessionsModel::handle(ctx)
|
||||
.as_ref(ctx)
|
||||
.session(terminal_view_id)
|
||||
.is_some_and(|s| s.status == CLIAgentSessionStatus::InProgress)
|
||||
})
|
||||
.await
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
running
|
||||
}
|
||||
|
||||
/// Create a [`NamedTempFile`] with the given prefix and write `content` into it.
|
||||
///
|
||||
/// Used by third-party harnesses to stage prompts / system prompts on disk
|
||||
/// before launching the CLI, avoiding shell-quoting issues with complex input.
|
||||
pub(super) fn write_temp_file(
|
||||
prefix: &str,
|
||||
content: &str,
|
||||
) -> Result<NamedTempFile, AgentDriverError> {
|
||||
let mut file = tempfile::Builder::new()
|
||||
.prefix(prefix)
|
||||
.suffix(".txt")
|
||||
.tempfile()
|
||||
.map_err(|e| {
|
||||
AgentDriverError::ConfigBuildFailed(anyhow::anyhow!(
|
||||
"Failed to create temp file '{prefix}': {e}"
|
||||
))
|
||||
})?;
|
||||
file.write_all(content.as_bytes()).map_err(|e| {
|
||||
AgentDriverError::ConfigBuildFailed(anyhow::anyhow!(
|
||||
"Failed to write temp file '{prefix}': {e}"
|
||||
))
|
||||
})?;
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
/// Upload a [`SerializedBlock`] as the JSON block snapshot for a third-party harness conversation.
|
||||
pub(crate) async fn upload_block_snapshot(
|
||||
client: &dyn HarnessSupportClient,
|
||||
conversation_id: AIConversationId,
|
||||
block: SerializedBlock,
|
||||
) -> Result<()> {
|
||||
log::info!("Uploading block snapshot for CLI agent to conversation {conversation_id}");
|
||||
let target = client
|
||||
.get_block_snapshot_upload_target(&conversation_id)
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!("Unable to get block upload slot for conversation {conversation_id}")
|
||||
})?;
|
||||
|
||||
let body = block
|
||||
.to_json()
|
||||
.with_context(|| format!("Unable to serialize block for conversation {conversation_id}"))?;
|
||||
|
||||
upload_to_target(client.http_client(), &target, body).await
|
||||
}
|
||||
|
||||
/// Fetch the current block snapshot for `block_id` and upload it to the server.
|
||||
///
|
||||
/// If the snapshot cannot be fetched, logs a warning and returns `Ok(())`.
|
||||
pub(super) async fn upload_current_block_snapshot(
|
||||
foreground: &ModelSpawner<AgentDriver>,
|
||||
terminal_driver: &ModelHandle<TerminalDriver>,
|
||||
client: &dyn HarnessSupportClient,
|
||||
conversation_id: AIConversationId,
|
||||
block_id: BlockId,
|
||||
) -> Result<()> {
|
||||
let td = terminal_driver.clone();
|
||||
let snapshot = foreground
|
||||
.spawn(move |_, ctx| td.as_ref(ctx).block_snapshot(&block_id, ctx))
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Agent driver dropped"))?;
|
||||
match snapshot {
|
||||
Some(block) => upload_block_snapshot(client, conversation_id, block).await,
|
||||
None => {
|
||||
log::warn!("No block snapshot found for harness command");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "mod_test.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,33 @@
|
||||
use super::validate_cli_installed;
|
||||
use crate::ai::agent_sdk::driver::AgentDriverError;
|
||||
|
||||
fn assert_harness_setup_failed(err: &AgentDriverError) -> (&str, &str) {
|
||||
match err {
|
||||
AgentDriverError::HarnessSetupFailed { harness, reason } => (harness, reason),
|
||||
other => panic!("expected HarnessSetupFailed, got: {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
#[test]
|
||||
fn validate_cli_installed_succeeds_for_known_binary() {
|
||||
assert!(validate_cli_installed("ls", None).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_cli_installed_fails_for_missing_binary() {
|
||||
let err = validate_cli_installed("__nonexistent_cli_abc123__", None).unwrap_err();
|
||||
let (harness, reason) = assert_harness_setup_failed(&err);
|
||||
assert_eq!(harness, "__nonexistent_cli_abc123__");
|
||||
assert!(reason.contains("not found"));
|
||||
assert!(!reason.contains("Install it first"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_cli_installed_includes_docs_url_in_error() {
|
||||
let url = "https://example.com/install";
|
||||
let err = validate_cli_installed("__nonexistent_cli_abc123__", Some(url)).unwrap_err();
|
||||
let (_, reason) = assert_harness_setup_failed(&err);
|
||||
assert!(reason.contains(url));
|
||||
assert!(reason.contains("Install it first"));
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,648 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
ffi::OsString,
|
||||
future::Future,
|
||||
path::PathBuf,
|
||||
pin::Pin,
|
||||
sync::Arc,
|
||||
task::{Context, Poll},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use futures::channel::oneshot;
|
||||
use session_sharing_protocol::common::{Role, SessionId};
|
||||
use session_sharing_protocol::sharer::SessionSourceType;
|
||||
use warp_cli::share::{ShareAccessLevel, ShareRequest, ShareSubject};
|
||||
use warp_completer::completer::CommandOutput;
|
||||
use warp_core::command::ExitCode;
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warp_util::path::ShellFamily;
|
||||
use warpui::{
|
||||
r#async::FutureExt, AppContext, Entity, ModelContext, ModelHandle, SingletonEntity as _,
|
||||
ViewHandle,
|
||||
};
|
||||
|
||||
use crate::terminal::model::session::ExecuteCommandOptions;
|
||||
|
||||
use crate::{
|
||||
ai::ambient_agents::AmbientAgentTaskId,
|
||||
pane_group::NewTerminalOptions,
|
||||
root_view::{open_new_with_workspace_source, NewWorkspaceSource},
|
||||
terminal::{
|
||||
model::block::{BlockId, SerializedBlock},
|
||||
shared_session::{self, IsSharedSessionCreator},
|
||||
shell::ShellType,
|
||||
view::ConversationRestorationInNewPaneType,
|
||||
TerminalView,
|
||||
},
|
||||
util::sync::Condition,
|
||||
workspaces::user_workspaces::UserWorkspaces,
|
||||
};
|
||||
|
||||
use crate::ai::attachment_utils::attachments_download_dir;
|
||||
|
||||
use super::AgentDriverError;
|
||||
|
||||
/// Describes why an agent's session-sharing request failed.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum ShareSessionError {
|
||||
/// Connection to the session-sharing server failed.
|
||||
#[error("Internal error")]
|
||||
Internal(#[source] Arc<anyhow::Error>),
|
||||
/// The server rejected the session-sharing request.
|
||||
#[error("{0}")]
|
||||
Failed(String),
|
||||
/// Session sharing is disabled for this user or team.
|
||||
#[error(
|
||||
"Session sharing is not enabled. This is likely because an administrator has disabled session sharing for your team."
|
||||
)]
|
||||
Disabled,
|
||||
/// The session-sharing request timed out.
|
||||
#[error("Timed out waiting for session sharing to start")]
|
||||
Timeout,
|
||||
/// The session-sharing channel was dropped before completing.
|
||||
#[error("Session sharing was interrupted")]
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
const TERMINAL_SESSION_BOOTSTRAP_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
const TERMINAL_SESSION_SHARE_DELAY: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Options for creating the terminal view before constructing a [`TerminalDriver`].
|
||||
pub(crate) struct TerminalDriverOptions {
|
||||
pub working_dir: PathBuf,
|
||||
pub env_vars: HashMap<OsString, OsString>,
|
||||
pub should_share: bool,
|
||||
pub task_id: Option<AmbientAgentTaskId>,
|
||||
pub conversation_restoration: Option<ConversationRestorationInNewPaneType>,
|
||||
}
|
||||
|
||||
/// Events emitted by [`TerminalDriver`] for [`super::AgentDriver`] to react to.
|
||||
pub(crate) enum TerminalDriverEvent {
|
||||
/// Terminal bootstrap is taking unusually long.
|
||||
SlowBootstrap,
|
||||
/// The terminal session has established a shared session.
|
||||
EstablishedSharedSession {
|
||||
session_id: session_sharing_protocol::common::SessionId,
|
||||
join_url: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Manages the terminal session lifecycle for the agent driver.
|
||||
///
|
||||
/// Responsibilities:
|
||||
/// - Monitoring for terminal bootstrapping to be done
|
||||
/// - Configuring session sharing and applying guest requests
|
||||
/// - Executing commands in the session
|
||||
/// - Detecting block completion
|
||||
pub(crate) struct TerminalDriver {
|
||||
terminal_view: ViewHandle<TerminalView>,
|
||||
session_bootstrapped: Condition,
|
||||
/// The session ID once sharing has been established.
|
||||
shared_session_id: Option<SessionId>,
|
||||
/// Receiver for the session sharing result. Present when sharing is expected
|
||||
/// and `wait_for_session_shared` has not yet been called.
|
||||
session_share_rx: Option<oneshot::Receiver<Result<(), ShareSessionError>>>,
|
||||
pending_share_requests: Vec<ShareRequest>,
|
||||
waiting_command: Option<oneshot::Sender<ExitCode>>,
|
||||
|
||||
/// State for the pending command we're expecting to start executing.
|
||||
/// The `String` is the expected command text, and the sender is used
|
||||
/// to send the block ID to the waiting caller.
|
||||
pending_command_start: Option<(String, oneshot::Sender<BlockId>)>,
|
||||
}
|
||||
|
||||
impl Entity for TerminalDriver {
|
||||
type Event = TerminalDriverEvent;
|
||||
}
|
||||
|
||||
/// Create the terminal window and extract the [`ViewHandle<TerminalView>`].
|
||||
///
|
||||
/// This is separate from [`TerminalDriver::new`] because [`AppContext::add_model`]
|
||||
/// requires an infallible constructor; the fallible window/view creation must happen first.
|
||||
fn create_terminal_view(
|
||||
options: TerminalDriverOptions,
|
||||
ctx: &mut AppContext,
|
||||
) -> Result<ViewHandle<TerminalView>, AgentDriverError> {
|
||||
let is_shared_session_creator = if options.should_share {
|
||||
IsSharedSessionCreator::Yes {
|
||||
source_type: SessionSourceType::AmbientAgent {
|
||||
task_id: options.task_id.map(|t| t.to_string()),
|
||||
},
|
||||
}
|
||||
} else {
|
||||
IsSharedSessionCreator::No
|
||||
};
|
||||
|
||||
let (_, root_view) = open_new_with_workspace_source(
|
||||
NewWorkspaceSource::Session {
|
||||
options: Box::new(NewTerminalOptions {
|
||||
is_shared_session_creator,
|
||||
initial_directory: Some(options.working_dir),
|
||||
env_vars: options.env_vars,
|
||||
conversation_restoration: options.conversation_restoration,
|
||||
..Default::default()
|
||||
}),
|
||||
},
|
||||
ctx,
|
||||
);
|
||||
|
||||
root_view
|
||||
.as_ref(ctx)
|
||||
.workspace_view()
|
||||
.ok_or(AgentDriverError::TerminalUnavailable)?
|
||||
.as_ref(ctx)
|
||||
.active_tab_pane_group()
|
||||
.as_ref(ctx)
|
||||
.active_session_view(ctx)
|
||||
.ok_or(AgentDriverError::TerminalUnavailable)
|
||||
}
|
||||
|
||||
impl TerminalDriver {
|
||||
/// Create a terminal view from the given options and wrap it in a new `TerminalDriver` model.
|
||||
pub(crate) fn create(
|
||||
options: TerminalDriverOptions,
|
||||
ctx: &mut AppContext,
|
||||
) -> Result<ModelHandle<Self>, AgentDriverError> {
|
||||
let should_share = options.should_share;
|
||||
let task_id = options.task_id;
|
||||
let working_dir = options.working_dir.clone();
|
||||
let terminal_view = create_terminal_view(options, ctx)?;
|
||||
Ok(ctx.add_model(|ctx| Self::new(terminal_view, should_share, task_id, working_dir, ctx)))
|
||||
}
|
||||
|
||||
/// Wrap an already-created terminal view in a new `TerminalDriver` model.
|
||||
///
|
||||
/// Unlike [`Self::create`], this does not open a new window — it reuses an
|
||||
/// existing view (e.g. a docker sandbox pane). Session sharing is disabled
|
||||
/// and no task ID is associated.
|
||||
pub(crate) fn create_from_existing_view(
|
||||
terminal_view: ViewHandle<TerminalView>,
|
||||
ctx: &mut AppContext,
|
||||
) -> ModelHandle<Self> {
|
||||
ctx.add_model(|ctx| Self::new(terminal_view, false, None, PathBuf::default(), ctx))
|
||||
}
|
||||
|
||||
/// Set up event subscriptions and session-sharing conditions for an
|
||||
/// already-created terminal view.
|
||||
fn new(
|
||||
terminal_view: ViewHandle<TerminalView>,
|
||||
should_share: bool,
|
||||
task_id: Option<AmbientAgentTaskId>,
|
||||
working_dir: PathBuf,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Self {
|
||||
let session_bootstrapped = Condition::new();
|
||||
|
||||
// Create a oneshot channel for session sharing when sharing is expected.
|
||||
// When sharing is disabled (or running against ngrok), leave both halves
|
||||
// as None so that `wait_for_session_shared` returns immediately.
|
||||
let sharing_expected =
|
||||
should_share && !warp_core::channel::ChannelState::server_root_url().contains("ngrok");
|
||||
let (mut session_share_tx, session_share_rx) = if sharing_expected {
|
||||
if !FeatureFlag::CreatingSharedSessions.is_enabled() {
|
||||
// Session sharing was requested but the feature is not enabled for this
|
||||
// user/team (typically an enterprise/admin setting). Fail immediately
|
||||
// with a clear error rather than waiting for a timeout.
|
||||
log::warn!(
|
||||
"Session sharing requested but the CreatingSharedSessions feature flag \
|
||||
is not enabled. This is likely due to a team administrator disabling \
|
||||
session sharing."
|
||||
);
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let _ = tx.send(Err(ShareSessionError::Disabled));
|
||||
(None, Some(rx))
|
||||
} else {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
(Some(tx), Some(rx))
|
||||
}
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
// Set the task_id and attachments download dir on the AI controller right away
|
||||
// so they're available for session sharing and file downloads.
|
||||
// Only set the download dir when a task_id is present (cloud mode),
|
||||
// since attachments require a task to fetch presigned URLs.
|
||||
if let Some(tid) = task_id {
|
||||
let attachments_dir = attachments_download_dir(&working_dir);
|
||||
terminal_view.update(ctx, |terminal, ctx| {
|
||||
terminal.ai_controller().update(ctx, |controller, ctx| {
|
||||
controller.set_ambient_agent_task_id(Some(tid), ctx);
|
||||
controller.set_attachments_download_dir(attachments_dir);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
ctx.subscribe_to_view(&terminal_view, move |me, event, ctx| {
|
||||
me.handle_terminal_view_event(event, &mut session_share_tx, ctx);
|
||||
});
|
||||
|
||||
// If the session already bootstrapped before we subscribed, set the
|
||||
// condition immediately so callers of `wait_for_session_bootstrapped`
|
||||
// don't block forever.
|
||||
let already_bootstrapped = terminal_view.read(ctx, |terminal, _| {
|
||||
terminal
|
||||
.model
|
||||
.lock()
|
||||
.block_list()
|
||||
.is_bootstrapping_precmd_done()
|
||||
});
|
||||
if already_bootstrapped {
|
||||
session_bootstrapped.set();
|
||||
}
|
||||
|
||||
Self {
|
||||
terminal_view,
|
||||
session_bootstrapped,
|
||||
shared_session_id: None,
|
||||
session_share_rx,
|
||||
pending_share_requests: Vec::new(),
|
||||
waiting_command: None,
|
||||
pending_command_start: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a handle to the backing terminal view.
|
||||
pub fn terminal_view(&self) -> &ViewHandle<TerminalView> {
|
||||
&self.terminal_view
|
||||
}
|
||||
|
||||
/// Provide mutable access to the terminal view through a closure.
|
||||
pub fn with_terminal_view(
|
||||
&self,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
f: impl FnOnce(&mut TerminalView, &mut warpui::ViewContext<TerminalView>),
|
||||
) {
|
||||
self.terminal_view.update(ctx, f);
|
||||
}
|
||||
|
||||
/// Request that the terminal session be shared with the given participants.
|
||||
///
|
||||
/// This has no effect if the session is not being shared.
|
||||
pub fn add_share_requests(
|
||||
&mut self,
|
||||
share_requests: impl IntoIterator<Item = ShareRequest>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
self.pending_share_requests.extend(share_requests);
|
||||
if self.shared_session_id.is_some() {
|
||||
self.apply_share_requests(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply pending session-sharing guest requests.
|
||||
fn apply_share_requests(&mut self, ctx: &mut ModelContext<Self>) {
|
||||
if self.pending_share_requests.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let share_requests = std::mem::take(&mut self.pending_share_requests);
|
||||
self.terminal_view.update(ctx, |terminal_view, ctx| {
|
||||
let mut viewer_emails = Vec::new();
|
||||
let mut editor_emails = Vec::new();
|
||||
|
||||
for request in share_requests {
|
||||
let role = match request.access_level {
|
||||
ShareAccessLevel::View => Role::Reader,
|
||||
ShareAccessLevel::Edit => Role::Executor,
|
||||
};
|
||||
|
||||
match request.subject {
|
||||
ShareSubject::Team => {
|
||||
if let Some(team_uid) = UserWorkspaces::as_ref(ctx).current_team_uid() {
|
||||
terminal_view.update_session_team_permissions(
|
||||
Some(role),
|
||||
team_uid.to_string(),
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
}
|
||||
ShareSubject::Public => {
|
||||
// Apply an anyone-with-link ACL at the requested role.
|
||||
// This uses the same path as the share modal's
|
||||
// "anyone with link" toggle. The workspace-level
|
||||
// anyone-with-link setting on the server still gates
|
||||
// whether the ACL write succeeds.
|
||||
terminal_view.update_session_link_permissions(Some(role), ctx);
|
||||
}
|
||||
ShareSubject::User { email } => match request.access_level {
|
||||
ShareAccessLevel::View => viewer_emails.push(email),
|
||||
ShareAccessLevel::Edit => editor_emails.push(email),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
if !viewer_emails.is_empty() {
|
||||
terminal_view.add_guests(viewer_emails, Role::Reader, ctx);
|
||||
}
|
||||
if !editor_emails.is_empty() {
|
||||
terminal_view.add_guests(editor_emails, Role::Executor, ctx);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Submit `text` to the active CLI agent on the terminal PTY using the
|
||||
/// agent-specific submission strategy.
|
||||
///
|
||||
/// Used to send exit commands to third-party harnesses.
|
||||
pub(super) fn send_text_to_cli(&self, text: String, ctx: &mut ModelContext<Self>) {
|
||||
self.terminal_view.update(ctx, |terminal, ctx| {
|
||||
terminal.submit_text_to_cli_agent_pty(text, ctx);
|
||||
});
|
||||
}
|
||||
|
||||
/// Return a snapshot of the block with the given ID.
|
||||
pub fn block_snapshot(&self, block_id: &BlockId, ctx: &AppContext) -> Option<SerializedBlock> {
|
||||
let terminal = self.terminal_view.as_ref(ctx);
|
||||
let model = terminal.model.lock();
|
||||
model
|
||||
.block_list()
|
||||
.block_with_id(block_id)
|
||||
.map(SerializedBlock::from)
|
||||
}
|
||||
|
||||
/// Execute a command in the terminal and return a future that resolves to a
|
||||
/// [`CommandHandle`] once the command starts executing.
|
||||
pub fn execute_command(
|
||||
&mut self,
|
||||
command: &str,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Result<impl Future<Output = Result<CommandHandle, AgentDriverError>>, AgentDriverError>
|
||||
{
|
||||
let (exit_tx, exit_rx) = oneshot::channel::<ExitCode>();
|
||||
let (start_tx, start_rx) = oneshot::channel::<BlockId>();
|
||||
|
||||
// We should not be able to execute a command while we are still waiting on another one.
|
||||
// This is enforced by the caller by waiting on rx before continuing.
|
||||
if self.waiting_command.is_some() || self.pending_command_start.is_some() {
|
||||
return Err(AgentDriverError::InvalidRuntimeState);
|
||||
}
|
||||
|
||||
let command_string = command.to_string();
|
||||
self.terminal_view.update(ctx, |terminal, ctx| {
|
||||
self.waiting_command = Some(exit_tx);
|
||||
self.pending_command_start = Some((command_string, start_tx));
|
||||
terminal.execute_command_or_set_pending(command, ctx);
|
||||
});
|
||||
|
||||
Ok(async move {
|
||||
let block_id = start_rx
|
||||
.await
|
||||
.map_err(|_| AgentDriverError::InvalidRuntimeState)?;
|
||||
Ok(CommandHandle {
|
||||
exit_status_rx: exit_rx,
|
||||
block_id,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Execute a command through the active session's in-band command
|
||||
/// executor, without adding a block to the user-visible blocklist.
|
||||
///
|
||||
/// Intended for silent probes (e.g. `test -d`) that the agent needs to
|
||||
/// drive through the terminal session (so they run against the correct
|
||||
/// filesystem, including inside a Docker sandbox) but should not clutter
|
||||
/// the user's command history.
|
||||
pub fn execute_silent_command(
|
||||
&self,
|
||||
command: String,
|
||||
ctx: &ModelContext<Self>,
|
||||
) -> impl Future<Output = Result<CommandOutput, AgentDriverError>> {
|
||||
let session = self.terminal_view.read(ctx, |terminal, app| {
|
||||
terminal
|
||||
.active_block_session_id()
|
||||
.and_then(|id| terminal.sessions_model().as_ref(app).get(id))
|
||||
});
|
||||
async move {
|
||||
let session = session.ok_or(AgentDriverError::InvalidRuntimeState)?;
|
||||
session
|
||||
.execute_command(&command, None, None, ExecuteCommandOptions::default())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
log::warn!("silent command failed: {e:#}");
|
||||
AgentDriverError::InvalidRuntimeState
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the shell type of the active terminal session, if known.
|
||||
pub fn active_session_shell_type(&self, ctx: &AppContext) -> Option<ShellType> {
|
||||
self.terminal_view
|
||||
.read(ctx, |terminal, app| terminal.active_session_shell_type(app))
|
||||
}
|
||||
|
||||
/// Build the shell-aware `cd <escaped>` command for the active session.
|
||||
///
|
||||
/// Shared between [`Self::cd`] and [`Self::cd_silent`] so both paths use
|
||||
/// the same [`ShellFamily::shell_escape`] logic (posix single-quoting,
|
||||
/// fish backslash, pwsh double-quote doubling) and don't drift.
|
||||
fn build_cd_command(&self, target: &str, ctx: &AppContext) -> String {
|
||||
let shell_family = self.terminal_view.read(ctx, |terminal, app| {
|
||||
terminal
|
||||
.active_session_shell_type(app)
|
||||
.map(ShellFamily::from)
|
||||
.unwrap_or(ShellFamily::Posix)
|
||||
});
|
||||
let escaped_target = shell_family.shell_escape(target);
|
||||
format!("cd {escaped_target}")
|
||||
}
|
||||
|
||||
/// Change directory within the active terminal session.
|
||||
pub fn cd(
|
||||
&mut self,
|
||||
target: &str,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Result<impl Future<Output = Result<CommandHandle, AgentDriverError>>, AgentDriverError>
|
||||
{
|
||||
let cd_command = self.build_cd_command(target, ctx);
|
||||
self.execute_command(&cd_command, ctx)
|
||||
}
|
||||
|
||||
/// Change directory within the active terminal session, silently — no
|
||||
/// visible block is added to the user-facing blocklist.
|
||||
///
|
||||
/// Uses the same shell-aware escaping as [`Self::cd`] but dispatches
|
||||
/// through [`Self::execute_silent_command`]. Intended for callers that
|
||||
/// need to position the session's CWD as an implementation detail of a
|
||||
/// larger setup step (e.g. positioning the session before running
|
||||
/// silent probes).
|
||||
pub fn cd_silent(
|
||||
&self,
|
||||
target: &str,
|
||||
ctx: &ModelContext<Self>,
|
||||
) -> impl Future<Output = Result<CommandOutput, AgentDriverError>> {
|
||||
let cd_command = self.build_cd_command(target, ctx);
|
||||
self.execute_silent_command(cd_command, ctx)
|
||||
}
|
||||
|
||||
/// The current working directory of the active terminal session, if known.
|
||||
#[allow(dead_code)]
|
||||
pub fn current_directory(&self, ctx: &AppContext) -> Option<PathBuf> {
|
||||
// TODO(ben): This should handle non-local paths.
|
||||
self.terminal_view
|
||||
.as_ref(ctx)
|
||||
.active_session_path_if_local(ctx)
|
||||
}
|
||||
|
||||
/// Returns a future that resolves when the session has bootstrapped.
|
||||
///
|
||||
/// This only waits for the `SessionBootstrapped` terminal view event.
|
||||
pub fn wait_for_session_bootstrapped(
|
||||
&self,
|
||||
) -> impl Future<Output = Result<(), AgentDriverError>> {
|
||||
let session_bootstrapped = self.session_bootstrapped.clone();
|
||||
|
||||
async move {
|
||||
session_bootstrapped
|
||||
.wait()
|
||||
.with_timeout(TERMINAL_SESSION_BOOTSTRAP_TIMEOUT)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
log::error!("Timed out waiting for session bootstrap");
|
||||
AgentDriverError::BootstrapFailed
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a future that resolves when (optional) session sharing has started.
|
||||
///
|
||||
/// This is separate from `wait_for_session_bootstrapped` so that callers can:
|
||||
/// - wait for terminal bootstrap early (e.g. before starting MCP servers)
|
||||
/// - wait for session sharing later (e.g. right before running visible commands)
|
||||
pub fn wait_for_session_shared(
|
||||
&mut self,
|
||||
) -> impl Future<Output = Result<(), AgentDriverError>> {
|
||||
let rx = self.session_share_rx.take();
|
||||
|
||||
async move {
|
||||
let Some(rx) = rx else {
|
||||
// Sharing is disabled or already resolved.
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
match rx.with_timeout(TERMINAL_SESSION_SHARE_DELAY).await {
|
||||
Ok(Ok(Ok(()))) => Ok(()),
|
||||
Ok(Ok(Err(error))) => {
|
||||
log::error!("Session sharing failed: {error}");
|
||||
Err(AgentDriverError::ShareSessionFailed { error })
|
||||
}
|
||||
Ok(Err(_canceled)) => {
|
||||
log::error!("Session sharing channel dropped");
|
||||
Err(AgentDriverError::ShareSessionFailed {
|
||||
error: ShareSessionError::Interrupted,
|
||||
})
|
||||
}
|
||||
Err(_timeout) => {
|
||||
log::error!(
|
||||
"Timed out waiting for session sharing to start after {}s",
|
||||
TERMINAL_SESSION_SHARE_DELAY.as_secs()
|
||||
);
|
||||
Err(AgentDriverError::ShareSessionFailed {
|
||||
error: ShareSessionError::Timeout,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A handle to a running terminal command.
|
||||
///
|
||||
/// Resolves to the command's [`ExitCode`] when the block completes.
|
||||
/// Also carries the [`BlockId`] so callers can retrieve the block snapshot
|
||||
/// after completion.
|
||||
pub(crate) struct CommandHandle {
|
||||
exit_status_rx: oneshot::Receiver<ExitCode>,
|
||||
block_id: BlockId,
|
||||
}
|
||||
|
||||
impl CommandHandle {
|
||||
/// The block ID of the command that was executed.
|
||||
pub fn block_id(&self) -> &BlockId {
|
||||
&self.block_id
|
||||
}
|
||||
}
|
||||
|
||||
impl Future for CommandHandle {
|
||||
type Output = Result<ExitCode, AgentDriverError>;
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
Pin::new(&mut self.exit_status_rx)
|
||||
.poll(cx)
|
||||
.map(|result| result.map_err(|_| AgentDriverError::InvalidRuntimeState))
|
||||
}
|
||||
}
|
||||
|
||||
impl TerminalDriver {
|
||||
/// Handle terminal view events.
|
||||
fn handle_terminal_view_event(
|
||||
&mut self,
|
||||
event: &crate::terminal::view::Event,
|
||||
session_share_tx: &mut Option<oneshot::Sender<Result<(), ShareSessionError>>>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
match event {
|
||||
crate::terminal::view::Event::SessionBootstrapped => {
|
||||
self.session_bootstrapped.set();
|
||||
}
|
||||
crate::terminal::view::Event::SlowBootstrap => {
|
||||
ctx.emit(TerminalDriverEvent::SlowBootstrap);
|
||||
}
|
||||
crate::terminal::view::Event::EstablishedSharedSession { session_id } => {
|
||||
self.shared_session_id = Some(*session_id);
|
||||
if let Some(tx) = session_share_tx.take() {
|
||||
let _ = tx.send(Ok(()));
|
||||
}
|
||||
|
||||
// Apply any pending share requests now that the session is established.
|
||||
self.apply_share_requests(ctx);
|
||||
|
||||
ctx.emit(TerminalDriverEvent::EstablishedSharedSession {
|
||||
session_id: *session_id,
|
||||
join_url: shared_session::join_link(session_id),
|
||||
});
|
||||
}
|
||||
crate::terminal::view::Event::FailedToShareSession { reason, cause } => {
|
||||
if let Some(tx) = session_share_tx.take() {
|
||||
let error = match cause {
|
||||
Some(cause) => ShareSessionError::Internal(cause.clone()),
|
||||
None => ShareSessionError::Failed(reason.clone()),
|
||||
};
|
||||
let _ = tx.send(Err(error));
|
||||
}
|
||||
}
|
||||
crate::terminal::view::Event::ExecuteCommand(event) => {
|
||||
if let Some((_expected_command, sender)) = self
|
||||
.pending_command_start
|
||||
.take_if(|(cmd, _)| *cmd == event.command)
|
||||
{
|
||||
let block_id = self.terminal_view.read(ctx, |terminal, _| {
|
||||
terminal.model.lock().block_list().active_block_id().clone()
|
||||
});
|
||||
let _ = sender.send(block_id);
|
||||
}
|
||||
}
|
||||
crate::terminal::view::Event::BlockCompleted { block, .. } => {
|
||||
if let Some(sender) = self.waiting_command.take_if(|_| {
|
||||
let bootstrapping_done = self.terminal_view.read(ctx, |terminal, _| {
|
||||
terminal
|
||||
.model
|
||||
.lock()
|
||||
.block_list()
|
||||
.is_bootstrapping_precmd_done()
|
||||
});
|
||||
// This was originally checking `bootstrapping_done && block.did_execute`.
|
||||
// Oddly, we've seen cases where we missed the preexec hook, so
|
||||
// `block.did_execute` is false even though the command actually did run.
|
||||
// To hedge against this while we're still figuring out the root cause,
|
||||
// we instead simply make sure it was not a background block.
|
||||
bootstrapping_done && !block.is_background
|
||||
}) {
|
||||
let _ = sender.send(block.exit_code);
|
||||
}
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,388 @@
|
||||
use std::{ffi::OsString, sync::Arc, time::Duration};
|
||||
|
||||
use futures::channel::oneshot;
|
||||
use warp_cli::agent::Harness;
|
||||
use warp_cli::{
|
||||
OZ_CLI_ENV, OZ_HARNESS_ENV, OZ_PARENT_RUN_ID_ENV, OZ_RUN_ID_ENV, SERVER_ROOT_URL_OVERRIDE_ENV,
|
||||
SESSION_SHARING_SERVER_URL_OVERRIDE_ENV, WS_SERVER_URL_OVERRIDE_ENV,
|
||||
};
|
||||
use warp_core::channel::ChannelState;
|
||||
|
||||
use super::{
|
||||
IdleTimeoutSender, LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV,
|
||||
LEGACY_OZ_PARENT_STATE_ROOT_ENV, OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV,
|
||||
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
};
|
||||
use crate::ai::agent::{
|
||||
task::TaskId, AIAgentActionResult, AIAgentActionResultType, AIAgentInput, AIAgentOutput,
|
||||
AIAgentOutputMessage, ArtifactCreatedData, MessageId, UploadArtifactResult,
|
||||
};
|
||||
use crate::ai::mcp::parsing::normalize_mcp_json;
|
||||
use crate::ai::{agent_sdk::task_env_vars, ambient_agents::AmbientAgentTaskId};
|
||||
|
||||
#[test]
|
||||
fn test_normalize_single_cli_server() {
|
||||
let input = r#"{"command": "npx", "args": ["-y", "mcp-server"]}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
// Should wrap with a generated name
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
let parsed = parsed.as_object().unwrap();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
let (_name, server) = parsed.iter().next().unwrap();
|
||||
assert_eq!(server["command"].as_str().unwrap(), "npx");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_single_sse_server() {
|
||||
let input = r#"{"url": "http://localhost:3000/mcp", "headers": {"API_KEY": "value"}}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
// Should wrap with a generated name
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
let parsed = parsed.as_object().unwrap();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
let (_name, server) = parsed.iter().next().unwrap();
|
||||
assert_eq!(server["url"].as_str().unwrap(), "http://localhost:3000/mcp");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_already_wrapped_server() {
|
||||
let input = r#"{"my-server": {"command": "npx", "args": []}}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
// Should return as-is (no command/url at top level)
|
||||
assert_eq!(result, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_mcp_servers_wrapper() {
|
||||
let input = r#"{"mcpServers": {"server-name": {"command": "npx", "args": []}}}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
// Should return as-is (no command/url at top level)
|
||||
assert_eq!(result, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_servers_wrapper() {
|
||||
let input = r#"{"servers": {"server-name": {"url": "http://example.com"}}}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
// Should return as-is (no command/url at top level)
|
||||
assert_eq!(result, input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_invalid_json() {
|
||||
let input = "not valid json";
|
||||
let result = normalize_mcp_json(input);
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_cli_server_with_env() {
|
||||
let input = r#"{"command": "npx", "args": ["-y", "mcp-server"], "env": {"API_KEY": "secret"}}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
let parsed = parsed.as_object().unwrap();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
let (_name, server) = parsed.iter().next().unwrap();
|
||||
assert_eq!(server["env"]["API_KEY"].as_str().unwrap(), "secret");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_sse_server_with_headers() {
|
||||
let input =
|
||||
r#"{"url": "http://localhost:5000/mcp", "headers": {"Authorization": "Bearer token"}}"#;
|
||||
let result = normalize_mcp_json(input).unwrap();
|
||||
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
let parsed = parsed.as_object().unwrap();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
let (_name, server) = parsed.iter().next().unwrap();
|
||||
assert_eq!(
|
||||
server["headers"]["Authorization"].as_str().unwrap(),
|
||||
"Bearer token"
|
||||
);
|
||||
}
|
||||
|
||||
// ── IdleTimeoutSender tests ──────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_sender_send_now_delivers_value() {
|
||||
let (tx, mut rx) = oneshot::channel::<i32>();
|
||||
let idle_timeout = IdleTimeoutSender::new(tx);
|
||||
idle_timeout.end_run_now(42);
|
||||
assert_eq!(rx.try_recv().unwrap(), Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_sender_send_now_only_delivers_once() {
|
||||
let (tx, mut rx) = oneshot::channel::<i32>();
|
||||
let idle_timeout = IdleTimeoutSender::new(tx);
|
||||
idle_timeout.end_run_now(1);
|
||||
idle_timeout.end_run_now(2);
|
||||
assert_eq!(rx.try_recv().unwrap(), Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_sender_send_after_delivers_after_timeout() {
|
||||
let (tx, mut rx) = oneshot::channel::<i32>();
|
||||
let idle_timeout = IdleTimeoutSender::new(tx);
|
||||
idle_timeout.end_run_after(Duration::from_millis(50), 99);
|
||||
|
||||
// Not yet delivered.
|
||||
assert_eq!(rx.try_recv().unwrap(), None);
|
||||
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
assert_eq!(rx.try_recv().unwrap(), Some(99));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_sender_cancel_prevents_delivery() {
|
||||
let (tx, mut rx) = oneshot::channel::<i32>();
|
||||
let idle_timeout = IdleTimeoutSender::new(tx);
|
||||
idle_timeout.end_run_after(Duration::from_millis(50), 99);
|
||||
idle_timeout.cancel_idle_timeout();
|
||||
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
// Sender was not consumed, so the channel is still open but empty.
|
||||
assert_eq!(rx.try_recv().unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_sender_cancel_then_send_now_delivers() {
|
||||
let (tx, mut rx) = oneshot::channel::<i32>();
|
||||
let idle_timeout = IdleTimeoutSender::new(tx);
|
||||
idle_timeout.end_run_after(Duration::from_millis(50), 1);
|
||||
idle_timeout.cancel_idle_timeout();
|
||||
idle_timeout.end_run_now(2);
|
||||
|
||||
assert_eq!(rx.try_recv().unwrap(), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_sender_later_send_after_supersedes_earlier() {
|
||||
let (tx, mut rx) = oneshot::channel::<i32>();
|
||||
let idle_timeout = IdleTimeoutSender::new(tx);
|
||||
// First timer: long timeout.
|
||||
idle_timeout.end_run_after(Duration::from_secs(10), 1);
|
||||
// Second timer: short timeout. The first is implicitly cancelled.
|
||||
idle_timeout.end_run_after(Duration::from_millis(50), 2);
|
||||
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
assert_eq!(rx.try_recv().unwrap(), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_env_vars_include_parent_run_id_when_present() {
|
||||
let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440000".parse().unwrap();
|
||||
let env_vars = task_env_vars(Some(&task_id), Some("parent-run-123"), Harness::Claude);
|
||||
let overrides_allowed = ChannelState::channel().allows_server_url_overrides();
|
||||
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_RUN_ID_ENV)),
|
||||
Some(&OsString::from(task_id.to_string()))
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_PARENT_RUN_ID_ENV)),
|
||||
Some(&OsString::from("parent-run-123"))
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_HARNESS_ENV)),
|
||||
Some(&OsString::from("claude"))
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV)),
|
||||
Some(&OsString::from("1"))
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(
|
||||
LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV
|
||||
)),
|
||||
Some(&OsString::from("1"))
|
||||
);
|
||||
assert!(env_vars
|
||||
.get(&OsString::from(OZ_CLI_ENV))
|
||||
.is_some_and(|value| !value.is_empty()));
|
||||
|
||||
let server_root_url = ChannelState::server_root_url().into_owned();
|
||||
if overrides_allowed && !server_root_url.is_empty() {
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(SERVER_ROOT_URL_OVERRIDE_ENV)),
|
||||
Some(&OsString::from(server_root_url))
|
||||
);
|
||||
} else {
|
||||
assert!(!env_vars.contains_key(&OsString::from(SERVER_ROOT_URL_OVERRIDE_ENV)));
|
||||
}
|
||||
|
||||
let ws_server_url = ChannelState::ws_server_url().into_owned();
|
||||
if overrides_allowed && !ws_server_url.is_empty() {
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(WS_SERVER_URL_OVERRIDE_ENV)),
|
||||
Some(&OsString::from(ws_server_url))
|
||||
);
|
||||
} else {
|
||||
assert!(!env_vars.contains_key(&OsString::from(WS_SERVER_URL_OVERRIDE_ENV)));
|
||||
}
|
||||
|
||||
if overrides_allowed {
|
||||
match ChannelState::session_sharing_server_url() {
|
||||
Some(url) if !url.is_empty() => assert_eq!(
|
||||
env_vars.get(&OsString::from(SESSION_SHARING_SERVER_URL_OVERRIDE_ENV)),
|
||||
Some(&OsString::from(url.into_owned()))
|
||||
),
|
||||
_ => {
|
||||
assert!(!env_vars
|
||||
.contains_key(&OsString::from(SESSION_SHARING_SERVER_URL_OVERRIDE_ENV)))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
assert!(!env_vars.contains_key(&OsString::from(SESSION_SHARING_SERVER_URL_OVERRIDE_ENV)));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_env_vars_omit_parent_run_id_when_absent() {
|
||||
let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440001".parse().unwrap();
|
||||
let env_vars = task_env_vars(Some(&task_id), None, Harness::Oz);
|
||||
let overrides_allowed = ChannelState::channel().allows_server_url_overrides();
|
||||
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_RUN_ID_ENV)),
|
||||
Some(&OsString::from(task_id.to_string()))
|
||||
);
|
||||
assert!(!env_vars.contains_key(&OsString::from(OZ_PARENT_RUN_ID_ENV)));
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_HARNESS_ENV)),
|
||||
Some(&OsString::from("oz"))
|
||||
);
|
||||
assert!(!env_vars.contains_key(&OsString::from(OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV)));
|
||||
assert!(!env_vars.contains_key(&OsString::from(
|
||||
LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV
|
||||
)));
|
||||
assert_eq!(
|
||||
env_vars.contains_key(&OsString::from(SERVER_ROOT_URL_OVERRIDE_ENV)),
|
||||
overrides_allowed && !ChannelState::server_root_url().is_empty()
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.contains_key(&OsString::from(WS_SERVER_URL_OVERRIDE_ENV)),
|
||||
overrides_allowed && !ChannelState::ws_server_url().is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_env_vars_enable_external_parent_listener_for_claude_runs_without_parent_run_id() {
|
||||
let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440002".parse().unwrap();
|
||||
let env_vars = task_env_vars(Some(&task_id), None, Harness::Claude);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV)),
|
||||
Some(&OsString::from("1"))
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(
|
||||
LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV
|
||||
)),
|
||||
Some(&OsString::from("1"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn task_env_vars_propagate_message_listener_state_root_with_legacy_alias() {
|
||||
let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440003".parse().unwrap();
|
||||
std::env::set_var(
|
||||
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
|
||||
"/tmp/message-listener-root",
|
||||
);
|
||||
let env_vars = task_env_vars(Some(&task_id), None, Harness::Claude);
|
||||
std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV);
|
||||
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV)),
|
||||
Some(&OsString::from("/tmp/message-listener-root"))
|
||||
);
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(LEGACY_OZ_PARENT_STATE_ROOT_ENV)),
|
||||
Some(&OsString::from("/tmp/message-listener-root"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_env_vars_can_use_opencode_harness() {
|
||||
let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440004".parse().unwrap();
|
||||
let env_vars = task_env_vars(Some(&task_id), Some("parent-run-456"), Harness::OpenCode);
|
||||
|
||||
assert_eq!(
|
||||
env_vars.get(&OsString::from(OZ_HARNESS_ENV)),
|
||||
Some(&OsString::from("opencode"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_format_output_includes_filename_for_file_artifact_created_event() {
|
||||
let output = AIAgentOutput {
|
||||
messages: vec![AIAgentOutputMessage::artifact_created(
|
||||
MessageId::new("message-1".to_string()),
|
||||
ArtifactCreatedData::File {
|
||||
artifact_uid: "artifact-uid".to_string(),
|
||||
filepath: "outputs/report.txt".to_string(),
|
||||
filename: "report.txt".to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
description: Some("Build output for the latest run".to_string()),
|
||||
size_bytes: 42,
|
||||
},
|
||||
)],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
super::output::json::format_output(&output, &mut bytes).expect("json formatting should work");
|
||||
|
||||
let value: serde_json::Value =
|
||||
serde_json::from_slice(&bytes).expect("output should be valid json");
|
||||
|
||||
assert_eq!(value["type"], "artifact_created");
|
||||
assert_eq!(value["artifact_type"], "file");
|
||||
assert_eq!(value["artifact_uid"], "artifact-uid");
|
||||
assert_eq!(value["filepath"], "outputs/report.txt");
|
||||
assert_eq!(value["filename"], "report.txt");
|
||||
assert_eq!(value["mime_type"], "text/plain");
|
||||
assert_eq!(value["description"], "Build output for the latest run");
|
||||
assert_eq!(value["size_bytes"], 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_format_input_omits_filepath_and_description_for_proto_upload_result() {
|
||||
let input = AIAgentInput::ActionResult {
|
||||
result: AIAgentActionResult {
|
||||
id: "tool-call-1".to_string().into(),
|
||||
task_id: TaskId::new("task-1".to_string()),
|
||||
result: AIAgentActionResultType::UploadArtifact(UploadArtifactResult::Success {
|
||||
artifact_uid: "artifact-123".to_string(),
|
||||
filepath: None,
|
||||
mime_type: "text/plain".to_string(),
|
||||
description: None,
|
||||
size_bytes: 42,
|
||||
}),
|
||||
},
|
||||
context: Arc::from([]),
|
||||
};
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
super::output::json::format_input(&input, &mut bytes).expect("json formatting should work");
|
||||
|
||||
let value: serde_json::Value =
|
||||
serde_json::from_slice(&bytes).expect("output should be valid json");
|
||||
|
||||
assert_eq!(value["type"], "tool_result");
|
||||
assert_eq!(value["tool"], "upload_artifact");
|
||||
assert_eq!(value["artifact_uid"], "artifact-123");
|
||||
assert_eq!(value["mime_type"], "text/plain");
|
||||
assert_eq!(value["size_bytes"], 42);
|
||||
assert!(value.get("filepath").is_none());
|
||||
assert!(value.get("description").is_none());
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,128 @@
|
||||
use std::process;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde_json::json;
|
||||
use warp_cli::federate::{FederateCommand, IssueGcpTokenArgs, IssueTokenArgs};
|
||||
use warp_cli::{agent::OutputFormat, GlobalOptions};
|
||||
use warp_core::{features::FeatureFlag, report_error};
|
||||
use warp_managed_secrets::ManagedSecretManager;
|
||||
use warpui::{platform::TerminationMode, AppContext, SingletonEntity as _};
|
||||
|
||||
use super::common::set_ambient_task_context_from_run_id;
|
||||
|
||||
/// Run identity federation commands.
|
||||
pub fn run(
|
||||
ctx: &mut AppContext,
|
||||
global_options: GlobalOptions,
|
||||
command: FederateCommand,
|
||||
) -> Result<()> {
|
||||
if !FeatureFlag::OzIdentityFederation.is_enabled() {
|
||||
return Err(anyhow::anyhow!("This feature is not enabled"));
|
||||
}
|
||||
match command {
|
||||
FederateCommand::IssueToken(args) => issue_token(ctx, args, global_options.output_format),
|
||||
FederateCommand::IssueGcpToken(args) => issue_gcp_token(ctx, args),
|
||||
}
|
||||
}
|
||||
|
||||
fn issue_token(
|
||||
ctx: &mut AppContext,
|
||||
args: IssueTokenArgs,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
// Set the task ID so the ambient workload token header is sent.
|
||||
set_ambient_task_context_from_run_id(ctx, &args.run_id)?;
|
||||
|
||||
let duration: std::time::Duration = args.duration.into();
|
||||
let audience = args.audience;
|
||||
let subject_template = match args.subject_template {
|
||||
Some(template) => vec1::Vec1::try_from_vec(template)
|
||||
.map_err(|_| anyhow::anyhow!("--subject-template requires at least one value"))?,
|
||||
None => vec1::vec1!["principal".to_owned()],
|
||||
};
|
||||
|
||||
ManagedSecretManager::handle(ctx).update(ctx, move |manager, ctx| {
|
||||
let future =
|
||||
manager.issue_task_identity_token(warp_managed_secrets::client::IdentityTokenOptions {
|
||||
audience,
|
||||
requested_duration: duration,
|
||||
subject_template,
|
||||
});
|
||||
ctx.spawn(future, move |_, result, ctx| match result {
|
||||
Ok(token) => {
|
||||
let token_value = token.token;
|
||||
let expires_at = token.expires_at.to_rfc3339();
|
||||
let issuer = token.issuer;
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
let output = json!({
|
||||
"token": token_value,
|
||||
"expires_at": expires_at,
|
||||
"issuer": issuer,
|
||||
});
|
||||
let output =
|
||||
serde_json::to_string(&output).expect("token output should serialize");
|
||||
println!("{output}");
|
||||
}
|
||||
OutputFormat::Text => {
|
||||
println!("{token_value}");
|
||||
}
|
||||
OutputFormat::Pretty => {
|
||||
println!("Token: {token_value}");
|
||||
println!("Expires at: {expires_at}");
|
||||
println!("Issuer: {issuer}");
|
||||
}
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(err, ctx);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
fn issue_gcp_token(ctx: &mut AppContext, args: IssueGcpTokenArgs) -> Result<()> {
|
||||
// Set the task ID so the ambient workload token header is sent.
|
||||
set_ambient_task_context_from_run_id(ctx, &args.run_id)?;
|
||||
|
||||
let duration: std::time::Duration = args.duration.into();
|
||||
let audience = args.audience;
|
||||
let token_type = args.token_type;
|
||||
let output_file = args.output_file;
|
||||
|
||||
ManagedSecretManager::handle(ctx).update(ctx, move |manager, ctx| {
|
||||
let future =
|
||||
manager.issue_gcp_workload_identity_federation_token(audience, token_type, duration);
|
||||
ctx.spawn(future, move |_, result, ctx| match result {
|
||||
Ok(token) => {
|
||||
let output =
|
||||
serde_json::to_string(&token).expect("gcp token output should serialize");
|
||||
|
||||
// If we can't cache the token, report an error but don't fail the command.
|
||||
if let Some(output_path) = output_file {
|
||||
if let Err(err) = std::fs::write(&output_path, &output) {
|
||||
report_error!(anyhow!(err)
|
||||
.context(format!("Error writing GCP token to {output_path}")));
|
||||
}
|
||||
}
|
||||
|
||||
println!("{output}");
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
// The GCP SDK requires the executable to print a JSON error to stderr
|
||||
// and exit with a non-zero status code. Because of this, we exit
|
||||
// directly instead of via `ctx.terminate_app` (which would print a
|
||||
// non-JSON error).
|
||||
let output =
|
||||
serde_json::to_string(&err).expect("gcp error output should serialize");
|
||||
eprintln!("{output}");
|
||||
process::exit(1);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
//! `warp harness-support` CLI dispatch and the singleton model all subcommands run async work on.
|
||||
//!
|
||||
//! Subcommands:
|
||||
//! - [`ping`] — fetches the current run by task ID and prints its info.
|
||||
//! - [`report_artifact`] — reports an artifact (e.g. a PR) back to the Oz platform.
|
||||
use anyhow::Result;
|
||||
use warp_cli::agent::OutputFormat;
|
||||
use warp_cli::harness_support::{
|
||||
FinishTaskArgs, HarnessSupportArgs, HarnessSupportCommand, NotifyUserArgs, ReportArtifactArgs,
|
||||
ReportArtifactCommand, TaskStatus,
|
||||
};
|
||||
use warp_cli::GlobalOptions;
|
||||
use warp_core::features::FeatureFlag;
|
||||
use warpui::{platform::TerminationMode, AppContext, ModelHandle, SingletonEntity};
|
||||
|
||||
use super::common::set_ambient_task_context_from_run_id;
|
||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||
use crate::ai::artifacts::Artifact;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
|
||||
/// Run harness-support commands.
|
||||
pub fn run(
|
||||
ctx: &mut AppContext,
|
||||
global_options: GlobalOptions,
|
||||
args: HarnessSupportArgs,
|
||||
) -> Result<()> {
|
||||
if !FeatureFlag::AgentHarness.is_enabled() {
|
||||
return Err(anyhow::anyhow!("This feature is not enabled"));
|
||||
}
|
||||
|
||||
// Store the run ID so that it's included on all server requests, along with a workload token.
|
||||
let task_id = set_ambient_task_context_from_run_id(ctx, &args.run_id)?;
|
||||
let runner = ctx.add_singleton_model(|_| HarnessSupportRunner);
|
||||
|
||||
match args.command {
|
||||
HarnessSupportCommand::Ping => ping(ctx, runner, task_id, global_options.output_format),
|
||||
HarnessSupportCommand::ReportArtifact(report_args) => {
|
||||
report_artifact(ctx, runner, report_args, global_options.output_format)
|
||||
}
|
||||
HarnessSupportCommand::NotifyUser(notify_args) => {
|
||||
notify_user(ctx, runner, notify_args, global_options.output_format)
|
||||
}
|
||||
HarnessSupportCommand::FinishTask(finish_args) => {
|
||||
finish_task(ctx, runner, finish_args, global_options.output_format)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch the current run by ID and print its info.
|
||||
fn ping(
|
||||
ctx: &mut AppContext,
|
||||
runner: ModelHandle<HarnessSupportRunner>,
|
||||
task_id: AmbientAgentTaskId,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
runner.update(ctx, |_, ctx| {
|
||||
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
|
||||
|
||||
ctx.spawn(
|
||||
async move {
|
||||
let task = ai_client.get_ambient_agent_task(&task_id).await?;
|
||||
Ok(task)
|
||||
},
|
||||
move |_, result, ctx| match result {
|
||||
Ok(task) => {
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
let json = serde_json::to_string(&task).unwrap_or_else(|e| {
|
||||
serde_json::json!({"error": e.to_string()}).to_string()
|
||||
});
|
||||
println!("{json}");
|
||||
}
|
||||
OutputFormat::Pretty | OutputFormat::Text => {
|
||||
super::ambient::print_tasks(&[task]);
|
||||
}
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(err, ctx);
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Report an artifact back to the Oz platform.
|
||||
fn report_artifact(
|
||||
ctx: &mut AppContext,
|
||||
runner: ModelHandle<HarnessSupportRunner>,
|
||||
args: ReportArtifactArgs,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
runner.update(ctx, |_, ctx| {
|
||||
let client = ServerApiProvider::as_ref(ctx).get_harness_support_client();
|
||||
|
||||
let artifact = match args.command {
|
||||
ReportArtifactCommand::PullRequest(pr_args) => Artifact::PullRequest {
|
||||
url: pr_args.url,
|
||||
branch: pr_args.branch,
|
||||
repo: None,
|
||||
number: None,
|
||||
},
|
||||
};
|
||||
|
||||
ctx.spawn(
|
||||
async move { client.report_artifact(&artifact).await },
|
||||
move |_, result, ctx| match result {
|
||||
Ok(response) => {
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
|
||||
serde_json::json!({"error": e.to_string()}).to_string()
|
||||
});
|
||||
println!("{json}");
|
||||
}
|
||||
OutputFormat::Pretty | OutputFormat::Text => {
|
||||
println!("Artifact reported: {}", response.artifact_uid);
|
||||
}
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(err, ctx);
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a progress notification to the task's originating platform.
|
||||
fn notify_user(
|
||||
ctx: &mut AppContext,
|
||||
runner: ModelHandle<HarnessSupportRunner>,
|
||||
args: NotifyUserArgs,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
runner.update(ctx, |_, ctx| {
|
||||
let client = ServerApiProvider::as_ref(ctx).get_harness_support_client();
|
||||
|
||||
ctx.spawn(
|
||||
async move { client.notify_user(&args.message).await },
|
||||
move |_, result, ctx| match result {
|
||||
Ok(()) => {
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
println!("{{}}");
|
||||
}
|
||||
OutputFormat::Pretty | OutputFormat::Text => {
|
||||
println!("Notification sent.");
|
||||
}
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(err, ctx);
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Report task completion or failure.
|
||||
fn finish_task(
|
||||
ctx: &mut AppContext,
|
||||
runner: ModelHandle<HarnessSupportRunner>,
|
||||
args: FinishTaskArgs,
|
||||
output_format: OutputFormat,
|
||||
) -> Result<()> {
|
||||
runner.update(ctx, |_, ctx| {
|
||||
let client = ServerApiProvider::as_ref(ctx).get_harness_support_client();
|
||||
|
||||
ctx.spawn(
|
||||
async move {
|
||||
let success = args.status == TaskStatus::Success;
|
||||
client.finish_task(success, &args.summary).await
|
||||
},
|
||||
move |_, result, ctx| match result {
|
||||
Ok(()) => {
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
println!("{{}}");
|
||||
}
|
||||
OutputFormat::Pretty | OutputFormat::Text => {
|
||||
println!("Task finished.");
|
||||
}
|
||||
}
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(err, ctx);
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Singleton model for running async harness-support operations.
|
||||
struct HarnessSupportRunner;
|
||||
|
||||
impl warpui::Entity for HarnessSupportRunner {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for HarnessSupportRunner {}
|
||||
@@ -0,0 +1,524 @@
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
use futures::future;
|
||||
use warp_cli::{
|
||||
integration::{CreateIntegrationArgs, IntegrationCommand, UpdateIntegrationArgs},
|
||||
provider::ProviderType,
|
||||
GlobalOptions,
|
||||
};
|
||||
use warp_graphql::mutations::create_simple_integration::CreateSimpleIntegrationOutput;
|
||||
use warp_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus;
|
||||
use warp_graphql::queries::get_simple_integrations::SimpleIntegrationsOutput;
|
||||
use warpui::{platform::TerminationMode, AppContext, ModelContext, SingletonEntity};
|
||||
|
||||
use super::common::{EnvironmentChoice, ResolveConfigurationError};
|
||||
use super::integration_output;
|
||||
use super::oauth_flow::poll_oauth_until_terminal;
|
||||
|
||||
pub fn run(
|
||||
ctx: &mut AppContext,
|
||||
global_options: GlobalOptions,
|
||||
command: IntegrationCommand,
|
||||
) -> anyhow::Result<()> {
|
||||
let runner = ctx.add_singleton_model(|_ctx| IntegrationCommandRunner);
|
||||
match command {
|
||||
IntegrationCommand::Create(args) => {
|
||||
runner.update(ctx, |runner, ctx| runner.create(args, ctx));
|
||||
}
|
||||
IntegrationCommand::Update(args) => {
|
||||
runner.update(ctx, |runner, ctx| runner.update(args, ctx));
|
||||
}
|
||||
IntegrationCommand::List => {
|
||||
runner.update(ctx, |runner, ctx| runner.list(global_options, ctx));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct IntegrationCommandRunner;
|
||||
|
||||
impl IntegrationCommandRunner {
|
||||
fn list(&self, global_options: GlobalOptions, ctx: &mut ModelContext<Self>) {
|
||||
// Hardcoded set of providers that this client knows how to render.
|
||||
let providers = vec![ProviderType::Linear, ProviderType::Slack];
|
||||
let provider_slugs: Vec<String> = providers.into_iter().map(|p| p.slug()).collect();
|
||||
|
||||
let integrations_client = ServerApiProvider::as_ref(ctx).get_integrations_client();
|
||||
|
||||
let list_future = async move {
|
||||
integrations_client
|
||||
.list_simple_integrations(provider_slugs)
|
||||
.await
|
||||
};
|
||||
|
||||
ctx.spawn(
|
||||
list_future,
|
||||
move |_, result: anyhow::Result<SimpleIntegrationsOutput>, ctx| match result {
|
||||
Ok(output) => {
|
||||
integration_output::print_integrations(&output, global_options.output_format);
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
}
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn create(&self, args: CreateIntegrationArgs, ctx: &mut ModelContext<Self>) {
|
||||
let refresh_future = super::common::refresh_workspace_metadata(ctx);
|
||||
let warp_drive_sync_future = super::common::refresh_warp_drive(ctx);
|
||||
let setup_future = future::try_join(refresh_future, warp_drive_sync_future);
|
||||
|
||||
ctx.spawn(setup_future, move |runner, setup_result, ctx| {
|
||||
if let Err(err) = setup_result {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
|
||||
let loaded_file = match args.config_file.file.as_deref() {
|
||||
Some(path) => match super::config_file::load_config_file(path) {
|
||||
Ok(file) => Some(file),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
let integration_type = args.provider.slug();
|
||||
let enabled = true;
|
||||
let is_update = false;
|
||||
|
||||
let cli_mcp_servers =
|
||||
match super::mcp_config::build_mcp_servers_from_specs(&args.mcp_specs) {
|
||||
Ok(mcp_servers) => mcp_servers,
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut merged_config = super::config_file::merge_with_precedence(
|
||||
loaded_file.as_ref(),
|
||||
crate::ai::ambient_agents::AgentConfigSnapshot {
|
||||
name: None,
|
||||
environment_id: args.environment.environment.clone(),
|
||||
model_id: args.model.model.clone(),
|
||||
base_prompt: args.prompt.clone(),
|
||||
mcp_servers: cli_mcp_servers,
|
||||
profile_id: None,
|
||||
worker_host: args.worker_host.clone(),
|
||||
skill_spec: None,
|
||||
// TODO(QUALITY-295): Support computer use flag in integrations.
|
||||
computer_use_enabled: None,
|
||||
// TODO(REMOTE-1134): Support harness selection for integrations.
|
||||
harness: None,
|
||||
harness_auth_secrets: None,
|
||||
},
|
||||
);
|
||||
|
||||
// We must wait until after workspace metadata is refreshed to check available LLMs.
|
||||
let model_id = match merged_config
|
||||
.model_id
|
||||
.as_deref()
|
||||
.map(|model_id| super::common::validate_agent_mode_base_model_id(model_id, ctx))
|
||||
.transpose()
|
||||
{
|
||||
Ok(model_id) => model_id.map(|model_id| model_id.to_string()),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let base_prompt = merged_config.base_prompt.take();
|
||||
let worker_host = merged_config.worker_host.take();
|
||||
|
||||
let mcp_servers_json = match merged_config.mcp_servers.take() {
|
||||
Some(map) => match serde_json::to_string(&map) {
|
||||
Ok(json) => Some(json),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err.into())));
|
||||
return;
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
//If the user didn't explicitly request no environment, load environment from the config
|
||||
let mut environment_args = args.environment;
|
||||
if environment_args.environment.is_none() && !environment_args.no_environment {
|
||||
environment_args.environment = merged_config.environment_id.take();
|
||||
}
|
||||
|
||||
let environment_uid = match EnvironmentChoice::resolve_for_create(environment_args, ctx)
|
||||
{
|
||||
Ok(EnvironmentChoice::None) => {
|
||||
eprintln!("Creating integration without an environment.");
|
||||
None
|
||||
}
|
||||
Ok(EnvironmentChoice::Environment { id, .. }) => {
|
||||
eprintln!("Creating integration with environment {id}.");
|
||||
Some(id)
|
||||
}
|
||||
Err(ResolveConfigurationError::Canceled) => {
|
||||
eprintln!("Integration creation canceled.");
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, None);
|
||||
return;
|
||||
}
|
||||
Err(err) => {
|
||||
super::report_fatal_error(anyhow::anyhow!(err), ctx);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
runner.start_create_or_update_flow(
|
||||
ctx,
|
||||
integration_type,
|
||||
environment_uid,
|
||||
base_prompt,
|
||||
model_id,
|
||||
mcp_servers_json,
|
||||
None,
|
||||
worker_host,
|
||||
enabled,
|
||||
is_update,
|
||||
1,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn start_create_or_update_flow(
|
||||
&self,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
integration_type: String,
|
||||
environment_uid: Option<String>,
|
||||
base_prompt: Option<String>,
|
||||
model_id: Option<String>,
|
||||
mcp_servers_json: Option<String>,
|
||||
remove_mcp_server_names: Option<Vec<String>>,
|
||||
worker_host: Option<String>,
|
||||
enabled: bool,
|
||||
is_update: bool,
|
||||
attempt: u32,
|
||||
) {
|
||||
const MAX_CREATE_ATTEMPTS: u32 = 8;
|
||||
let action = if is_update { "update" } else { "creation" };
|
||||
|
||||
if attempt > MAX_CREATE_ATTEMPTS {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!(
|
||||
"Exceeded maximum number of integration creation attempts ({}). Retry.",
|
||||
MAX_CREATE_ATTEMPTS
|
||||
))),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let integrations_client = ServerApiProvider::as_ref(ctx).get_integrations_client();
|
||||
|
||||
let future_integration_type = integration_type.clone();
|
||||
let future_environment_uid = environment_uid.clone();
|
||||
let future_base_prompt = base_prompt.clone();
|
||||
let future_model_id = model_id.clone();
|
||||
let future_mcp_servers_json = mcp_servers_json.clone();
|
||||
let future_remove_mcp_server_names = remove_mcp_server_names.clone();
|
||||
let future_worker_host = worker_host.clone();
|
||||
let future_is_update = is_update;
|
||||
|
||||
let create_future = async move {
|
||||
integrations_client
|
||||
.create_or_update_simple_integration(
|
||||
future_integration_type,
|
||||
future_is_update,
|
||||
future_environment_uid,
|
||||
future_base_prompt,
|
||||
future_model_id,
|
||||
future_mcp_servers_json,
|
||||
future_remove_mcp_server_names,
|
||||
future_worker_host,
|
||||
enabled,
|
||||
)
|
||||
.await
|
||||
};
|
||||
|
||||
ctx.spawn(
|
||||
create_future,
|
||||
move |_runner, result: anyhow::Result<CreateSimpleIntegrationOutput>, ctx| {
|
||||
match result {
|
||||
Ok(output) => {
|
||||
println!("{}", output.message);
|
||||
|
||||
let auth_url = output.auth_url;
|
||||
let tx_id = output.tx_id;
|
||||
|
||||
match (auth_url, tx_id) {
|
||||
(Some(auth_url), Some(tx_id)) => {
|
||||
// We have another auth step: open URL and poll txId.
|
||||
println!("Authorize the provider here: {auth_url}\n");
|
||||
ctx.open_url(&auth_url);
|
||||
|
||||
let integrations_client = ServerApiProvider::as_ref(ctx)
|
||||
.get_integrations_client();
|
||||
let tx_id = tx_id.into_inner();
|
||||
|
||||
let poll_future =
|
||||
poll_oauth_until_terminal(integrations_client, tx_id);
|
||||
|
||||
let next_integration_type = integration_type.clone();
|
||||
let next_environment_uid = environment_uid.clone();
|
||||
let next_base_prompt = base_prompt.clone();
|
||||
let next_model_id = model_id.clone();
|
||||
let next_mcp_servers_json = mcp_servers_json.clone();
|
||||
let next_remove_mcp_server_names = remove_mcp_server_names.clone();
|
||||
let next_worker_host = worker_host.clone();
|
||||
let next_enabled = enabled;
|
||||
let next_is_update = is_update;
|
||||
let next_attempt = attempt + 1;
|
||||
|
||||
ctx.spawn(
|
||||
poll_future,
|
||||
move |runner, poll_result, ctx| {
|
||||
match poll_result {
|
||||
Ok(OauthConnectTxStatus::Completed) => {
|
||||
// Inner loop done; try create or update again (outer loop).
|
||||
// This may happen multiple times if the user needs to authorize multiple services.
|
||||
runner.start_create_or_update_flow(
|
||||
ctx,
|
||||
next_integration_type,
|
||||
next_environment_uid,
|
||||
next_base_prompt,
|
||||
next_model_id,
|
||||
next_mcp_servers_json,
|
||||
next_remove_mcp_server_names,
|
||||
next_worker_host,
|
||||
next_enabled,
|
||||
next_is_update,
|
||||
next_attempt,
|
||||
);
|
||||
}
|
||||
Ok(OauthConnectTxStatus::Failed) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!("OAuth authorization failed."))),
|
||||
);
|
||||
}
|
||||
Ok(OauthConnectTxStatus::Expired) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!("OAuth authorization expired."))),
|
||||
);
|
||||
}
|
||||
Ok(OauthConnectTxStatus::Pending)
|
||||
| Ok(OauthConnectTxStatus::InProgress) => {
|
||||
// Should not be returned by poll_oauth_until_terminal.
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!("Unexpected non-terminal OAuth status returned"))),
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!("Error polling OAuth status: {err}"))),
|
||||
);
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
(Some(auth_url), None) => {
|
||||
println!("Authorize the provider here: {auth_url}\n");
|
||||
ctx.open_url(&auth_url);
|
||||
println!(
|
||||
"After authorizing, re-run the command to continue the integration {action} process.",
|
||||
);
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
None,
|
||||
);
|
||||
}
|
||||
(None, Some(_)) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!("Server did not return an authURL for the integration creation process."))),
|
||||
);
|
||||
}
|
||||
(None, None) => {
|
||||
// No more auth steps; finalize.
|
||||
if output.success {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
None,
|
||||
);
|
||||
} else {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(anyhow::anyhow!("Integration creation reported failure: {}", output.message))),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(err)),
|
||||
);
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn update(&self, args: UpdateIntegrationArgs, ctx: &mut ModelContext<Self>) {
|
||||
let refresh_future = super::common::refresh_workspace_metadata(ctx);
|
||||
let warp_drive_sync_future = super::common::refresh_warp_drive(ctx);
|
||||
let setup_future = future::try_join(refresh_future, warp_drive_sync_future);
|
||||
|
||||
ctx.spawn(setup_future, move |runner, setup_result, ctx| {
|
||||
if let Err(err) = setup_result {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
|
||||
let loaded_file = match args.config_file.file.as_deref() {
|
||||
Some(path) => match super::config_file::load_config_file(path) {
|
||||
Ok(file) => Some(file),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
let remove_mcp = args.remove_mcp.clone();
|
||||
|
||||
let integration_type = args.provider.slug();
|
||||
let enabled = true;
|
||||
let is_update = true;
|
||||
|
||||
let cli_mcp_servers =
|
||||
match super::mcp_config::build_mcp_servers_from_specs(&args.mcp_specs) {
|
||||
Ok(mcp_servers) => mcp_servers,
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut merged_config = super::config_file::merge_with_precedence(
|
||||
loaded_file.as_ref(),
|
||||
crate::ai::ambient_agents::AgentConfigSnapshot {
|
||||
name: None,
|
||||
environment_id: args.environment.environment.clone(),
|
||||
model_id: args.model.model.clone(),
|
||||
base_prompt: args.prompt.clone(),
|
||||
mcp_servers: cli_mcp_servers,
|
||||
profile_id: None,
|
||||
worker_host: args.worker_host.clone(),
|
||||
skill_spec: None,
|
||||
// TODO(QUALITY-295): Support computer use flag in integrations.
|
||||
computer_use_enabled: None,
|
||||
// TODO(REMOTE-1134): Support harness selection for integrations.
|
||||
harness: None,
|
||||
harness_auth_secrets: None,
|
||||
},
|
||||
);
|
||||
|
||||
// We must wait until after workspace metadata is refreshed to check available LLMs.
|
||||
let model_id = match merged_config
|
||||
.model_id
|
||||
.as_deref()
|
||||
.map(|model_id| super::common::validate_agent_mode_base_model_id(model_id, ctx))
|
||||
.transpose()
|
||||
{
|
||||
Ok(model_id) => model_id.map(|model_id| model_id.to_string()),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(TerminationMode::ForceTerminate, Some(Err(err)));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let base_prompt = merged_config.base_prompt.take();
|
||||
let worker_host = merged_config.worker_host.take();
|
||||
|
||||
// MCP update semantics are patch-only:
|
||||
// - `mcp_servers_json` adds/overwrites MCP servers.
|
||||
// - `remove_mcp_server_names` removes MCP servers.
|
||||
// If both are present, removals win by filtering removed names out of the JSON payload.
|
||||
let mcp_servers_json = match merged_config.mcp_servers.take() {
|
||||
Some(mut map) => {
|
||||
for name in &remove_mcp {
|
||||
map.remove(name);
|
||||
}
|
||||
|
||||
if map.is_empty() {
|
||||
None
|
||||
} else {
|
||||
match serde_json::to_string(&map) {
|
||||
Ok(json) => Some(json),
|
||||
Err(err) => {
|
||||
ctx.terminate_app(
|
||||
TerminationMode::ForceTerminate,
|
||||
Some(Err(err.into())),
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
||||
let remove_mcp_server_names = if args.remove_mcp.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(args.remove_mcp)
|
||||
};
|
||||
|
||||
if args.environment.remove_environment {
|
||||
// Explicitly requested to update without an environment.
|
||||
runner.start_create_or_update_flow(
|
||||
ctx,
|
||||
integration_type,
|
||||
Some(String::new()),
|
||||
base_prompt,
|
||||
model_id,
|
||||
mcp_servers_json,
|
||||
remove_mcp_server_names,
|
||||
worker_host,
|
||||
enabled,
|
||||
is_update,
|
||||
1,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let environment_uid = merged_config.environment_id.take();
|
||||
|
||||
runner.start_create_or_update_flow(
|
||||
ctx,
|
||||
integration_type,
|
||||
environment_uid,
|
||||
base_prompt,
|
||||
model_id,
|
||||
mcp_servers_json,
|
||||
remove_mcp_server_names,
|
||||
worker_host,
|
||||
enabled,
|
||||
is_update,
|
||||
1,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl warpui::Entity for IntegrationCommandRunner {
|
||||
type Event = ();
|
||||
}
|
||||
impl SingletonEntity for IntegrationCommandRunner {}
|
||||
@@ -0,0 +1,355 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use comfy_table::{presets::UTF8_FULL, Cell, Table};
|
||||
use serde::Serialize;
|
||||
use serde_json::{Map, Value};
|
||||
use warp_cli::agent::OutputFormat;
|
||||
|
||||
use crate::ai::agent_sdk::output::{self, TableFormat};
|
||||
use crate::util::time_format::format_approx_duration_from_now_utc;
|
||||
use warp_graphql::queries::get_simple_integrations::{
|
||||
ListedSimpleIntegrationConfig, SimpleIntegration, SimpleIntegrationConnectionStatus,
|
||||
SimpleIntegrationsOutput,
|
||||
};
|
||||
|
||||
const MAX_LINE_WIDTH: usize = 90;
|
||||
|
||||
/// Print simple integrations.
|
||||
pub fn print_integrations(graphql_output: &SimpleIntegrationsOutput, output_format: OutputFormat) {
|
||||
if let Some(message) = &graphql_output.message {
|
||||
eprintln!("{message}");
|
||||
return;
|
||||
}
|
||||
|
||||
let integrations = &graphql_output.integrations;
|
||||
|
||||
if integrations.is_empty() {
|
||||
println!("No integrations found.");
|
||||
return;
|
||||
}
|
||||
|
||||
match output_format {
|
||||
OutputFormat::Json | OutputFormat::Ndjson => {
|
||||
// Convert to serializable format and use common output utilities
|
||||
let integration_infos: Vec<IntegrationInfo> = integrations
|
||||
.iter()
|
||||
.map(IntegrationInfo::from_graphql)
|
||||
.collect();
|
||||
output::print_list(integration_infos, output_format);
|
||||
}
|
||||
OutputFormat::Pretty | OutputFormat::Text => {
|
||||
// Use the existing card-style layout for pretty/text output
|
||||
if integrations.len() == 1 {
|
||||
println!("\nIntegration:");
|
||||
} else {
|
||||
println!("\nIntegrations:");
|
||||
}
|
||||
|
||||
for integration in integrations {
|
||||
print_integration_card(integration);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn render_labeled_wrapped_lines(label: &str, lines: &[String], width: usize) -> String {
|
||||
let indent = " ".repeat(label.len() + 2); // align under "{label}: "
|
||||
let mut out = String::new();
|
||||
|
||||
for (idx, line) in lines.iter().enumerate() {
|
||||
let wrapped = crate::ai::agent_sdk::text_layout::word_wrap(line, width);
|
||||
for (widx, wline) in wrapped.iter().enumerate() {
|
||||
if !out.is_empty() {
|
||||
out.push('\n');
|
||||
}
|
||||
if idx == 0 && widx == 0 {
|
||||
out.push_str(&format!("{label}: {wline}"));
|
||||
} else {
|
||||
out.push_str(&indent);
|
||||
out.push_str(wline);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
fn format_mcp_server_display(name: &str, config: &Value) -> String {
|
||||
let Some(obj) = config.as_object() else {
|
||||
return name.to_string();
|
||||
};
|
||||
|
||||
if let Some(url) = obj
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
return format!("{name}: {url}");
|
||||
}
|
||||
|
||||
if let Some(command) = obj
|
||||
.get("command")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
let args = obj
|
||||
.get("args")
|
||||
.and_then(Value::as_array)
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::to_string)
|
||||
.collect::<Vec<String>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
if args.is_empty() {
|
||||
return format!("{name}: {command}");
|
||||
}
|
||||
|
||||
return format!("{name}: {command} {}", args.join(" "));
|
||||
}
|
||||
|
||||
if let Some(warp_id) = obj
|
||||
.get("warp_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
return format!("{name}: warp_id={warp_id}");
|
||||
}
|
||||
|
||||
name.to_string()
|
||||
}
|
||||
|
||||
fn mcp_server_display_lines(config: &ListedSimpleIntegrationConfig) -> Vec<String> {
|
||||
let json = config.mcp_servers_json.trim();
|
||||
if json.is_empty() || json == "{}" {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let Ok(map) = serde_json::from_str::<Map<String, Value>>(json) else {
|
||||
return Vec::new();
|
||||
};
|
||||
|
||||
let mut entries: Vec<(String, Value)> = map.into_iter().collect();
|
||||
entries.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(name, cfg)| format_mcp_server_display(&name, &cfg))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn print_integration_card(integration: &SimpleIntegration) {
|
||||
let mut table = Table::new();
|
||||
table.load_preset(UTF8_FULL);
|
||||
|
||||
// Row 1: provider name (title-cased slug) and description, no label
|
||||
let provider_name =
|
||||
crate::ai::agent_sdk::text_layout::title_case_identifier(&integration.provider_slug);
|
||||
let title_row = crate::ai::agent_sdk::text_layout::render_labeled_wrapped_field(
|
||||
&provider_name,
|
||||
&integration.description,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![title_row]);
|
||||
|
||||
// Row 2: Status: <emoji> Status description
|
||||
let emoji = status_emoji(integration.connection_status);
|
||||
let explanation = status_explanation(integration.connection_status);
|
||||
let status_text = format!("{emoji} {explanation}");
|
||||
let status_row = crate::ai::agent_sdk::text_layout::render_labeled_wrapped_field(
|
||||
"Status",
|
||||
&status_text,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![status_row]);
|
||||
|
||||
// Environment row.
|
||||
let env_value = match &integration.integration_config {
|
||||
Some(ListedSimpleIntegrationConfig {
|
||||
environment_uid, ..
|
||||
}) if !environment_uid.is_empty() => environment_uid.clone(),
|
||||
_ => "(none)".to_string(),
|
||||
};
|
||||
let env_row = crate::ai::agent_sdk::text_layout::render_labeled_wrapped_field(
|
||||
"Environment",
|
||||
&env_value,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![env_row]);
|
||||
|
||||
// Model row (only if present).
|
||||
if let Some(ListedSimpleIntegrationConfig { model_id, .. }) = &integration.integration_config {
|
||||
if !model_id.is_empty() {
|
||||
let model_row = crate::ai::agent_sdk::text_layout::render_labeled_wrapped_field(
|
||||
"Model",
|
||||
model_id,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![model_row]);
|
||||
}
|
||||
}
|
||||
|
||||
// Base prompt row (only if present).
|
||||
if let Some(ListedSimpleIntegrationConfig { base_prompt, .. }) = &integration.integration_config
|
||||
{
|
||||
if !base_prompt.is_empty() {
|
||||
let base_prompt_row = crate::ai::agent_sdk::text_layout::render_labeled_wrapped_field(
|
||||
"Base prompt",
|
||||
base_prompt,
|
||||
MAX_LINE_WIDTH,
|
||||
);
|
||||
table.add_row(vec![base_prompt_row]);
|
||||
}
|
||||
}
|
||||
|
||||
// MCP servers row (only if present).
|
||||
if let Some(config) = &integration.integration_config {
|
||||
let lines = mcp_server_display_lines(config);
|
||||
if !lines.is_empty() {
|
||||
let row = render_labeled_wrapped_lines("MCP servers", &lines, MAX_LINE_WIDTH);
|
||||
table.add_row(vec![row]);
|
||||
}
|
||||
}
|
||||
|
||||
// Timestamps: keep created/updated in a single row, no label.
|
||||
let mut created_updated = String::new();
|
||||
if let Some(created) = integration.created_at {
|
||||
let dt = created.utc();
|
||||
let formatted = format_approx_duration_from_now_utc(dt);
|
||||
created_updated.push_str(&format!("Created: {formatted}"));
|
||||
}
|
||||
if let Some(updated) = integration.updated_at {
|
||||
let dt = updated.utc();
|
||||
let formatted = format_approx_duration_from_now_utc(dt);
|
||||
if !created_updated.is_empty() {
|
||||
created_updated.push_str(" | ");
|
||||
}
|
||||
created_updated.push_str(&format!("Updated: {formatted}"));
|
||||
}
|
||||
if !created_updated.is_empty() {
|
||||
let wrapped =
|
||||
crate::ai::agent_sdk::text_layout::word_wrap(&created_updated, MAX_LINE_WIDTH);
|
||||
let ts_cell = wrapped.join("\n");
|
||||
table.add_row(vec![ts_cell]);
|
||||
}
|
||||
|
||||
println!("{table}");
|
||||
}
|
||||
|
||||
fn status_emoji(status: SimpleIntegrationConnectionStatus) -> &'static str {
|
||||
match status {
|
||||
SimpleIntegrationConnectionStatus::NotConnected => "❌",
|
||||
// TODO(bens): these warning emojis render weirdly, maybe switch?
|
||||
SimpleIntegrationConnectionStatus::ConnectionError => "⚠️",
|
||||
SimpleIntegrationConnectionStatus::IntegrationNotConfigured => "⚠️",
|
||||
SimpleIntegrationConnectionStatus::NotEnabled => "⚠️",
|
||||
SimpleIntegrationConnectionStatus::Active => "✅",
|
||||
}
|
||||
}
|
||||
|
||||
fn status_explanation(status: SimpleIntegrationConnectionStatus) -> &'static str {
|
||||
match status {
|
||||
SimpleIntegrationConnectionStatus::NotConnected => "This integration is not connected.",
|
||||
SimpleIntegrationConnectionStatus::ConnectionError => {
|
||||
"This provider is connected but there is an error."
|
||||
}
|
||||
SimpleIntegrationConnectionStatus::IntegrationNotConfigured => {
|
||||
"Connection is active, but the agent integration has not been configured yet."
|
||||
}
|
||||
SimpleIntegrationConnectionStatus::NotEnabled => {
|
||||
"Integration is configured but currently disabled."
|
||||
}
|
||||
SimpleIntegrationConnectionStatus::Active => "Integration is connected and enabled.",
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializable integration info for output.
|
||||
#[derive(Serialize)]
|
||||
struct IntegrationInfo {
|
||||
provider: String,
|
||||
description: String,
|
||||
status: String,
|
||||
environment_uid: Option<String>,
|
||||
base_prompt: Option<String>,
|
||||
created_at: Option<DateTime<Utc>>,
|
||||
updated_at: Option<DateTime<Utc>>,
|
||||
#[serde(skip_serializing)]
|
||||
created_at_formatted: String,
|
||||
#[serde(skip_serializing)]
|
||||
updated_at_formatted: String,
|
||||
}
|
||||
|
||||
impl IntegrationInfo {
|
||||
fn from_graphql(integration: &SimpleIntegration) -> Self {
|
||||
let provider =
|
||||
crate::ai::agent_sdk::text_layout::title_case_identifier(&integration.provider_slug);
|
||||
let status = status_explanation(integration.connection_status).to_string();
|
||||
|
||||
let environment_uid = integration.integration_config.as_ref().and_then(|config| {
|
||||
if config.environment_uid.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(config.environment_uid.clone())
|
||||
}
|
||||
});
|
||||
|
||||
let base_prompt = integration.integration_config.as_ref().and_then(|config| {
|
||||
if config.base_prompt.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(config.base_prompt.clone())
|
||||
}
|
||||
});
|
||||
|
||||
let created_at = integration.created_at.map(|t| t.utc());
|
||||
let updated_at = integration.updated_at.map(|t| t.utc());
|
||||
|
||||
let created_at_formatted = created_at
|
||||
.map(format_approx_duration_from_now_utc)
|
||||
.unwrap_or_else(|| "Unknown".to_string());
|
||||
|
||||
let updated_at_formatted = updated_at
|
||||
.map(format_approx_duration_from_now_utc)
|
||||
.unwrap_or_else(|| "Unknown".to_string());
|
||||
|
||||
Self {
|
||||
provider,
|
||||
description: integration.description.clone(),
|
||||
status,
|
||||
environment_uid,
|
||||
base_prompt,
|
||||
created_at,
|
||||
updated_at,
|
||||
created_at_formatted,
|
||||
updated_at_formatted,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TableFormat for IntegrationInfo {
|
||||
fn header() -> Vec<Cell> {
|
||||
vec![
|
||||
Cell::new("Provider"),
|
||||
Cell::new("Description"),
|
||||
Cell::new("Status"),
|
||||
Cell::new("Environment"),
|
||||
Cell::new("Created"),
|
||||
Cell::new("Updated"),
|
||||
]
|
||||
}
|
||||
|
||||
fn row(&self) -> Vec<Cell> {
|
||||
vec![
|
||||
Cell::new(&self.provider),
|
||||
Cell::new(&self.description),
|
||||
Cell::new(&self.status),
|
||||
Cell::new(self.environment_uid.as_deref().unwrap_or("(none)")),
|
||||
Cell::new(&self.created_at_formatted),
|
||||
Cell::new(&self.updated_at_formatted),
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
use comfy_table::Cell;
|
||||
use serde::Serialize;
|
||||
use warp_cli::{mcp::MCPCommand, GlobalOptions};
|
||||
use warpui::{AppContext, ModelContext, SingletonEntity};
|
||||
|
||||
use crate::ai::agent_sdk::output::{self, TableFormat};
|
||||
use crate::ai::mcp::TemplatableMCPServerManager;
|
||||
use crate::server::cloud_objects::update_manager::UpdateManager;
|
||||
|
||||
/// Handle MCP-related CLI commands.
|
||||
pub fn run(
|
||||
ctx: &mut AppContext,
|
||||
global_options: GlobalOptions,
|
||||
command: MCPCommand,
|
||||
) -> anyhow::Result<()> {
|
||||
let runner = ctx.add_singleton_model(|_ctx| MCPCommandRunner);
|
||||
match command {
|
||||
MCPCommand::List => {
|
||||
runner.update(ctx, |runner, ctx| runner.list(global_options, ctx));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Singleton model for running async work as part of MCP CLI commands.
|
||||
struct MCPCommandRunner;
|
||||
|
||||
impl MCPCommandRunner {
|
||||
fn list(&self, global_options: GlobalOptions, ctx: &mut ModelContext<Self>) {
|
||||
let initial_sync = UpdateManager::as_ref(ctx).initial_load_complete();
|
||||
|
||||
ctx.spawn(initial_sync, move |_, _, ctx| {
|
||||
let mut servers = TemplatableMCPServerManager::get_all_runnable_mcp_servers(ctx);
|
||||
servers.sort_by_key(|(uuid, _)| *uuid);
|
||||
|
||||
output::print_list(
|
||||
servers
|
||||
.into_iter()
|
||||
.map(|(uuid, name)| MCPServerInfo { uuid, name }),
|
||||
global_options.output_format,
|
||||
);
|
||||
|
||||
ctx.terminate_app(warpui::platform::TerminationMode::ForceTerminate, None);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl warpui::Entity for MCPCommandRunner {
|
||||
type Event = ();
|
||||
}
|
||||
impl SingletonEntity for MCPCommandRunner {}
|
||||
|
||||
/// MCP server information that's shown in the `list` command.
|
||||
#[derive(Serialize)]
|
||||
struct MCPServerInfo {
|
||||
uuid: uuid::Uuid,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl TableFormat for MCPServerInfo {
|
||||
fn header() -> Vec<Cell> {
|
||||
vec![Cell::new("UUID"), Cell::new("Name")]
|
||||
}
|
||||
|
||||
fn row(&self) -> Vec<Cell> {
|
||||
vec![Cell::new(self.uuid), Cell::new(&self.name)]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
use anyhow::Context as _;
|
||||
use serde_json::{Map, Value};
|
||||
use warp_cli::mcp::MCPSpec;
|
||||
|
||||
use crate::ai::mcp::TemplatableMCPServer;
|
||||
|
||||
/// Build the `mcp_servers` map to send to the public ambient-agent API.
|
||||
///
|
||||
/// Returns the unwrapped server map (`{ <server_name>: <server_config>, ... }`).
|
||||
/// If user input includes wrapper shapes like `{ "mcpServers": { ... } }`, we unpack them.
|
||||
///
|
||||
/// Notes:
|
||||
/// - UUID specs are coerced into `{"<uuid>": {"warp_id": "<uuid>"}}`.
|
||||
/// - We do light validation to catch obvious config errors before sending the request.
|
||||
pub(super) fn build_mcp_servers_from_specs(
|
||||
specs: &[MCPSpec],
|
||||
) -> anyhow::Result<Option<Map<String, Value>>> {
|
||||
if specs.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut merged = Map::new();
|
||||
|
||||
for spec in specs {
|
||||
match spec {
|
||||
MCPSpec::Uuid(uuid) => {
|
||||
// TODO: Look up and use the real MCP server name from MCP managers instead of using the UUID.
|
||||
let name = uuid.to_string();
|
||||
insert_unique(
|
||||
&mut merged,
|
||||
name.clone(),
|
||||
Value::Object({
|
||||
let mut obj = Map::new();
|
||||
obj.insert("warp_id".to_string(), Value::String(name));
|
||||
obj
|
||||
}),
|
||||
)?;
|
||||
}
|
||||
MCPSpec::Json(json_str) => {
|
||||
let json_str = normalize_mcp_json_for_single_server(json_str)?;
|
||||
let value = parse_json_with_optional_braces(&json_str)?;
|
||||
|
||||
let server_map = TemplatableMCPServer::find_template_map(value)
|
||||
.context("Failed to parse MCP server map")?;
|
||||
|
||||
for (name, config) in server_map {
|
||||
insert_unique(&mut merged, name, config)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
validate_mcp_servers(&merged)?;
|
||||
|
||||
if merged.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(merged))
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_unique(map: &mut Map<String, Value>, name: String, config: Value) -> anyhow::Result<()> {
|
||||
if map.contains_key(&name) {
|
||||
anyhow::bail!("Duplicate MCP server name '{name}' specified multiple times");
|
||||
}
|
||||
|
||||
map.insert(name, config);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_json_with_optional_braces(input: &str) -> anyhow::Result<Value> {
|
||||
// Some docs don't show curly braces around the json object, so add them if necessary.
|
||||
let json = input.trim();
|
||||
let json = if json.starts_with('{') {
|
||||
json.to_owned()
|
||||
} else {
|
||||
format!("{{{json}}}")
|
||||
};
|
||||
|
||||
serde_json::from_str(&json).with_context(|| "Invalid MCP JSON".to_string())
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn normalize_mcp_json_for_single_server(input: &str) -> anyhow::Result<String> {
|
||||
crate::ai::mcp::parsing::normalize_mcp_json(input)
|
||||
.map_err(|e| anyhow::anyhow!(e))
|
||||
.context("Failed to normalize MCP JSON")
|
||||
}
|
||||
|
||||
// The CLI + ambient-agent API isn’t used in WASM builds, but this module still needs to compile.
|
||||
// Implement the same normalization behavior (single-server shorthand wrap) locally.
|
||||
#[cfg(target_family = "wasm")]
|
||||
fn normalize_mcp_json_for_single_server(input: &str) -> anyhow::Result<String> {
|
||||
let json = input.trim();
|
||||
let json_for_parsing = if json.starts_with('{') {
|
||||
json.to_owned()
|
||||
} else {
|
||||
format!("{{{json}}}")
|
||||
};
|
||||
|
||||
let value: Value =
|
||||
serde_json::from_str(&json_for_parsing).with_context(|| "Invalid MCP JSON".to_string())?;
|
||||
|
||||
let is_single_server = value.get("command").is_some() || value.get("url").is_some();
|
||||
if is_single_server {
|
||||
let name = uuid::Uuid::new_v4().to_string();
|
||||
let mut map = Map::new();
|
||||
map.insert(name, value);
|
||||
Ok(Value::Object(map).to_string())
|
||||
} else {
|
||||
Ok(input.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn validate_mcp_servers(mcp_servers: &Map<String, Value>) -> anyhow::Result<()> {
|
||||
for (name, config) in mcp_servers {
|
||||
validate_server_config(name, config)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_server_config(server_name: &str, config: &Value) -> anyhow::Result<()> {
|
||||
let obj = config.as_object().ok_or_else(|| {
|
||||
anyhow::anyhow!("MCP server '{server_name}' config must be a JSON object")
|
||||
})?;
|
||||
|
||||
let has_warp_id = obj.contains_key("warp_id");
|
||||
let has_command = obj.contains_key("command");
|
||||
let has_url = obj.contains_key("url");
|
||||
|
||||
let kind_count = usize::from(has_warp_id) + usize::from(has_command) + usize::from(has_url);
|
||||
if kind_count != 1 {
|
||||
anyhow::bail!(
|
||||
"MCP server '{server_name}' must have exactly one of: 'warp_id', 'command', or 'url'"
|
||||
);
|
||||
}
|
||||
|
||||
if has_warp_id {
|
||||
let warp_id = obj.get("warp_id").and_then(Value::as_str).ok_or_else(|| {
|
||||
anyhow::anyhow!("MCP server '{server_name}' field 'warp_id' must be a string")
|
||||
})?;
|
||||
|
||||
uuid::Uuid::parse_str(warp_id).with_context(|| {
|
||||
format!("MCP server '{server_name}' field 'warp_id' must be a UUID")
|
||||
})?;
|
||||
}
|
||||
|
||||
if has_command {
|
||||
let command = obj.get("command").and_then(Value::as_str).ok_or_else(|| {
|
||||
anyhow::anyhow!("MCP server '{server_name}' field 'command' must be a string")
|
||||
})?;
|
||||
|
||||
if command.is_empty() {
|
||||
anyhow::bail!("MCP server '{server_name}' field 'command' must be non-empty");
|
||||
}
|
||||
|
||||
if let Some(args) = obj.get("args") {
|
||||
let args = args.as_array().ok_or_else(|| {
|
||||
anyhow::anyhow!("MCP server '{server_name}' field 'args' must be an array")
|
||||
})?;
|
||||
|
||||
for (idx, arg) in args.iter().enumerate() {
|
||||
if !arg.is_string() {
|
||||
anyhow::bail!(
|
||||
"MCP server '{server_name}' field 'args[{idx}]' must be a string"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if has_url {
|
||||
let url = obj.get("url").and_then(Value::as_str).ok_or_else(|| {
|
||||
anyhow::anyhow!("MCP server '{server_name}' field 'url' must be a string")
|
||||
})?;
|
||||
|
||||
if url.is_empty() {
|
||||
anyhow::bail!("MCP server '{server_name}' field 'url' must be non-empty");
|
||||
}
|
||||
}
|
||||
|
||||
validate_string_map_field(obj, server_name, "env")?;
|
||||
validate_string_map_field(obj, server_name, "headers")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_string_map_field(
|
||||
obj: &Map<String, Value>,
|
||||
server_name: &str,
|
||||
field: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let Some(value) = obj.get(field) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let map = value.as_object().ok_or_else(|| {
|
||||
anyhow::anyhow!("MCP server '{server_name}' field '{field}' must be an object")
|
||||
})?;
|
||||
|
||||
for (key, value) in map {
|
||||
if !value.is_string() {
|
||||
anyhow::bail!("MCP server '{server_name}' field '{field}.{key}' must be a string");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "mcp_config_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,300 @@
|
||||
use serde_json::{json, Map, Value};
|
||||
use warp_cli::mcp::MCPSpec;
|
||||
|
||||
use super::build_mcp_servers_from_specs;
|
||||
|
||||
fn build(specs: Vec<MCPSpec>) -> Map<String, Value> {
|
||||
build_mcp_servers_from_specs(&specs)
|
||||
.expect("builder should not error")
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_specs_returns_none() {
|
||||
assert!(build_mcp_servers_from_specs(&[]).unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uuid_spec_is_coerced_to_warp_id() {
|
||||
let uuid = uuid::Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap();
|
||||
let servers = build(vec![MCPSpec::Uuid(uuid)]);
|
||||
|
||||
let entry = servers.get(&uuid.to_string()).unwrap();
|
||||
assert_eq!(
|
||||
entry["warp_id"].as_str(),
|
||||
Some("550e8400-e29b-41d4-a716-446655440000")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrapper_mcp_servers_is_unpacked() {
|
||||
let spec = json!({
|
||||
"mcpServers": {
|
||||
"github": { "command": "npx", "args": ["-y", "server"] }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("github"));
|
||||
assert_eq!(servers["github"]["command"].as_str(), Some("npx"));
|
||||
assert!(servers.get("mcpServers").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrapper_mcp_servers_snake_case_is_unpacked() {
|
||||
let spec = json!({
|
||||
"mcp_servers": {
|
||||
"s": { "url": "https://example.com/mcp" }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("s"));
|
||||
assert_eq!(
|
||||
servers["s"]["url"].as_str(),
|
||||
Some("https://example.com/mcp")
|
||||
);
|
||||
assert!(servers.get("mcp_servers").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrapper_servers_is_unpacked() {
|
||||
let spec = json!({
|
||||
"servers": {
|
||||
"s": { "command": "python", "args": ["mcp.py"] }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("s"));
|
||||
assert_eq!(servers["s"]["command"].as_str(), Some("python"));
|
||||
assert!(servers.get("servers").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrapper_mcp_dot_servers_is_unpacked() {
|
||||
let spec = json!({
|
||||
"mcp": {
|
||||
"servers": {
|
||||
"s": { "url": "https://example.com/mcp" }
|
||||
}
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("s"));
|
||||
assert_eq!(
|
||||
servers["s"]["url"].as_str(),
|
||||
Some("https://example.com/mcp")
|
||||
);
|
||||
assert!(servers.get("mcp").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plain_map_is_accepted() {
|
||||
let spec = json!({
|
||||
"github": { "command": "npx", "args": [] },
|
||||
"remote": { "url": "https://example.com/mcp" }
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("github"));
|
||||
assert!(servers.contains_key("remote"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_outer_braces_is_accepted() {
|
||||
// Emulate copying docs that omit the top-level `{}`.
|
||||
let full = json!({
|
||||
"mcpServers": {
|
||||
"s": { "command": "npx", "args": [] }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let inner = &full[1..full.len() - 1];
|
||||
let spec = format!(" {inner} ");
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("s"));
|
||||
assert_eq!(servers["s"]["command"].as_str(), Some("npx"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_server_shorthand_command_is_wrapped() {
|
||||
let spec = json!({ "command": "npx", "args": [] }).to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert_eq!(servers.len(), 1);
|
||||
let (_name, config) = servers.iter().next().unwrap();
|
||||
assert_eq!(config["command"].as_str(), Some("npx"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn command_without_args_is_accepted() {
|
||||
// args should be optional for command-based MCP servers
|
||||
let spec = json!({
|
||||
"mcpServers": {
|
||||
"s": { "command": "uvx" }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert!(servers.contains_key("s"));
|
||||
assert_eq!(servers["s"]["command"].as_str(), Some("uvx"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_server_shorthand_url_is_wrapped() {
|
||||
let spec = json!({ "url": "https://example.com/mcp" }).to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
|
||||
assert_eq!(servers.len(), 1);
|
||||
let (_name, config) = servers.iter().next().unwrap();
|
||||
assert_eq!(config["url"].as_str(), Some("https://example.com/mcp"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_multiple_specs_and_duplicate_name_errors() {
|
||||
let s1 = json!({ "mcpServers": { "a": { "command": "npx", "args": [] } } }).to_string();
|
||||
let s2 = json!({ "mcpServers": { "b": { "url": "https://example.com/mcp" } } }).to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(s1.clone()), MCPSpec::Json(s2)]);
|
||||
assert!(servers.contains_key("a"));
|
||||
assert!(servers.contains_key("b"));
|
||||
|
||||
let err =
|
||||
build_mcp_servers_from_specs(&[MCPSpec::Json(s1.clone()), MCPSpec::Json(s1)]).unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("Duplicate MCP server name 'a'"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_escaped_strings_in_env_values() {
|
||||
let spec = json!({
|
||||
"mcpServers": {
|
||||
"s": {
|
||||
"command": "npx",
|
||||
"args": [],
|
||||
"env": {
|
||||
"TOKEN": "a\"b\\c\n"
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let servers = build(vec![MCPSpec::Json(spec)]);
|
||||
let token = servers["s"]["env"]["TOKEN"].as_str().unwrap();
|
||||
|
||||
// `serde_json` will decode escapes.
|
||||
assert_eq!(token, "a\"b\\c\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validation_rejects_invalid_entries() {
|
||||
// Both command and url.
|
||||
let spec = json!({
|
||||
"mcpServers": {
|
||||
"bad": { "command": "npx", "url": "https://example.com" }
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err
|
||||
.to_string()
|
||||
.contains("must have exactly one of: 'warp_id', 'command', or 'url'"));
|
||||
|
||||
// warp_id must be a UUID string.
|
||||
let spec = json!({ "mcpServers": { "bad": { "warp_id": "not-a-uuid" } } }).to_string();
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err.to_string().contains("field 'warp_id' must be a UUID"));
|
||||
|
||||
// args must be array.
|
||||
let spec = json!({ "mcpServers": { "bad": { "command": "npx", "args": "nope" } } }).to_string();
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err.to_string().contains("field 'args' must be an array"));
|
||||
|
||||
// args entries must be strings.
|
||||
let spec = json!({ "mcpServers": { "bad": { "command": "npx", "args": [1] } } }).to_string();
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err.to_string().contains("args[0]"));
|
||||
|
||||
// env values must be strings.
|
||||
let spec = json!({
|
||||
"mcpServers": { "bad": { "command": "npx", "args": [], "env": { "X": 1 } } }
|
||||
})
|
||||
.to_string();
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err.to_string().contains("env.X"));
|
||||
|
||||
// headers values must be strings.
|
||||
let spec = json!({
|
||||
"mcpServers": { "bad": { "url": "https://example.com", "headers": { "X": 1 } } }
|
||||
})
|
||||
.to_string();
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err.to_string().contains("headers.X"));
|
||||
|
||||
// server config must be an object.
|
||||
let spec = json!({ "mcpServers": { "bad": 1 } }).to_string();
|
||||
let err = build_mcp_servers_from_specs(&[MCPSpec::Json(spec)]).unwrap_err();
|
||||
assert!(err.to_string().contains("config must be a JSON object"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serializes_mcp_servers_as_object_not_string() {
|
||||
use crate::ai::ambient_agents::AgentConfigSnapshot;
|
||||
use crate::server::server_api::ai::SpawnAgentRequest;
|
||||
|
||||
let uuid = uuid::Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap();
|
||||
let mcp_servers = build_mcp_servers_from_specs(&[MCPSpec::Uuid(uuid)])
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let request = SpawnAgentRequest {
|
||||
prompt: "hello".to_string(),
|
||||
config: Some(AgentConfigSnapshot {
|
||||
mcp_servers: Some(mcp_servers),
|
||||
..Default::default()
|
||||
}),
|
||||
title: None,
|
||||
team: None,
|
||||
skill: None,
|
||||
attachments: vec![],
|
||||
interactive: None,
|
||||
parent_run_id: None,
|
||||
runtime_skills: vec![],
|
||||
referenced_attachments: vec![],
|
||||
};
|
||||
|
||||
let value = serde_json::to_value(&request).unwrap();
|
||||
|
||||
let config = value.get("config").unwrap();
|
||||
let mcp_servers = config.get("mcp_servers").unwrap();
|
||||
|
||||
assert!(mcp_servers.is_object());
|
||||
assert!(mcp_servers.get("mcpServers").is_none());
|
||||
|
||||
let server = mcp_servers.get(uuid.to_string()).unwrap();
|
||||
assert_eq!(
|
||||
server.get("warp_id").unwrap(),
|
||||
&Value::String(uuid.to_string())
|
||||
);
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user