public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/active_agent_views_model.rs
+118 -31
View File
@@ -1,17 +1,21 @@
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 galaxyui::{
AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity, WeakModelHandle,
WindowId,
};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_conversations_model::{AgentConversationEntry, AgentConversationEntryId};
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::blocklist::agent_view::{AgentViewController, AgentViewControllerEvent};
use crate::ai::blocklist::orchestration_event_streamer::{
register_agent_event_consumer, unregister_agent_event_consumer,
};
use crate::ai::blocklist::BlocklistAIHistoryModel;
use crate::terminal::model::session::active_session::ActiveSession;
/// Contains the handles needed to track an active agent view.
struct ActiveAgentViewHandles {
controller: WeakModelHandle<AgentViewController>,
@@ -44,6 +48,17 @@ pub enum ConversationOrTaskId {
TaskId(AmbientAgentTaskId),
}
impl From<AgentConversationEntryId> for ConversationOrTaskId {
fn from(id: AgentConversationEntryId) -> Self {
match id {
AgentConversationEntryId::Conversation(conversation_id) => {
ConversationOrTaskId::ConversationId(conversation_id)
}
AgentConversationEntryId::AmbientRun(task_id) => ConversationOrTaskId::TaskId(task_id),
}
}
}
impl ConversationOrTaskId {
pub fn conversation_id(&self) -> Option<AIConversationId> {
match self {
@@ -94,6 +109,34 @@ impl ActiveAgentViewsModel {
}
}
fn update_focused_conversation_for_terminal(
&mut self,
terminal_view_id: EntityId,
conversation_id: Option<ConversationOrTaskId>,
) {
for focused_terminal_state in self.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 = conversation_id;
}
}
if let Some(last_focused_terminal_state) = &mut self.last_focused_terminal_state {
if last_focused_terminal_state.focused_terminal_id == terminal_view_id
&& !matches!(
last_focused_terminal_state.active_conversation_id,
Some(ConversationOrTaskId::TaskId(_))
)
{
last_focused_terminal_state.active_conversation_id = conversation_id;
}
}
}
/// 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(
@@ -122,7 +165,19 @@ impl ActiveAgentViewsModel {
},
);
ctx.subscribe_to_model(controller, move |model, event, ctx| match event {
// On pane re-attach the controller's `agent_view_state` is still
// `Active` while the unregister path has already torn down the
// streamer consumer. Re-register here; the `EnteredAgentView`
// subscription only fires on subsequent state transitions.
if let Some(conversation_id) = controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
{
register_agent_event_consumer(conversation_id, terminal_view_id, ctx);
}
ctx.subscribe_to_model(controller, move |model, _, event, ctx| match event {
AgentViewControllerEvent::EnteredAgentView {
conversation_id, ..
} => {
@@ -132,37 +187,31 @@ impl ActiveAgentViewsModel {
// 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);
}
}
model.update_focused_conversation_for_terminal(terminal_view_id, Some(conv_id));
// Bridge the controller's lifecycle into the streamer's
// per-conversation consumer registry.
register_agent_event_consumer(*conversation_id, terminal_view_id, ctx);
// Emit so subscribers can move this conversation to the Active section.
ctx.emit(ActiveAgentViewsEvent::TerminalViewFocused);
}
AgentViewControllerEvent::ExitedAgentView {
conversation_id, ..
conversation_id,
is_exit_before_new_entrance,
..
} => {
// Skip if this exit is part of an in-place switch — the follow-up
// entrance will register the new conversation's consumer.
if *is_exit_before_new_entrance {
return;
}
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;
}
}
model.update_focused_conversation_for_terminal(terminal_view_id, None);
unregister_agent_event_consumer(*conversation_id, terminal_view_id, ctx);
// Emit so subscribers can move this conversation to the Past section.
ctx.emit(ActiveAgentViewsEvent::ConversationClosed {
conversation_id: *conversation_id,
@@ -198,6 +247,9 @@ impl ActiveAgentViewsModel {
}
if let Some(conversation_id) = closed_conversation_id {
// The pane-close path bypasses exit_agent_view_internal, so
// unregister the streamer consumer here.
unregister_agent_event_consumer(conversation_id, terminal_pane_id, ctx);
ctx.emit(ActiveAgentViewsEvent::ConversationClosed { conversation_id });
}
}
@@ -252,6 +304,12 @@ impl ActiveAgentViewsModel {
.and_then(|state| state.active_conversation_id)
}
pub fn get_focused_terminal_view_id(&self, window_id: WindowId) -> Option<EntityId> {
self.focused_terminal_states
.get(&window_id)
.map(|state| state.focused_terminal_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
@@ -259,6 +317,14 @@ impl ActiveAgentViewsModel {
.map(|state| state.focused_terminal_id)
}
/// Get the most recent focused conversation or ambient task ID, persisted
/// across non-terminal focus changes.
pub fn get_last_focused_conversation(&self) -> Option<ConversationOrTaskId> {
self.last_focused_terminal_state
.as_ref()
.and_then(|state| state.active_conversation_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).
@@ -307,6 +373,8 @@ impl ActiveAgentViewsModel {
task_id: AmbientAgentTaskId,
ctx: &mut ModelContext<Self>,
) {
self.ambient_sessions
.retain(|view_id, id| *view_id == terminal_view_id || *id != task_id);
let existing = self.ambient_sessions.insert(terminal_view_id, task_id);
if existing != Some(task_id) {
self.last_opened_times
@@ -323,9 +391,11 @@ impl ActiveAgentViewsModel {
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 });
if !self.ambient_sessions.values().any(|id| *id == task_id) {
self.last_opened_times
.remove(&ConversationOrTaskId::TaskId(task_id));
ctx.emit(ActiveAgentViewsEvent::AmbientSessionClosed { task_id });
}
}
}
@@ -443,6 +513,23 @@ impl ActiveAgentViewsModel {
None
}
pub fn get_terminal_view_id_for_entry(
&self,
entry: &AgentConversationEntry,
ctx: &AppContext,
) -> Option<EntityId> {
if let Some(task_id) = entry.identity.ambient_agent_task_id {
if let Some(terminal_view_id) = self.get_terminal_view_id_for_ambient_task(task_id) {
return Some(terminal_view_id);
}
}
if let Some(conversation_id) = entry.identity.local_conversation_id {
return self.get_terminal_view_id_for_conversation(conversation_id, ctx);
}
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.
app/src/ai/active_agent_views_model_tests.rs
+118
View File
@@ -12,6 +12,46 @@ fn new_task_id() -> AmbientAgentTaskId {
AmbientAgentTaskId::from_str(&uuid::Uuid::new_v4().to_string()).unwrap()
}
#[test]
fn conversation_switch_updates_last_focused_terminal_state() {
App::test((), |mut app| async move {
let model = setup_model(&mut app);
let window = WindowId::new();
let terminal = EntityId::new();
let conversation_1 = AIConversationId::new();
let conversation_2 = AIConversationId::new();
model.update(&mut app, |model, ctx| {
model.handle_pane_focus_change(window, Some(terminal), None, ctx);
model.update_focused_conversation_for_terminal(
terminal,
Some(ConversationOrTaskId::ConversationId(conversation_1)),
);
});
model.read(&app, |model, _| {
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal));
assert_eq!(
model.get_last_focused_conversation(),
Some(ConversationOrTaskId::ConversationId(conversation_1))
);
});
model.update(&mut app, |model, _| {
model.update_focused_conversation_for_terminal(
terminal,
Some(ConversationOrTaskId::ConversationId(conversation_2)),
);
});
model.read(&app, |model, _| {
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal));
assert_eq!(
model.get_last_focused_conversation(),
Some(ConversationOrTaskId::ConversationId(conversation_2))
);
});
});
}
#[test]
fn per_window_focused_state_is_independent() {
App::test((), |mut app| async move {
@@ -88,6 +128,10 @@ fn last_focused_terminal_tracks_most_recent_globally() {
});
model.read(&app, |model, _| {
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal_a));
assert_eq!(
model.get_last_focused_conversation(),
Some(ConversationOrTaskId::TaskId(task_a))
);
});
model.update(&mut app, |model, ctx| {
@@ -95,6 +139,10 @@ fn last_focused_terminal_tracks_most_recent_globally() {
});
model.read(&app, |model, _| {
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal_b));
assert_eq!(
model.get_last_focused_conversation(),
Some(ConversationOrTaskId::TaskId(task_b))
);
});
// Clearing window B's focus should NOT clear last_focused (it persists).
@@ -103,6 +151,10 @@ fn last_focused_terminal_tracks_most_recent_globally() {
});
model.read(&app, |model, _| {
assert_eq!(model.get_last_focused_terminal_id(), Some(terminal_b));
assert_eq!(
model.get_last_focused_conversation(),
Some(ConversationOrTaskId::TaskId(task_b))
);
});
});
}
@@ -144,6 +196,72 @@ fn focus_change_without_task_id_has_no_conversation() {
});
}
#[test]
fn ambient_session_registration_replaces_stale_terminal_for_same_task() {
App::test((), |mut app| async move {
let model = setup_model(&mut app);
let terminal_a = EntityId::new();
let terminal_b = EntityId::new();
let task = new_task_id();
model.update(&mut app, |model, ctx| {
model.register_ambient_session(terminal_a, task, ctx);
model.register_ambient_session(terminal_b, task, ctx);
});
model.read(&app, |model, _| {
assert_eq!(
model.get_terminal_view_id_for_ambient_task(task),
Some(terminal_b)
);
assert_eq!(model.ambient_sessions.len(), 1);
});
});
}
#[test]
fn ambient_session_unregister_keeps_task_open_until_last_terminal_is_removed() {
App::test((), |mut app| async move {
let model = setup_model(&mut app);
let terminal_a = EntityId::new();
let terminal_b = EntityId::new();
let task = new_task_id();
model.update(&mut app, |model, _| {
model.ambient_sessions.insert(terminal_a, task);
model.ambient_sessions.insert(terminal_b, task);
model
.last_opened_times
.insert(ConversationOrTaskId::TaskId(task), Utc::now());
});
model.update(&mut app, |model, ctx| {
model.unregister_ambient_session(terminal_a, ctx);
});
model.read(&app, |model, _| {
assert_eq!(
model.get_terminal_view_id_for_ambient_task(task),
Some(terminal_b)
);
assert!(model
.last_opened_times
.contains_key(&ConversationOrTaskId::TaskId(task)));
});
model.update(&mut app, |model, ctx| {
model.unregister_ambient_session(terminal_b, ctx);
});
model.read(&app, |model, _| {
assert_eq!(model.get_terminal_view_id_for_ambient_task(task), None);
assert!(!model
.last_opened_times
.contains_key(&ConversationOrTaskId::TaskId(task)));
});
});
}
#[test]
fn remove_focused_state_for_window_cleans_up() {
App::test((), |mut app| async move {
app/src/ai/agent/api.rs
+105 -25
View File
@@ -3,40 +3,38 @@ mod convert_from;
mod convert_to;
mod r#impl;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
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 mcp::TemplatableMCPServerInfo;
pub use r#impl::generate_multi_agent_output;
use serde::Serialize;
use galaxy_core::channel::ChannelState;
use galaxy_core::execution_mode::AppExecutionMode;
use galaxy_core::features::FeatureFlag;
use serde::Serialize;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::{
ai::{blocklist::SessionContext, llms::LLMId},
server::server_api::AIApiError,
};
use galaxy_core::user_preferences::GetUserPreferences;
use galaxyui::{AppContext, EntityId, SingletonEntity as _};
use super::{AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions};
use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput};
use crate::ai::mcp::templatable_manager::TemplatableMCPServerInfo;
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput, SessionContext};
use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel;
use crate::ai::execution_profiles::AIExecutionProfileAppExt;
use crate::ai::llms::{LLMId, LLMPreferences};
use crate::ai::mcp::TemplatableMCPServerManager;
use crate::server::server_api::AIApiError;
use crate::settings::AISettings;
use crate::terminal::safe_mode_settings::get_secret_obfuscation_mode;
use crate::workspaces::user_workspaces::UserWorkspaces;
use galaxy_core::user_preferences::GetUserPreferences;
use galaxyui::{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.
@@ -109,13 +107,21 @@ pub struct RequestParams {
pub computer_use_model: LLMId,
pub is_memory_enabled: bool,
pub warp_drive_context_enabled: bool,
pub context_window_limit: Option<u32>,
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,
/// User-provided custom model providers (BYOK endpoints).
pub custom_model_providers:
Option<warp_multi_agent_api::request::settings::CustomModelProviders>,
/// User-defined custom model routers referenced by the current selection. Mirrors
/// `custom_model_providers`: the selected model's `config_key` indexes into this
/// registry. `None` when no custom router is selected.
pub custom_model_routers: Option<warp_multi_agent_api::request::settings::CustomModelRouters>,
pub allow_use_of_warp_credits: bool,
pub autonomy_level: warp_multi_agent_api::AutonomyLevel,
pub isolation_level: warp_multi_agent_api::IsolationLevel,
pub web_search_enabled: bool,
@@ -170,6 +176,44 @@ pub struct ConversationData {
}
impl RequestParams {
#[cfg(test)]
pub fn new_for_test() -> Self {
Self {
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: LLMId::from("test-model"),
coding_model: LLMId::from("test-model"),
cli_agent_model: LLMId::from("test-model"),
computer_use_model: LLMId::from("test-model"),
is_memory_enabled: false,
warp_drive_context_enabled: false,
context_window_limit: None,
mcp_context: None,
planning_enabled: false,
should_redact_secrets: false,
api_keys: None,
custom_model_providers: None,
custom_model_routers: None,
allow_use_of_warp_credits: false,
autonomy_level: Default::default(),
isolation_level: Default::default(),
web_search_enabled: false,
computer_use_enabled: false,
ask_user_question_enabled: false,
research_agent_enabled: false,
orchestration_enabled: false,
supported_tools_override: None,
parent_agent_id: None,
agent_name: None,
}
}
pub fn new(
terminal_view_id: Option<EntityId>,
session_context: SessionContext,
@@ -251,12 +295,31 @@ impl RequestParams {
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_bedrock_enabled(app),
let api_key_manager = ApiKeyManager::as_ref(app);
let is_byo_enabled = user_workspaces.is_byo_api_key_enabled(app);
#[cfg(not(target_family = "wasm"))]
let geap_binding = crate::ai::geap_credentials::current_geap_policy(app).mint_binding();
#[cfg(target_family = "wasm")]
let geap_binding: Option<::ai::api_keys::GeapMintBinding> = None;
let api_keys = api_key_manager.api_keys_for_request(
is_byo_enabled,
user_workspaces.is_aws_bedrock_credentials_enabled(app),
geap_binding,
);
let allow_use_of_warp_credits_with_byok =
*AISettings::as_ref(app).can_use_warp_credits_with_byok;
let is_custom_inference_enabled = user_workspaces.is_custom_inference_enabled(app);
let custom_model_providers = FeatureFlag::CustomInferenceEndpoints
.is_enabled()
.then(|| {
api_key_manager.custom_model_providers_for_request(is_custom_inference_enabled)
})
.flatten();
let custom_model_routers = FeatureFlag::CustomModelRouters.is_enabled().then(|| {
LLMPreferences::as_ref(app).custom_model_routers_for_request(
&request_input.model_id,
&request_input.coding_model_id,
)
});
let allow_use_of_warp_credits = *AISettings::as_ref(app).can_use_warp_credits_for_fallback;
let app_execution_mode = AppExecutionMode::as_ref(app);
let autonomy_level = if app_execution_mode.is_autonomous() {
@@ -292,11 +355,25 @@ impl RequestParams {
!= crate::ai::execution_profiles::AskUserQuestionPermission::Never;
let orchestration_enabled = ai_settings.is_orchestration_enabled(app)
&& BlocklistAIPermissions::as_ref(app)
.get_run_agents_setting(app, terminal_view_id)
.is_enabled()
&& session_context
.session_type()
.as_ref()
.is_none_or(|t| matches!(t, crate::terminal::model::session::SessionType::Local));
// Reconcile the persisted override against the active base model's
// current `LLMContextWindow` instead of trusting whatever was stored
// last. If the active model isn't configurable or has been removed
// server-side, drop the override; otherwise clamp it to the model's
// current `[min, max]` range. This closes the window between an
// in-flight model metadata refresh and the next request.
let context_window_limit = AIExecutionProfilesModel::as_ref(app)
.active_profile(terminal_view_id, app)
.data()
.context_window_limit_for_request(app);
Self {
input: request_input.all_inputs().cloned().collect(),
conversation_token: conversation.server_conversation_token,
@@ -304,6 +381,7 @@ impl RequestParams {
ambient_agent_task_id: conversation.ambient_agent_task_id,
tasks: conversation.tasks,
existing_suggestions: conversation.existing_suggestions,
context_window_limit,
metadata,
session_context,
model: request_input.model_id.clone(),
@@ -316,7 +394,9 @@ impl RequestParams {
planning_enabled: true,
should_redact_secrets,
api_keys,
allow_use_of_warp_credits_with_byok,
custom_model_providers,
custom_model_routers,
allow_use_of_warp_credits,
autonomy_level,
isolation_level,
web_search_enabled,
app/src/ai/agent/api/convert_conversation.rs
+208 -62
View File
@@ -4,12 +4,28 @@
//! If some UI state is stored in the client, it needs to also be represented in the proto tasks somehow so it can be restored.
//! Some conversions may be lossy if it's not important to recover that UI state.
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use ai::agent::action_result::{
AskUserQuestionAnswerItem, AskUserQuestionResult, FetchConversationResult, ReadSkillResult,
RequestComputerUseResult, SendMessageToAgentResult, StartAgentResult, StartAgentVersion,
UseComputerResult,
};
use ai::skills::{ParsedSkill, SkillPathOrigin};
use chrono::{DateTime, Local, TimeZone};
use persistence::model::AgentConversationData;
use galaxy_core::command::ExitCode;
use warp_multi_agent_api as api;
use warp_multi_agent_api::ask_user_question_result::answer_item::Answer as AskUserQuestionAnswer;
use crate::ai::agent::api::convert_from::{
convert_user_query_mode, ConversionParams, ConvertAPIMessageToClientOutputMessage,
MaybeAIAgentOutputMessage,
};
use crate::ai::agent::conversation::update_todo_list_from_todo_op;
use crate::ai::agent::conversation::{AIConversation, AIConversationId};
use crate::ai::agent::conversation::{
update_todo_list_from_todo_op, AIConversation, AIConversationId, ServerAIConversationMetadata,
};
use crate::ai::agent::task::TaskId;
use crate::ai::agent::todos::AIAgentTodoList;
use crate::ai::agent::{
@@ -24,7 +40,7 @@ use crate::ai::agent::{
RequestFileEditsResult, SearchCodebaseFailureReason, SearchCodebaseResult, ServerOutputId,
Shared, ShellCommandCompletedTrigger, ShellCommandError, SuggestNewConversationResult,
SuggestPromptResult, TransferShellCommandControlToUserResult, UpdatedFileContext,
UploadArtifactResult, WriteToLongRunningShellCommandResult,
UploadArtifactResult, UserQueryMode, WriteToLongRunningShellCommandResult,
};
use crate::ai::block_context::BlockContext;
use crate::ai::document::ai_document_model::{AIDocumentId, AIDocumentVersion};
@@ -32,22 +48,6 @@ use crate::ai::llms::LLMId;
use crate::ai_assistant::execution_context::{WarpAiExecutionContext, WarpAiOsContext};
use crate::terminal::model::block::BlockId;
use crate::terminal::model::terminal_model::BlockIndex;
use ai::agent::action_result::{
AskUserQuestionAnswerItem, AskUserQuestionResult, FetchConversationResult, ReadSkillResult,
RequestComputerUseResult, SendMessageToAgentResult, StartAgentResult, StartAgentVersion,
UseComputerResult,
};
use ai::skills::ParsedSkill;
use chrono::{DateTime, Local, TimeZone};
use galaxy_core::command::ExitCode;
use persistence::model::AgentConversationData;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use warp_multi_agent_api as api;
use warp_multi_agent_api::ask_user_question_result::answer_item::Answer as AskUserQuestionAnswer;
use crate::ai::agent::conversation::ServerAIConversationMetadata;
use crate::ai::agent::UserQueryMode;
/// How to restore a conversation from the cloud.
pub enum RestorationMode {
@@ -81,12 +81,14 @@ pub fn convert_conversation_data_to_ai_conversation(
artifacts_json: None,
parent_agent_id: None,
agent_name: None,
orchestration_harness_type: None,
parent_conversation_id: None,
is_remote_child: false,
root_task_is_optimistic: None,
run_id: None,
autoexecute_override: None,
last_event_sequence: None,
progressive_summary: None,
messages_summarized_up_to: 0,
pinned: false,
},
RestorationMode::Continue => AgentConversationData {
server_conversation_token: Some(
@@ -98,16 +100,16 @@ pub fn convert_conversation_data_to_ai_conversation(
artifacts_json: serde_json::to_string(&metadata.artifacts).ok(),
parent_agent_id: None,
agent_name: None,
orchestration_harness_type: None,
parent_conversation_id: None,
// TODO: Populate run_id from server metadata once it is exposed
// in ServerAIConversationMetadata. For cloud conversations that
// were spawned via the server API, the run_id is created at task
// dispatch time; adding it here would avoid a round-trip to StreamInit.
run_id: None,
is_remote_child: false,
root_task_is_optimistic: None,
run_id: metadata
.ambient_agent_task_id
.map(|task_id| task_id.to_string()),
autoexecute_override: None,
last_event_sequence: None,
progressive_summary: None,
messages_summarized_up_to: 0,
pinned: false,
},
};
@@ -239,7 +241,8 @@ pub(crate) fn convert_input_context(context: Option<&api::InputContext>) -> Arc<
};
// Convert binary data to base64
use base64::{engine::general_purpose, Engine};
use base64::engine::general_purpose;
use base64::Engine;
let data = general_purpose::STANDARD.encode(&image.data);
result.push(AIAgentContext::Image(ImageContext {
@@ -436,7 +439,10 @@ impl ConvertToExchanges for &api::Task {
api::message::system_query::Type::ResumeConversation(_)
| api::message::system_query::Type::GeneratePassiveSuggestions(_)
// TODO: Implement this for real. ZB adding this to bump proto version for unrelated API changes.
| api::message::system_query::Type::SummarizeConversation(_)=> false,
| api::message::system_query::Type::SummarizeConversation(_)
// HandoffRehydration is injected by the server for agent-only
// context; the client must never render it as user input.
| api::message::system_query::Type::HandoffRehydration(_) => false,
}
}
api::message::Message::ToolCallResult(tool_call_result) => {
@@ -462,7 +468,10 @@ impl ConvertToExchanges for &api::Task {
}
api::message::Message::InvokeSkill(invoke_skill) => {
if let Some(api_skill) = invoke_skill.skill.clone() {
if let Ok(parsed_skill) = ParsedSkill::try_from(api_skill) {
if let Ok(parsed_skill) = ParsedSkill::try_from_api_with_origin(
api_skill,
&SkillPathOrigin::RestoredDisplayOnly,
) {
let user_query = invoke_skill
.user_query
.clone()
@@ -513,7 +522,8 @@ impl ConvertToExchanges for &api::Task {
| api::message::Message::DebugOutput(_)
| api::message::Message::ArtifactEvent(_)
| api::message::Message::MessagesReceivedFromAgents(_)
| api::message::Message::ModelUsed(_) => false,
| api::message::Message::ModelUsed(_)
| api::message::Message::OrchestrationConfigSnapshot(_) => false,
};
if !added_message_as_exchange_input {
@@ -524,6 +534,7 @@ impl ConvertToExchanges for &api::Task {
// TODO(alokedesai): Support persistence for the code review state.
active_code_review: None,
task_id: &TaskId::new(api_message.task_id.clone()),
skill_path_origin: &SkillPathOrigin::Unavailable,
})
{
current_outputs.push(output_msg);
@@ -581,6 +592,14 @@ pub(crate) fn convert_tool_call_result_to_input(
command: result.command.clone(),
output: finished.output.clone(),
exit_code: ExitCode::from(finished.exit_code),
start_ts: finished
.start_ts
.as_ref()
.map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
completed_ts: finished
.finish_ts
.as_ref()
.map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
}
}
Some(api::run_shell_command_result::Result::LongRunningCommandSnapshot(
@@ -627,6 +646,8 @@ pub(crate) fn convert_tool_call_result_to_input(
block_id: finished.command_id.clone().into(),
output: finished.output.clone(),
exit_code: ExitCode::from(finished.exit_code),
start_ts: finished.start_ts.as_ref().map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
completed_ts: finished.finish_ts.as_ref().map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
},
Some(api::write_to_long_running_shell_command_result::Result::Error(api::ShellCommandError{
r#type: Some(api::shell_command_error::Type::CommandNotFound(()))
@@ -1217,6 +1238,14 @@ pub(crate) fn convert_tool_call_result_to_input(
block_id: finished.command_id.clone().into(),
output: finished.output.clone(),
exit_code: ExitCode::from(finished.exit_code),
start_ts: finished
.start_ts
.as_ref()
.map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
completed_ts: finished
.finish_ts
.as_ref()
.map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
}
}
Some(
@@ -1267,6 +1296,8 @@ pub(crate) fn convert_tool_call_result_to_input(
block_id: finished.command_id.clone().into(),
output: finished.output.clone(),
exit_code: ExitCode::from(finished.exit_code),
start_ts: finished.start_ts.as_ref().map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
completed_ts: finished.finish_ts.as_ref().map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
},
Some(api::transfer_shell_command_control_to_user_result::Result::Error(
api::ShellCommandError {
@@ -1315,34 +1346,52 @@ pub(crate) fn convert_tool_call_result_to_input(
})
}
Some(ToolCallResultType::UseComputer(result)) => {
let use_computer_result = match &result.result {
Some(api::use_computer_result::Result::Success(success)) => {
let screenshot = success.screenshot.as_ref().map(|s| {
// The original dimensions are not preserved through the API, so we use
// the current dimensions for both.
computer_use::Screenshot {
width: s.width as usize,
height: s.height as usize,
original_width: s.width as usize,
original_height: s.height as usize,
data: s.data.clone(),
mime_type: s.mime_type.clone().into(),
}
});
let cursor_position = success
.cursor_position
.as_ref()
.map(|c| computer_use::Vector2I::new(c.x, c.y));
UseComputerResult::Success(computer_use::ActionResult {
screenshot,
cursor_position,
})
}
Some(api::use_computer_result::Result::Error(error)) => {
UseComputerResult::Error(error.message.clone())
}
None => UseComputerResult::Cancelled,
};
let use_computer_result =
match &result.result {
Some(api::use_computer_result::Result::Success(success)) => {
let screenshot = success.screenshot.as_ref().map(|s| {
// The original dimensions are not preserved through the API, so we use
// the current dimensions for both.
computer_use::Screenshot {
width: s.width as usize,
height: s.height as usize,
original_width: s.width as usize,
original_height: s.height as usize,
data: s.data.clone(),
mime_type: s.mime_type.clone().into(),
}
});
let cursor_position = success
.cursor_position
.as_ref()
.map(|c| computer_use::Vector2I::new(c.x, c.y));
let windows = success
.windows
.iter()
.map(convert_api_window_info)
.collect();
// A present captured-window message indicates a window screenshot was taken.
// The window id is an opaque string on the wire; on macOS it is a CGWindowID,
// so parse it back to a u32, defaulting to 0 when it is not parseable.
let captured_window = success.captured_window.as_ref().map(|c| {
computer_use::CapturedWindow {
window_id: c.window_id.parse().unwrap_or(0),
width_px: c.width_px,
height_px: c.height_px,
}
});
UseComputerResult::Success(computer_use::ActionResult {
screenshot,
cursor_position,
windows,
captured_window,
})
}
Some(api::use_computer_result::Result::Error(error)) => {
UseComputerResult::Error(error.message.clone())
}
None => UseComputerResult::Cancelled,
};
Some(AIAgentInput::ActionResult {
result: AIAgentActionResult {
@@ -1361,6 +1410,7 @@ pub(crate) fn convert_tool_call_result_to_input(
api::request_computer_use_result::Approved {
screen_dimensions: Some(screen_dimensions),
initial_screenshot: Some(initial_screenshot),
windows,
..
},
Some(platform),
@@ -1374,6 +1424,7 @@ pub(crate) fn convert_tool_call_result_to_input(
mime_type: initial_screenshot.mime_type.clone().into(),
},
platform,
windows: windows.iter().map(convert_api_window_info).collect(),
},
_ => RequestComputerUseResult::Error(
"Missing screen dimensions, initial screenshot, or valid platform"
@@ -1544,6 +1595,76 @@ pub(crate) fn convert_tool_call_result_to_input(
context,
})
}
Some(ToolCallResultType::RunAgentsResult(result)) => {
use ai::agent::action_result::{
RunAgentsAgentOutcome, RunAgentsAgentOutcomeKind, RunAgentsLaunchedExecutionMode,
RunAgentsResult,
};
let run_agents_result = match &result.outcome {
Some(api::run_agents_result::Outcome::Launched(launched)) => {
let execution_mode = match &launched.resolved_execution_mode {
Some(api::run_agents_result::launched::ResolvedExecutionMode::Remote(
remote,
)) => RunAgentsLaunchedExecutionMode::Remote {
environment_id: remote.environment_id.clone(),
worker_host: remote.worker_host.clone(),
computer_use_enabled: remote.computer_use_enabled,
},
Some(api::run_agents_result::launched::ResolvedExecutionMode::Local(_))
| None => RunAgentsLaunchedExecutionMode::Local,
};
let agents = launched
.agents
.iter()
.map(|outcome| RunAgentsAgentOutcome {
name: outcome.name.clone(),
kind: match &outcome.result {
Some(api::run_agents_result::agent_outcome::Result::Launched(
launched_agent,
)) => RunAgentsAgentOutcomeKind::Launched {
agent_id: launched_agent.agent_id.clone(),
},
Some(api::run_agents_result::agent_outcome::Result::Failed(
failed,
)) => RunAgentsAgentOutcomeKind::Failed {
error: failed.error.clone(),
},
None => RunAgentsAgentOutcomeKind::Failed {
error: String::new(),
},
},
})
.collect();
RunAgentsResult::Launched {
model_id: launched.resolved_model_id.clone(),
harness_type:
crate::ai::agent::api::convert_from::convert_run_agents_harness(
launched.resolved_harness.as_ref(),
)
.unwrap_or_default(),
execution_mode,
agents,
}
}
Some(api::run_agents_result::Outcome::Denied(denied)) => RunAgentsResult::Denied {
reason: denied.reason.clone(),
},
Some(api::run_agents_result::Outcome::Failure(failure)) => {
RunAgentsResult::Failure {
error: failure.error.clone(),
}
}
None => RunAgentsResult::Cancelled,
};
Some(AIAgentInput::ActionResult {
result: AIAgentActionResult {
id: tool_call_id.into(),
task_id: task_id.clone(),
result: AIAgentActionResultType::RunAgents(run_agents_result),
},
context,
})
}
// Deprecated/unused result types or absent result.
Some(ToolCallResultType::SuggestCreatePlan(..))
| Some(ToolCallResultType::SuggestPlan(..))
@@ -1551,6 +1672,7 @@ pub(crate) fn convert_tool_call_result_to_input(
log::warn!("No result present for tool call ID: {tool_call_id}");
None
}
Some(ToolCallResultType::WaitForEvents(_)) => None,
}
}
@@ -1678,8 +1800,14 @@ fn create_cancelled_result_for_tool_call(
ToolType::SendMessageToAgent(_) => {
AIAgentActionResultType::SendMessageToAgent(SendMessageToAgentResult::Cancelled)
}
ToolType::RunAgents(_) => {
AIAgentActionResultType::RunAgents(ai::agent::action_result::RunAgentsResult::Cancelled)
}
// These tools are deprecated.
ToolType::SuggestCreatePlan(_) | ToolType::SuggestPlan(_) => return None,
ToolType::WaitForEvents(_) => {
return None;
}
};
Some(AIAgentInput::ActionResult {
@@ -1777,6 +1905,10 @@ fn create_exchange_from_messages(
model_id: model.model_id.clone().into(),
display_name: model.model_display_name.clone(),
is_fallback: model.is_fallback,
prompt_cache_expires_at: model
.prompt_cache_expires_at
.as_ref()
.map(|ts| proto_timestamp_to_local_datetime(ts.seconds, ts.nanos)),
}),
request_cost: None,
};
@@ -1879,7 +2011,8 @@ where
| api::message::Message::DebugOutput(_)
| api::message::Message::ArtifactEvent(_)
| api::message::Message::InvokeSkill(_)
| api::message::Message::ModelUsed(_) => {
| api::message::Message::ModelUsed(_)
| api::message::Message::OrchestrationConfigSnapshot(_) => {
message.timestamp.as_ref().map(|timestamp| {
proto_timestamp_to_local_datetime(timestamp.seconds, timestamp.nanos)
})
@@ -1985,7 +2118,7 @@ fn convert_passive_suggestion_result_to_input(
context,
})
}
fn proto_timestamp_to_local_datetime(seconds: i64, nanos: i32) -> DateTime<Local> {
pub(crate) fn proto_timestamp_to_local_datetime(seconds: i64, nanos: i32) -> DateTime<Local> {
let nanos = if nanos < 0 { 0 } else { nanos as u32 };
Local
@@ -2014,6 +2147,19 @@ fn convert_api_platform(platform: i32) -> Option<computer_use::Platform> {
}
}
/// Reconstructs the internal computer_use window record from the API `WindowInfo` message.
fn convert_api_window_info(window: &api::WindowInfo) -> computer_use::WindowInfo {
computer_use::WindowInfo {
// The window id arrives as an opaque string; on macOS it is a CGWindowID (u32). Default to
// 0 when it is not parseable.
window_id: window.window_id.parse().unwrap_or(0),
pid: window.pid,
app_name: window.app_name.clone(),
title: window.title.clone(),
layer: window.layer,
}
}
#[cfg(test)]
#[path = "convert_conversation_tests.rs"]
mod tests;
app/src/ai/agent/api/convert_conversation_tests.rs
+224 -5
View File
@@ -1,7 +1,56 @@
use crate::ai::agent::api::convert_conversation::*;
use crate::ai::agent::{AIAgentInput, UserQueryMode};
use std::collections::HashMap;
use chrono::Utc;
use warp_multi_agent_api as api;
use crate::ai::agent::api::convert_conversation::*;
use crate::ai::agent::api::ServerConversationToken;
use crate::ai::agent::conversation::{
AIAgentHarness, AIConversationId, ServerAIConversationMetadata,
};
use crate::ai::agent::{AIAgentInput, UserQueryMode};
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::cloud_object::{Revision, ServerMetadata, ServerPermissions};
use crate::persistence::model::ConversationUsageMetadata;
use crate::server::ids::ServerId;
fn test_server_metadata(
server_token: &str,
ambient_agent_task_id: Option<AmbientAgentTaskId>,
) -> ServerAIConversationMetadata {
ServerAIConversationMetadata {
title: "test conversation".to_string(),
working_directory: None,
harness: AIAgentHarness::Oz,
usage: ConversationUsageMetadata {
was_summarized: false,
context_window_usage: 0.0,
credits_spent: 0.0,
platform_credits_spent: 0.0,
credits_spent_for_last_block: None,
token_usage: vec![],
tool_usage_metadata: Default::default(),
context_window_segments: Vec::new(),
},
metadata: 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,
},
permissions: ServerPermissions::mock_personal(),
creator: None,
ambient_agent_task_id,
server_conversation_token: ServerConversationToken::new(server_token.to_string()),
artifacts: vec![],
}
}
fn test_skill() -> api::Skill {
api::Skill {
descriptor: Some(api::SkillDescriptor {
@@ -25,6 +74,40 @@ fn test_skill() -> api::Skill {
}
}
#[test]
#[allow(deprecated)]
fn test_convert_conversation_data_to_ai_conversation_sets_restored_run_id() {
let conversation_id = AIConversationId::new();
let ambient_agent_task_id: AmbientAgentTaskId =
"550e8400-e29b-41d4-a716-446655440000".parse().unwrap();
let conversation_data = api::ConversationData {
tasks: vec![api::Task {
id: "root".to_string(),
messages: vec![],
dependencies: None,
description: String::new(),
summary: String::new(),
server_data: String::new(),
}],
ordered_message_ids: vec![],
};
let conversation = convert_conversation_data_to_ai_conversation(
conversation_id,
&conversation_data,
test_server_metadata("server-token", Some(ambient_agent_task_id)),
RestorationMode::Continue,
)
.expect("conversation should restore");
assert_eq!(conversation.id(), conversation_id);
assert_eq!(conversation.task_id(), Some(ambient_agent_task_id));
assert_eq!(
conversation.run_id(),
Some(ambient_agent_task_id.to_string())
);
}
#[test]
fn test_convert_tool_call_result_to_input_transfer_control_snapshot() {
let task_id = crate::ai::agent::task::TaskId::new("task".to_string());
@@ -277,6 +360,7 @@ fn test_into_exchanges_basic() {
// Create minimal test data
let messages = vec![
api::Message {
fetched_memories: vec![],
id: "user_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -292,6 +376,7 @@ fn test_into_exchanges_basic() {
timestamp: None,
},
api::Message {
fetched_memories: vec![],
id: "agent_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -305,6 +390,7 @@ fn test_into_exchanges_basic() {
timestamp: None,
},
api::Message {
fetched_memories: vec![],
id: "user_msg2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -320,6 +406,7 @@ fn test_into_exchanges_basic() {
timestamp: None,
},
api::Message {
fetched_memories: vec![],
id: "agent_msg2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -333,6 +420,7 @@ fn test_into_exchanges_basic() {
timestamp: None,
},
api::Message {
fetched_memories: vec![],
id: "user_msg3".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -348,6 +436,7 @@ fn test_into_exchanges_basic() {
timestamp: None,
},
api::Message {
fetched_memories: vec![],
id: "agent_msg3".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -392,6 +481,7 @@ fn test_invoke_skill_arguments_round_trip() {
let query = "arg1 arg2".to_string();
let messages = vec![
api::Message {
fetched_memories: vec![],
id: "invoke_skill_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -412,6 +502,7 @@ fn test_invoke_skill_arguments_round_trip() {
timestamp: None,
},
api::Message {
fetched_memories: vec![],
id: "agent_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -448,7 +539,7 @@ fn test_invoke_skill_arguments_round_trip() {
Some("arg1 arg2")
);
assert_eq!(
exchanges[0].input[0].user_query().as_deref(),
exchanges[0].input[0].display_query().as_deref(),
Some("/test-skill arg1 arg2")
);
}
@@ -459,6 +550,7 @@ fn test_invoke_skill_arguments_round_trip() {
#[test]
fn test_invoke_skill_missing_user_query_maps_to_none() {
let messages = vec![api::Message {
fetched_memories: vec![],
id: "invoke_skill_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -492,7 +584,7 @@ fn test_invoke_skill_missing_user_query_maps_to_none() {
assert_eq!(skill.name, "test-skill");
assert_eq!(user_query, &None);
assert_eq!(
exchanges[0].input[0].user_query().as_deref(),
exchanges[0].input[0].display_query().as_deref(),
Some("/test-skill")
);
}
@@ -505,6 +597,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
let messages = vec![
// User query
api::Message {
fetched_memories: vec![],
id: "user_query".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -521,6 +614,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Agent response
api::Message {
fetched_memories: vec![],
id: "agent_response".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -535,6 +629,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Tool call 1
api::Message {
fetched_memories: vec![],
id: "tool_call_1".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -558,6 +653,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Tool call 2
api::Message {
fetched_memories: vec![],
id: "tool_call_2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -581,6 +677,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Tool call 3
api::Message {
fetched_memories: vec![],
id: "tool_call_3".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -604,6 +701,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Tool call result - cancelled (call_2)
api::Message {
fetched_memories: vec![],
id: "result_cancelled".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -620,6 +718,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Tool call result - success (call_1)
api::Message {
fetched_memories: vec![],
id: "result_success_1".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -639,6 +738,8 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
command_id: "command_1".to_string(),
output: "1".to_string(),
exit_code: 0,
start_ts: None,
finish_ts: None,
},
)),
},
@@ -650,6 +751,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Tool call result - success (call_3)
api::Message {
fetched_memories: vec![],
id: "result_success_3".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -669,6 +771,8 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
command_id: "command_2".to_string(),
output: "3".to_string(),
exit_code: 0,
start_ts: None,
finish_ts: None,
},
)),
},
@@ -680,6 +784,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Final agent response
api::Message {
fetched_memories: vec![],
id: "final_response".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -694,6 +799,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Follow-up user query
api::Message {
fetched_memories: vec![],
id: "followup_query".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -710,6 +816,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() {
},
// Final agent response
api::Message {
fetched_memories: vec![],
id: "final_response2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -811,6 +918,7 @@ fn test_into_exchanges_with_code_diffs() {
let messages = vec![
// User query asking for code changes
api::Message {
fetched_memories: vec![],
id: "user_query".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -827,6 +935,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// Agent response
api::Message {
fetched_memories: vec![],
id: "agent_response".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -841,6 +950,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// File diff tool call
api::Message {
fetched_memories: vec![],
id: "diff_call".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -862,6 +972,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// User cancels the diff
api::Message {
fetched_memories: vec![],
id: "diff_cancelled".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -878,6 +989,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// User provides feedback
api::Message {
fetched_memories: vec![],
id: "user_feedback".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -894,6 +1006,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// Agent response
api::Message {
fetched_memories: vec![],
id: "agent_response_2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -908,6 +1021,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// Second file diff tool call
api::Message {
fetched_memories: vec![],
id: "diff_call_2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -929,6 +1043,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// User accepts the diff
api::Message {
fetched_memories: vec![],
id: "diff_accepted".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -956,6 +1071,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// Final agent response
api::Message {
fetched_memories: vec![],
id: "final_response".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -970,6 +1086,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// Follow-up user query
api::Message {
fetched_memories: vec![],
id: "followup".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -986,6 +1103,7 @@ fn test_into_exchanges_with_code_diffs() {
},
// Final agent response
api::Message {
fetched_memories: vec![],
id: "final_response_2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1083,6 +1201,7 @@ fn test_into_exchanges_with_code_diffs() {
fn test_user_query_mode_conversion() {
// Test conversion with Plan mode
let messages = vec![api::Message {
fetched_memories: vec![],
id: "user_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1129,6 +1248,7 @@ fn test_user_query_mode_conversion() {
// Test conversion with Normal mode (no type set)
let messages_normal = vec![api::Message {
fetched_memories: vec![],
id: "user_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1173,6 +1293,7 @@ fn test_user_query_mode_conversion() {
// Test conversion with no mode field (should default to Normal)
let messages_default = vec![api::Message {
fetched_memories: vec![],
id: "user_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1226,6 +1347,7 @@ fn test_exchanges_grouped_by_request_id() {
let messages = vec![
// Message 0: Server message (should be ignored or handled gracefully)
api::Message {
fetched_memories: vec![],
id: "2512077c-0ede-46b0-8f69-230c8792df07".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(),
@@ -1243,6 +1365,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 1: User query with request_id 78e236b8
api::Message {
fetched_memories: vec![],
id: "4d6c450d-3d54-446f-974c-5c414e6083e9".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(),
@@ -1259,6 +1382,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 2: Agent output with same request_id
api::Message {
fetched_memories: vec![],
id: "10210d1a-5298-45ef-90ba-df6367805080".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(),
@@ -1273,6 +1397,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 3: Tool call with same request_id
api::Message {
fetched_memories: vec![],
id: "936c7c86-eb4a-4edf-97c0-22f5c61b35a6".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(),
@@ -1296,6 +1421,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 4: Tool call result with NEW request_id 59a3947f (starts new exchange)
api::Message {
fetched_memories: vec![],
id: "cbebf5fb-4dd8-4aef-be45-bb916eff552c".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "59a3947f-fc7e-413a-96b5-baecd7e406dc".to_string(),
@@ -1330,6 +1456,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 5: Agent output with same request_id
api::Message {
fetched_memories: vec![],
id: "7a89857d-fa33-4d45-88e3-5fa9cbce3f20".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "59a3947f-fc7e-413a-96b5-baecd7e406dc".to_string(),
@@ -1344,6 +1471,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 6: Write to long running command with NEW request_id 9f85acb2 (starts new exchange)
api::Message {
fetched_memories: vec![],
id: "dac6d336-9fcb-4e34-bc2b-b06e70f52ec5".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "9f85acb2-0b1f-41b1-a0de-3623e131758a".to_string(),
@@ -1374,6 +1502,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 7: Final tool call result with same request_id
api::Message {
fetched_memories: vec![],
id: "ad319d66-fac0-4169-8bf1-e6004aca1619".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "9f85acb2-0b1f-41b1-a0de-3623e131758a".to_string(),
@@ -1395,6 +1524,8 @@ fn test_exchanges_grouped_by_request_id() {
command_id: "cmd1".to_string(),
output: "Done".to_string(),
exit_code: 0,
start_ts: None,
finish_ts: None,
},
)),
},
@@ -1404,6 +1535,7 @@ fn test_exchanges_grouped_by_request_id() {
},
// Message 8: Final agent output with same request_id
api::Message {
fetched_memories: vec![],
id: "f15f8a59-2e9c-416e-b216-83b3bd52d6be".to_string(),
task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(),
request_id: "9f85acb2-0b1f-41b1-a0de-3623e131758a".to_string(),
@@ -1492,6 +1624,7 @@ fn test_multiple_create_documents_get_default_version() {
let messages = vec![
// User query
api::Message {
fetched_memories: vec![],
id: "user_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1508,6 +1641,7 @@ fn test_multiple_create_documents_get_default_version() {
},
// Agent output
api::Message {
fetched_memories: vec![],
id: "agent_text".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1522,6 +1656,7 @@ fn test_multiple_create_documents_get_default_version() {
},
// First CreateDocuments tool call
api::Message {
fetched_memories: vec![],
id: "tool_call_create_a".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1544,6 +1679,7 @@ fn test_multiple_create_documents_get_default_version() {
},
// First CreateDocuments result
api::Message {
fetched_memories: vec![],
id: "result_create_a".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1572,6 +1708,7 @@ fn test_multiple_create_documents_get_default_version() {
},
// Agent output before second plan
api::Message {
fetched_memories: vec![],
id: "agent_text_2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1586,6 +1723,7 @@ fn test_multiple_create_documents_get_default_version() {
},
// Second CreateDocuments tool call
api::Message {
fetched_memories: vec![],
id: "tool_call_create_b".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1608,6 +1746,7 @@ fn test_multiple_create_documents_get_default_version() {
},
// Second CreateDocuments result
api::Message {
fetched_memories: vec![],
id: "result_create_b".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1694,7 +1833,6 @@ fn test_multiple_create_documents_get_default_version() {
#[test]
fn test_create_then_edit_then_create_version_tracking() {
use crate::ai::agent::{AIAgentActionResultType, CreateDocumentsResult, EditDocumentsResult};
use crate::ai::document::ai_document_model::AIDocumentVersion;
let doc_id_a = uuid::Uuid::new_v4().to_string();
let doc_id_b = uuid::Uuid::new_v4().to_string();
@@ -1702,6 +1840,7 @@ fn test_create_then_edit_then_create_version_tracking() {
let messages = vec![
// User query
api::Message {
fetched_memories: vec![],
id: "user_msg".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1718,6 +1857,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Agent output
api::Message {
fetched_memories: vec![],
id: "agent_text".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1732,6 +1872,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Create doc A tool call
api::Message {
fetched_memories: vec![],
id: "tool_call_create_a".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1754,6 +1895,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Create doc A result
api::Message {
fetched_memories: vec![],
id: "result_create_a".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1782,6 +1924,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Agent output before edit
api::Message {
fetched_memories: vec![],
id: "agent_text_2".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1796,6 +1939,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Edit doc A tool call
api::Message {
fetched_memories: vec![],
id: "tool_call_edit_a".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1817,6 +1961,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Edit doc A result
api::Message {
fetched_memories: vec![],
id: "result_edit_a".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1845,6 +1990,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Agent output before second create
api::Message {
fetched_memories: vec![],
id: "agent_text_3".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1859,6 +2005,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Create doc B tool call
api::Message {
fetched_memories: vec![],
id: "tool_call_create_b".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1881,6 +2028,7 @@ fn test_create_then_edit_then_create_version_tracking() {
},
// Create doc B result
api::Message {
fetched_memories: vec![],
id: "result_create_b".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
@@ -1979,3 +2127,74 @@ fn test_create_then_edit_then_create_version_tracking() {
"Created doc B should have default version (v1), independent of doc A"
);
}
/// Verify that a `SystemQuery::HandoffRehydration` message does not produce
/// a displayed input when restoring a conversation. It must be treated as
/// hidden, so the exchange should have zero user-visible inputs.
#[test]
fn test_handoff_rehydration_system_query_is_hidden() {
let messages = vec![
// HandoffRehydration system query – should be hidden
api::Message {
fetched_memories: vec![],
id: "msg_handoff".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
citations: vec![],
message: Some(api::message::Message::SystemQuery(
api::message::SystemQuery {
r#type: Some(api::message::system_query::Type::HandoffRehydration(
api::message::HandoffRehydration {
instructions: "restore handoff state".to_string(),
},
)),
context: None,
},
)),
request_id: "req1".to_string(),
timestamp: None,
},
// Agent output that follows the hidden system query
api::Message {
fetched_memories: vec![],
id: "msg_output".to_string(),
task_id: "task1".to_string(),
server_message_data: "".to_string(),
citations: vec![],
message: Some(api::message::Message::AgentOutput(
api::message::AgentOutput {
text: "I have restored the handoff state.".to_string(),
},
)),
request_id: "req1".to_string(),
timestamp: None,
},
];
let task = api::Task {
id: "task1".to_string(),
messages,
dependencies: None,
description: "".to_string(),
summary: "".to_string(),
server_data: "".to_string(),
};
let exchanges = task.into_exchanges();
assert_eq!(exchanges.len(), 1, "Should produce exactly one exchange");
let exchange = &exchanges[0];
// The HandoffRehydration should NOT appear as input
assert!(
exchange.input.is_empty(),
"HandoffRehydration must not produce a displayed input, got: {:?}",
exchange.input
);
// The agent output should still be present
let output = exchange.output_status.output().expect("should have output");
assert!(
!output.get().messages.is_empty(),
"Agent output should still be rendered"
);
}
app/src/ai/agent/api/convert_from.rs
+137 -36
View File
@@ -2,34 +2,33 @@
use std::collections::HashMap;
use std::time::Duration;
use ai::agent::action::{LifecycleEventType as StartAgentLifecycleEventType, ReadSkillRequest};
use ai::agent::action_result::StartAgentVersion;
use ai::agent::convert::ToolToAIAgentActionError;
use ai::agent::UnknownCitationTypeError;
use ai::skills::{
skill_reference_from_api_skill_ref, skill_reference_from_read_skill_ref, SkillPathOrigin,
};
use api::ask_user_question::question::QuestionType;
use galaxy_core::channel::ChannelState;
use warp_multi_agent_api as api;
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;
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,
AIAgentAction, AIAgentActionType, AIAgentAttachment, AIAgentCitation, AIAgentInput,
AIAgentOutputMessage, AIAgentText, AIAgentTodo, ArtifactCreatedData, CloneRepositoryURL,
MessageId, RunAgentsAgentRunConfig, RunAgentsExecutionMode, RunAgentsRequest,
StartAgentExecutionMode, SubagentCall, SubagentType, SuggestedAgentModeWorkflow, SuggestedRule,
Suggestions, SummarizationType, TodoOperation, UserQueryMode, 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 galaxy_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;
@@ -53,6 +52,18 @@ impl TryFrom<api::Attachment> for AIAgentAttachment {
}
}
fn convert_read_skill(
read_skill: api::message::tool_call::ReadSkill,
skill_path_origin: &SkillPathOrigin,
) -> Result<AIAgentActionType, ToolToAIAgentActionError> {
let Some(reference) = read_skill.skill_reference else {
return Err(ToolToAIAgentActionError::MissingSkillReference);
};
let skill = skill_reference_from_read_skill_ref(reference, skill_path_origin)
.map_err(|_| ToolToAIAgentActionError::MissingSkillReference)?;
Ok(AIAgentActionType::ReadSkill(ReadSkillRequest { skill }))
}
/// 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 {
@@ -81,6 +92,22 @@ fn convert_start_agent_v2_harness_type(
.filter(|harness_type| !harness_type.trim().is_empty())
}
/// Maps the proto `Harness` oneof to a client-side string identifier
/// (e.g. "oz", "claude"). Returns `None` for an unset variant.
pub(crate) fn convert_run_agents_harness(harness: Option<&api::Harness>) -> Option<String> {
let variant = harness?.variant.as_ref()?;
Some(
match variant {
api::harness::Variant::Oz(_) => "oz",
api::harness::Variant::ClaudeCode(_) => "claude",
api::harness::Variant::OpenCode(_) => "opencode",
api::harness::Variant::Gemini(_) => "gemini",
api::harness::Variant::Codex(_) => "codex",
}
.to_string(),
)
}
fn convert_start_agent_execution_mode(
execution_mode: Option<api::start_agent::ExecutionMode>,
) -> StartAgentExecutionMode {
@@ -94,8 +121,62 @@ fn convert_start_agent_execution_mode(
}
}
fn convert_run_agents_execution_mode(
execution_mode: Option<api::run_agents::ExecutionMode>,
) -> RunAgentsExecutionMode {
match execution_mode {
Some(api::run_agents::ExecutionMode::Remote(remote)) => RunAgentsExecutionMode::Remote {
environment_id: remote.environment_id,
worker_host: remote.worker_host,
computer_use_enabled: remote.computer_use_enabled,
},
Some(api::run_agents::ExecutionMode::Local(_)) | None => RunAgentsExecutionMode::Local,
}
}
fn convert_run_agents(
run_agents: api::RunAgents,
skill_path_origin: &SkillPathOrigin,
) -> AIAgentActionType {
let api::RunAgents {
summary,
base_prompt,
skills,
model_id,
harness,
agent_run_configs,
execution_mode,
plan_id,
} = run_agents;
AIAgentActionType::RunAgents(RunAgentsRequest {
summary,
base_prompt,
skills: skills
.into_iter()
.filter_map(|skill| skill_reference_from_api_skill_ref(skill, skill_path_origin))
.collect(),
model_id,
harness_type: convert_run_agents_harness(harness.as_ref()).unwrap_or_default(),
execution_mode: convert_run_agents_execution_mode(execution_mode),
agent_run_configs: agent_run_configs
.into_iter()
.map(|config| RunAgentsAgentRunConfig {
name: config.name,
prompt: config.prompt,
title: config.title,
})
.collect(),
plan_id,
// Auth secret is a client-side dispatch concern populated by the
// confirmation card from `CloudAgentSettings.last_selected_auth_secret`
// before Accept. The proto does not carry it.
harness_auth_secret_name: None,
})
}
fn convert_start_agent_v2_execution_mode(
execution_mode: Option<api::start_agent_v2::ExecutionMode>,
skill_path_origin: &SkillPathOrigin,
) -> StartAgentExecutionMode {
match execution_mode.and_then(|execution_mode| execution_mode.mode) {
Some(api::start_agent_v2::execution_mode::Mode::Remote(remote)) => {
@@ -104,7 +185,9 @@ fn convert_start_agent_v2_execution_mode(
skill_references: remote
.skills
.into_iter()
.filter_map(convert_skill_reference)
.filter_map(|skill| {
skill_reference_from_api_skill_ref(skill, skill_path_origin)
})
.collect(),
model_id: remote.model_id,
computer_use_enabled: remote.computer_use_enabled,
@@ -112,6 +195,9 @@ fn convert_start_agent_v2_execution_mode(
harness_type: convert_start_agent_v2_harness_type(remote.harness)
.unwrap_or_default(),
title: remote.title,
// Auth secret is plumbed client-side via `RunAgentsRequest`;
// StartAgentV2 from the server never carries it.
auth_secret_name: None,
}
}
Some(api::start_agent_v2::execution_mode::Mode::Local(local)) => {
@@ -123,16 +209,6 @@ fn convert_start_agent_v2_execution_mode(
}
}
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 {
@@ -167,6 +243,7 @@ pub struct ConversionParams<'a> {
pub task_id: &'a TaskId,
pub current_todo_list: Option<&'a AIAgentTodoList>,
pub active_code_review: Option<&'a CodeReview>,
pub skill_path_origin: &'a SkillPathOrigin,
}
/// Trait for converting an [`api::Message`] to an [`AIAgentOutputMessage`].
@@ -569,7 +646,11 @@ impl ConvertAPIMessageToClientOutputMessage for api::Message {
| api::message::Message::CodeReview(_)
| api::message::Message::ServerEvent(_)
| api::message::Message::InvokeSkill(_)
| api::message::Message::PassiveSuggestionResult(_) => {
| api::message::Message::PassiveSuggestionResult(_)
// Stage 2 plan-card config snapshot: hydrated separately by the
// plan card's `AIDocumentModel` subscription, not via the
// exchange/output stream. No client output message representation.
| api::message::Message::OrchestrationConfigSnapshot(_) => {
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation)
}
}
@@ -600,7 +681,7 @@ trait ConvertAPIToolCallToAIAgentAction {
) -> Result<MaybeAIAgentAction, ToolToAIAgentActionError>;
}
/// Trys to convert an [`api::message::ToolCall`] to an [`AIAgentAction`].
/// Tries 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.
@@ -700,6 +781,7 @@ impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
create_standard_action(request_computer_use.into())
}
api::message::tool_call::Tool::Subagent(subagent) => {
use api::message::tool_call::subagent::conversation_search_metadata::Target;
use api::message::tool_call::subagent::Metadata;
let subagent_type = match subagent.metadata {
Some(Metadata::Cli(_)) => SubagentType::Cli,
@@ -713,14 +795,23 @@ impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
} else {
Some(cs_meta.query)
};
let conversation_id = if cs_meta.conversation_id.is_empty() {
None
} else {
Some(cs_meta.conversation_id)
let (conversation_id, agent_run_id) = match cs_meta.target {
Some(Target::ConversationId(conversation_id))
if !conversation_id.is_empty() =>
{
(Some(conversation_id), None)
}
Some(Target::AgentRunId(agent_run_id)) if !agent_run_id.is_empty() => {
(None, Some(agent_run_id))
}
Some(Target::ConversationId(_))
| Some(Target::AgentRunId(_))
| None => (None, None),
};
SubagentType::ConversationSearch {
query,
conversation_id,
agent_run_id,
}
}
Some(Metadata::WarpDocumentationSearch(_)) => {
@@ -757,6 +848,7 @@ impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
prompt: start_agent.prompt,
execution_mode: convert_start_agent_v2_execution_mode(
start_agent.execution_mode,
params.skill_path_origin,
),
lifecycle_subscription: start_agent.lifecycle_subscription.map(
|subscription| {
@@ -769,6 +861,9 @@ impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
),
})
}
api::message::tool_call::Tool::RunAgents(orchestrate) => {
create_standard_action(convert_run_agents(orchestrate, params.skill_path_origin))
}
api::message::tool_call::Tool::SendMessageToAgent(send_message) => {
create_standard_action(AIAgentActionType::SendMessageToAgent {
addresses: send_message.addresses,
@@ -780,7 +875,7 @@ impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
create_standard_action(insert_review_comments.into())
}
api::message::tool_call::Tool::ReadSkill(read_skill) => {
create_standard_action(read_skill.try_into()?)
create_standard_action(convert_read_skill(read_skill, params.skill_path_origin)?)
}
api::message::tool_call::Tool::FetchConversation(fetch_conversation) => {
create_standard_action(fetch_conversation.into())
@@ -798,6 +893,12 @@ impl ConvertAPIToolCallToAIAgentAction for api::message::ToolCall {
api::message::tool_call::Tool::Server(_) => {
Ok(MaybeAIAgentAction::NoClientRepresentation)
}
api::message::tool_call::Tool::WaitForEvents(payload) => {
create_standard_action(AIAgentActionType::WaitForEvents {
tool_call_id: self.tool_call_id.clone(),
idle_timeout_seconds: payload.idle_timeout_seconds,
})
}
_ => Err(ToolToAIAgentActionError::UnexpectedTool),
}
}
app/src/ai/agent/api/convert_from_tests.rs
+27 -4
View File
@@ -1,3 +1,10 @@
use std::path::PathBuf;
use ai::agent::action::AskUserQuestionType;
use ai::skills::{SkillPathOrigin, SkillReference};
use warp_multi_agent_api as api;
use warp_util::local_or_remote_path::LocalOrRemotePath;
use super::{
convert_api_question, ConversionParams, ConvertAPIMessageToClientOutputMessage,
MaybeAIAgentOutputMessage,
@@ -6,9 +13,6 @@ 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,
@@ -17,6 +21,7 @@ fn start_agent_tool_call_message(
lifecycle_subscription_event_types: Option<Vec<i32>>,
) -> api::Message {
api::Message {
fetched_memories: vec![],
id: "message-id".to_string(),
task_id: "task-id".to_string(),
server_message_data: String::new(),
@@ -63,6 +68,7 @@ fn start_agent_v2_tool_call_message(
lifecycle_subscription_event_types: Option<Vec<i32>>,
) -> api::Message {
api::Message {
fetched_memories: vec![],
id: "message-id".to_string(),
task_id: "task-id".to_string(),
server_message_data: String::new(),
@@ -87,6 +93,7 @@ fn start_agent_v2_tool_call_message(
fn upload_artifact_tool_call_message(path: &str, description: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: "message-id".to_string(),
task_id: "task-id".to_string(),
server_message_data: String::new(),
@@ -140,6 +147,7 @@ fn remote_start_agent_v2_execution_mode(
fn file_artifact_created_message(filepath: &str, description: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: "message-id".to_string(),
task_id: "task-id".to_string(),
server_message_data: String::new(),
@@ -298,6 +306,7 @@ fn converts_start_agent_tool_call_to_action_with_prompt() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -327,6 +336,7 @@ fn converts_local_start_agent_v2_without_harness_type_to_defaults() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -354,6 +364,7 @@ fn converts_upload_artifact_tool_call_to_action() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -377,6 +388,7 @@ fn converts_file_artifact_created_message_with_filename() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -407,6 +419,7 @@ fn converts_start_agent_tool_calls_with_different_prompt_lengths() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("partial conversion should succeed");
let updated_output = updated_message
@@ -414,6 +427,7 @@ fn converts_start_agent_tool_calls_with_different_prompt_lengths() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("updated conversion should succeed");
@@ -442,6 +456,7 @@ fn converts_start_agent_with_explicit_empty_lifecycle_subscription() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -471,6 +486,7 @@ fn converts_start_agent_with_cancelled_and_blocked_lifecycle_subscription() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -503,6 +519,7 @@ fn converts_remote_start_agent_with_environment_id() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -520,6 +537,7 @@ fn converts_remote_start_agent_with_environment_id() {
worker_host: String::new(),
harness_type: String::new(),
title: String::new(),
auth_secret_name: None,
}
);
assert_eq!(lifecycle_subscription, None);
@@ -540,6 +558,7 @@ fn converts_remote_start_agent_v2_with_skill_references() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -552,7 +571,7 @@ fn converts_remote_start_agent_v2_with_skill_references() {
StartAgentExecutionMode::Remote {
environment_id: "env-123".to_string(),
skill_references: vec![
SkillReference::Path("/tmp/SKILL.md".into()),
SkillReference::Path(LocalOrRemotePath::Local(PathBuf::from("/tmp/SKILL.md",))),
SkillReference::BundledSkillId("review-comments".to_string()),
],
model_id: "gpt-test".to_string(),
@@ -560,6 +579,7 @@ fn converts_remote_start_agent_v2_with_skill_references() {
worker_host: "worker-host".to_string(),
harness_type: "claude-code".to_string(),
title: "Remote child".to_string(),
auth_secret_name: None,
}
);
assert_eq!(lifecycle_subscription, None);
@@ -580,6 +600,7 @@ fn converts_local_start_agent_v2_with_harness_type() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("conversion should succeed");
@@ -599,6 +620,7 @@ 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 {
fetched_memories: vec![],
id: "message".to_string(),
task_id: "task".to_string(),
server_message_data: String::new(),
@@ -622,6 +644,7 @@ fn transfer_control_tool_call_converts_to_action_message() {
task_id: &task_id,
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
})
.expect("transfer-control conversion should succeed");
app/src/ai/agent/api/convert_to.rs
+87 -19
View File
@@ -5,14 +5,12 @@ 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,
use crate::ai::agent::{
AIAgentActionResult, AIAgentActionResultType, AIAgentAttachment, AIAgentContext, AIAgentInput,
DriveObjectPayload, MCPContext, PassiveSuggestionResultType, PassiveSuggestionTrigger,
RunningCommand, StaticQueryType, Suggestions, UserQueryMode,
};
use crate::ai::block_context::BlockContext;
fn local_datetime_to_timestamp(timestamp: DateTime<Local>) -> prost_types::Timestamp {
prost_types::Timestamp {
@@ -46,11 +44,6 @@ impl TryFrom<StaticQueryType> for api::request::input::query_with_canned_respons
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())
}
@@ -209,9 +202,9 @@ pub(super) fn convert_input(
)),
});
}
AIAgentInput::SummarizeConversation { prompt } => {
AIAgentInput::SummarizeConversation { prompt, context } => {
return Ok(api::request::Input {
context: None,
context: Some(convert_context(context.as_ref())),
r#type: Some(api::request::input::Type::SummarizeConversation(
api::request::input::SummarizeConversation {
prompt: prompt.unwrap_or_default(),
@@ -435,6 +428,19 @@ fn convert_input_to_user_input(
),
)
}
AIAgentInput::OrchestrationConfigUpdate {
plan_id,
config,
status,
} => Ok(
api::request::input::user_inputs::user_input::Input::OrchestrationConfigUpdate(
api::OrchestrationConfigUpdate {
plan_id,
config: Some(config.to_proto()),
status: status.to_proto(),
},
),
),
AIAgentInput::ResumeConversation { .. } => Err(ConvertToAPITypeError::Ignore),
AIAgentInput::InitProjectRules { .. } => Err(ConvertToAPITypeError::Ignore),
AIAgentInput::CodeReview { .. } => Err(ConvertToAPITypeError::Ignore),
@@ -692,6 +698,12 @@ impl TryFrom<AIAgentActionResult> for api::request::input::user_inputs::user_inp
AIAgentActionResultType::AskUserQuestion(ask_user_question_result) => {
Some(ask_user_question_result.into())
}
AIAgentActionResultType::RunAgents(orchestrate_result) => {
Some(orchestrate_result.try_into()?)
}
AIAgentActionResultType::WaitForEvents(wait_for_events_result) => {
Some(wait_for_events_result.try_into()?)
}
};
Ok(
api::request::input::user_inputs::user_input::Input::ToolCallResult(
@@ -706,6 +718,7 @@ impl TryFrom<AIAgentActionResult> for api::request::input::user_inputs::user_inp
fn convert_context(context: &[AIAgentContext]) -> api::InputContext {
let mut api_context = api::InputContext::default();
let mut git_context = None;
for context in context.iter().cloned() {
match context {
AIAgentContext::Block(block) => {
@@ -789,11 +802,40 @@ fn convert_context(context: &[AIAgentContext]) -> api::InputContext {
}
}
AIAgentContext::Git { head, branch } => {
api_context.git = Some(api::input_context::Git {
head,
branch: branch.unwrap_or_default(),
let api_git_context =
git_context.get_or_insert_with(api::input_context::Git::default);
api_git_context.head = head;
api_git_context.branch = branch.unwrap_or_default();
}
AIAgentContext::Repository { name, owner } => {
let api_git_context =
git_context.get_or_insert_with(api::input_context::Git::default);
api_git_context.repository = Some(api::input_context::git::Repository {
name,
owner: owner.unwrap_or_default(),
});
}
AIAgentContext::PullRequest {
number,
state,
draft,
base_branch,
} => {
if number <= 0 {
continue;
}
let Some(state) = api_pull_request_state(&state, draft) else {
continue;
};
let pull_request = api::input_context::git::PullRequest {
number,
state: state as i32,
base_branch,
};
let api_git_context =
git_context.get_or_insert_with(api::input_context::Git::default);
api_git_context.pull_request = Some(pull_request);
}
AIAgentContext::Skills { skills } => {
api_context.updated_skills_context = Some(api::input_context::SkillsContext {
available_skills: skills
@@ -810,9 +852,34 @@ fn convert_context(context: &[AIAgentContext]) -> api::InputContext {
}
}
}
api_context.git = git_context;
api_context
}
/// Maps a GitHub PR state plus draft flag to the proto `State` enum.
///
/// Returns `None` for unknown states so the caller can skip emitting a
/// `pull_request` sub-message rather than sending `STATE_UNSPECIFIED` to the
/// server.
fn api_pull_request_state(
state: &str,
draft: bool,
) -> Option<api::input_context::git::pull_request::State> {
use api::input_context::git::pull_request::State;
match state.to_ascii_uppercase().as_str() {
"OPEN" => {
if draft {
Some(State::OpenDraft)
} else {
Some(State::Open)
}
}
"CLOSED" => Some(State::Closed),
"MERGED" => Some(State::Merged),
_ => None,
}
}
impl From<Suggestions> for api::Suggestions {
fn from(value: Suggestions) -> Self {
Self {
@@ -940,12 +1007,13 @@ impl From<BlockContext> for api::ExecutedShellCommand {
}
}
/// Trys to convert a [`serde_json::Value`] to a [`prost_types::Value`].
/// Tries 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 std::collections::BTreeMap;
use prost_types::value::Kind::*;
use serde_json::Value::*;
use std::collections::BTreeMap;
Ok(prost_types::Value {
kind: Some(match value {
app/src/ai/agent/api/convert_to_tests.rs
+111 -5
View File
@@ -1,11 +1,111 @@
use crate::ai::agent::task::TaskId;
use crate::ai::agent::{
AIAgentActionResult, AIAgentActionResultType, TransferShellCommandControlToUserResult,
};
use crate::terminal::model::block::BlockId;
use chrono::{DateTime, Utc};
use galaxy_core::command::ExitCode;
use warp_multi_agent_api as api;
use crate::ai::agent::task::TaskId;
use crate::ai::agent::{
AIAgentActionResult, AIAgentActionResultType, AIAgentContext,
TransferShellCommandControlToUserResult,
};
use crate::terminal::model::block::BlockId;
#[test]
fn git_context_converts_repository_and_pull_request_metadata() {
let context = vec![
AIAgentContext::Git {
head: "abc123".to_string(),
branch: Some("feature/repo-pr".to_string()),
},
AIAgentContext::Repository {
name: "warp-internal".to_string(),
owner: Some("warpdotdev".to_string()),
},
AIAgentContext::PullRequest {
number: 42,
state: "OPEN".to_string(),
draft: true,
base_branch: "main".to_string(),
},
];
let api_context = super::convert_context(&context);
let git = api_context.git.expect("expected git context");
assert_eq!(git.head, "abc123");
assert_eq!(git.branch, "feature/repo-pr");
let repository = git.repository.expect("expected repository context");
assert_eq!(repository.name, "warp-internal");
assert_eq!(repository.owner, "warpdotdev");
let pull_request = git.pull_request.expect("expected pull request context");
assert_eq!(pull_request.number, 42);
assert_eq!(
pull_request.state,
api::input_context::git::pull_request::State::OpenDraft as i32
);
assert_eq!(pull_request.base_branch, "main");
}
#[test]
fn git_context_skips_pull_request_metadata_with_invalid_number() {
for number in [0, -1] {
let context = vec![
AIAgentContext::Git {
head: "abc123".to_string(),
branch: Some("feature/repo-pr".to_string()),
},
AIAgentContext::PullRequest {
number,
state: "OPEN".to_string(),
draft: false,
base_branch: "main".to_string(),
},
];
let api_context = super::convert_context(&context);
let git = api_context.git.expect("expected git context");
assert_eq!(git.head, "abc123");
assert_eq!(git.branch, "feature/repo-pr");
assert_eq!(git.pull_request, None);
}
}
#[test]
fn git_context_skips_pull_request_metadata_with_unknown_state() {
let context = vec![
AIAgentContext::Git {
head: "abc123".to_string(),
branch: Some("feature/repo-pr".to_string()),
},
AIAgentContext::PullRequest {
number: 42,
state: "SOMETHING_ELSE".to_string(),
draft: false,
base_branch: "main".to_string(),
},
];
let api_context = super::convert_context(&context);
let git = api_context.git.expect("expected git context");
assert_eq!(git.pull_request, None);
}
#[test]
fn git_context_deserializes_legacy_string_pull_request_number() {
let pull_request = serde_json::from_str::<AIAgentContext>(
r#"{"PullRequest":{"number":"42","state":"OPEN","draft":false,"base_branch":"main"}}"#,
)
.expect("expected legacy serialized pull request context");
let api_context = super::convert_context(&[pull_request]);
let pull_request = api_context
.git
.expect("expected git context")
.pull_request
.expect("expected pull request context");
assert_eq!(pull_request.number, 42);
}
#[test]
fn transfer_control_snapshot_result_converts_to_tool_call_result_input() {
let block_id = BlockId::default();
@@ -51,6 +151,8 @@ fn transfer_control_snapshot_result_converts_to_tool_call_result_input() {
#[test]
fn transfer_control_finished_result_converts_to_tool_call_result_input() {
let block_id = BlockId::default();
let start_ts = DateTime::from(Utc::now());
let completed_ts = DateTime::from(Utc::now());
let input =
api::request::input::user_inputs::user_input::Input::try_from(AIAgentActionResult {
id: "tool_call".to_string().into(),
@@ -60,6 +162,8 @@ fn transfer_control_finished_result_converts_to_tool_call_result_input() {
block_id: block_id.clone(),
output: "done".to_string(),
exit_code: ExitCode::from(17),
start_ts: Some(start_ts),
completed_ts: Some(completed_ts),
},
),
})
@@ -78,6 +182,8 @@ fn transfer_control_finished_result_converts_to_tool_call_result_input() {
assert_eq!(finished.command_id, block_id.to_string());
assert_eq!(finished.output, "done");
assert_eq!(finished.exit_code, 17);
assert_eq!(finished.start_ts, Some(super::local_datetime_to_timestamp(start_ts)));
assert_eq!(finished.finish_ts, Some(super::local_datetime_to_timestamp(completed_ts)));
}
other => panic!("Expected command-finished result, got {other:?}"),
},
app/src/ai/agent/api/impl.rs
+88 -54
View File
@@ -1,15 +1,15 @@
use std::{collections::HashMap, sync::Arc};
use std::collections::HashMap;
use std::sync::Arc;
use crate::{ai::agent::redaction, terminal::model::session::SessionType};
use futures_util::StreamExt;
use galaxy_core::features::FeatureFlag;
use warp_multi_agent_api as api;
use crate::ai::bedrock::translator::{self, TranslatorRequest};
use crate::ai::openai::translator as openai_translator;
use crate::ai::provider::ProviderConfig;
use super::{convert_to::convert_input, ConvertToAPITypeError, RequestParams, ResponseStream};
use super::convert_to::convert_input;
use super::{ConvertToAPITypeError, RequestParams, ResponseStream};
use crate::ai::agent::redaction;
use crate::server::server_api::{AIApiError, ServerApi};
use crate::terminal::model::session::SessionType;
pub async fn generate_multi_agent_output(
provider_config: ProviderConfig,
@@ -53,10 +53,10 @@ pub async fn generate_multi_agent_output(
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 api_keys = api_keys_with_warp_credit_fallback_setting(
params.api_keys,
params.allow_use_of_warp_credits,
);
let mut request = api::Request {
task_context: Some(api::request::TaskContext {
@@ -68,6 +68,7 @@ pub async fn generate_multi_agent_output(
base: params.model.into(),
cli_agent: params.cli_agent_model.into(),
computer_use_agent: params.computer_use_model.into(),
base_model_context_window_limit: params.context_window_limit.unwrap_or(0),
..Default::default()
}),
rules_enabled: params.is_memory_enabled,
@@ -99,7 +100,11 @@ pub async fn generate_multi_agent_output(
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(),
supports_orchestration_v2: supports_orchestration_v2(params.orchestration_enabled),
supports_background_computer_use: FeatureFlag::BackgroundComputerUse.is_enabled()
&& computer_use::background_supported(),
custom_model_providers: params.custom_model_providers,
custom_model_routers: params.custom_model_routers,
}),
metadata: Some(api::request::Metadata {
logging: logging_metadata,
@@ -113,6 +118,8 @@ pub async fn generate_multi_agent_output(
.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 assigned a new id yet).
params
.forked_from_conversation_token
.map(|token| token.as_str().to_string())
@@ -129,41 +136,19 @@ pub async fn generate_multi_agent_output(
mcp_context: params.mcp_context.map(Into::into),
};
let model_id = request
.settings
.as_ref()
.and_then(|s| s.model_config.as_ref())
.map(|mc| mc.base.clone())
.unwrap_or_default();
match provider_config {
ProviderConfig::Bedrock(config) => {
let translator_request = TranslatorRequest {
config,
model_id,
root_task_id: params.root_task_id.clone(),
bedrock_message_history: params.bedrock_message_history.clone(),
bedrock_tool_result_archive: params.bedrock_tool_result_archive.clone(),
bedrock_progressive_summary: params.bedrock_progressive_summary.clone(),
bedrock_messages_sent: params.bedrock_messages_sent.clone(),
};
match translator::execute(translator_request, &mut request).await {
Ok(stream) => {
let output_stream = stream.take_until(cancellation_rx);
Ok(Box::pin(output_stream))
}
Err(e) => {
log::error!("[bedrock] Translator error: {e}");
let err = Arc::new(crate::server::server_api::AIApiError::Stream {
stream_type: "bedrock_converse",
source: anyhow::anyhow!("{e}"),
});
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
let response_stream =
warp_multi_agent_client::generate_multi_agent_output(server_api.as_ref(), &request).await;
match response_stream {
Ok(stream) => {
let output_stream = stream
.then(|result| async {
match result {
Ok(event) => Ok(event),
Err(error) => Err(convert_multi_agent_client_error(error).await),
}
})
.take_until(cancellation_rx);
Ok(Box::pin(output_stream))
}
ProviderConfig::OpenAI(config) => {
let translator_request = openai_translator::TranslatorRequest {
@@ -202,12 +187,53 @@ pub async fn generate_multi_agent_output(
),
});
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
let _ = tx
.send(Err(convert_multi_agent_client_error(e).await))
.await;
Ok(Box::pin(rx))
}
}
}
async fn convert_multi_agent_client_error(
error: warp_multi_agent_client::Error,
) -> Arc<AIApiError> {
let error = match error {
warp_multi_agent_client::Error::Authentication(error)
| warp_multi_agent_client::Error::AmbientHeaders(error) => AIApiError::Other(error),
warp_multi_agent_client::Error::Base64Decode(error) => {
AIApiError::Other(anyhow::Error::from(error))
}
warp_multi_agent_client::Error::ProtobufDecode(error) => {
AIApiError::Other(anyhow::Error::from(error))
}
warp_multi_agent_client::Error::EventSource(error) => {
AIApiError::from_stream_error("GenerateMultiAgentOutput", *error).await
}
};
Arc::new(error)
}
fn api_keys_with_warp_credit_fallback_setting(
api_keys: Option<api::request::settings::ApiKeys>,
allow_use_of_warp_credits: bool,
) -> Option<api::request::settings::ApiKeys> {
match api_keys {
Some(mut api_keys) => {
api_keys.allow_use_of_warp_credits = allow_use_of_warp_credits;
Some(api_keys)
}
None if allow_use_of_warp_credits => Some(api::request::settings::ApiKeys {
allow_use_of_warp_credits: true,
..Default::default()
}),
None => None,
}
}
fn supports_orchestration_v2(orchestration_enabled: bool) -> bool {
orchestration_enabled
}
fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
let mut supported_tools = vec![
api::ToolType::Grep,
@@ -245,7 +271,14 @@ fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
}
}
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.
supported_tools.extend(&[api::ToolType::ReadFiles, api::ToolType::ApplyFileDiffs]);
if FeatureFlag::RemoteCodebaseIndexing.is_enabled() {
supported_tools.push(api::ToolType::SearchCodebase);
}
}
Some(SessionType::WarpifiedRemote { host_id: None }) => {}
}
@@ -264,12 +297,10 @@ fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
}
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);
supported_tools.extend([api::ToolType::RunAgents, api::ToolType::SendMessageToAgent]);
// Declare client-handled wait_for_events so the server doesn't
// fall back to the legacy server-handled form.
supported_tools.push(api::ToolType::WaitForEvents);
}
if FeatureFlag::AskUserQuestion.is_enabled() && params.ask_user_question_enabled {
@@ -299,6 +330,9 @@ fn get_supported_cli_agent_tools(params: &RequestParams) -> Vec<api::ToolType> {
}
Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
supported_cli_agent_tools.push(api::ToolType::ReadFiles);
if FeatureFlag::RemoteCodebaseIndexing.is_enabled() {
supported_cli_agent_tools.push(api::ToolType::SearchCodebase);
}
}
Some(SessionType::WarpifiedRemote { host_id: None }) => {}
}
app/src/ai/agent/api/impl_tests.rs
+124 -5
View File
@@ -1,10 +1,15 @@
use galaxy_core::features::FeatureFlag;
use galaxy_core::HostId;
use warp_multi_agent_api as api;
use super::{
api_keys_with_warp_credit_fallback_setting, get_supported_cli_agent_tools, get_supported_tools,
supports_orchestration_v2,
};
use crate::ai::agent::api::RequestParams;
use crate::ai::blocklist::SessionContext;
use crate::ai::llms::LLMId;
use galaxy_core::features::FeatureFlag;
use warp_multi_agent_api as api;
use super::get_supported_tools;
use crate::terminal::model::session::SessionType;
fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool) -> RequestParams {
let model = LLMId::from("test-model");
@@ -24,11 +29,14 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool
computer_use_model: model,
is_memory_enabled: false,
warp_drive_context_enabled: false,
context_window_limit: None,
mcp_context: None,
planning_enabled: true,
should_redact_secrets: false,
api_keys: None,
allow_use_of_warp_credits_with_byok: false,
custom_model_providers: None,
custom_model_routers: None,
allow_use_of_warp_credits: false,
autonomy_level: api::AutonomyLevel::Supervised,
isolation_level: api::IsolationLevel::None,
web_search_enabled: false,
@@ -47,6 +55,85 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool
}
}
fn request_params_for_remote(host_id: Option<HostId>) -> RequestParams {
let mut params = request_params_with_ask_user_question_enabled(false);
params.session_context =
SessionContext::new_with_session_type_for_test(Some(SessionType::WarpifiedRemote {
host_id,
}));
params
}
#[test]
fn api_keys_with_warp_credit_fallback_setting_returns_none_without_keys_or_fallback() {
let api_keys = api_keys_with_warp_credit_fallback_setting(None, false);
assert!(api_keys.is_none());
}
#[test]
fn api_keys_with_warp_credit_fallback_setting_creates_fallback_only_api_keys() {
let api_keys = api_keys_with_warp_credit_fallback_setting(None, true)
.expect("fallback setting should create ApiKeys");
assert!(api_keys.allow_use_of_warp_credits);
assert!(api_keys.anthropic.is_empty());
assert!(api_keys.openai.is_empty());
assert!(api_keys.google.is_empty());
assert!(api_keys.open_router.is_empty());
assert!(api_keys.aws_credentials.is_none());
}
#[test]
fn api_keys_with_warp_credit_fallback_setting_preserves_existing_keys() {
let api_keys = api_keys_with_warp_credit_fallback_setting(
Some(api::request::settings::ApiKeys {
anthropic: "anthropic-key".to_string(),
openai: String::new(),
google: String::new(),
open_router: String::new(),
grok_oauth_access_token: String::new(),
allow_use_of_warp_credits: false,
aws_credentials: None,
google_cloud_credentials: None,
}),
true,
)
.expect("existing ApiKeys should be preserved");
assert_eq!(api_keys.anthropic, "anthropic-key");
assert!(api_keys.allow_use_of_warp_credits);
}
#[test]
fn supports_orchestration_v2_matches_request_orchestration_setting() {
assert!(supports_orchestration_v2(true));
assert!(!supports_orchestration_v2(false));
}
#[test]
fn supported_tools_include_orchestration_tools_when_orchestration_enabled() {
let mut params = request_params_with_ask_user_question_enabled(false);
params.orchestration_enabled = true;
let supported_tools = get_supported_tools(&params);
assert!(supported_tools.contains(&api::ToolType::RunAgents));
assert!(supported_tools.contains(&api::ToolType::SendMessageToAgent));
assert!(!supported_tools.contains(&api::ToolType::StartAgent));
assert!(!supported_tools.contains(&api::ToolType::StartAgentV2));
}
#[test]
fn supported_tools_omit_orchestration_tools_when_orchestration_disabled() {
let params = request_params_with_ask_user_question_enabled(false);
let supported_tools = get_supported_tools(&params);
assert!(!supported_tools.contains(&api::ToolType::RunAgents));
assert!(!supported_tools.contains(&api::ToolType::SendMessageToAgent));
assert!(!supported_tools.contains(&api::ToolType::StartAgent));
assert!(!supported_tools.contains(&api::ToolType::StartAgentV2));
}
#[test]
fn supported_tools_omits_ask_user_question_when_disabled() {
let params = request_params_with_ask_user_question_enabled(false);
@@ -84,3 +171,35 @@ fn supported_tools_omit_upload_artifact_when_feature_flag_is_disabled() {
assert!(!supported_tools.contains(&api::ToolType::UploadFileArtifact));
}
#[test]
fn remote_supported_tools_include_search_codebase_when_connected_and_feature_flag_is_enabled() {
let _flag = FeatureFlag::RemoteCodebaseIndexing.override_enabled(true);
let params = request_params_for_remote(Some(HostId::new("host".to_string())));
let supported_tools = get_supported_tools(&params);
let supported_cli_agent_tools = get_supported_cli_agent_tools(&params);
assert!(supported_tools.contains(&api::ToolType::SearchCodebase));
assert!(supported_cli_agent_tools.contains(&api::ToolType::SearchCodebase));
}
#[test]
fn remote_supported_tools_omit_search_codebase_when_feature_flag_is_disabled() {
let _flag = FeatureFlag::RemoteCodebaseIndexing.override_enabled(false);
let params = request_params_for_remote(Some(HostId::new("host".to_string())));
let supported_tools = get_supported_tools(&params);
let supported_cli_agent_tools = get_supported_cli_agent_tools(&params);
assert!(!supported_tools.contains(&api::ToolType::SearchCodebase));
assert!(!supported_cli_agent_tools.contains(&api::ToolType::SearchCodebase));
}
#[test]
fn remote_supported_tools_omit_search_codebase_when_remote_is_not_connected() {
let _flag = FeatureFlag::RemoteCodebaseIndexing.override_enabled(true);
let params = request_params_for_remote(None);
let supported_tools = get_supported_tools(&params);
let supported_cli_agent_tools = get_supported_cli_agent_tools(&params);
assert!(!supported_tools.contains(&api::ToolType::SearchCodebase));
assert!(!supported_cli_agent_tools.contains(&api::ToolType::SearchCodebase));
}
app/src/ai/agent/comment.rs
+6 -10
View File
@@ -1,5 +1,5 @@
use crate::code::buffer_location::LocalOrRemotePath;
use crate::code_review::comments::CommentId;
use std::path::PathBuf;
/// The current state of a code review.
#[derive(Debug, Clone, Default)]
@@ -31,18 +31,14 @@ 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 path_component = file_path.path_component();
let file_name = path_component.file_name().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");
let path_component = file_path.path_component();
let file_name = path_component.file_name().unwrap_or("Invalid File Name");
file_name.to_string()
}
(None, _) => self
@@ -101,6 +97,6 @@ impl From<crate::code_review::comments::AttachedReviewCommentTarget> for ReviewD
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ReviewDiff {
pub file_path: Option<PathBuf>,
pub file_path: Option<LocalOrRemotePath>,
pub line_number: Option<usize>,
}
app/src/ai/agent/conversation.rs
+1071 -459
View File
File diff suppressed because it is too large Load Diff
app/src/ai/agent/conversation_tests.rs
+692 -5
View File
@@ -1,12 +1,23 @@
use std::collections::HashMap;
use super::{
artifact_from_fork_proto, AIConversation, AIConversationAutoexecuteMode, AIConversationId,
};
use crate::ai::artifacts::Artifact;
use crate::persistence::model::AgentConversationData;
use ai::api_keys::ApiKeyManager;
use galaxy_core::features::FeatureFlag;
use warp_multi_agent_api as api;
use warpui::{App, SingletonEntity};
use super::{
artifact_from_fork_proto, footer_model_token_usage, AIConversation,
AIConversationAutoexecuteMode, AIConversationId, ConversationStatus, RestoreConversationError,
};
use crate::ai::artifacts::Artifact;
use crate::ai::llms::LLMPreferences;
use crate::auth::auth_manager::AuthManager;
use crate::auth::AuthStateProvider;
use crate::network::NetworkStatus;
use crate::persistence::model::AgentConversationData;
use crate::server::server_api::ServerApiProvider;
use crate::test_util::settings::initialize_settings_for_tests;
use crate::workspaces::user_workspaces::UserWorkspaces;
fn restored_conversation(conversation_data: Option<AgentConversationData>) -> AIConversation {
AIConversation::new_restored(
@@ -24,8 +35,25 @@ fn restored_conversation(conversation_data: Option<AgentConversationData>) -> AI
.unwrap()
}
fn restored_conversation_with_root_description(description: &str) -> AIConversation {
AIConversation::new_restored(
AIConversationId::new(),
vec![api::Task {
id: "root-task".to_string(),
messages: vec![],
dependencies: None,
description: description.to_string(),
summary: String::new(),
server_data: String::new(),
}],
None,
)
.unwrap()
}
fn user_query_message(id: &str, request_id: &str, query: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: "root-task".to_string(),
server_message_data: String::new(),
@@ -44,6 +72,7 @@ fn user_query_message(id: &str, request_id: &str, query: &str) -> api::Message {
fn agent_output_message(id: &str, request_id: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: "root-task".to_string(),
server_message_data: String::new(),
@@ -86,6 +115,43 @@ fn restored_conversation_with_queries(queries: &[&str]) -> AIConversation {
.unwrap()
}
fn initialize_custom_endpoint_usage_test_app(app: &mut App) {
initialize_settings_for_tests(app);
app.add_singleton_model(|_| ServerApiProvider::new_for_test());
app.add_singleton_model(|_| NetworkStatus::new());
app.add_singleton_model(UserWorkspaces::default_mock);
app.add_singleton_model(|_| AuthStateProvider::new_for_test());
app.add_singleton_model(AuthManager::new_for_test);
}
#[allow(deprecated)]
fn custom_endpoint_usage_metadata(
config_key: &str,
total_tokens: u32,
) -> api::response_event::stream_finished::ConversationUsageMetadata {
let category = "primary_agent".to_string();
api::response_event::stream_finished::ConversationUsageMetadata {
context_window_usage: 0.0,
credits_spent: 0.0,
platform_credits_spent: 0.0,
summarized: false,
token_usage: vec![],
tool_usage_metadata: None,
total_input_tokens: 0,
warp_token_usage: HashMap::new(),
byok_token_usage: HashMap::new(),
context_window_segments: Vec::new(),
custom_endpoint_token_usage: HashMap::from([(
config_key.to_string(),
api::response_event::stream_finished::ModelTokenUsage {
model_id: config_key.to_string(),
total_tokens,
token_usage_by_category: HashMap::from([(category, total_tokens)]),
},
)]),
}
}
#[test]
fn latest_user_query_returns_latest_non_empty_user_query() {
let conversation =
@@ -107,6 +173,49 @@ fn latest_user_query_trims_and_skips_empty_queries() {
);
}
#[test]
fn title_uses_root_task_description() {
let conversation = restored_conversation_with_root_description("Root task title");
assert_eq!(conversation.title().as_deref(), Some("Root task title"));
}
#[test]
fn title_falls_back_to_initial_query_when_root_description_is_empty() {
let conversation = restored_conversation_with_queries(&["Initial query"]);
assert_eq!(conversation.title().as_deref(), Some("Initial query"));
}
#[test]
fn reassign_exchange_ids_keeps_exchange_lookup_consistent() {
let mut conversation = restored_conversation_with_queries(&["one", "two"]);
let old_ids: Vec<_> = conversation.all_exchanges().iter().map(|e| e.id).collect();
assert!(!old_ids.is_empty());
// Pre-condition: every original id resolves via the exchange-id index.
for id in &old_ids {
assert!(conversation.exchange_with_id(*id).is_some());
}
conversation.reassign_exchange_ids();
// Reassigning regenerates ids without changing the exchange count, so
// `modify_task` does not rebuild the index; correctness relies on the
// explicit `rebuild_exchange_id_index()` call. The stale ids must be gone.
for id in &old_ids {
assert!(conversation.exchange_with_id(*id).is_none());
}
// Every current id resolves via the rebuilt index.
let new_ids: Vec<_> = conversation.all_exchanges().iter().map(|e| e.id).collect();
assert_eq!(new_ids.len(), old_ids.len());
for id in &new_ids {
assert!(conversation.exchange_with_id(*id).is_some());
}
}
#[test]
fn restored_conversation_defaults_autoexecute_override_when_not_persisted() {
let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
@@ -121,6 +230,416 @@ fn restored_conversation_defaults_autoexecute_override_when_not_persisted() {
);
}
#[test]
fn restored_conversation_uses_persisted_last_event_sequence() {
let conversation_data: AgentConversationData =
serde_json::from_str(r#"{"server_conversation_token":null,"last_event_sequence":42}"#)
.unwrap();
let conversation = restored_conversation(Some(conversation_data));
assert_eq!(conversation.last_event_sequence(), Some(42));
}
#[test]
fn restored_conversation_uses_persisted_remote_child_marker() {
let conversation_data: AgentConversationData =
serde_json::from_str(r#"{"server_conversation_token":null,"is_remote_child":true}"#)
.unwrap();
let conversation = restored_conversation(Some(conversation_data));
assert!(conversation.is_remote_child());
}
#[test]
fn child_conversation_detection_uses_parent_agent_id() {
let conversation_data: AgentConversationData = serde_json::from_str(
r#"{"server_conversation_token":null,"parent_agent_id":"parent-run-id"}"#,
)
.unwrap();
let conversation = restored_conversation(Some(conversation_data));
assert!(conversation.is_child_agent_conversation());
assert_eq!(conversation.parent_conversation_id(), None);
}
/// When the persisted task list is empty (e.g. a child conversation persisted
/// before any server response), restoring via `new_restored_synthesizing_on_empty`
/// must produce a fresh in-progress optimistic root, mirroring
/// `AIConversation::new()`.
#[test]
fn restored_conversation_with_empty_task_list_creates_in_progress_optimistic_root() {
let conversation =
AIConversation::new_restored_synthesizing_on_empty(AIConversationId::new(), vec![], None)
.expect("empty task list must synthesize an optimistic root");
let root_task = conversation
.get_root_task()
.expect("synthesized root task should exist");
assert!(root_task.is_root_task());
assert!(
root_task.source().is_none(),
"synthesized root is optimistic and has no api::Task source"
);
assert!(
!root_task.id().to_string().is_empty(),
"synthesized optimistic root must have a non-empty UUID id"
);
assert_eq!(conversation.status(), &ConversationStatus::InProgress);
assert!(conversation.status_error_message().is_none());
}
#[test]
fn update_cost_and_usage_resolves_custom_endpoint_alias_for_footer_usage() {
App::test((), |mut app| async move {
initialize_custom_endpoint_usage_test_app(&mut app);
ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| {
manager.add_custom_endpoint(
"Endpoint".to_string(),
"https://custom.example".to_string(),
"key".to_string(),
vec![(
"raw-model".to_string(),
Some("Friendly alias".to_string()),
Some("config-key".to_string()),
)],
ctx,
);
});
app.add_singleton_model(LLMPreferences::new);
let mut conversation = AIConversation::new(false, false);
app.read(|ctx| {
conversation
.update_cost_and_usage_for_request(
None,
vec![],
Some(custom_endpoint_usage_metadata("config-key", 6)),
false,
ctx,
)
.expect("custom endpoint usage should update");
});
let usage = conversation
.token_usage()
.iter()
.find(|usage| usage.model_id == "Friendly alias")
.expect("custom endpoint alias should resolve into footer usage");
assert_eq!(usage.custom_endpoint_tokens, 6);
assert_eq!(usage.byok_tokens, 0);
assert_eq!(
usage
.custom_endpoint_token_usage_by_category
.get("primary_agent"),
Some(&6)
);
});
}
#[test]
fn update_cost_and_usage_uses_fallback_label_for_unknown_custom_endpoint() {
App::test((), |mut app| async move {
initialize_custom_endpoint_usage_test_app(&mut app);
app.add_singleton_model(LLMPreferences::new);
let mut conversation = AIConversation::new(false, false);
app.read(|ctx| {
conversation
.update_cost_and_usage_for_request(
None,
vec![],
Some(custom_endpoint_usage_metadata("missing-config-key", 9)),
false,
ctx,
)
.expect("fallback custom endpoint usage should update");
});
let usage = conversation
.token_usage()
.iter()
.find(|usage| usage.model_id == "Custom endpoint")
.expect("unknown custom endpoint usage should use the fallback label");
assert_eq!(usage.custom_endpoint_tokens, 9);
assert_eq!(usage.byok_tokens, 0);
assert_eq!(
usage
.custom_endpoint_token_usage_by_category
.get("primary_agent"),
Some(&9)
);
});
}
#[allow(deprecated)]
#[test]
fn footer_model_token_usage_keeps_custom_endpoint_usage_distinct_from_same_labeled_models() {
App::test((), |mut app| async move {
initialize_custom_endpoint_usage_test_app(&mut app);
ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| {
manager.add_custom_endpoint(
"Endpoint".to_string(),
"https://custom.example".to_string(),
"key".to_string(),
vec![(
"raw-model".to_string(),
Some("Resolved custom".to_string()),
Some("config-key".to_string()),
)],
ctx,
);
});
app.add_singleton_model(LLMPreferences::new);
let category = "primary_agent".to_string();
let usage_metadata = api::response_event::stream_finished::ConversationUsageMetadata {
context_window_usage: 0.0,
credits_spent: 0.0,
platform_credits_spent: 0.0,
summarized: false,
#[allow(deprecated)]
token_usage: vec![],
tool_usage_metadata: None,
total_input_tokens: 0,
warp_token_usage: HashMap::new(),
byok_token_usage: HashMap::from([(
"Resolved custom".to_string(),
api::response_event::stream_finished::ModelTokenUsage {
model_id: "Resolved custom".to_string(),
total_tokens: 4,
token_usage_by_category: HashMap::from([(category.clone(), 4)]),
},
)]),
custom_endpoint_token_usage: HashMap::from([(
"config-key".to_string(),
api::response_event::stream_finished::ModelTokenUsage {
model_id: "config-key".to_string(),
total_tokens: 6,
token_usage_by_category: HashMap::from([(category.clone(), 6)]),
},
)]),
context_window_segments: Vec::new(),
};
let model_usage =
app.read(|ctx| footer_model_token_usage(&usage_metadata, LLMPreferences::as_ref(ctx)));
let byok_usage = model_usage
.iter()
.find(|usage| usage.model_id == "Resolved custom" && usage.byok_tokens == 4)
.expect("existing model usage should be present");
let custom_usage = model_usage
.iter()
.find(|usage| usage.model_id == "Resolved custom" && usage.custom_endpoint_tokens == 6)
.expect("custom endpoint usage should remain distinct");
assert_eq!(model_usage.len(), 2);
assert_eq!(
byok_usage.byok_token_usage_by_category.get(&category),
Some(&4)
);
assert_eq!(
custom_usage
.custom_endpoint_token_usage_by_category
.get(&category),
Some(&6)
);
assert_eq!(byok_usage.warp_tokens, 0);
assert_eq!(custom_usage.warp_tokens, 0);
assert_eq!(custom_usage.byok_tokens, 0);
});
}
#[allow(deprecated)]
#[test]
fn footer_model_token_usage_preserves_unresolved_custom_endpoint_usage_with_fallback_label() {
App::test((), |mut app| async move {
initialize_custom_endpoint_usage_test_app(&mut app);
app.add_singleton_model(LLMPreferences::new);
let category = "primary_agent".to_string();
let usage_metadata = api::response_event::stream_finished::ConversationUsageMetadata {
context_window_usage: 0.0,
credits_spent: 0.0,
platform_credits_spent: 0.0,
summarized: false,
#[allow(deprecated)]
token_usage: vec![],
tool_usage_metadata: None,
total_input_tokens: 0,
warp_token_usage: HashMap::new(),
byok_token_usage: HashMap::new(),
custom_endpoint_token_usage: HashMap::from([(
"missing-config-key".to_string(),
api::response_event::stream_finished::ModelTokenUsage {
model_id: "missing-config-key".to_string(),
total_tokens: 9,
token_usage_by_category: HashMap::from([(category.clone(), 9)]),
},
)]),
context_window_segments: Vec::new(),
};
let model_usage =
app.read(|ctx| footer_model_token_usage(&usage_metadata, LLMPreferences::as_ref(ctx)));
let custom_usage = model_usage
.iter()
.find(|usage| usage.model_id == "Custom endpoint")
.expect("fallback custom endpoint usage should be present");
assert_eq!(model_usage.len(), 1);
assert_eq!(custom_usage.custom_endpoint_tokens, 9);
assert_eq!(custom_usage.byok_tokens, 0);
assert_eq!(
custom_usage
.custom_endpoint_token_usage_by_category
.get(&category),
Some(&9)
);
assert_eq!(custom_usage.warp_tokens, 0);
});
}
/// The legacy `AgentConversationData.root_task_is_optimistic` flag must be
/// ignored on restore. A non-empty task list always produces a real
/// server-backed root regardless of whether the flag is set.
#[test]
fn restored_conversation_ignores_legacy_root_task_is_optimistic_flag_with_non_empty_tasks() {
let conversation_data: AgentConversationData = serde_json::from_str(
r#"{"server_conversation_token":null,"root_task_is_optimistic":true}"#,
)
.unwrap();
let conversation = restored_conversation(Some(conversation_data));
let root_task = conversation
.get_root_task()
.expect("root task should exist");
assert_eq!(root_task.id().to_string(), "root-task");
assert!(root_task.is_root_task());
assert!(
root_task.source().is_some(),
"with a real task list, the legacy optimistic flag must be ignored",
);
}
/// The legacy `root_task_is_optimistic` flag is ignored when restoring an
/// empty task list via `new_restored_synthesizing_on_empty`.
#[test]
fn restored_conversation_ignores_legacy_root_task_is_optimistic_flag_with_empty_tasks() {
let conversation_data: AgentConversationData = serde_json::from_str(
r#"{"server_conversation_token":null,"root_task_is_optimistic":true}"#,
)
.unwrap();
let conversation = AIConversation::new_restored_synthesizing_on_empty(
AIConversationId::new(),
vec![],
Some(conversation_data),
)
.expect("empty task list must synthesize an optimistic root regardless of legacy flag");
let root_task = conversation
.get_root_task()
.expect("synthesized root task should exist");
assert!(root_task.is_root_task());
assert!(root_task.source().is_none());
assert_eq!(conversation.status(), &ConversationStatus::InProgress);
}
/// Strict `new_restored` returns `NoRootTask` for an empty task list.
#[test]
fn new_restored_with_empty_task_list_returns_no_root_task_error() {
let result = AIConversation::new_restored(AIConversationId::new(), vec![], None);
assert!(
matches!(result, Err(RestoreConversationError::NoRootTask)),
"empty task list via strict new_restored must return NoRootTask; got {result:?}",
);
}
/// When multiple parentless tasks exist (e.g. a legacy orphan optimistic
/// stub alongside the real server root), `new_restored` must prefer the
/// candidate whose `messages` is non-empty. Each ordering runs in a loop to
/// surface any nondeterminism in candidate selection.
#[test]
fn test_new_restored_prefers_parentless_task_with_messages_over_empty_stub() {
let stub = api::Task {
id: "optimistic-stub-uuid".to_string(),
messages: vec![],
dependencies: None,
description: String::new(),
summary: String::new(),
server_data: String::new(),
};
let real = api::Task {
id: "server-root-id".to_string(),
messages: vec![user_query_message("user-msg", "request-1", "real query")],
dependencies: None,
description: String::new(),
summary: String::new(),
server_data: String::new(),
};
// Stub appears first in the vec.
for _ in 0..50 {
let conversation = AIConversation::new_restored(
AIConversationId::new(),
vec![stub.clone(), real.clone()],
None,
)
.expect("restore with stub + real parentless tasks must succeed");
let root_task = conversation
.get_root_task()
.expect("restored conversation must have a root task");
assert_eq!(
root_task.id().to_string(),
"server-root-id",
"expected the real (non-empty) parentless task to win when stub is first",
);
let source = root_task
.source()
.expect("chosen root must have api::Task source");
assert!(
!source.messages.is_empty(),
"chosen root must have non-empty messages",
);
}
// Real appears first in the vec.
for _ in 0..50 {
let conversation = AIConversation::new_restored(
AIConversationId::new(),
vec![real.clone(), stub.clone()],
None,
)
.expect("restore with real + stub parentless tasks must succeed");
let root_task = conversation
.get_root_task()
.expect("restored conversation must have a root task");
assert_eq!(
root_task.id().to_string(),
"server-root-id",
"expected the real (non-empty) parentless task to win when real is first",
);
let source = root_task
.source()
.expect("chosen root must have api::Task source");
assert!(
!source.messages.is_empty(),
"chosen root must have non-empty messages",
);
}
}
#[test]
fn cli_agent_transcript_vehicle_is_excluded_from_navigation() {
let conversation = AIConversation::new(false, true);
assert!(conversation.should_exclude_from_navigation());
}
#[test]
fn restored_conversation_defaults_unknown_persisted_autoexecute_override() {
let _flag = FeatureFlag::RememberFastForwardState.override_enabled(true);
@@ -197,3 +716,171 @@ fn fork_artifacts_adds_file_artifacts_to_conversation() {
})
);
}
#[test]
fn waiting_for_events_display_label_is_waiting() {
assert_eq!(
format!("{}", ConversationStatus::WaitingForEvents),
"Waiting"
);
}
/// `is_done` returns true only for `Success | Error | Cancelled`;
/// `WaitingForEvents` and `Blocked` are not done because the run can still
/// resume on its own.
#[test]
fn is_done_only_includes_success_error_cancelled() {
assert!(ConversationStatus::Success.is_done());
assert!(ConversationStatus::Error.is_done());
assert!(ConversationStatus::Cancelled.is_done());
assert!(!ConversationStatus::InProgress.is_done());
assert!(!ConversationStatus::Blocked {
blocked_action: "approve".to_string()
}
.is_done());
assert!(!ConversationStatus::WaitingForEvents.is_done());
}
/// `is_waiting_for_events` is true only for the new variant.
#[test]
fn is_waiting_for_events_returns_true_only_for_waiting_for_events_variant() {
assert!(ConversationStatus::WaitingForEvents.is_waiting_for_events());
assert!(!ConversationStatus::InProgress.is_waiting_for_events());
assert!(!ConversationStatus::Success.is_waiting_for_events());
assert!(!ConversationStatus::Error.is_waiting_for_events());
assert!(!ConversationStatus::Cancelled.is_waiting_for_events());
assert!(!ConversationStatus::Blocked {
blocked_action: "approve".to_string()
}
.is_waiting_for_events());
}
/// A conversation that was yielded via `wait_for_events` at shutdown
/// restores as whatever `derive_status_from_root_task` returns (Success
/// for a cleanly-streamed last exchange). The unresolved tool call stays
/// in the transcript as an orphan; the next outbound request triggers
/// the server's existing supersede mechanism to synthesize the matching
/// `Cancel`. The waiting state itself is not durable across restart.
#[test]
fn restored_conversation_does_not_re_enter_waiting_for_events() {
let conversation_data: AgentConversationData =
serde_json::from_str(r#"{"server_conversation_token":null}"#).unwrap();
let conversation = restored_conversation(Some(conversation_data));
assert_eq!(conversation.status(), &ConversationStatus::Success);
}
fn fetched_memory(
memory_id: &str,
content: &str,
memory_store_id: &str,
source: Option<api::message::fetched_memory::Source>,
) -> api::message::FetchedMemory {
api::message::FetchedMemory {
memory_id: memory_id.to_string(),
content: content.to_string(),
memory_store_id: memory_store_id.to_string(),
source,
}
}
fn conversation_source(conversation_id: &str) -> Option<api::message::fetched_memory::Source> {
Some(api::message::fetched_memory::Source::Conversation(
api::message::fetched_memory::Conversation {
conversation_id: conversation_id.to_string(),
},
))
}
fn restored_conversation_with_memories_per_query(
memories_per_query: Vec<Vec<api::message::FetchedMemory>>,
) -> AIConversation {
let messages = memories_per_query
.into_iter()
.enumerate()
.flat_map(|(index, memories)| {
let request_id = format!("request-{index}");
let query = api::Message {
fetched_memories: memories,
..user_query_message(&format!("user-{index}"), &request_id, "query")
};
[
query,
agent_output_message(&format!("agent-{index}"), &request_id),
]
})
.collect();
AIConversation::new_restored(
AIConversationId::new(),
vec![api::Task {
id: "root-task".to_string(),
messages,
..Default::default()
}],
None,
)
.unwrap()
}
#[test]
fn fetched_memories_is_empty_when_no_message_has_memories() {
let conversation = restored_conversation_with_memories_per_query(vec![vec![]]);
assert_eq!(conversation.fetched_memories(), vec![]);
}
#[test]
fn fetched_memories_preserves_order_across_and_within_messages() {
let conversation = restored_conversation_with_memories_per_query(vec![
vec![
fetched_memory("m1", "first", "store-1", None),
fetched_memory("m2", "second", "store-1", None),
],
vec![fetched_memory("m3", "third", "store-2", None)],
]);
let ids: Vec<String> = conversation
.fetched_memories()
.into_iter()
.map(|memory| memory.memory_id)
.collect();
assert_eq!(ids, vec!["m1", "m2", "m3"]);
}
#[test]
fn fetched_memories_dedupes_keeping_first_position_and_latest_data() {
let conversation = restored_conversation_with_memories_per_query(vec![
vec![
fetched_memory("m1", "old content", "store-1", None),
fetched_memory("m2", "other", "store-1", None),
],
vec![
fetched_memory(
"m1",
"new content",
"store-1",
conversation_source("conversation-1"),
),
fetched_memory("m1", "same memory id different store", "store-2", None),
],
]);
let memories = conversation.fetched_memories();
assert_eq!(
memories,
vec![
fetched_memory(
"m1",
"new content",
"store-1",
conversation_source("conversation-1"),
),
fetched_memory("m2", "other", "store-1", None),
fetched_memory("m1", "same memory id different store", "store-2", None),
]
);
}
app/src/ai/agent/conversation_yaml.rs
+70 -4
View File
@@ -9,11 +9,10 @@ use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use warp_multi_agent_api as api;
use api::message::tool_call::Tool;
use api::message::tool_call_result::Result as ToolCallResultType;
use api::message::Message;
use warp_multi_agent_api as api;
use super::task::helper::{SubagentExt, ToolExt};
@@ -243,7 +242,8 @@ fn write_task_messages(
| Message::SystemQuery(_)
| Message::CodeReview(_)
| Message::ServerEvent(_)
| Message::InvokeSkill(_) => {}
| Message::InvokeSkill(_)
| Message::OrchestrationConfigSnapshot(_) => {}
}
}
Ok(())
@@ -507,6 +507,19 @@ fn write_tool_call_args(out: &mut String, tool: &Tool) {
out.push_str(&format!(" - deleted: {}\n", df.file_path));
}
}
Tool::RunAgents(o) => {
out.push_str(&format!(
"summary: \"{}\"\n",
escape_yaml_string(&o.summary)
));
out.push_str("agents:\n");
for cfg in &o.agent_run_configs {
out.push_str(&format!(
" - name: \"{}\"\n",
escape_yaml_string(&cfg.name)
));
}
}
// No additional args worth serializing.
Tool::ReadShellCommandOutput(_)
| Tool::UseComputer(_)
@@ -519,13 +532,65 @@ fn write_tool_call_args(out: &mut String, tool: &Tool) {
| Tool::InitProject(_)
| Tool::Server(_)
| Tool::Subagent(_)
| Tool::TransferShellCommandControlToUser(_) => {}
| Tool::TransferShellCommandControlToUser(_)
| Tool::WaitForEvents(_) => {}
}
}
/// Writes content from structured tool call results.
fn write_tool_call_result_content(out: &mut String, result: &ToolCallResultType) {
match result {
ToolCallResultType::RunAgentsResult(r) => match &r.outcome {
Some(api::run_agents_result::Outcome::Launched(launched)) => {
out.push_str("status: launched\n");
out.push_str(&format!("agent_count: {}\n", launched.agents.len()));
if !launched.agents.is_empty() {
out.push_str("agents:\n");
for agent in &launched.agents {
out.push_str(&format!(
" - name: \"{}\"\n",
escape_yaml_string(&agent.name)
));
match &agent.result {
Some(api::run_agents_result::agent_outcome::Result::Launched(
launched,
)) => {
out.push_str(" status: launched\n");
out.push_str(&format!(" agent_id: {}\n", launched.agent_id));
}
Some(api::run_agents_result::agent_outcome::Result::Failed(failed)) => {
out.push_str(" status: failed\n");
out.push_str(&format!(
" error: \"{}\"\n",
escape_yaml_string(&failed.error)
));
}
None => {
out.push_str(" status: unknown\n");
}
}
}
out.push_str(
"next_step: \"Use send_message_to_agent with the existing agent_id instead of running agents again.\"\n",
);
}
}
Some(api::run_agents_result::Outcome::Denied(denied)) => {
out.push_str("status: launch_denied\n");
out.push_str(&format!(
"reason: \"{}\"\n",
escape_yaml_string(&denied.reason)
));
}
Some(api::run_agents_result::Outcome::Failure(failure)) => {
out.push_str("status: failure\n");
out.push_str(&format!(
"error: \"{}\"\n",
escape_yaml_string(&failure.error)
));
}
None => {}
},
ToolCallResultType::StartAgentV2(r) => match &r.result {
Some(api::start_agent_v2_result::Result::Success(s)) => {
out.push_str(&format!("agent_id: {}\n", s.agent_id));
@@ -535,6 +600,7 @@ fn write_tool_call_result_content(out: &mut String, result: &ToolCallResultType)
}
None => {}
},
ToolCallResultType::WaitForEvents(_) => {}
ToolCallResultType::RunShellCommand(r) => {
if let Some(res) = &r.result {
use api::run_shell_command_result::Result;
app/src/ai/agent/conversation_yaml_tests.rs
+68 -2
View File
@@ -3,12 +3,11 @@ use std::path::Path;
use warp_multi_agent_api as api;
use super::{base_dir, materialize_tasks_to_yaml};
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)
@@ -22,6 +21,7 @@ fn list_dir_sorted(dir: &Path) -> Vec<String> {
fn make_user_query_message(id: &str, task_id: &str, query: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: String::new(),
@@ -45,6 +45,7 @@ fn make_tool_call_message(
tool: api::message::tool_call::Tool,
) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: String::new(),
@@ -65,6 +66,7 @@ fn make_tool_call_result_message(
result: api::message::tool_call_result::Result,
) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: String::new(),
@@ -133,6 +135,70 @@ fn mixed_message_types_produce_sequentially_indexed_files() {
cleanup_dir(&dir);
}
#[test]
fn run_agents_result_serializes_agent_ids() {
let task_id = "root";
let tasks = vec![create_api_task(
task_id,
vec![make_tool_call_result_message(
"m1",
task_id,
"tc_run_agents",
api::message::tool_call_result::Result::RunAgentsResult(api::RunAgentsResult {
outcome: Some(api::run_agents_result::Outcome::Launched(
api::run_agents_result::Launched {
resolved_model_id: "auto".to_string(),
resolved_harness: Some(api::Harness {
variant: Some(api::harness::Variant::Oz(api::harness::Oz {})),
}),
resolved_execution_mode: Some(
api::run_agents_result::launched::ResolvedExecutionMode::Local(
api::run_agents::Local {},
),
),
agents: vec![
api::run_agents_result::AgentOutcome {
name: "child".to_string(),
result: Some(
api::run_agents_result::agent_outcome::Result::Launched(
api::run_agents_result::LaunchedAgent {
agent_id: "agent-123".to_string(),
},
),
),
},
api::run_agents_result::AgentOutcome {
name: "other".to_string(),
result: Some(
api::run_agents_result::agent_outcome::Result::Failed(
api::run_agents_result::FailedAgent {
error: "failed to start".to_string(),
},
),
),
},
],
},
)),
}),
)],
)];
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("status: launched"));
assert!(content.contains("agent_count: 2"));
assert!(content.contains("name: \"child\""));
assert!(content.contains("agent_id: agent-123"));
assert!(content.contains("name: \"other\""));
assert!(content.contains("error: \"failed to start\""));
assert!(content.contains("Use send_message_to_agent with the existing agent_id"));
cleanup_dir(&dir);
}
#[test]
fn subagent_file_and_subdirectory_share_same_index() {
let root_id = "root";
app/src/ai/agent/icons.rs
+4 -2
View File
@@ -1,6 +1,8 @@
use galaxy_core::ui::{appearance::Appearance, theme::AnsiColorIdentifier};
use galaxy_core::ui::appearance::Appearance;
use galaxy_core::ui::theme::AnsiColorIdentifier;
use crate::ui_components::{blended_colors, icons::Icon};
use crate::ui_components::blended_colors;
use crate::ui_components::icons::Icon;
pub fn todo_list_icon(appearance: &Appearance) -> galaxyui::elements::Icon {
galaxyui::elements::Icon::new(
app/src/ai/agent/linearization_tests.rs
+5 -1
View File
@@ -1,11 +1,13 @@
use std::collections::HashMap;
use super::*;
use warp_multi_agent_api as api;
use super::*;
// Helper function to create a basic message
fn create_message(id: &str, task_id: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: "server_data".to_string(),
@@ -22,6 +24,7 @@ fn create_message(id: &str, task_id: &str) -> api::Message {
fn create_subagent_tool_call_message(id: &str, task_id: &str, subtask_id: &str) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: "server_data".to_string(),
@@ -44,6 +47,7 @@ fn create_subagent_tool_call_message(id: &str, task_id: &str, subtask_id: &str)
// 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 {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: "server_data".to_string(),
app/src/ai/agent/mod.rs
+397 -96
View File
@@ -12,50 +12,48 @@ mod task_store;
pub(super) mod telemetry;
pub(super) mod util;
// Re-export types that were moved to the ai crate.
pub use ai::agent::{action::*, action_result::*, AIAgentCitation, FileLocations};
use galaxy_core::features::FeatureFlag;
#[cfg(test)]
mod suggestion_test;
use crate::ai::block_context::BlockContext;
use crate::ai::blocklist::block::view_impl::output::are_all_text_sections_empty;
use crate::ai::skills::SkillDescriptor;
use crate::code::editor_management::CodeSource;
use crate::code_review::comments::{
AttachedReviewComment as CodeReviewComment, ReviewCommentBatch,
};
use crate::search::slash_command_menu::static_commands::commands;
use crate::server::server_api::AIApiError;
use ai::skills::ParsedSkill;
use chrono::{DateTime, Local, TimeDelta};
use comment::ReviewComment;
use task::TaskId;
pub use telemetry::AIIdentifiers;
use galaxy_editor::render::model::LineCount;
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::fmt::Display;
use std::ops::{AddAssign, Deref, DerefMut, Range};
use std::sync::Arc;
use std::time::Duration;
// Re-export types that were moved to the ai crate.
pub use ai::agent::action::*;
pub use ai::agent::action_result::*;
use ai::agent::orchestration_config::{OrchestrationConfig, OrchestrationConfigStatus};
pub use ai::agent::{AIAgentCitation, FileLocations};
use ai::skills::ParsedSkill;
use chrono::{DateTime, Local, TimeDelta};
use comment::ReviewComment;
use derivative::Derivative;
use markdown_parser::{parse_markdown, FormattedTable, FormattedText, FormattedTextInline};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use session_sharing_protocol::common::ParticipantId;
use task::TaskId;
pub use telemetry::AIIdentifiers;
use uuid::Uuid;
use galaxy_core::channel::ChannelState;
use galaxy_core::features::FeatureFlag;
use warp_editor::render::model::LineCount;
use warp_multi_agent_api::{diff_hunk as diff_hunk_api, AgentEvent, AgentType};
pub use self::api::{MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError};
use super::llms::LLMId;
use crate::ai::block_context::BlockContext;
use crate::ai::blocklist::block::view_impl::output::are_all_text_sections_empty;
use crate::ai::skills::SkillDescriptor;
use crate::ai_assistant::execution_context::WarpAiExecutionContext;
use crate::code::editor_management::CodeSource;
use crate::code_review::comments::{
AttachedReviewComment as CodeReviewComment, ReviewCommentBatch,
};
use crate::search::slash_command_menu::static_commands::commands;
use crate::server::server_api::{AIApiError, DeserializationError};
use crate::terminal::model::block::BlockId;
use crate::terminal::shell::ShellType;
use crate::terminal::view::block_onboarding::onboarding_agentic_suggestions_block::OnboardingChipType;
use crate::TelemetryEvent;
use derivative::Derivative;
use markdown_parser::{parse_markdown, FormattedTable, FormattedText, FormattedTextInline};
use serde::{Deserialize, Serialize};
use session_sharing_protocol::common::ParticipantId;
use super::llms::LLMId;
/// A server supplied ID for a specific AI generated output.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
@@ -84,6 +82,8 @@ impl ServerOutputId {
pub enum CancellationReason {
/// The user explicitly cancelled without providing a follow-up.
ManuallyCancelled,
/// Warp automatically cancelled the local run so it could continue in Cloud Mode.
AutomaticCloudHandoff,
/// The user submitted a follow-up query during streaming which implicitly cancelled the current one.
FollowUpSubmitted {
@@ -99,22 +99,65 @@ pub enum CancellationReason {
// The user deleted the conversation while it was in progress.
Deleted,
/// The long-running command completed while the agent was still streaming.
/// The long-running command completed while the agent was still streaming a response started via inline agent view.
/// This should be treated as a successful completion, not a cancellation.
OptimisticCLISubagentCompletion,
/// Note this is only used for inline agent view (user starting an agent to monitor an already running command),
/// not when CLI subagent monitors a requested command.
CommandFinishedDuringInlineAgentView,
/// The user manually took control of a long-running command away from the agent.
/// The agent conversation is still in progress — it will resume after the command
/// finishes or once the user hands control back. The stream is cancelled only to
/// stop the CLI subagent monitoring loop, not to end the conversation.
CLISubagentUserTakeover,
/// An agent-issued command caused the shell process to exit (e.g. it ran
/// `exit`, or ran a failing command after enabling `set -e`). The in-flight
/// stream/actions are cancelled to stop work, but the conversation is
/// finalized as a terminal `Error` (with a shell-exit message) by the
/// controller rather than reported as a user cancellation.
AgentExitedShell,
}
/// How a [`CancellationReason`] maps to the conversation's resulting status.
/// This is the single source of truth consumed by the stream- and
/// action-cancellation machinery; see [`CancellationReason::conversation_outcome`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CancellationOutcome {
/// Leave the conversation `InProgress`; it will continue on its own (a
/// follow-up request or a resumed long-running command) without further user
/// input.
KeepInProgress,
/// Finalize the conversation as a successful completion (`Success`).
Succeeded,
/// Finalize the conversation as a user cancellation (`Cancelled`).
Cancelled,
/// Terminal, but a dedicated path (not the cancellation machinery) writes the
/// status — the cancellation is only a stop signal and must not stamp a status.
/// Currently used for shell exit, which is finalized as `Error` by
/// `fail_conversation_due_to_shell_exit`. Unlike `KeepInProgress`, the
/// conversation is ending; only the status write is suppressed.
FinalizedExternally,
}
impl Display for CancellationReason {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CancellationReason::ManuallyCancelled => write!(f, "manual cancellation"),
CancellationReason::AutomaticCloudHandoff => write!(f, "automatic cloud handoff"),
CancellationReason::FollowUpSubmitted { .. } => write!(f, "follow-up submission"),
CancellationReason::UserCommandExecuted => write!(f, "user command execution"),
CancellationReason::Reverted => write!(f, "revert"),
CancellationReason::Deleted => write!(f, "deleted"),
CancellationReason::OptimisticCLISubagentCompletion => {
CancellationReason::CommandFinishedDuringInlineAgentView => {
write!(f, "LRC command completed")
}
CancellationReason::CLISubagentUserTakeover => {
write!(f, "CLI subagent user takeover")
}
CancellationReason::AgentExitedShell => {
write!(f, "agent command exited the shell")
}
}
}
}
@@ -139,8 +182,36 @@ impl CancellationReason {
matches!(self, CancellationReason::Reverted)
}
pub fn is_lrc_command_completed(&self) -> bool {
matches!(self, CancellationReason::OptimisticCLISubagentCompletion)
/// How a cancellation reason maps to the
/// conversation's resulting status. Every site that finalizes a cancelled
/// stream or action consults this instead of re-deriving the disposition,
/// so the reason -> status mapping lives in one exhaustive place.
/// Note that sometimes the action result is treated as authoritative for determining
/// conversation status even when there is a cancellation reason (taking priority over this)
pub fn conversation_outcome(&self) -> CancellationOutcome {
match self {
// The conversation continues without further user input (a follow-up
// request or a resumed long-running command drives it forward), so
// its status must stay InProgress.
CancellationReason::FollowUpSubmitted {
is_for_same_conversation: true,
}
| CancellationReason::CLISubagentUserTakeover => CancellationOutcome::KeepInProgress,
// A long-running command finishing (optimistically) or a revert are
// successful completions rather than cancellations.
CancellationReason::CommandFinishedDuringInlineAgentView
| CancellationReason::Reverted => CancellationOutcome::Succeeded,
// The shell died under the agent; a dedicated path finalizes this as a
// terminal `Error`, so the cancellation machinery must not stamp a status.
CancellationReason::AgentExitedShell => CancellationOutcome::FinalizedExternally,
CancellationReason::ManuallyCancelled
| CancellationReason::AutomaticCloudHandoff
| CancellationReason::UserCommandExecuted
| CancellationReason::Deleted
| CancellationReason::FollowUpSubmitted {
is_for_same_conversation: false,
} => CancellationOutcome::Cancelled,
}
}
}
@@ -401,6 +472,9 @@ pub struct OutputModelInfo {
pub model_id: LLMId,
pub display_name: String,
pub is_fallback: bool,
/// When the provider-side prompt cache for this request is expected to
/// expire. `None` means unknown / no cache-expiry info.
pub prompt_cache_expires_at: Option<DateTime<Local>>,
}
impl Display for AIAgentOutput {
@@ -615,9 +689,11 @@ impl AIAgentOutput {
}
/// Represents user visible errors.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[derive(Clone, Debug)]
pub enum RenderableAIError {
QuotaLimit,
QuotaLimit {
user_display_message: Option<String>,
},
ServerOverloaded,
InternalWarpError,
ContextWindowExceeded(String),
@@ -628,16 +704,59 @@ pub enum RenderableAIError {
AwsBedrockCredentialsExpiredOrInvalid {
model_name: String,
},
/// A transient network failure (lost connection or truncated response stream). Carries its
/// own complete user-facing copy; `kind` preserves the structured cause (including the raw
/// API error) so user reports can disambiguate the different causes behind the shared message.
TransientNetworkError {
kind: TransientNetworkErrorKind,
will_attempt_resume: bool,
/// When `will_attempt_resume` is true, this indicates whether we're waiting for network
/// connectivity before attempting the resume.
waiting_for_network: bool,
},
Other {
error_message: String,
will_attempt_resume: bool,
/// When `will_attempt_resume` is true, this indicates whether we're waiting for network
/// connectivity before attempting the resume.
waiting_for_network: bool,
/// True when the error originates from a user-side issue (e.g., model not allowed,
/// blocked due to fraud, plan restriction). Maps the task to FAILED state instead of ERROR.
is_user_error: bool,
},
/// An agent-issued command caused the shell process to exit, so the run
/// cannot continue. Surfaced as a terminal failure (FAILED) rather than a
/// user cancellation.
AgentExitedShell,
}
impl RenderableAIError {
const TRANSIENT_NETWORK_ERROR_MESSAGE: &'static str =
"Warp lost connection while receiving the agent response. This is usually temporary.";
/// User-facing message shown when an agent-issued command exits the shell.
pub const AGENT_EXITED_SHELL_MESSAGE: &'static str =
"The shell exited while the agent was running a command, so the run could not continue. Ensure the agent is not asked to run commands or source scripts that can exit the shell.";
/// Creates a transient network error. `kind` is the structured cause (including the raw API
/// error where one exists), preserved so user reports can disambiguate the different causes
/// behind the shared user-facing copy.
pub fn transient_network_error(
will_attempt_resume: bool,
waiting_for_network: bool,
kind: TransientNetworkErrorKind,
) -> Self {
Self::TransientNetworkError {
kind,
will_attempt_resume,
waiting_for_network,
}
}
fn is_transient_network_transport_error(error: &reqwest::Error) -> bool {
// If reqwest has an HTTP status, the server responded. Preserve the existing generic
// rendering for those failures rather than calling them lost connections.
error.status().is_none()
}
pub fn is_invalid_api_key(&self) -> bool {
matches!(self, Self::InvalidApiKey { .. })
}
@@ -653,20 +772,99 @@ impl RenderableAIError {
Self::Other {
will_attempt_resume: true,
..
} | Self::TransientNetworkError {
will_attempt_resume: true,
..
}
)
}
/// Whether the failed-output UI should be suppressed while an automatic resume is in
/// flight. Release builds stay quiet so transient blips that recover on their own
/// don't surface an alarming error; dogfood builds (Local/Dev) keep the old, more
/// aggressive behavior so developers still see every transport failure.
pub fn should_suppress_during_recovery(&self) -> bool {
self.will_attempt_resume() && !ChannelState::channel().is_dogfood()
}
/// Constructs a generic [`RenderableAIError::Other`] from a message.
/// `is_user_error` selects the task classification (true → FAILED, false →
/// ERROR). The resume/network flags are false: this is for terminal,
/// out-of-band errors that are not auto-resumed.
pub fn other(error_message: impl Into<String>, is_user_error: bool) -> Self {
Self::Other {
error_message: error_message.into(),
will_attempt_resume: false,
waiting_for_network: false,
is_user_error,
}
}
}
impl From<&AIApiError> for RenderableAIError {
fn from(value: &AIApiError) -> Self {
match value {
AIApiError::QuotaLimit => Self::QuotaLimit,
/// The cause behind a [`RenderableAIError::TransientNetworkError`]. Kept structured (rather than
/// collapsed to a free-form string) so user reports preserve the raw error; rendered to text only
/// at display time.
#[derive(Clone, Debug, thiserror::Error)]
pub enum TransientNetworkErrorKind {
/// A lost connection or truncated response stream — the raw underlying API error. Rendered via
/// `Debug` so reports preserve the full structured error rather than its terse `Display`.
#[error("{0:?}")]
Api(Arc<AIApiError>),
/// The response stream completed with an unfinished exchange and no error event.
#[error("stream completed with an unfinished exchange and no error event")]
UnfinishedExchange,
/// The conversation was left in a transient-error state but the last exchange carried no
/// structured error to surface.
#[error("no structured error on the last exchange")]
MissingExchangeError,
}
impl From<&Arc<AIApiError>> for RenderableAIError {
fn from(value: &Arc<AIApiError>) -> Self {
// Non-retryable 4xx errors (403 fraud block, 400 model/plan restriction, etc.)
// are user-originating — map them to a user error so the task reaches FAILED
// state rather than ERROR state.
let is_user_error = !value.is_recoverable();
match value.as_ref() {
AIApiError::QuotaLimit {
user_display_message,
} => Self::QuotaLimit {
user_display_message: user_display_message.clone(),
},
AIApiError::ServerOverloaded => Self::ServerOverloaded,
_ => Self::Other {
AIApiError::Transport(error)
| AIApiError::Deserialization(DeserializationError::Transport(error)) => {
// A transport error with no HTTP status is a lost-connection failure; one that
// carries a status means the server responded, so it gets generic rendering.
if Self::is_transient_network_transport_error(error) {
Self::transient_network_error(
false,
false,
TransientNetworkErrorKind::Api(value.clone()),
)
} else {
Self::Other {
error_message: format!("Request failed with error: {value:?}"),
will_attempt_resume: false,
waiting_for_network: false,
is_user_error,
}
}
}
AIApiError::UnexpectedEof => Self::transient_network_error(
false,
false,
TransientNetworkErrorKind::Api(value.clone()),
),
AIApiError::Deserialization(DeserializationError::Json(_))
| AIApiError::NoContextFound
| AIApiError::ErrorStatus(_, _)
| AIApiError::Other(_)
| AIApiError::Stream { .. } => Self::Other {
error_message: format!("Request failed with error: {value:?}"),
will_attempt_resume: false,
waiting_for_network: false,
is_user_error,
},
}
}
@@ -675,7 +873,15 @@ impl From<&AIApiError> for RenderableAIError {
impl Display for RenderableAIError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::QuotaLimit => write!(f, "Quota limit reached."),
Self::QuotaLimit {
user_display_message,
} => {
if let Some(message) = user_display_message {
write!(f, "{message}")
} else {
write!(f, "Quota limit reached.")
}
}
Self::ServerOverloaded => {
write!(f, "Warp is currently overloaded. Please try again later.")
}
@@ -692,7 +898,15 @@ impl Display for RenderableAIError {
"AWS Bedrock credentials expired or invalid for {model_name}"
)
}
Self::TransientNetworkError { kind, .. } => {
write!(
f,
"{}\n\nDebug info: {kind}",
Self::TRANSIENT_NETWORK_ERROR_MESSAGE
)
}
Self::Other { error_message, .. } => write!(f, "{error_message}"),
Self::AgentExitedShell => write!(f, "{}", Self::AGENT_EXITED_SHELL_MESSAGE),
}
}
}
@@ -717,13 +931,19 @@ impl ProgrammingLanguage {
#[cfg_attr(target_family = "wasm", allow(unused))]
pub fn to_extension(&self) -> Option<&str> {
match self {
// The arms below cover both canonical language names emitted by the agent (e.g.
// "rust", "kotlin") and common markdown code-fence aliases (e.g. "rs", "kt") to keep
// syntax highlighting working when the model uses either. The set of recognized
// languages here is kept in sync with `SUPPORTED_LANGUAGES` in the `languages` crate.
Self::Other(language) => match language.to_lowercase().as_str() {
"rust" => Some("rs"),
"go" => Some("go"),
"python" => Some("py"),
"javascript" => Some("js"),
"typescript" => Some("ts"),
"yaml" => Some("yaml"),
"rust" | "rs" => Some("rs"),
"go" | "golang" => Some("go"),
"python" | "py" => Some("py"),
"javascript" | "js" => Some("js"),
"typescript" | "ts" => Some("ts"),
"jsx" => Some("jsx"),
"tsx" => Some("tsx"),
"yaml" | "yml" => Some("yaml"),
"cpp" | "c++" => Some("cpp"),
"java" => Some("java"),
"groovy" => Some("java"),
@@ -733,17 +953,23 @@ impl ProgrammingLanguage {
"css" => Some("css"),
"c" => Some("c"),
"json" => Some("json"),
"hcl" => Some("hcl"),
"jq" => Some("jq"),
"hcl" | "terraform" | "tf" => Some("hcl"),
"lua" => Some("lua"),
"ruby" => Some("rb"),
"ruby" | "rb" => Some("rb"),
"php" => Some("php"),
"toml" => Some("toml"),
"swift" => Some("swift"),
"kotlin" => Some("kt"),
"kotlin" | "kt" => Some("kt"),
"powershell" => Some("ps1"),
"elixir" => Some("exs"),
"scala" => Some("scala"),
"sql" => Some("sql"),
"objective-c" | "objc" => Some("m"),
"starlark" => Some("bzl"),
"xml" => Some("xml"),
"vue" => Some("vue"),
"dockerfile" | "docker" | "containerfile" => Some("dockerfile"),
_ => None,
},
Self::Shell(ShellType::PowerShell) => Some("ps1"),
@@ -1513,9 +1739,12 @@ pub enum SubagentType {
Summarization,
ConversationSearch {
query: Option<String>,
/// The ID of the conversation being searched. None when searching the
/// current conversation.
/// Search targets are mutually exclusive; at most one of `conversation_id` or
/// `agent_run_id` should be populated for a single conversation search subagent.
/// The ID of the conversation being searched.
conversation_id: Option<String>,
/// The ID of the agent run being searched.
agent_run_id: Option<String>,
},
WarpDocumentationSearch,
Unknown,
@@ -2010,6 +2239,30 @@ pub enum AIAgentContext {
branch: Option<String>,
},
/// Information about the git repository in the current working directory.
Repository {
/// The repository name (e.g. "warp-internal").
name: String,
/// The repository owner/organization (e.g. "warpdotdev"), if determinable from the remote URL.
owner: Option<String>,
},
/// Information about the GitHub pull request associated with the current branch.
PullRequest {
/// The pull request number.
#[serde(default, deserialize_with = "deserialize_pull_request_number")]
number: i32,
/// The pull request state (for example, `OPEN`, `MERGED`, or `CLOSED`).
#[serde(default)]
state: String,
/// Whether the pull request is marked as draft.
#[serde(default)]
draft: bool,
/// The pull request's base branch.
#[serde(default)]
base_branch: String,
},
/// List of available skills is provided to the agent during initialization
/// or when updated.
Skills {
@@ -2020,6 +2273,37 @@ pub enum AIAgentContext {
Block(Box<BlockContext>),
}
fn deserialize_pull_request_number<'de, D>(deserializer: D) -> Result<i32, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = serde_json::Value::deserialize(deserializer)?;
match value {
serde_json::Value::Null => Ok(0),
serde_json::Value::String(s) => {
if !s.chars().all(|c| c.is_ascii_digit()) {
return Ok(0);
}
Ok(s.parse()
.ok()
.filter(|number| *number > 0)
.unwrap_or_default())
}
serde_json::Value::Number(n) => {
let Some(number) = n.as_i64() else {
return Ok(0);
};
Ok(i32::try_from(number)
.ok()
.filter(|number| *number > 0)
.unwrap_or_default())
}
value => Err(serde::de::Error::custom(format!(
"expected string or number for pull request number, got {value}"
))),
}
}
#[derive(Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct ImageContext {
/// Base64-encoded image data.
@@ -2229,16 +2513,12 @@ pub enum StaticQueryType {
Code,
Deploy,
SomethingElse,
CustomOnboardingRequest,
EvaluationSuite,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[allow(clippy::enum_variant_names)]
pub enum EntrypointType {
Onboarding {
chip_type: OnboardingChipType,
},
PromptSuggestion {
is_static: bool,
is_coding: bool,
@@ -2327,6 +2607,31 @@ pub enum UserQueryMode {
Orchestrate,
}
pub fn extract_user_query_mode(query: String) -> (String, UserQueryMode) {
if let Some(query) = commands::strip_command_prefix(&query, commands::PLAN_NAME) {
(query, UserQueryMode::Plan)
} else if let Some(query) = commands::strip_command_prefix(&query, commands::ORCHESTRATE_NAME) {
(query, UserQueryMode::Orchestrate)
} else {
(query, UserQueryMode::Normal)
}
}
/// Reconstructs the display form of a user query that has been stripped via
/// [`extract_user_query_mode`], by re-prepending the slash-command prefix
/// associated with [`UserQueryMode`].
///
/// This is the inverse of [`extract_user_query_mode`] and the canonical way
/// for UI to render a stored `(mode, query)` pair so the displayed prompt
/// always matches what the user originally submitted.
pub fn display_user_query_with_mode(mode: UserQueryMode, query: &str) -> String {
match mode {
UserQueryMode::Normal => query.to_owned(),
UserQueryMode::Plan => format!("{} {query}", commands::PLAN.name),
UserQueryMode::Orchestrate => format!("{} {query}", commands::ORCHESTRATE.name),
}
}
// TODO(zachbai): Refactor this to consolidate with `LongRunningCommandSnapshot` and `Snapshot`
// variants of `ReadShellCommandOutputResult` and `WriteToLongRunningShellCommandResult`.
#[derive(Clone, Debug, PartialEq)]
@@ -2422,6 +2727,7 @@ pub enum AIAgentInput {
SummarizeConversation {
prompt: Option<String>,
context: Arc<[AIAgentContext]>,
},
/// Invoke a skill. The skill content is passed as instructions to the agent.
@@ -2468,6 +2774,15 @@ pub enum AIAgentInput {
suggestion: PassiveSuggestionResultType,
context: Arc<[AIAgentContext]>,
},
/// Piggybacked orchestration config update from the plan card.
/// Sent on the next outbound request after the user edits the
/// config block or toggles approval.
OrchestrationConfigUpdate {
plan_id: String,
config: OrchestrationConfig,
status: OrchestrationConfigStatus,
},
}
/// Data for a single message received by an agent from another agent.
@@ -2525,7 +2840,7 @@ impl Display for AIAgentInput {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::UserQuery { .. } => {
write!(f, "UserQuery: {}", self.user_query().unwrap_or_default())
write!(f, "UserQuery: {}", self.display_query().unwrap_or_default())
}
Self::AutoCodeDiffQuery { query, .. } => {
write!(f, "AutoCodeDiffQuery: {query}")
@@ -2561,24 +2876,23 @@ impl Display for AIAgentInput {
write!(f, "EventsFromAgents({} events)", events.len())
}
Self::PassiveSuggestionResult { .. } => write!(f, "PassiveSuggestionResult"),
Self::OrchestrationConfigUpdate { .. } => write!(f, "OrchestrationConfigUpdate"),
}
}
}
impl AIAgentInput {
pub fn user_query(&self) -> Option<String> {
/// Display text for any input that surfaces a prompt-like query in the UI
/// (typed queries, slash commands, skill invocations, etc.). Unlike
/// [`Self::is_user_query`], which strictly matches the `UserQuery` variant,
/// this returns `Some` for several input variants.
pub fn display_query(&self) -> Option<String> {
match self {
Self::UserQuery {
query,
user_query_mode,
..
} => match user_query_mode {
UserQueryMode::Plan => Some(format!("{} {query}", commands::PLAN.name)),
UserQueryMode::Orchestrate => {
Some(format!("{} {query}", commands::ORCHESTRATE.name))
}
UserQueryMode::Normal => Some(query.clone()),
},
} => Some(display_user_query_with_mode(*user_query_mode, query)),
Self::CreateNewProject { query, .. } => Some(query.clone()),
Self::CloneRepository {
clone_repo_url: url,
@@ -2624,7 +2938,8 @@ impl AIAgentInput {
| Self::StartFromAmbientRunPrompt { .. }
| Self::MessagesReceivedFromAgents { .. }
| Self::EventsFromAgents { .. }
| Self::PassiveSuggestionResult { .. } => None,
| Self::PassiveSuggestionResult { .. }
| Self::OrchestrationConfigUpdate { .. } => None,
}
}
@@ -2634,7 +2949,7 @@ impl AIAgentInput {
&self,
initial_conversation_query: Option<&String>,
) -> Option<String> {
let mut query = self.user_query()?;
let mut query = self.display_query()?;
if self
.user_query_mode()
.is_none_or(|mode| matches!(mode, UserQueryMode::Normal))
@@ -2719,9 +3034,10 @@ impl AIAgentInput {
| Self::InvokeSkill { context, .. }
| Self::StartFromAmbientRunPrompt { context, .. }
| Self::PassiveSuggestionResult { context, .. } => Some(context),
Self::SummarizeConversation { .. }
| Self::MessagesReceivedFromAgents { .. }
| Self::EventsFromAgents { .. } => None,
Self::SummarizeConversation { context, .. } => Some(context),
Self::MessagesReceivedFromAgents { .. }
| Self::EventsFromAgents { .. }
| Self::OrchestrationConfigUpdate { .. } => None,
}
}
@@ -2752,7 +3068,8 @@ impl AIAgentInput {
| Self::StartFromAmbientRunPrompt { .. }
| Self::MessagesReceivedFromAgents { .. }
| Self::EventsFromAgents { .. }
| Self::PassiveSuggestionResult { .. } => None,
| Self::PassiveSuggestionResult { .. }
| Self::OrchestrationConfigUpdate { .. } => None,
}
}
@@ -2859,7 +3176,7 @@ impl AIAgentExchange {
let user_queries: Vec<String> = self
.input
.iter()
.filter_map(|input| input.user_query())
.filter_map(|input| input.display_query())
.collect();
user_queries.join("\n")
}
@@ -2909,7 +3226,9 @@ impl AIAgentExchange {
}
pub fn has_user_query(&self) -> bool {
self.input.iter().any(|input| input.user_query().is_some())
self.input
.iter()
.any(|input| input.display_query().is_some())
}
pub fn has_accepted_file_edit(&self) -> bool {
@@ -2970,25 +3289,7 @@ pub struct RequestMetadata {
pub is_auto_resume_after_error: bool,
}
/// A globally unique ID for a suggested objects.
///
/// This is used for telemetry purposes to track and connect both:
/// - Suggested objects generated by the AI agent
/// - The corresponding objects stored in the cloud (if the suggestion was accepted)
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
pub struct SuggestedLoggingId(String);
impl Display for SuggestedLoggingId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<String> for SuggestedLoggingId {
fn from(value: String) -> Self {
Self(value)
}
}
pub use cloud_object_models::SuggestedLoggingId;
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct SuggestedRule {
@@ -3034,5 +3335,5 @@ impl Suggestions {
}
#[cfg(test)]
#[path = "mod_test.rs"]
#[path = "mod_tests.rs"]
mod tests;
app/src/ai/agent/mod_test.rs → app/src/ai/agent/mod_tests.rs
+164 -4
View File
@@ -1,15 +1,18 @@
use std::ops::Range;
use std::sync::Arc;
use anyhow::anyhow;
use markdown_parser::{FormattedText, FormattedTextFragment, FormattedTextLine};
use warp_multi_agent_api::{FileContent, FileContentLineRange};
use crate::ai::agent::{
AIAgentOutput, AIAgentOutputMessage, AIAgentOutputMessageType, AIAgentText, AIAgentTextSection,
AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram, AnyFileContent,
FileContext, FormattedTextWrapper, MessageId, ProgrammingLanguage,
AIAgentContext, AIAgentOutput, AIAgentOutputMessage, AIAgentOutputMessageType, AIAgentText,
AIAgentTextSection, AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram,
AnyFileContent, FileContext, FormattedTextWrapper, MessageId, ProgrammingLanguage,
RenderableAIError, TransientNetworkErrorKind,
};
use crate::server::server_api::AIApiError;
use crate::terminal::shell::ShellType;
use markdown_parser::{FormattedText, FormattedTextFragment, FormattedTextLine};
fn to_range(range: Range<u32>) -> Option<FileContentLineRange> {
Some(FileContentLineRange {
@@ -46,6 +49,80 @@ fn formatted_text_wrapper_preserves_content() {
assert_eq!(ft.lines.len(), 2);
}
fn deserialize_pull_request_number_from_json(number_json: &str) -> serde_json::Result<i32> {
let context = serde_json::from_str::<AIAgentContext>(&format!(
r#"{{"PullRequest":{{"number":{number_json}}}}}"#
))?;
match context {
AIAgentContext::PullRequest { number, .. } => Ok(number),
other => panic!("expected pull request context, got {other:?}"),
}
}
#[test]
fn pull_request_number_deserializer_accepts_positive_number_and_string() {
assert_eq!(deserialize_pull_request_number_from_json("42").unwrap(), 42);
assert_eq!(
deserialize_pull_request_number_from_json(r#""42""#).unwrap(),
42
);
}
#[test]
fn pull_request_number_deserializer_defaults_invalid_numbers() {
for number_json in ["null", "0", "-1", "1.5", "2147483648", r#""""#, r#""abc""#] {
assert_eq!(
deserialize_pull_request_number_from_json(number_json).unwrap(),
0,
"expected {number_json} to deserialize to default pull request number",
);
}
}
#[test]
fn pull_request_number_deserializer_rejects_unsupported_json_types() {
for number_json in ["true", "[]", "{}"] {
assert!(
deserialize_pull_request_number_from_json(number_json).is_err(),
"expected {number_json} to fail deserialization",
);
}
}
#[test]
fn transient_network_error_includes_user_facing_message_and_debug_details() {
let error = RenderableAIError::transient_network_error(
false,
false,
TransientNetworkErrorKind::Api(Arc::new(AIApiError::Other(anyhow!("connection reset")))),
);
let rendered = error.to_string();
assert!(
rendered.starts_with(
"Warp lost connection while receiving the agent response. This is usually temporary.\n\nDebug info: "
),
"unexpected rendering: {rendered}"
);
// The raw underlying API error must survive into the debug section.
assert!(
rendered.contains("connection reset"),
"raw error detail should surface in debug info: {rendered}"
);
assert!(!error.will_attempt_resume());
}
#[test]
fn transient_network_error_reports_pending_resume() {
let error = RenderableAIError::transient_network_error(
true,
false,
TransientNetworkErrorKind::Api(Arc::new(AIApiError::Other(anyhow!("connection reset")))),
);
assert!(error.will_attempt_resume());
}
#[test]
fn test_convert_files() {
let a = FileContext::new(
@@ -152,6 +229,86 @@ fn test_programming_language_from_string() {
);
}
#[test]
fn test_programming_language_to_extension() {
// Each entry is (markdown language token, expected extension). The expected extension
// must resolve back to a recognized language via `languages::language_by_filename` so that
// syntax highlighting is applied to the AI block.
let cases: &[(&str, &str)] = &[
// Canonical names.
("rust", "rs"),
("go", "go"),
("python", "py"),
("javascript", "js"),
("typescript", "ts"),
("yaml", "yaml"),
("cpp", "cpp"),
("java", "java"),
("c#", "cs"),
("csharp", "cs"),
("html", "html"),
("css", "css"),
("c", "c"),
("json", "json"),
("hcl", "hcl"),
("lua", "lua"),
("ruby", "rb"),
("php", "php"),
("toml", "toml"),
("swift", "swift"),
("kotlin", "kt"),
("powershell", "ps1"),
("elixir", "exs"),
("scala", "scala"),
("sql", "sql"),
// Languages newly covered by this fix — previously fell through to None and rendered
// without syntax highlighting in AI blocks even though the `languages` crate supports them.
("jsx", "jsx"),
("tsx", "tsx"),
("xml", "xml"),
("vue", "vue"),
("dockerfile", "dockerfile"),
("starlark", "bzl"),
("objective-c", "m"),
("objc", "m"),
// Common markdown code-fence aliases.
("rs", "rs"),
("golang", "go"),
("py", "py"),
("js", "js"),
("ts", "ts"),
("yml", "yaml"),
("c++", "cpp"),
("rb", "rb"),
("kt", "kt"),
("terraform", "hcl"),
("tf", "hcl"),
("docker", "dockerfile"),
("containerfile", "dockerfile"),
];
for (token, expected_extension) in cases {
let language = ProgrammingLanguage::from((*token).to_string());
assert_eq!(
language.to_extension(),
Some(*expected_extension),
"expected to_extension({token:?}) to be Some({expected_extension:?})",
);
}
// PowerShell remains the only Shell variant whose extension is exposed; this preserves
// existing behavior for the other Shell variants which are intentionally not extended here.
assert_eq!(
ProgrammingLanguage::Shell(ShellType::PowerShell).to_extension(),
Some("ps1"),
);
// Unrecognized tokens still return None.
assert_eq!(
ProgrammingLanguage::Other("definitely-not-a-language".to_string()).to_extension(),
None,
);
}
#[test]
fn format_for_copy_preserves_visual_markdown_sections() {
let output = AIAgentOutput {
@@ -189,3 +346,6 @@ fn format_for_copy_preserves_visual_markdown_sections() {
"Intro\n![Diagram](./diagram.png)\n```mermaid\ngraph TD\nA --> B\n```"
);
}
#[path = "suggestions_tests.rs"]
mod suggestions;
app/src/ai/agent/redaction.rs
+13 -6
View File
@@ -1,15 +1,14 @@
use std::sync::Arc;
use super::super::blocklist::block::secret_redaction::{
find_secrets_in_text, SECRET_REDACTION_REPLACEMENT_CHARACTER,
};
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)
@@ -53,10 +52,11 @@ pub(crate) fn redact_inputs(inputs: &mut [AIAgentInput]) {
| AIAgentInput::StartFromAmbientRunPrompt { context, .. } => {
redact_context(Arc::make_mut(context));
}
AIAgentInput::SummarizeConversation { prompt } => {
AIAgentInput::SummarizeConversation { prompt, context } => {
if let Some(p) = prompt {
redact_secrets(p);
}
redact_context(Arc::make_mut(context));
}
AIAgentInput::CreateEnvironment { context, .. } => {
redact_context(Arc::make_mut(context));
@@ -105,7 +105,8 @@ pub(crate) fn redact_inputs(inputs: &mut [AIAgentInput]) {
}
// No user-provided text to redact in inter-agent relay inputs.
AIAgentInput::MessagesReceivedFromAgents { .. }
| AIAgentInput::EventsFromAgents { .. } => {}
| AIAgentInput::EventsFromAgents { .. }
| AIAgentInput::OrchestrationConfigUpdate { .. } => {}
AIAgentInput::ActionResult { result, context } => {
redact_context(Arc::make_mut(context));
match &mut result.result {
@@ -260,6 +261,10 @@ pub(crate) fn redact_inputs(inputs: &mut [AIAgentInput]) {
AIAgentActionResultType::AskUserQuestion(result) => {
redact_ask_user_question_result(result);
}
// Orchestrate results contain agent IDs / canonical error
// strings only; no user-provided text to redact.
AIAgentActionResultType::RunAgents(_)
| AIAgentActionResultType::WaitForEvents(_) => {}
}
}
AIAgentInput::FetchReviewComments { repo_path, context } => {
@@ -344,6 +349,8 @@ fn redact_context(context: &mut [AIAgentContext]) {
| AIAgentContext::Codebase { .. }
| AIAgentContext::ProjectRules { .. }
| AIAgentContext::Git { .. }
| AIAgentContext::Repository { .. }
| AIAgentContext::PullRequest { .. }
| AIAgentContext::File(_)
| AIAgentContext::Skills { .. } => {}
}
app/src/ai/agent/suggestion_test.rs → app/src/ai/agent/suggestions_tests.rs
View File
app/src/ai/agent/task.rs
+73 -60
View File
@@ -1,44 +1,35 @@
pub mod helper;
pub mod transaction;
use std::{
collections::{HashMap, HashSet},
fmt::Display,
ops::Deref,
};
use std::collections::{HashMap, HashSet};
use std::fmt::Display;
use std::ops::Deref;
use chrono::DateTime;
use ai::skills::SkillPathOrigin;
use field_mask::{FieldMaskError, FieldMaskOperation};
use helper::{MessageExt, SubagentExt, ToolCallExt};
use itertools::Itertools;
use prost_types::FieldMask;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use warp_multi_agent_api::{
self as api,
message::{tool_call::subagent::Metadata, Message},
};
use warp_multi_agent_api::message::tool_call::subagent::Metadata;
use warp_multi_agent_api::message::Message;
use warp_multi_agent_api::{self as api};
use crate::{
ai::{
agent::comment::CodeReview,
document::ai_document_model::{AIDocumentId, AIDocumentVersion},
},
server::datetime_ext::DateTimeExt,
terminal::model::block::BlockId,
AIAgentTodoList,
use super::api::convert_conversation::convert_tool_call_result_to_input;
use super::api::{
user_inputs_from_messages, ConversionParams, ConvertAPIMessageToClientOutputMessage,
};
use super::comment::CodeReview;
use super::conversation::{context_in_exchanges, update_todo_list_from_todo_op};
use super::{
api::{
convert_conversation::convert_tool_call_result_to_input, user_inputs_from_messages,
ConversionParams, ConvertAPIMessageToClientOutputMessage,
},
conversation::{context_in_exchanges, update_todo_list_from_todo_op},
AIAgentContext, AIAgentExchange, AIAgentExchangeId, AIAgentOutput, AIAgentOutputMessage,
AIAgentOutputStatus, MaybeAIAgentOutputMessage, MessageId, MessageToAIAgentOutputMessageError,
Shared,
};
use crate::ai::document::ai_document_model::{AIDocumentId, AIDocumentVersion};
use crate::terminal::model::block::BlockId;
use crate::AIAgentTodoList;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TaskId(String);
@@ -128,7 +119,6 @@ struct ServerTask {
}
mod optimistic {
use crate::terminal::model::block::BlockId;
#[derive(Debug, Clone)]
pub(super) struct CLIAgentSubtask {
@@ -176,6 +166,12 @@ pub struct Task {
/// List of `AIAgentExchange`s corresponding to messages contained in this task.
exchanges: Vec<AIAgentExchange>,
}
#[derive(Clone, Copy)]
pub(super) struct TaskMessageContext<'a> {
pub(super) current_todo_list: Option<&'a AIAgentTodoList>,
pub(super) active_code_review: Option<&'a CodeReview>,
pub(super) skill_path_origin: &'a SkillPathOrigin,
}
impl Task {
pub(super) fn new_optimistic_root() -> Self {
@@ -236,6 +232,7 @@ impl Task {
parent_task: Option<&api::Task>,
current_todo_list: Option<&AIAgentTodoList>,
active_code_review: Option<&CodeReview>,
skill_path_origin: &SkillPathOrigin,
) -> Result<Self, UpgradeOptimisticTaskError> {
match self.data {
TaskImpl::Optimistic(optimistic::Task::Root) => {
@@ -281,8 +278,11 @@ impl Task {
if let Err(e) = self.update_exchange_from_messages(
messages,
exchange_id,
current_todo_list,
active_code_review,
TaskMessageContext {
current_todo_list,
active_code_review,
skill_path_origin,
},
false,
) {
log::error!(
@@ -315,7 +315,8 @@ impl Task {
parent_task: &api::Task,
existing_exchange: &AIAgentExchange,
current_todo_list: Option<&AIAgentTodoList>,
current_comment_state: Option<&CodeReview>,
active_code_review: Option<&CodeReview>,
skill_path_origin: &SkillPathOrigin,
should_convert_input_messages: bool,
) -> Self {
let subagent_call_and_id = parent_task.messages.iter().find_map(|message| {
@@ -330,7 +331,7 @@ impl Task {
input: vec![],
output_status: AIAgentOutputStatus::Streaming { output: None },
added_message_ids: Default::default(),
start_time: DateTime::now().into(),
start_time: chrono::Local::now(),
finish_time: None,
time_to_first_token_ms: None,
working_directory: existing_exchange.working_directory.clone(),
@@ -368,8 +369,11 @@ impl Task {
me.update_exchange_from_messages(
messages_clone,
new_exchange_id,
current_todo_list,
current_comment_state,
TaskMessageContext {
current_todo_list,
active_code_review,
skill_path_origin,
},
should_convert_input_messages,
)
.expect("Exchange exists and output is in 'streaming' state.");
@@ -462,7 +466,7 @@ impl Task {
input: vec![],
output_status: AIAgentOutputStatus::Streaming { output: None },
added_message_ids: Default::default(),
start_time: DateTime::now().into(),
start_time: chrono::Local::now(),
finish_time: None,
time_to_first_token_ms: None,
working_directory: existing_exchange.working_directory.clone(),
@@ -603,6 +607,19 @@ impl Task {
self.try_get_source().ok()
}
pub(super) fn source_for_persistence(&self) -> Option<api::Task> {
match &self.data {
TaskImpl::Server(server_data) => Some(server_data.source.clone()),
// Optimistic root tasks have a client-generated UUID and no
// server-side identity yet. Persisting a stub `api::Task` for them
// produces an orphan row in `agent_tasks` that survives the later
// server-side upgrade and breaks restore by competing with the
// real server root for parentless-task selection. See QUALITY-774.
TaskImpl::Optimistic(optimistic::Task::Root) => None,
TaskImpl::Optimistic(optimistic::Task::CLIAgent(_)) => None,
}
}
pub fn messages(&self) -> impl Iterator<Item = &api::Message> {
self.source()
.into_iter()
@@ -677,8 +694,7 @@ impl Task {
&mut self,
messages: Vec<api::Message>,
exchange_id: AIAgentExchangeId,
current_todo_list: Option<&AIAgentTodoList>,
current_comments: Option<&CodeReview>,
message_context: TaskMessageContext<'_>,
should_convert_input_messages: bool,
) -> Result<(), UpdateTaskError> {
if self.source().is_none() {
@@ -687,8 +703,7 @@ impl Task {
self.update_exchange_from_messages(
messages.clone(),
exchange_id,
current_todo_list,
current_comments,
message_context,
should_convert_input_messages,
)?;
self.try_get_source_mut()?.messages.extend(messages);
@@ -699,8 +714,7 @@ impl Task {
&mut self,
message: api::Message,
exchange_id: AIAgentExchangeId,
current_todo_list: Option<&AIAgentTodoList>,
current_comments: Option<&CodeReview>,
message_context: TaskMessageContext<'_>,
mask: FieldMask,
should_convert_input_messages: bool,
) -> Result<&api::Message, UpdateTaskError> {
@@ -714,8 +728,7 @@ impl Task {
self.add_messages(
vec![message.clone()],
exchange_id,
current_todo_list,
current_comments,
message_context,
should_convert_input_messages,
)?;
return self
@@ -734,10 +747,13 @@ impl Task {
.exchange_mut(exchange_id)
.ok_or(UpdateTaskError::ExchangeNotFound)?;
exchange_to_update.upsert_output_for_message(
&id,
&updated_message,
current_todo_list,
current_comments,
ConversionParams {
task_id: &id,
current_todo_list: message_context.current_todo_list,
active_code_review: message_context.active_code_review,
skill_path_origin: message_context.skill_path_origin,
},
)?;
// Task message updates can carry tool call result updates with them,
@@ -782,8 +798,7 @@ impl Task {
&mut self,
message: api::Message,
exchange_id: AIAgentExchangeId,
current_todo_list: Option<&AIAgentTodoList>,
current_comments: Option<&CodeReview>,
message_context: TaskMessageContext<'_>,
mask: FieldMask,
) -> Result<&api::Message, UpdateTaskError> {
let Some((idx, existing_message)) = self
@@ -836,10 +851,13 @@ impl Task {
.exchange_mut(exchange_id)
.ok_or(UpdateTaskError::ExchangeNotFound)?;
exchange_to_update.upsert_output_for_message(
&id,
&updated_message,
current_todo_list,
current_comments,
ConversionParams {
task_id: &id,
current_todo_list: message_context.current_todo_list,
active_code_review: message_context.active_code_review,
skill_path_origin: message_context.skill_path_origin,
},
)?;
let source = self.try_get_source_mut()?;
@@ -964,8 +982,7 @@ impl Task {
&mut self,
messages: Vec<api::Message>,
exchange_id: AIAgentExchangeId,
current_todo_list: Option<&AIAgentTodoList>,
active_code_review: Option<&CodeReview>,
message_context: TaskMessageContext<'_>,
should_convert_input_messages: bool,
) -> Result<(), UpdateTaskError> {
let exchange = self
@@ -1007,8 +1024,9 @@ impl Task {
.filter_map(|m| {
match m.to_client_output_message(ConversionParams {
task_id: &self.id,
current_todo_list,
active_code_review,
current_todo_list: message_context.current_todo_list,
active_code_review: message_context.active_code_review,
skill_path_origin: message_context.skill_path_origin,
}) {
Ok(MaybeAIAgentOutputMessage::Message(m)) => Some(Ok(m)),
Ok(MaybeAIAgentOutputMessage::NoClientRepresentation) => None,
@@ -1043,10 +1061,8 @@ impl AIAgentExchange {
/// Note: this means updates will insert a new entry after previously added entries.
fn upsert_output_for_message(
&self,
task_id: &TaskId,
task_message: &api::Message,
todo_list: Option<&AIAgentTodoList>,
comments: Option<&CodeReview>,
conversion_params: super::api::ConversionParams<'_>,
) -> Result<(), UpdateTaskError> {
if let AIAgentOutputStatus::Streaming {
output: Some(output),
@@ -1079,11 +1095,8 @@ impl AIAgentExchange {
match task_message
.clone()
.to_client_output_message(ConversionParams {
current_todo_list: todo_list,
active_code_review: comments,
task_id,
})? {
.to_client_output_message(conversion_params)?
{
MaybeAIAgentOutputMessage::Message(m) => {
log::info!(
"[bedrock-debug] upsert_output_for_message: client_message_type={:?}",
app/src/ai/agent/task/helper.rs
+4
View File
@@ -139,6 +139,10 @@ impl ToolExt for api::message::tool_call::Tool {
Tool::AskUserQuestion(_) => "ask_user_question",
Tool::SendMessageToAgent(_) => "send_message_to_agent",
Tool::TransferShellCommandControlToUser(_) => "transfer_shell_command_control",
Tool::RunAgents(_) => "orchestrate",
// Matches the legacy server-handled name so analytics don't
// double-count the rollout.
Tool::WaitForEvents(_) => "wait_for_events",
}
}
}
app/src/ai/agent/task/transaction.rs
+1 -2
View File
@@ -1,8 +1,7 @@
use std::collections::HashMap;
use crate::ai::agent::task::TaskId;
use super::Task;
use crate::ai::agent::task::TaskId;
/// Keeps track of the state of tasks before they are modified.
/// Messages are assumed to be only updated during the same transaction
app/src/ai/agent/task_store.rs
+54 -19
View File
@@ -2,18 +2,11 @@ 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,
};
use super::task::helper::{MessageExt, ToolCallExt};
use super::task::{Task, TaskId};
use super::{AIAgentExchange, AIAgentExchangeId, AIAgentOutputMessageType};
use crate::ai::agent::{AIAgentContext, AIAgentInput};
use crate::ai::skills::SkillDescriptor;
#[derive(Debug, Clone)]
struct ExchangeRef {
@@ -27,6 +20,12 @@ pub struct TaskStore {
root_task_id: TaskId,
tasks: HashMap<TaskId, Task>,
linearized_refs: Vec<ExchangeRef>,
exchange_id_index: HashMap<AIAgentExchangeId, ExchangeRef>,
/// If the root task was upgraded from an optimistic (client-generated) ID
/// to a server-assigned ID, stores the original optimistic ID so that
/// deferred event handlers referencing the stale ID can still resolve
/// the task via `root_task_id`.
optimistic_root_task_id: Option<TaskId>,
}
impl TaskStore {
@@ -35,7 +34,9 @@ impl TaskStore {
let mut store = Self {
tasks: HashMap::new(),
linearized_refs: Vec::new(),
exchange_id_index: HashMap::new(),
root_task_id: root_task_id.clone(),
optimistic_root_task_id: None,
};
store.tasks.insert(root_task_id, root_task);
store.rebuild_linearized_refs_index();
@@ -48,7 +49,9 @@ impl TaskStore {
let mut store = Self {
tasks,
linearized_refs: Vec::new(),
exchange_id_index: HashMap::new(),
root_task_id,
optimistic_root_task_id: None,
};
store.rebuild_linearized_refs_index();
store
@@ -59,7 +62,10 @@ impl TaskStore {
}
pub fn get(&self, task_id: &TaskId) -> Option<&Task> {
self.tasks.get(task_id)
self.tasks.get(task_id).or_else(|| {
let old_id = self.optimistic_root_task_id.as_ref()?;
(old_id == task_id).then(|| self.tasks.get(&self.root_task_id))?
})
}
pub fn tasks(&self) -> impl Iterator<Item = &Task> {
@@ -104,15 +110,15 @@ impl TaskStore {
None
}
/// Modifies a task via the provided closure and rebuilds the exchange index
/// if exchanges changed.
/// Modifies a task via the provided closure and rebuilds the exchange index if the exchange
/// count changes.
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 exchange_count_before = task.exchanges_len();
let result = f(task);
let exchange_count_after = self
.tasks
@@ -142,10 +148,39 @@ impl TaskStore {
self.remove(&old_root_id);
let new_root_id = root_task.id().clone();
if old_root_id != new_root_id {
self.optimistic_root_task_id = Some(old_root_id);
}
self.root_task_id = new_root_id;
self.insert(root_task);
}
pub fn exchange_by_id(&self, exchange_id: AIAgentExchangeId) -> Option<&AIAgentExchange> {
let exchange_ref = self.exchange_id_index.get(&exchange_id)?;
self.lookup_exchange(exchange_ref)
}
pub(super) fn rebuild_exchange_id_index(&mut self) {
self.exchange_id_index = self
.tasks
.values()
.flat_map(|task| {
let task_id = task.id().clone();
task.exchanges()
.enumerate()
.map(move |(exchange_index, exchange)| {
(
exchange.id,
ExchangeRef {
task_id: task_id.clone(),
exchange_index,
},
)
})
})
.collect();
}
pub fn first_exchange(&self) -> Option<&AIAgentExchange> {
self.linearized_refs
.first()
@@ -266,7 +301,7 @@ impl TaskStore {
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);
self.rebuild_linearized_refs_index();
Some(task)
}
@@ -280,6 +315,7 @@ impl TaskStore {
/// 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);
self.rebuild_exchange_id_index();
}
/// Builds linearized exchange refs via DFS traversal without mutating self.
@@ -330,9 +366,8 @@ impl TaskStore {
#[cfg(test)]
mod testing {
use crate::ai::agent::task::TaskId;
use super::TaskStore;
use crate::ai::agent::task::TaskId;
impl TaskStore {
pub fn contains(&self, task_id: &TaskId) -> bool {
app/src/ai/agent/task_store_tests.rs
+125 -12
View File
@@ -3,17 +3,14 @@ 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;
use crate::ai::agent::task::{Task, TaskId};
use crate::ai::agent::{
AIAgentExchange, AIAgentExchangeId, AIAgentOutput, AIAgentOutputMessage,
AIAgentOutputMessageType, AIAgentOutputStatus, FinishedAIAgentOutput, MessageId, Shared,
SubagentCall,
};
use crate::ai::llms::LLMId;
fn create_test_exchange() -> AIAgentExchange {
AIAgentExchange {
@@ -216,7 +213,8 @@ fn test_set_root_task_replaces_old() {
assert_eq!(store.task_count(), 1);
assert_eq!(store.exchange_count(), 3);
assert!(store.get(&task1_id).is_none());
// task1_id is now aliased to the new root task via optimistic_root_task_id
assert!(store.get(&task1_id).is_some());
assert!(store.get(&task2_id).is_some());
assert_eq!(store.root_task_id(), &task2_id);
@@ -485,7 +483,6 @@ fn test_linearization_nested_subtasks() {
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(), None);
let child_id = child_subtask.id().clone();
@@ -623,3 +620,119 @@ fn test_all_exchanges_by_task_with_subtasks() {
assert_eq!(by_task[2].1.len(), 1);
assert_eq!(by_task[2].1[0].id, root_exchange3_id);
}
#[test]
fn test_exchange_by_id_resolves_subtask_exchange() {
// The index spans all tasks, not just the root, so exchanges that live in a
// subtask must be resolvable by id.
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);
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();
let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
store.append_exchange(&root_task_id, subagent_exchange);
store.insert(subtask);
for id in &subtask_exchange_ids {
assert_eq!(store.exchange_by_id(*id).map(|e| e.id), Some(*id));
}
}
#[test]
fn test_exchange_by_id_after_remove_task() {
let root_task = create_test_task_with_exchanges(1);
let root_task_id = root_task.id().clone();
let root_exchange_id = root_task.exchanges().next().unwrap().id;
let mut store = TaskStore::with_root_task(root_task);
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();
let subagent_exchange = create_exchange_with_subagent_call(&subtask_id);
store.append_exchange(&root_task_id, subagent_exchange);
store.insert(subtask);
// Sanity: the subtask's exchanges resolve before removal.
assert!(store.exchange_by_id(subtask_exchange_ids[0]).is_some());
store.remove(&subtask_id);
// The removed task's exchanges are no longer resolvable.
for id in &subtask_exchange_ids {
assert!(store.exchange_by_id(*id).is_none());
}
// The surviving root exchange still resolves.
assert_eq!(
store.exchange_by_id(root_exchange_id).map(|e| e.id),
Some(root_exchange_id)
);
}
#[test]
fn test_exchange_by_id_after_remove_task_exchange_index_shift() {
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);
// Remove the middle exchange, which shifts the index of everything after it.
store.remove_task_exchange(&task_id, exchange_ids[1]);
// The removed id no longer resolves.
assert!(store.exchange_by_id(exchange_ids[1]).is_none());
// The exchange that followed it (now at a different index) still resolves to itself.
assert_eq!(
store.exchange_by_id(exchange_ids[2]).map(|e| e.id),
Some(exchange_ids[2])
);
// The exchange before it is unaffected.
assert_eq!(
store.exchange_by_id(exchange_ids[0]).map(|e| e.id),
Some(exchange_ids[0])
);
}
#[test]
fn test_exchange_by_id_after_modify_task_append() {
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;
store.modify_task(&task_id, |task| {
task.append_exchange(new_exchange);
});
// The newly appended exchange is found via the rebuilt index.
assert_eq!(
store.exchange_by_id(new_exchange_id).map(|e| e.id),
Some(new_exchange_id)
);
}
#[test]
fn test_exchange_by_id_after_set_root_task() {
let task1 = create_test_task_with_exchanges(2);
let task1_exchange_ids: Vec<_> = task1.exchanges().map(|e| e.id).collect();
let mut store = TaskStore::with_root_task(task1);
let task2 = create_test_task_with_exchanges(3);
let task2_exchange_ids: Vec<_> = task2.exchanges().map(|e| e.id).collect();
store.set_root_task(task2);
// The old root's exchanges no longer resolve.
for id in &task1_exchange_ids {
assert!(store.exchange_by_id(*id).is_none());
}
// The new root's exchanges resolve.
for id in &task2_exchange_ids {
assert_eq!(store.exchange_by_id(*id).map(|e| e.id), Some(*id));
}
}
app/src/ai/agent/task_tests.rs
+12 -7
View File
@@ -1,5 +1,11 @@
use std::collections::HashSet;
use ai::skills::SkillPathOrigin;
use chrono::Local;
use prost_types::FieldMask;
use warp_multi_agent_api as api;
use super::{ExtractMessagesError, Task, TaskMessageContext};
use crate::ai::agent::{
AIAgentActionType, AIAgentExchange, AIAgentOutput, AIAgentOutputMessageType,
AIAgentOutputStatus, MessageId, Shared,
@@ -8,11 +14,6 @@ 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 {
@@ -47,6 +48,7 @@ fn create_start_agent_tool_call_message(
prompt: &str,
) -> api::Message {
api::Message {
fetched_memories: vec![],
id: id.to_string(),
task_id: task_id.to_string(),
server_message_data: String::new(),
@@ -113,8 +115,11 @@ fn test_upsert_message_adds_start_agent_prompt_to_output() {
"run tests",
),
exchange_id,
None,
None,
TaskMessageContext {
current_todo_list: None,
active_code_review: None,
skill_path_origin: &SkillPathOrigin::Local,
},
FieldMask {
paths: vec!["message.tool_call".to_string()],
},
app/src/ai/agent/telemetry.rs
+11 -21
View File
@@ -1,18 +1,14 @@
use galaxyui::{AppContext, SingletonEntity};
use serde::Serialize;
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,
};
use crate::ai::llms::LLMId;
use crate::server::telemetry::AgentModeCitation as CitationForTelemetry;
use crate::CloudModel;
pub trait ForTelemetry {
type Output;
@@ -37,6 +33,14 @@ impl ForTelemetry for AIAgentCitation {
Some(CitationForTelemetry::WarpDocs { page: path.clone() })
}
Self::WebPage { url } => Some(CitationForTelemetry::WebPage { url: url.clone() }),
Self::AgentMemory {
memory_store_id,
memory_id,
..
} => Some(CitationForTelemetry::AgentMemory {
memory_store_id: memory_store_id.clone(),
memory_id: memory_id.clone(),
}),
}
}
}
@@ -44,20 +48,6 @@ impl ForTelemetry for AIAgentCitation {
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,
app/src/ai/agent/todos/mod.rs
+1 -2
View File
@@ -1,6 +1,5 @@
use crate::ai::agent::AIAgentTodo;
use super::AIAgentTodoId;
use crate::ai::agent::AIAgentTodo;
pub(crate) mod popup;
#[derive(Debug, Clone, PartialEq, Eq, Default)]
app/src/ai/agent/todos/popup.rs
+17 -17
View File
@@ -1,27 +1,24 @@
use crate::ai::blocklist::{BlocklistAIContextEvent, BlocklistAIContextModel};
use pathfinder_color::ColorU;
use galaxy_core::ui::appearance::Appearance;
use galaxy_core::ui::theme::Fill;
use galaxyui::elements::{
ClippedScrollStateHandle, ClippedScrollable, Dismiss, Empty, Expanded, ParentElement,
SavePosition, ScrollTarget, ScrollToPositionMode, ScrollbarWidth, Shrinkable,
Border, ClippedScrollStateHandle, ClippedScrollable, ConstrainedBox, Container, CornerRadius,
CrossAxisAlignment, Dismiss, DropShadow, Empty, Expanded, Flex, MainAxisSize, ParentElement,
Radius, SavePosition, ScrollTarget, ScrollToPositionMode, ScrollbarWidth, Shrinkable, Text,
};
use galaxyui::fonts::FamilyId;
use galaxyui::ModelHandle;
use galaxyui::SingletonEntity;
use galaxyui::fonts::{FamilyId, Properties, Weight};
use galaxyui::keymap::FixedBinding;
use galaxyui::{
elements::{
Border, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, DropShadow, Flex,
MainAxisSize, Radius, Text,
},
fonts::{Properties, Weight},
keymap::FixedBinding,
AppContext, Element, Entity, EntityId, TypedActionView, View, ViewContext,
AppContext, Element, Entity, EntityId, ModelHandle, SingletonEntity, TypedActionView, View,
ViewContext,
};
use pathfinder_color::ColorU;
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::ai::blocklist::{
BlocklistAIContextEvent, BlocklistAIContextModel, BlocklistAIHistoryEvent,
BlocklistAIHistoryModel,
};
use crate::ui_components::blended_colors;
pub struct AgentTodosPopupView {
@@ -86,8 +83,11 @@ impl AgentTodosPopupView {
event: &BlocklistAIHistoryEvent,
ctx: &mut ViewContext<Self>,
) {
if let BlocklistAIHistoryEvent::UpdatedTodoList { terminal_view_id } = event {
if *terminal_view_id == self.terminal_view_id {
if let BlocklistAIHistoryEvent::UpdatedTodoList {
terminal_surface_id,
} = event
{
if *terminal_surface_id == self.terminal_view_id {
ctx.notify();
}
}
app/src/ai/agent/util.rs
+10 -7
View File
@@ -1,9 +1,6 @@
use super::{
AIAgentTextSection, AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram,
AgentOutputTable, ProgrammingLanguage,
};
use crate::code::editor_management::CodeSource;
use crate::features::FeatureFlag;
use std::collections::HashMap;
use std::path::PathBuf;
use ai::gfm_table::{format_gfm_table, maybe_collect_gfm_table_lines};
use galaxy_util::path::LineAndColumnArg;
use itertools::Itertools;
@@ -13,7 +10,13 @@ use markdown_parser::{
};
use mermaid_to_svg::is_mermaid_diagram;
use regex::Regex;
use std::{collections::HashMap, path::PathBuf};
use super::{
AIAgentTextSection, AgentOutputImage, AgentOutputImageLayout, AgentOutputMermaidDiagram,
AgentOutputTable, ProgrammingLanguage,
};
use crate::code::editor_management::CodeSource;
use crate::features::FeatureFlag;
lazy_static! {
/// Markdown prefix for code blocks. Matches on triple backticks followed by a language.
app/src/ai/agent_conversations_model.rs
+675 -646
View File
File diff suppressed because it is too large Load Diff
app/src/ai/agent_conversations_model/entry.rs
+683
View File
@@ -0,0 +1,683 @@
use chrono::{DateTime, Utc};
use session_sharing_protocol::common::SessionId;
use warp_cli::agent::Harness;
use galaxy_core::features::FeatureFlag;
use warpui::{AppContext, SingletonEntity};
use super::{
artifacts_match_filter, AgentManagementFilters, AgentRunDisplayStatus, ArtifactFilter,
ConversationMetadata, CreatedOnFilter, CreatorFilter, EnvironmentFilter, HarnessFilter,
OwnerFilter, SessionStatus, SourceFilter, StatusFilter,
};
use crate::ai::active_agent_views_model::{ActiveAgentViewsModel, ConversationOrTaskId};
use crate::ai::agent::api::ServerConversationToken;
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::ambient_agents::{AgentSource, AmbientAgentTask, AmbientAgentTaskId};
use crate::ai::artifacts::Artifact;
use crate::ai::blocklist::history_model::{AIConversationMetadata, BlocklistAIHistoryModel};
use crate::ai::conversation_navigation::ConversationNavigationData;
use crate::auth::{AuthStateProvider, UserUid};
use crate::util::time_format::human_readable_precise_duration;
use crate::workspace::RestoreConversationLayout;
use crate::workspaces::user_profiles::{UserProfileWithUID, UserProfiles};
const SESSION_EXPIRATION_TIME: chrono::Duration = chrono::Duration::weeks(1);
/// Stable projection identity used by list and navigation surfaces.
///
/// Task-backed rows use the ambient run ID even when they are attached to a local
/// conversation, so task-specific affordances do not disappear when local data is present.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum AgentConversationEntryId {
AmbientRun(AmbientAgentTaskId),
Conversation(AIConversationId),
}
impl AgentConversationEntryId {
pub fn as_key(&self) -> String {
match self {
AgentConversationEntryId::AmbientRun(id) => format!("task_{id}"),
AgentConversationEntryId::Conversation(id) => format!("conv_{id}"),
}
}
}
impl From<ConversationOrTaskId> for AgentConversationEntryId {
fn from(id: ConversationOrTaskId) -> Self {
match id {
ConversationOrTaskId::ConversationId(conversation_id) => {
AgentConversationEntryId::Conversation(conversation_id)
}
ConversationOrTaskId::TaskId(task_id) => AgentConversationEntryId::AmbientRun(task_id),
}
}
}
/// Navigation request input for resolving an entry or server-token handle at action time.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AgentConversationNavigationSubject {
Entry(AgentConversationEntryId),
#[allow(dead_code)]
ServerToken(ServerConversationToken),
}
/// Normalized row data for agent conversation list, management, and navigation surfaces.
///
/// The entry keeps local conversation identity, ambient run identity, cloud token identity,
/// display fields, and available actions together so callers do not recompute navigation
/// policy from stale partial sources.
#[derive(Clone, Debug, PartialEq)]
pub struct AgentConversationEntry {
pub id: AgentConversationEntryId,
pub identity: AgentConversationIdentity,
pub provenance: AgentConversationProvenance,
pub display: AgentConversationDisplayData,
pub backing: AgentConversationBackingData,
pub capabilities: AgentConversationCapabilities,
}
/// Cross-system identifiers that may refer to the same underlying conversation/run.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AgentConversationIdentity {
pub local_conversation_id: Option<AIConversationId>,
pub ambient_agent_task_id: Option<AmbientAgentTaskId>,
pub server_conversation_token: Option<ServerConversationToken>,
pub session_id: Option<SessionId>,
}
/// Display-only fields for rendering a conversation entry without consulting source models.
#[derive(Clone, Debug, PartialEq)]
pub struct AgentConversationDisplayData {
pub title: String,
pub initial_query: Option<String>,
pub created_at: DateTime<Utc>,
pub last_updated: DateTime<Utc>,
pub status: AgentRunDisplayStatus,
pub creator: AgentConversationPrincipal,
pub executor: Option<AgentConversationPrincipal>,
pub request_usage: Option<f32>,
pub run_time: Option<String>,
pub session_status: Option<SessionStatus>,
pub source: Option<AgentSource>,
pub working_directory: Option<String>,
pub environment_id: Option<String>,
pub harness: Option<Harness>,
pub artifacts: Vec<Artifact>,
}
/// Type of principal that created or executed a run.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PrincipalType {
User,
ServiceAccount,
}
impl PrincipalType {
/// Parse from the wire-format string sent by the server.
pub fn parse(s: &str) -> Option<Self> {
if s.eq_ignore_ascii_case("user") {
Some(PrincipalType::User)
} else if s.eq_ignore_ascii_case("service_account") || s.eq_ignore_ascii_case("agent") {
Some(PrincipalType::ServiceAccount)
} else {
None
}
}
pub fn is_service_account(self) -> bool {
self == PrincipalType::ServiceAccount
}
}
/// Principal information normalized across local conversations and ambient runs.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct AgentConversationPrincipal {
pub name: Option<String>,
pub uid: Option<String>,
pub principal_type: Option<PrincipalType>,
}
/// Source category that explains why an entry exists and which backing systems can refresh it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AgentConversationProvenance {
LocalInteractive,
AmbientRun,
CloudSyncedConversation,
}
/// Availability flags for the source data that contributed to an entry.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AgentConversationBackingData {
pub has_loaded_conversation: bool,
pub has_local_persisted_data: bool,
pub has_cloud_data: bool,
pub has_ambient_run: bool,
}
/// Actions that should be exposed for an entry after applying current navigation policy.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AgentConversationCapabilities {
pub can_open: bool,
pub can_copy_link: bool,
pub can_share: bool,
pub can_delete: bool,
pub can_fork_locally: bool,
pub can_cancel: bool,
}
impl AgentConversationEntry {
pub(super) fn matches_filters(
&self,
filters: &AgentManagementFilters,
app: &AppContext,
) -> bool {
self.matches_owner_and_creator(&filters.owners, &filters.creator, app)
&& self.matches_status(&filters.status)
&& self.matches_source(&filters.source)
&& self.matches_created_on(&filters.created_on)
&& self.matches_artifact(&filters.artifact)
&& self.matches_environment(&filters.environment)
&& self.matches_harness(&filters.harness)
}
fn matches_owner_and_creator(
&self,
owner_filter: &OwnerFilter,
creator_filter: &CreatorFilter,
app: &AppContext,
) -> bool {
let current_user_id = AuthStateProvider::as_ref(app)
.get()
.user_id()
.map(|uid| uid.as_string());
let passes_owner = match owner_filter {
OwnerFilter::All => true,
OwnerFilter::PersonalOnly => {
if self.backing.has_ambient_run {
self.display.creator.uid == current_user_id
} else {
true
}
}
};
if !passes_owner || matches!(owner_filter, OwnerFilter::PersonalOnly) {
return passes_owner;
}
match creator_filter {
CreatorFilter::All => true,
CreatorFilter::Specific { name, .. } => {
self.display.creator.name.as_ref() == Some(name)
}
}
}
fn matches_status(&self, status_filter: &StatusFilter) -> bool {
match status_filter {
StatusFilter::All => true,
StatusFilter::Working | StatusFilter::Done | StatusFilter::Failed => {
self.display.status.status_filter() == *status_filter
}
}
}
fn matches_source(&self, source_filter: &SourceFilter) -> bool {
match source_filter {
SourceFilter::All => true,
SourceFilter::Specific(source) => self.display.source.as_ref() == Some(source),
}
}
fn matches_created_on(&self, created_on_filter: &CreatedOnFilter) -> bool {
let now = Utc::now();
let created_cutoff = match created_on_filter {
CreatedOnFilter::All => None,
CreatedOnFilter::Last24Hours => Some(now - chrono::Duration::hours(24)),
CreatedOnFilter::Past3Days => Some(now - chrono::Duration::days(3)),
CreatedOnFilter::LastWeek => Some(now - chrono::Duration::days(7)),
};
match created_cutoff {
Some(cutoff) => self.display.created_at >= cutoff,
None => true,
}
}
fn matches_artifact(&self, artifact_filter: &ArtifactFilter) -> bool {
artifacts_match_filter(&self.display.artifacts, artifact_filter)
}
fn matches_environment(&self, environment_filter: &EnvironmentFilter) -> bool {
match environment_filter {
EnvironmentFilter::All => true,
EnvironmentFilter::NoEnvironment => self.display.environment_id.is_none(),
EnvironmentFilter::Specific(id) => self.display.environment_id.as_ref() == Some(id),
}
}
fn matches_harness(&self, harness_filter: &HarnessFilter) -> bool {
match harness_filter {
HarnessFilter::All => true,
HarnessFilter::Specific(harness) => self.display.harness == Some(*harness),
}
}
pub fn has_open_action(
&self,
restore_layout: Option<RestoreConversationLayout>,
app: &AppContext,
) -> bool {
super::AgentConversationsModel::resolve_open_action(
AgentConversationNavigationSubject::Entry(self.id),
restore_layout,
app,
)
.is_some()
}
}
/// Returns the local conversation ID represented by the given task, if this task and a
/// conversation entry both point at the same underlying local run.
///
/// We first match using the orchestration agent ID (task ID / run ID under v2), and fall back
/// to the server conversation token for cases where the task only carries conversation identity
/// through `conversation_id`.
pub(super) fn conversation_id_shadowed_by_task(
task: &AmbientAgentTask,
history_model: &BlocklistAIHistoryModel,
) -> Option<AIConversationId> {
history_model
.conversation_id_for_agent_id(&task.run_id().to_string())
.or_else(|| {
task.conversation_id().and_then(|conversation_id| {
history_model.find_conversation_id_by_server_token(&ServerConversationToken::new(
conversation_id.to_string(),
))
})
})
}
pub(super) fn task_creator_name(task: &AmbientAgentTask, app: &AppContext) -> Option<String> {
task.creator_display_name().or_else(|| {
let uid = task.creator.as_ref().map(|creator| &creator.uid)?;
UserProfiles::as_ref(app).displayable_identifier_for_uid(UserUid::new(uid))
})
}
pub(super) fn task_creator_uid(task: &AmbientAgentTask) -> Option<String> {
task.creator.as_ref().map(|creator| creator.uid.clone())
}
fn current_user_name(app: &AppContext) -> Option<String> {
AuthStateProvider::as_ref(app).get().username_for_display()
}
fn current_user_uid(app: &AppContext) -> Option<String> {
AuthStateProvider::as_ref(app)
.get()
.user_id()
.map(|uid| uid.to_string())
}
fn task_session_id(task: &AmbientAgentTask) -> Option<SessionId> {
task.session_id.as_deref().and_then(parse_session_id)
}
fn task_session_status(task: &AmbientAgentTask) -> SessionStatus {
if FeatureFlag::CloudConversations.is_enabled() {
return if task.active_run_execution().session_link.is_some() {
SessionStatus::Available
} else {
SessionStatus::Unavailable
};
}
if task.active_run_execution().session_id.is_some() {
SessionStatus::Available
} else if (Utc::now() - task.created_at) > SESSION_EXPIRATION_TIME {
SessionStatus::Expired
} else {
SessionStatus::Unavailable
}
}
fn task_run_time(task: &AmbientAgentTask) -> Option<String> {
task.run_time().map(human_readable_precise_duration)
}
fn task_harness(task: &AmbientAgentTask) -> Option<Harness> {
task.agent_config_snapshot.as_ref().and_then(|config| {
config
.harness
.as_ref()
.map(|harness| harness.harness_type)
.or(Some(Harness::Oz))
})
}
fn conversation_title(
metadata: &ConversationMetadata,
history_model: &BlocklistAIHistoryModel,
) -> String {
history_model
.conversation(&metadata.nav_data.id)
.and_then(|conversation| conversation.title().clone())
.unwrap_or(metadata.nav_data.title.clone())
}
fn conversation_display_status(
metadata: &ConversationMetadata,
history_model: &BlocklistAIHistoryModel,
) -> AgentRunDisplayStatus {
history_model
.conversation(&metadata.nav_data.id)
.map(|conversation| AgentRunDisplayStatus::from_conversation_status(conversation.status()))
.unwrap_or(AgentRunDisplayStatus::ConversationSucceeded)
}
fn conversation_request_usage(
metadata: &ConversationMetadata,
history_model: &BlocklistAIHistoryModel,
) -> Option<f32> {
history_model
.conversation(&metadata.nav_data.id)
.map(|conversation| conversation.credits_spent())
.or_else(|| {
history_model
.get_conversation_metadata(&metadata.nav_data.id)
.and_then(|metadata| metadata.credits_spent)
})
}
fn conversation_artifacts(
metadata: &ConversationMetadata,
history_model: &BlocklistAIHistoryModel,
) -> Vec<Artifact> {
history_model
.conversation(&metadata.nav_data.id)
.map(|conversation| conversation.artifacts().to_vec())
.or_else(|| {
history_model
.get_conversation_metadata(&metadata.nav_data.id)
.map(|metadata| metadata.artifacts.clone())
})
.unwrap_or_default()
}
fn principal_from_user_profile(profile: &UserProfileWithUID) -> AgentConversationPrincipal {
let name = profile
.display_name
.as_ref()
.filter(|name| !name.is_empty())
.or_else(|| (!profile.email.is_empty()).then_some(&profile.email))
.cloned()
.or_else(|| Some(profile.firebase_uid.to_string()));
AgentConversationPrincipal {
name,
uid: Some(profile.firebase_uid.to_string()),
principal_type: Some(PrincipalType::User),
}
}
fn conversation_creator(
metadata: &ConversationMetadata,
history_model: &BlocklistAIHistoryModel,
app: &AppContext,
) -> AgentConversationPrincipal {
let server_metadata = history_model.get_server_conversation_metadata(&metadata.nav_data.id);
if let Some(profile) = server_metadata.and_then(|metadata| metadata.creator.as_ref()) {
return principal_from_user_profile(profile);
}
if let Some(uid) = server_metadata.and_then(|metadata| metadata.metadata.creator_uid.as_ref()) {
return AgentConversationPrincipal {
name: UserProfiles::as_ref(app).displayable_identifier_for_uid(UserUid::new(uid)),
uid: Some(uid.clone()),
principal_type: Some(PrincipalType::User),
};
}
AgentConversationPrincipal {
name: current_user_name(app),
uid: current_user_uid(app),
principal_type: Some(PrincipalType::User),
}
}
pub(super) fn entry_for_task(
task: &AmbientAgentTask,
history_model: &BlocklistAIHistoryModel,
app: &AppContext,
) -> AgentConversationEntry {
let local_conversation_id = conversation_id_shadowed_by_task(task, history_model);
let conversation_metadata =
local_conversation_id.and_then(|id| history_model.get_conversation_metadata(&id));
let server_conversation_token = task
.conversation_id()
.map(|id| ServerConversationToken::new(id.to_string()))
.or_else(|| {
local_conversation_id.and_then(|conversation_id| {
server_conversation_token_for_conversation(conversation_id, None, history_model)
})
});
let status = AgentRunDisplayStatus::from_task(task, app);
let has_active_session_id = task
.active_execution_session_id()
.and_then(parse_session_id)
.is_some();
let has_open_ambient_session = ActiveAgentViewsModel::as_ref(app)
.get_terminal_view_id_for_ambient_task(task.task_id)
.is_some();
let can_open = has_open_ambient_session
|| has_active_session_id
|| local_conversation_id.is_some()
|| server_conversation_token.is_some();
let can_copy_link = task.has_active_execution()
&& task.active_run_execution().session_link.is_some()
|| server_conversation_token.is_some();
AgentConversationEntry {
id: AgentConversationEntryId::AmbientRun(task.task_id),
identity: AgentConversationIdentity {
local_conversation_id,
ambient_agent_task_id: Some(task.task_id),
server_conversation_token,
session_id: task_session_id(task),
},
provenance: AgentConversationProvenance::AmbientRun,
display: AgentConversationDisplayData {
title: task.title.clone(),
initial_query: Some(task.prompt.clone()),
created_at: task.created_at,
last_updated: task.updated_at,
status: status.clone(),
creator: AgentConversationPrincipal {
name: task_creator_name(task, app),
uid: task_creator_uid(task),
principal_type: task
.creator
.as_ref()
.and_then(|c| PrincipalType::parse(&c.creator_type)),
},
executor: task
.executor
.as_ref()
.map(|executor| AgentConversationPrincipal {
name: executor.display_name.clone(),
uid: Some(executor.uid.clone()),
principal_type: PrincipalType::parse(&executor.creator_type),
}),
request_usage: task.credits_used(),
run_time: task_run_time(task),
session_status: Some(task_session_status(task)),
source: task.source.clone(),
working_directory: conversation_metadata
.and_then(|metadata| metadata.initial_working_directory.clone()),
environment_id: task
.agent_config_snapshot
.as_ref()
.and_then(|snapshot| snapshot.environment_id.clone()),
harness: task_harness(task),
artifacts: task.artifacts.clone(),
},
backing: AgentConversationBackingData {
has_loaded_conversation: local_conversation_id
.is_some_and(|id| history_model.conversation(&id).is_some()),
has_local_persisted_data: conversation_metadata
.is_some_and(|metadata| metadata.has_local_data),
has_cloud_data: conversation_metadata.is_some_and(|metadata| metadata.has_cloud_data)
|| task.conversation_id().is_some(),
has_ambient_run: true,
},
capabilities: AgentConversationCapabilities {
can_open,
can_copy_link,
can_share: task.conversation_id().is_some()
|| local_conversation_id
.is_some_and(|id| history_model.can_conversation_be_shared(&id)),
can_delete: false,
can_fork_locally: local_conversation_id.is_some(),
can_cancel: status.is_cancellable(),
},
}
}
pub(super) fn entry_for_conversation(
metadata: &ConversationMetadata,
history_model: &BlocklistAIHistoryModel,
app: &AppContext,
) -> AgentConversationEntry {
let conversation_metadata = history_model.get_conversation_metadata(&metadata.nav_data.id);
entry_for_conversation_parts(
metadata.nav_data.clone(),
conversation_metadata,
history_model,
app,
)
}
pub(super) fn entry_for_historical_metadata(
metadata: &AIConversationMetadata,
nav_data: ConversationNavigationData,
history_model: &BlocklistAIHistoryModel,
app: &AppContext,
) -> AgentConversationEntry {
entry_for_conversation_parts(nav_data, Some(metadata), history_model, app)
}
fn entry_for_conversation_parts(
nav_data: ConversationNavigationData,
conversation_metadata: Option<&AIConversationMetadata>,
history_model: &BlocklistAIHistoryModel,
app: &AppContext,
) -> AgentConversationEntry {
let metadata = ConversationMetadata { nav_data };
let conversation_id = metadata.nav_data.id;
let status = conversation_display_status(&metadata, history_model);
let has_loaded_conversation = history_model.conversation(&conversation_id).is_some();
let has_local_persisted_data = conversation_metadata
.is_some_and(|metadata| metadata.has_local_data)
|| has_loaded_conversation;
let has_cloud_data = conversation_metadata.is_some_and(|metadata| metadata.has_cloud_data)
|| server_conversation_token_for_conversation(
conversation_id,
Some(&metadata.nav_data),
history_model,
)
.is_some();
let provenance = if has_cloud_data {
AgentConversationProvenance::CloudSyncedConversation
} else {
AgentConversationProvenance::LocalInteractive
};
AgentConversationEntry {
id: AgentConversationEntryId::Conversation(conversation_id),
identity: AgentConversationIdentity {
local_conversation_id: Some(conversation_id),
ambient_agent_task_id: conversation_metadata
.and_then(|metadata| metadata.server_conversation_metadata.as_ref())
.and_then(|metadata| metadata.ambient_agent_task_id),
server_conversation_token: server_conversation_token_for_conversation(
conversation_id,
Some(&metadata.nav_data),
history_model,
),
session_id: None,
},
provenance,
display: AgentConversationDisplayData {
title: conversation_title(&metadata, history_model),
initial_query: metadata.nav_data.initial_query.clone(),
created_at: metadata.nav_data.last_updated.into(),
last_updated: metadata.nav_data.last_updated.into(),
status: status.clone(),
creator: conversation_creator(&metadata, history_model, app),
executor: None,
request_usage: conversation_request_usage(&metadata, history_model),
run_time: None,
session_status: None,
source: Some(AgentSource::Interactive),
working_directory: metadata
.nav_data
.latest_working_directory
.clone()
.or_else(|| metadata.nav_data.initial_working_directory.clone()),
environment_id: None,
harness: conversation_metadata
.and_then(|metadata| metadata.server_conversation_metadata.as_ref())
.map(|metadata| Harness::from(metadata.harness))
.or(Some(Harness::Oz)),
artifacts: conversation_artifacts(&metadata, history_model),
},
backing: AgentConversationBackingData {
has_loaded_conversation,
has_local_persisted_data,
has_cloud_data,
has_ambient_run: conversation_metadata
.is_some_and(AIConversationMetadata::is_ambient_agent_conversation),
},
capabilities: AgentConversationCapabilities {
can_open: has_local_persisted_data || has_cloud_data,
can_copy_link: server_conversation_token_for_conversation(
conversation_id,
Some(&metadata.nav_data),
history_model,
)
.is_some(),
can_share: history_model.can_conversation_be_shared(&conversation_id),
can_delete: has_local_persisted_data,
can_fork_locally: has_local_persisted_data,
can_cancel: status.is_cancellable(),
},
}
}
fn server_conversation_token_for_conversation(
conversation_id: AIConversationId,
nav_data: Option<&ConversationNavigationData>,
history_model: &BlocklistAIHistoryModel,
) -> Option<ServerConversationToken> {
history_model
.conversation(&conversation_id)
.and_then(|conversation| conversation.server_conversation_token())
.cloned()
.or_else(|| {
history_model
.get_conversation_metadata(&conversation_id)
.and_then(|metadata| metadata.server_conversation_token.clone())
})
.or_else(|| nav_data.and_then(|nav_data| nav_data.server_conversation_token.clone()))
}
pub(super) fn parse_session_id(session_id: &str) -> Option<SessionId> {
match session_id.parse::<SessionId>() {
Ok(session_id) => Some(session_id),
Err(e) => {
log::warn!("Failed to parse shared session ID: {e}");
None
}
}
}
app/src/ai/agent_conversations_model_tests.rs
+1435 -253
View File
File diff suppressed because it is too large Load Diff
app/src/ai/agent_events/driver.rs
+176 -36
View File
@@ -7,46 +7,102 @@ use futures::future::Either;
use futures::StreamExt;
use galaxyui::r#async::Timer;
use instant::Instant;
use galaxy_core::errors::AnyhowErrorExt as _;
use crate::server::retry_strategies::is_transient_http_error;
use crate::server::server_api::ai::AgentRunEvent;
use crate::server::server_api::presigned_upload::HttpStatusError;
use crate::server::server_api::ServerApi;
pub(crate) const DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
pub(crate) const DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS: &[u64] = &[30];
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;
/// Selects which server-side filter shape an [`AgentEventSource`] should use
/// when opening a stream.
///
/// `RunIds` maps to the `?run_ids[]=` query parameter on the SSE endpoint
/// and is used by child-only per-conversation streams and the dormant
/// Claude wake listener. `AncestorRunId` maps to the `?ancestor_run_id=`
/// shape: with `include_self=false` it streams events for every direct
/// child of the supplied parent run (the shared-session viewer's pill bar),
/// and with `include_self=true` it additionally streams the parent run's
/// own events so an owner-side orchestrator can receive child lifecycle
/// events plus its own inbox on one ordered stream.
#[derive(Clone, Debug)]
pub(crate) enum AgentEventFilter {
/// One stream per multiplexed set of run IDs. Matches today's
/// `?run_ids[]=` endpoint.
RunIds(Vec<String>),
/// Stream events for every direct child of the supplied parent run, and
/// (when `include_self` is true) the parent run itself. Matches the
/// `?ancestor_run_id=` endpoint.
AncestorRunId {
ancestor_run_id: String,
include_self: bool,
},
}
impl AgentEventFilter {
/// Returns a short debug label used in driver log lines so we don't have
/// to format the full `Vec<String>` payload on every retry.
pub(crate) fn log_label(&self) -> String {
match self {
AgentEventFilter::RunIds(ids) => format!("run_ids={ids:?}"),
AgentEventFilter::AncestorRunId {
ancestor_run_id,
include_self,
} => format!("ancestor_run_id={ancestor_run_id} include_self={include_self}"),
}
}
}
/// 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>,
/// Wire-level filter selecting which run IDs the stream serves. Either a
/// concrete multiplexed list of run IDs or an ancestor-scoped child set;
/// see [`AgentEventFilter`].
pub filter: AgentEventFilter,
/// 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.
/// Exponential-ish reconnect delays, in seconds, used after transient
/// stream failures (5xx, timeouts, connection resets).
pub reconnect_backoff_steps: &'static [u64],
/// Reconnect delays for permanent HTTP errors (4xx other than 408/429).
/// Typically much slower than transient backoff to reduce log spam while
/// still allowing recovery if the error was spurious.
pub permanent_error_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.
/// Failure count at which actionable reconnect failures are reported at Error level.
/// 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 {
/// orchestration and harness listeners, parameterised on the wire filter.
pub(crate) fn retry_forever(filter: AgentEventFilter, since_sequence: i64) -> Self {
Self {
run_ids,
filter,
since_sequence,
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
proactive_reconnect_after: Some(DEFAULT_AGENT_EVENT_PROACTIVE_RECONNECT),
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
}
}
/// Convenience: build a `retry_forever` config from a concrete list of
/// run IDs. Lets existing call sites keep their current ergonomics.
pub(crate) fn retry_forever_run_ids(run_ids: Vec<String>, since_sequence: i64) -> Self {
Self::retry_forever(AgentEventFilter::RunIds(run_ids), since_sequence)
}
}
/// Tells the shared driver whether to continue or stop after a handled event.
@@ -75,6 +131,27 @@ pub(crate) enum AgentEventDriverState {
ProactiveReconnect,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct AgentMessageEventMetadata {
pub sequence: i64,
pub message_id: String,
pub occurred_at: String,
}
impl AgentMessageEventMetadata {
pub(crate) fn from_event(event: &AgentRunEvent) -> Option<Self> {
if event.event_type != "new_message" {
return None;
}
Some(Self {
sequence: event.sequence,
message_id: event.ref_id.clone()?,
occurred_at: event.occurred_at.clone(),
})
}
}
/// Parsed items emitted by an [`AgentEventSource`].
pub(crate) enum AgentEventSourceItem {
Open,
@@ -91,13 +168,13 @@ cfg_if::cfg_if! {
}
}
/// Opens a stream of parsed agent events for one or more run IDs.
/// Opens a stream of parsed agent events for the supplied filter.
#[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],
filter: &AgentEventFilter,
since_sequence: i64,
) -> Result<AgentEventSourceStream>;
}
@@ -118,13 +195,28 @@ impl ServerApiAgentEventSource {
impl AgentEventSource for ServerApiAgentEventSource {
async fn open_stream(
&self,
run_ids: &[String],
filter: &AgentEventFilter,
since_sequence: i64,
) -> Result<AgentEventSourceStream> {
let stream = self
.server_api
.stream_agent_events(run_ids, since_sequence)
.await?;
let stream = match filter {
AgentEventFilter::RunIds(run_ids) => {
self.server_api
.stream_agent_events(run_ids, since_sequence)
.await?
}
AgentEventFilter::AncestorRunId {
ancestor_run_id,
include_self,
} => {
self.server_api
.stream_agent_events_for_ancestor(
ancestor_run_id,
*include_self,
since_sequence,
)
.await?
}
};
let stream = stream.filter_map(|event_result| async move {
match event_result {
@@ -138,7 +230,29 @@ impl AgentEventSource for ServerApiAgentEventSource {
}
}
}
Err(err) => Some(Err(anyhow!("SSE stream error: {err:?}"))),
Err(err) => {
let anyhow_err = match err {
reqwest_eventsource::Error::InvalidStatusCode(status_code, response) => {
let body = response
.text()
.await
.unwrap_or_else(|err| format!("(no response body: {err:#})"));
let status_err = HttpStatusError {
status: status_code.as_u16(),
body: body.clone(),
};
anyhow::Error::new(status_err).context(format!(
"SSE stream error: invalid status code {status_code}: {body}"
))
}
#[cfg(not(target_family = "wasm"))]
reqwest_eventsource::Error::Transport(err) => {
anyhow::Error::new(err).context("SSE stream error")
}
err => anyhow!("SSE stream error: {err:?}"),
};
Some(Err(anyhow_err))
}
}
});
@@ -188,18 +302,23 @@ where
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
}
// `open_stream` is lazy for the SSE-backed source: the TCP
// connect happens when the stream is first polled, not when
// this returns Ok. Wait for the `AgentEventSourceItem::Open`
// event below before declaring connectivity, so a server
// outage doesn't reset `failures` between every retry.
let mut stream = match source.open_stream(&config.filter, since_sequence).await {
Ok(stream) => stream,
Err(err) => {
failures += 1;
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
let backoff_steps = if is_transient_http_error(&err) {
config.reconnect_backoff_steps
} else {
config.permanent_error_backoff_steps
};
let backoff = agent_event_backoff(failures, backoff_steps);
log_stream_failure(
&config.run_ids,
&config.filter,
failures,
backoff,
&err,
@@ -249,7 +368,12 @@ where
}
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Open))) => {
failures = 0;
log::info!("Agent event stream opened for {:?}", config.run_ids);
has_connected_once = true;
notify_driver_state(consumer, AgentEventDriverState::Connected).await;
log::info!(
"Agent event stream opened for {}",
config.filter.log_label()
);
}
NextDriverItem::StreamItem(Some(Ok(AgentEventSourceItem::Event(event)))) => {
failures = 0;
@@ -273,9 +397,14 @@ where
}
NextDriverItem::StreamItem(Some(Err(err))) => {
failures += 1;
let backoff = agent_event_backoff(failures, config.reconnect_backoff_steps);
let backoff_steps = if is_transient_http_error(&err) {
config.reconnect_backoff_steps
} else {
config.permanent_error_backoff_steps
};
let backoff = agent_event_backoff(failures, backoff_steps);
log_stream_failure(
&config.run_ids,
&config.filter,
failures,
backoff,
&err,
@@ -293,12 +422,15 @@ where
Timer::after(backoff).await;
break;
}
// Clean stream closure (server-side close, not an HTTP
// error) — always use the transient backoff schedule since
// there is no HTTP status to classify.
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
"Agent event stream closed for {}, reconnecting in {backoff:?}",
config.filter.log_label()
);
notify_driver_state(
consumer,
@@ -332,21 +464,20 @@ async fn notify_driver_state<C: AgentEventConsumer>(
}
fn log_stream_failure(
run_ids: &[String],
filter: &AgentEventFilter,
failures: usize,
backoff: Duration,
err: &anyhow::Error,
failures_before_error_log: usize,
) {
if agent_event_failures_exceeded_threshold(failures, failures_before_error_log) {
let label = filter.log_label();
if agent_event_failure_should_log_error(err, failures, failures_before_error_log) {
log::error!(
"Agent event stream failed {failures} consecutive times for {:?}, retrying in {backoff:?}: {err:#}",
run_ids
"Agent event stream failed {failures} consecutive times for {label}, retrying in {backoff:?}: {err:#}"
);
} else {
log::warn!(
"Agent event stream failed for {:?}, retrying in {backoff:?}: {err:#}",
run_ids
"Agent event stream failed {failures} consecutive times for {label}, retrying in {backoff:?}: {err:#}"
);
}
}
@@ -361,6 +492,15 @@ pub(crate) fn agent_event_backoff(failures: usize, backoff_steps: &[u64]) -> Dur
Duration::from_secs(safe_steps[index])
}
#[cfg(test)]
pub(crate) fn agent_event_failures_exceeded_threshold(failures: usize, threshold: usize) -> bool {
failures >= threshold
}
pub(crate) fn agent_event_failure_should_log_error(
err: &anyhow::Error,
failures: usize,
threshold: usize,
) -> bool {
threshold > 0 && failures == threshold && err.is_actionable()
}
app/src/ai/agent_events/driver_tests.rs
+170 -6
View File
@@ -6,9 +6,12 @@ use anyhow::anyhow;
use async_trait::async_trait;
use futures::stream::{self, BoxStream};
use futures::StreamExt;
use galaxy_core::errors::AnyhowErrorExt as _;
use super::*;
use crate::ai::agent_events::driver::agent_event_failure_should_log_error;
use crate::server::server_api::ai::AgentRunEvent;
use crate::server::server_api::presigned_upload::HttpStatusError;
const ZERO_BACKOFF_STEPS: &[u64] = &[0];
@@ -28,7 +31,7 @@ impl FakeAgentEventSource {
impl AgentEventSource for FakeAgentEventSource {
async fn open_stream(
&self,
_run_ids: &[String],
_filter: &AgentEventFilter,
_since_sequence: i64,
) -> anyhow::Result<BoxStream<'static, anyhow::Result<AgentEventSourceItem>>> {
let response = self
@@ -136,9 +139,10 @@ async fn driver_skips_duplicate_sequences_and_persists_new_cursor() {
};
let config = AgentEventDriverConfig {
run_ids: vec!["child-run".to_string()],
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 2,
reconnect_backoff_steps: DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS,
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
proactive_reconnect_after: None,
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
};
@@ -181,9 +185,10 @@ async fn driver_resets_failures_after_successful_event_delivery() {
};
let config = AgentEventDriverConfig {
run_ids: vec!["child-run".to_string()],
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 0,
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
proactive_reconnect_after: None,
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
};
@@ -225,9 +230,10 @@ async fn driver_ignores_persist_cursor_errors() {
};
let config = AgentEventDriverConfig {
run_ids: vec!["child-run".to_string()],
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 0,
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
proactive_reconnect_after: None,
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
};
@@ -258,9 +264,10 @@ async fn driver_ignores_driver_state_errors() {
};
let config = AgentEventDriverConfig {
run_ids: vec!["child-run".to_string()],
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 0,
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
proactive_reconnect_after: None,
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
};
@@ -293,9 +300,10 @@ async fn driver_retries_initial_connection_until_stream_opens() {
};
let config = AgentEventDriverConfig {
run_ids: vec!["child-run".to_string()],
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 0,
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
permanent_error_backoff_steps: DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
proactive_reconnect_after: None,
failures_before_error_log: DEFAULT_AGENT_EVENT_FAILURES_BEFORE_ERROR_LOG,
};
@@ -350,3 +358,159 @@ fn failure_threshold_is_reached_at_and_above_limit() {
assert!(agent_event_failures_exceeded_threshold(5, 5));
assert!(agent_event_failures_exceeded_threshold(6, 5));
}
fn make_http_status_error(status: u16) -> anyhow::Error {
anyhow::Error::new(HttpStatusError {
status,
body: "not found".to_string(),
})
.context("SSE stream error")
}
#[test]
fn actionable_stream_status_reports_only_at_threshold_crossing() {
let err = make_http_status_error(400);
assert_eq!(
[
agent_event_failure_should_log_error(&err, 4, 5),
agent_event_failure_should_log_error(&err, 5, 5),
agent_event_failure_should_log_error(&err, 6, 5),
],
[false, true, false]
);
}
#[test]
fn zero_threshold_disables_stream_error_escalation() {
let err = make_http_status_error(400);
assert!(!agent_event_failure_should_log_error(&err, 1, 0));
}
#[test]
fn non_actionable_stream_statuses_do_not_report_at_threshold() {
for status in [408, 429] {
let err = make_http_status_error(status);
assert!(
!agent_event_failure_should_log_error(&err, 5, 5),
"status {status}"
);
}
}
#[test]
fn server_error_status_reports_at_threshold_crossing() {
let err = make_http_status_error(500);
assert!(agent_event_failure_should_log_error(&err, 5, 5));
}
#[test]
fn http_status_error_actionability_follows_status_classification() {
assert!(make_http_status_error(400).is_actionable());
assert!(make_http_status_error(500).is_actionable());
assert!(!make_http_status_error(429).is_actionable());
}
#[tokio::test]
async fn driver_uses_slow_backoff_on_permanent_http_error() {
let source = FakeAgentEventSource::new(vec![
// First attempt: stream opens then returns a 404-enriched error.
ok_stream(vec![
Ok(AgentEventSourceItem::Open),
Err(make_http_status_error(404)),
]),
// Second attempt: succeeds with a stoppable event.
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()
};
// Use asymmetric values: transient=9999s (would block if chosen),
// permanent=0s (instant). Proving backoff is 0s confirms the
// permanent schedule was selected.
let config = AgentEventDriverConfig {
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 0,
reconnect_backoff_steps: &[9999],
permanent_error_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]);
// The backoff must be from the permanent schedule (0s), not transient (9999s).
let retry_backoff = consumer
.driver_states
.iter()
.find_map(|s| match s {
AgentEventDriverState::RetryScheduled { backoff, .. } => Some(*backoff),
_ => None,
})
.unwrap();
assert_eq!(retry_backoff, Duration::from_secs(0));
}
#[tokio::test]
async fn driver_uses_fast_backoff_on_transient_http_error() {
let source = FakeAgentEventSource::new(vec![
// First attempt: stream opens then returns a 500-enriched error.
ok_stream(vec![
Ok(AgentEventSourceItem::Open),
Err(make_http_status_error(500)),
]),
// Second attempt: succeeds.
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()
};
// Set permanent backoff to something large so we can verify it was NOT used.
let config = AgentEventDriverConfig {
filter: AgentEventFilter::RunIds(vec!["child-run".to_string()]),
since_sequence: 0,
reconnect_backoff_steps: ZERO_BACKOFF_STEPS,
permanent_error_backoff_steps: &[9999],
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();
// Should have retried with the fast backoff (0s) and succeeded.
assert_eq!(consumer.handled_sequences, vec![1]);
// The backoff used should be from ZERO_BACKOFF_STEPS (0s), not permanent (9999s).
let retry_backoff = consumer
.driver_states
.iter()
.find_map(|s| match s {
AgentEventDriverState::RetryScheduled { backoff, .. } => Some(*backoff),
_ => None,
})
.unwrap();
assert_eq!(retry_backoff, Duration::from_secs(0));
}
app/src/ai/agent_events/message_hydrator.rs
+145 -16
View File
@@ -5,20 +5,33 @@ use anyhow::{anyhow, Context, Result};
#[cfg(not(target_family = "wasm"))]
use futures::future::Either;
#[cfg(not(target_family = "wasm"))]
use instant::Instant;
#[cfg(not(target_family = "wasm"))]
use reqwest::Error as ReqwestError;
#[cfg(not(target_family = "wasm"))]
use galaxyui::r#async::Timer;
use crate::ai::agent::ReceivedMessageInput;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::server::server_api::ai::{AIClient, AgentRunEvent, ReadAgentMessageResponse};
#[cfg(not(target_family = "wasm"))]
use crate::server::server_api::presigned_upload::HttpStatusError;
use crate::server::server_api::ServerApi;
pub(crate) const DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(5);
const DEFAULT_AGENT_MESSAGE_RETRY_DELAY: Duration = Duration::from_millis(50);
/// Hydrates `new_message` agent events into full message payloads and delivery
/// acknowledgements.
#[derive(Clone)]
pub(crate) struct MessageHydrator {
ai_client: Arc<dyn AIClient>,
task_scoped_server_api: Option<Arc<ServerApi>>,
task_id: Option<AmbientAgentTaskId>,
#[cfg_attr(target_family = "wasm", allow(dead_code))]
fetch_timeout: Duration,
#[cfg_attr(target_family = "wasm", allow(dead_code))]
retry_delay: Duration,
}
impl MessageHydrator {
@@ -26,16 +39,57 @@ impl MessageHydrator {
Self::with_fetch_timeout(ai_client, DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT)
}
pub(crate) fn for_task(server_api: Arc<ServerApi>, task_id: AmbientAgentTaskId) -> Self {
let ai_client: Arc<dyn AIClient> = server_api.clone();
Self {
ai_client,
task_scoped_server_api: Some(server_api),
task_id: Some(task_id),
fetch_timeout: DEFAULT_AGENT_MESSAGE_FETCH_TIMEOUT,
retry_delay: DEFAULT_AGENT_MESSAGE_RETRY_DELAY,
}
}
#[cfg(test)]
pub(crate) fn with_fetch_timing(
ai_client: Arc<dyn AIClient>,
fetch_timeout: Duration,
retry_delay: Duration,
) -> Self {
Self {
ai_client,
task_scoped_server_api: None,
task_id: None,
fetch_timeout,
retry_delay,
}
}
pub(crate) fn with_fetch_timeout(
ai_client: Arc<dyn AIClient>,
fetch_timeout: Duration,
) -> Self {
Self {
ai_client,
task_scoped_server_api: None,
task_id: None,
fetch_timeout,
retry_delay: DEFAULT_AGENT_MESSAGE_RETRY_DELAY,
}
}
async fn read_message(&self, message_id: &str) -> Result<ReadAgentMessageResponse> {
match (self.task_scoped_server_api.as_ref(), self.task_id) {
(Some(server_api), Some(task_id)) => {
server_api
.read_agent_message_for_task(&task_id, message_id)
.await
}
_ => self.ai_client.read_agent_message(message_id).await,
}
.with_context(|| format!("Failed to read agent message {message_id}"))
}
pub(crate) async fn hydrate_event_for_recipient(
&self,
event: &AgentRunEvent,
@@ -55,6 +109,17 @@ impl MessageHydrator {
return None;
}
};
if message.body.is_empty() {
log::warn!(
"Hydrated empty-body agent message: message_id={} event_sequence={} recipient_run_id={} sender_run_id={} subject={:?} task_id={:?}",
message.message_id,
event.sequence,
recipient_run_id,
message.sender_run_id,
message.subject,
self.task_id.map(|task_id| task_id.to_string())
);
}
Some(ReceivedMessageInput {
message_id: message.message_id,
@@ -70,16 +135,35 @@ impl MessageHydrator {
&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);
let deadline = Instant::now() + self.fetch_timeout;
let mut last_error = None;
match futures::future::select(read_message, timeout).await {
Either::Left((result, _)) => {
result.with_context(|| format!("Failed to read agent message {message_id}"))
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(message_read_timeout_error(message_id, last_error));
}
let read_message = self.read_message(message_id);
let timeout = Timer::after(remaining);
futures::pin_mut!(read_message);
futures::pin_mut!(timeout);
match futures::future::select(read_message, timeout).await {
Either::Left((Ok(message), _)) => return Ok(message),
Either::Left((Err(err), _)) if should_retry_message_read_error(&err) => {
last_error = Some(err);
let sleep_duration = self
.retry_delay
.min(deadline.saturating_duration_since(Instant::now()));
if sleep_duration.is_zero() {
return Err(message_read_timeout_error(message_id, last_error));
}
Timer::after(sleep_duration).await;
}
Either::Left((Err(err), _)) => return Err(err),
Either::Right(_) => return Err(message_read_timeout_error(message_id, last_error)),
}
Either::Right(_) => Err(anyhow!("Timed out reading agent message {message_id}")),
}
}
@@ -88,10 +172,7 @@ impl MessageHydrator {
&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}"))
self.read_message(message_id).await
}
pub(crate) async fn read_message_from_event_with_timeout(
@@ -105,10 +186,15 @@ impl MessageHydrator {
}
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"))
match (self.task_scoped_server_api.as_ref(), self.task_id) {
(Some(server_api), Some(task_id)) => {
server_api
.mark_message_delivered_for_task(&task_id, message_id)
.await
}
_ => 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>(
@@ -130,3 +216,46 @@ impl MessageHydrator {
failures
}
}
#[cfg(not(target_family = "wasm"))]
fn should_retry_message_read_error(err: &anyhow::Error) -> bool {
// Immediate read-after-event lag can surface as a short-lived 404 before
// the message row becomes readable. Restrict retries to status-preserving
// eventual-consistency/server responses plus clearly transient transport
// failures so permanent 4xxs fail fast instead of timing out.
if let Some(status) = message_read_error_status(err) {
return matches!(status, 404 | 408 | 429 | 500..=599);
}
err.chain().any(|cause| {
cause
.downcast_ref::<ReqwestError>()
.is_some_and(is_transient_message_read_transport_error)
})
}
#[cfg(not(target_family = "wasm"))]
fn message_read_error_status(err: &anyhow::Error) -> Option<u16> {
err.chain().find_map(|cause| {
cause
.downcast_ref::<HttpStatusError>()
.map(|http_err| http_err.status)
})
}
#[cfg(not(target_family = "wasm"))]
fn is_transient_message_read_transport_error(err: &ReqwestError) -> bool {
err.is_timeout() || err.is_connect()
}
#[cfg(not(target_family = "wasm"))]
fn message_read_timeout_error(
message_id: &str,
last_error: Option<anyhow::Error>,
) -> anyhow::Error {
let timeout_error = anyhow!("Timed out reading agent message {message_id}");
match last_error {
Some(err) => timeout_error.context(err),
None => timeout_error,
}
}
app/src/ai/agent_events/message_hydrator_tests.rs
+129 -11
View File
@@ -1,4 +1,6 @@
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use mockall::predicate::eq;
@@ -6,6 +8,7 @@ use super::*;
use crate::server::server_api::ai::{
AIClient, AgentRunEvent, MockAIClient, ReadAgentMessageResponse,
};
use crate::server::server_api::presigned_upload::HttpStatusError;
fn make_run_event(
sequence: i64,
@@ -23,6 +26,25 @@ fn make_run_event(
}
}
fn make_message_response(message_id: &str) -> ReadAgentMessageResponse {
ReadAgentMessageResponse {
message_id: message_id.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()),
}
}
fn http_status_read_error(status: u16) -> anyhow::Error {
anyhow::Error::new(HttpStatusError {
status,
body: format!("status {status} body"),
})
}
#[tokio::test]
async fn hydrator_reads_new_message_for_matching_run() {
let mut ai_client = MockAIClient::new();
@@ -30,17 +52,7 @@ async fn hydrator_reads_new_message_for_matching_run() {
.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()),
})
});
.returning(|_| Ok(make_message_response("msg-123")));
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
let hydrator = MessageHydrator::new(ai_client);
@@ -69,3 +81,109 @@ async fn hydrator_ignores_events_for_other_runs() {
.await
.is_none());
}
#[tokio::test]
async fn read_message_with_timeout_retries_transient_failures_until_success() {
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_clone = attempts.clone();
let mut ai_client = MockAIClient::new();
ai_client
.expect_read_agent_message()
.with(eq("msg-123"))
.times(2)
.returning(move |_| {
let attempt = attempts_clone.fetch_add(1, Ordering::SeqCst);
if attempt == 0 {
Err(http_status_read_error(404))
} else {
Ok(make_message_response("msg-123"))
}
});
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
let hydrator = MessageHydrator::with_fetch_timing(
ai_client,
Duration::from_millis(100),
Duration::from_millis(5),
);
let message = hydrator.read_message_with_timeout("msg-123").await.unwrap();
assert_eq!(attempts.load(Ordering::SeqCst), 2);
assert_eq!(message.message_id, "msg-123");
assert_eq!(message.body, "Switch to the failing test first.");
}
#[tokio::test]
async fn read_message_with_timeout_times_out_after_retrying_transient_failures() {
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_clone = attempts.clone();
let mut ai_client = MockAIClient::new();
ai_client
.expect_read_agent_message()
.with(eq("msg-123"))
.returning(move |_| {
attempts_clone.fetch_add(1, Ordering::SeqCst);
Err(http_status_read_error(404))
});
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
let hydrator = MessageHydrator::with_fetch_timing(
ai_client,
Duration::from_millis(120),
Duration::from_millis(20),
);
let err = hydrator
.read_message_with_timeout("msg-123")
.await
.expect_err("expected timeout after transient retries");
let err_chain = format!("{err:#}");
assert!(
err_chain.contains("Timed out reading agent message msg-123"),
"{err:#}"
);
assert!(
attempts.load(Ordering::SeqCst) >= 2,
"expected at least one retry before timeout"
);
}
#[tokio::test]
async fn read_message_with_timeout_does_not_retry_permanent_http_failures() {
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_clone = attempts.clone();
let mut ai_client = MockAIClient::new();
ai_client
.expect_read_agent_message()
.with(eq("msg-123"))
.times(1)
.returning(move |_| {
attempts_clone.fetch_add(1, Ordering::SeqCst);
Err(http_status_read_error(403))
});
let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
let hydrator = MessageHydrator::with_fetch_timing(
ai_client,
Duration::from_millis(100),
Duration::from_millis(5),
);
let err = hydrator
.read_message_with_timeout("msg-123")
.await
.expect_err("expected permanent failure to fail fast");
let err_chain = format!("{err:#}");
assert!(
err_chain.contains("HTTP request failed with status 403"),
"{err:#}"
);
assert_eq!(
attempts.load(Ordering::SeqCst),
1,
"permanent 4xx errors should not retry"
);
}
app/src/ai/agent_events/mod.rs
+2 -2
View File
@@ -8,11 +8,11 @@ mod message_hydrator;
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,
DEFAULT_AGENT_EVENT_RECONNECT_BACKOFF_STEPS, DEFAULT_PERMANENT_ERROR_BACKOFF_STEPS,
};
pub(crate) use driver::{
run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow,
AgentEventDriverConfig, ServerApiAgentEventSource,
AgentEventDriverConfig, AgentEventFilter, AgentMessageEventMetadata, ServerApiAgentEventSource,
};
pub(crate) use message_hydrator::MessageHydrator;
app/src/ai/agent_management/agent_management_model.rs
+106 -40
View File
@@ -1,9 +1,8 @@
use std::collections::HashMap;
use galaxy_core::features::FeatureFlag;
use galaxyui::{AppContext, Entity, EntityId, ModelContext, SingletonEntity, WindowId};
use crate::settings::AISettings;
use galaxy_core::send_telemetry_from_ctx;
use galaxyui::{AppContext, Entity, EntityId, ModelContext, SingletonEntity, ViewHandle, WindowId};
use crate::ai::active_agent_views_model::{ActiveAgentViewsEvent, ActiveAgentViewsModel};
use crate::ai::agent::conversation::{AIConversationId, ConversationStatus};
@@ -12,16 +11,16 @@ use crate::ai::agent_management::notifications::{
NotificationSourceAgent,
};
use crate::ai::artifacts::Artifact;
use crate::ai::blocklist::BlocklistAIHistoryEvent;
use crate::ai::blocklist::{BlocklistAIHistoryEvent, ConversationStatusUpdate, QueuedQueryModel};
use crate::server::telemetry::TelemetryEvent;
use crate::settings::AISettings;
use crate::terminal::cli_agent_sessions::{
CLIAgentSessionStatus, CLIAgentSessionsModel, CLIAgentSessionsModelEvent,
};
use crate::terminal::CLIAgent;
use crate::terminal::{CLIAgent, TerminalView};
use crate::workspace::util::is_terminal_view_in_same_tab;
use crate::workspace::{Workspace, WorkspaceRegistry};
use crate::BlocklistAIHistoryModel;
use galaxy_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.
@@ -43,17 +42,17 @@ 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| {
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| {
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| {
ctx.subscribe_to_model(&active_views_model, |me, _, event, ctx| {
me.handle_active_agent_views_changed(event, ctx);
});
@@ -165,14 +164,19 @@ impl AgentNotificationsModel {
CLIAgent::Codex => "Notification from Codex",
_ => "Task completed.",
};
let metadata = TerminalViewMetadata::lookup(*terminal_view_id, ctx);
self.add_notification(
title,
message.to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::CLI(*agent),
NotificationSourceAgent::CLI {
agent: *agent,
is_ambient: metadata.is_ambient,
},
NotificationOrigin::CLISession(*terminal_view_id),
*terminal_view_id,
vec![],
metadata.branch,
ctx,
);
}
@@ -180,16 +184,21 @@ impl AgentNotificationsModel {
let title = session_context
.display_title()
.unwrap_or_else(|| format!("{} needs attention", agent.display_name()));
let metadata = TerminalViewMetadata::lookup(*terminal_view_id, ctx);
self.add_notification(
title,
message
.clone()
.unwrap_or_else(|| "Waiting for input.".to_owned()),
NotificationCategory::Request,
NotificationSourceAgent::CLI(*agent),
NotificationSourceAgent::CLI {
agent: *agent,
is_ambient: metadata.is_ambient,
},
NotificationOrigin::CLISession(*terminal_view_id),
*terminal_view_id,
vec![],
metadata.branch,
ctx,
);
}
@@ -237,10 +246,11 @@ impl AgentNotificationsModel {
}
let BlocklistAIHistoryEvent::UpdatedConversationStatus {
terminal_view_id,
terminal_surface_id,
conversation_id,
// We shouldn't trigger toasts when restoring conversations on startup.
is_restored: false,
update: ConversationStatusUpdate::Changed { .. },
..
} = event
else {
return;
@@ -262,7 +272,7 @@ impl AgentNotificationsModel {
&status,
*conversation_id,
latest_query,
*terminal_view_id,
*terminal_surface_id,
ctx,
);
// The new mailbox path handled the event — skip the legacy toast path below.
@@ -273,7 +283,7 @@ impl AgentNotificationsModel {
return;
}
if is_terminal_view_visible(*terminal_view_id, ctx) {
if is_terminal_view_visible(*terminal_surface_id, ctx) {
return;
}
@@ -286,7 +296,7 @@ impl AgentNotificationsModel {
ctx.emit(AgentManagementEvent::ConversationNeedsAttention {
window_id,
tab_index,
terminal_view_id: *terminal_view_id,
terminal_view_id: *terminal_surface_id,
conversation_id: *conversation_id,
});
}
@@ -310,22 +320,36 @@ impl AgentNotificationsModel {
}
let title = latest_query.unwrap_or_else(|| "Agent task".to_owned());
let metadata = TerminalViewMetadata::lookup(terminal_view_id, ctx);
let oz_agent = NotificationSourceAgent::Oz {
is_ambient: metadata.is_ambient,
};
match status {
// When the agent resumes its work, clear stale notifications.
ConversationStatus::InProgress => {
// When the agent resumes its work (or is automatically recovering from a
// transient failure), clear stale notifications.
ConversationStatus::InProgress | ConversationStatus::TransientError => {
self.remove_notification_by_source(origin, ctx);
}
ConversationStatus::Success => {
// Suppress the completion notification when a queued follow-up prompt will
// auto-send as soon as this conversation finishes. The conversation isn't
// really in a stopped state, so the notification would be noisy. Pending
// artifacts are left intact so they roll into the notification fired when the
// conversation eventually finishes with an empty queue.
if QueuedQueryModel::as_ref(ctx).has_autofireable_prompt(conversation_id) {
return;
}
let artifacts = self.flush_pending_artifacts(conversation_id);
self.add_notification(
title,
"Task completed.".to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::Oz,
oz_agent,
origin,
terminal_view_id,
artifacts,
metadata.branch,
ctx,
);
}
@@ -335,10 +359,11 @@ impl AgentNotificationsModel {
title,
"Task was cancelled.".to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::Oz,
oz_agent,
origin,
terminal_view_id,
artifacts,
metadata.branch,
ctx,
);
}
@@ -347,10 +372,11 @@ impl AgentNotificationsModel {
title,
blocked_action.clone(),
NotificationCategory::Request,
NotificationSourceAgent::Oz,
oz_agent,
origin,
terminal_view_id,
vec![],
metadata.branch,
ctx,
);
}
@@ -360,13 +386,20 @@ impl AgentNotificationsModel {
title,
"Something went wrong.".to_owned(),
NotificationCategory::Error,
NotificationSourceAgent::Oz,
oz_agent,
origin,
terminal_view_id,
artifacts,
metadata.branch,
ctx,
);
}
// Yielded conversations are still active; mirror the
// InProgress arm and clear any stale notification for this
// origin.
ConversationStatus::WaitingForEvents => {
self.remove_notification_by_source(origin, ctx);
}
}
}
@@ -401,14 +434,12 @@ impl AgentNotificationsModel {
origin: NotificationOrigin,
terminal_view_id: EntityId,
artifacts: Vec<Artifact>,
branch: Option<String>,
ctx: &mut ModelContext<Self>,
) {
if !*AISettings::as_ref(ctx).show_agent_notifications {
return;
}
let show_agent_notifications = *AISettings::as_ref(ctx).show_agent_notifications;
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,
@@ -420,12 +451,14 @@ impl AgentNotificationsModel {
artifacts,
branch,
);
send_telemetry_from_ctx!(
TelemetryEvent::AgentNotificationShown {
agent_variant: agent.into(),
},
ctx
);
if show_agent_notifications {
send_telemetry_from_ctx!(
TelemetryEvent::AgentNotificationShown {
agent_variant: agent.into(),
},
ctx
);
}
let id = item.id;
self.notifications.push(item);
@@ -453,13 +486,22 @@ pub enum AgentManagementEvent {
impl ConversationStatus {
/// Returns true if the updating the conversation with this status should trigger some
/// notification to the user.
///
/// Exhaustive match so a new `ConversationStatus` variant forces a
/// deliberate decision about whether it should fire a notification.
pub fn should_trigger_notification(&self) -> bool {
matches!(
self,
match self {
ConversationStatus::Success
| ConversationStatus::Blocked { .. }
| ConversationStatus::Error
)
| ConversationStatus::Blocked { .. }
| ConversationStatus::Error => true,
// Streaming hasn't reached a notable state; a recovering or
// yielded conversation is still active; user-cancellations are
// self-evident.
ConversationStatus::InProgress
| ConversationStatus::TransientError
| ConversationStatus::WaitingForEvents
| ConversationStatus::Cancelled => false,
}
}
}
@@ -502,17 +544,41 @@ fn window_and_tab_idx_id_for_conversation(
})
}
fn resolve_git_branch_for_terminal_view(
/// Per-notification metadata derived from a single [`TerminalView`] lookup. Both fields
/// are read on the same emit path, so we resolve the view once and pass the projection
/// down rather than walking the workspace tree for each.
struct TerminalViewMetadata {
is_ambient: bool,
branch: Option<String>,
}
impl TerminalViewMetadata {
fn lookup(terminal_view_id: EntityId, app: &AppContext) -> Self {
let Some(terminal_view) = find_terminal_view_by_id(terminal_view_id, app) else {
return Self {
is_ambient: false,
branch: None,
};
};
let view = terminal_view.as_ref(app);
Self {
is_ambient: view.is_ambient_agent_session(app),
branch: view.current_git_branch(app),
}
}
}
fn find_terminal_view_by_id(
terminal_view_id: EntityId,
app: &AppContext,
) -> Option<String> {
) -> Option<ViewHandle<TerminalView>> {
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);
return Some(terminal_view);
}
}
}
app/src/ai/agent_management/agent_management_model_tests.rs
+241 -16
View File
@@ -1,14 +1,19 @@
use settings::Setting as _;
use galaxy_core::features::FeatureFlag;
use galaxyui::{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 galaxyui::{App, EntityId, ModelHandle, SingletonEntity};
use super::AgentNotificationsModel;
use crate::ai::active_agent_views_model::ActiveAgentViewsModel;
use crate::ai::agent::conversation::{AIConversation, AIConversationId, ConversationStatus};
use crate::ai::agent_management::notifications::{
NotificationCategory, NotificationFilter, NotificationOrigin, NotificationSourceAgent,
};
use crate::ai::artifacts::Artifact;
use crate::ai::blocklist::BlocklistAIHistoryEvent;
use crate::settings::AISettings;
use crate::terminal::cli_agent_sessions::CLIAgentSessionsModel;
use crate::test_util::settings::initialize_settings_for_tests;
use crate::{report_if_error, BlocklistAIHistoryModel};
fn setup_app(
app: &mut App,
@@ -16,7 +21,11 @@ fn setup_app(
ModelHandle<BlocklistAIHistoryModel>,
ModelHandle<AgentNotificationsModel>,
) {
initialize_settings_for_tests(app);
let history = app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[]));
// Registered after the history model since it subscribes to history events; the
// notifications model reads it to suppress completion notifications when a prompt is queued.
app.add_singleton_model(crate::ai::blocklist::QueuedQueryModel::new);
app.add_singleton_model(|_| CLIAgentSessionsModel::new());
app.add_singleton_model(|_| ActiveAgentViewsModel::new());
let notifications = app.add_singleton_model(AgentNotificationsModel::new);
@@ -51,7 +60,7 @@ fn artifact_event_accumulates_into_pending() {
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
terminal_surface_id: terminal_view_id,
conversation_id,
artifact: make_pr_artifact("https://github.com/org/repo/pull/42", "feature-branch"),
});
@@ -76,14 +85,14 @@ fn multiple_artifacts_accumulated_across_turns() {
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
terminal_surface_id: 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,
terminal_surface_id: terminal_view_id,
conversation_id,
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "main"),
});
@@ -98,6 +107,46 @@ fn multiple_artifacts_accumulated_across_turns() {
});
}
#[test]
fn add_notification_tracks_unread_activity_when_in_app_notifications_are_hidden() {
App::test((), |mut app| async move {
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
let (_history, notifications) = setup_app(&mut app);
AISettings::handle(&app).update(&mut app, |settings, ctx| {
report_if_error!(settings.show_agent_notifications.set_value(false, ctx));
});
let conversation_id = AIConversationId::new();
let terminal_view_id = EntityId::new();
notifications.update(&mut app, |model, ctx| {
model.add_notification(
"Agent task".to_owned(),
"Task completed.".to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::Oz { is_ambient: false },
NotificationOrigin::Conversation(conversation_id),
terminal_view_id,
vec![],
None,
ctx,
);
});
notifications.read(&app, |model, _| {
assert_eq!(
model
.notifications()
.filtered_count(NotificationFilter::All),
1
);
assert!(model
.notifications()
.has_unread_for_terminal_view(terminal_view_id));
});
});
}
#[test]
fn flush_drains_pending_artifacts() {
App::test((), |mut app| async move {
@@ -109,7 +158,7 @@ fn flush_drains_pending_artifacts() {
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
terminal_surface_id: terminal_view_id,
conversation_id,
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "branch-1"),
});
@@ -153,7 +202,7 @@ fn deletion_cleans_up_pending_artifacts() {
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
terminal_surface_id: terminal_view_id,
conversation_id,
artifact: make_pr_artifact("https://github.com/org/repo/pull/1", "branch-1"),
});
@@ -161,9 +210,10 @@ fn deletion_cleans_up_pending_artifacts() {
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
ctx.emit(BlocklistAIHistoryEvent::DeletedConversation {
terminal_view_id,
terminal_surface_id: terminal_view_id,
conversation_id,
conversation_title: None,
run_id: None,
});
});
@@ -185,14 +235,14 @@ fn separate_conversations_have_independent_pending_artifacts() {
history.update(&mut app, |_: &mut BlocklistAIHistoryModel, ctx| {
ctx.emit(BlocklistAIHistoryEvent::UpdatedConversationArtifacts {
terminal_view_id,
terminal_surface_id: 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,
terminal_surface_id: terminal_view_id,
conversation_id: conv_b,
artifact: make_plan_artifact("doc-b", "Plan B"),
});
@@ -209,3 +259,178 @@ fn separate_conversations_have_independent_pending_artifacts() {
});
});
}
// should_trigger_notification: pure-function tests pinning which statuses
// fire user-facing notifications. Terminal-error and blocked surface;
// in-progress, waiting-for-events, and user-cancelled do not.
#[test]
fn should_trigger_notification_returns_true_for_success() {
assert!(ConversationStatus::Success.should_trigger_notification());
}
#[test]
fn should_trigger_notification_returns_true_for_blocked() {
assert!(ConversationStatus::Blocked {
blocked_action: "approve diff".to_owned(),
}
.should_trigger_notification());
}
#[test]
fn should_trigger_notification_returns_true_for_error() {
assert!(ConversationStatus::Error.should_trigger_notification());
}
#[test]
fn should_trigger_notification_returns_false_for_in_progress() {
assert!(!ConversationStatus::InProgress.should_trigger_notification());
}
#[test]
fn should_trigger_notification_returns_false_for_waiting_for_events() {
assert!(!ConversationStatus::WaitingForEvents.should_trigger_notification());
}
#[test]
fn should_trigger_notification_returns_false_for_cancelled() {
assert!(!ConversationStatus::Cancelled.should_trigger_notification());
}
// Mailbox suppression for non-terminal status updates. In App::test the
// `is_conversation_open` gate always returns false, so the
// WaitingForEvents and InProgress arms both clear stale notifications
// regardless of status; this still pins the user-visible contract that
// no stale "Task completed" toast survives a non-terminal transition.
/// Disables `show_agent_notifications` so subsequent `add_notification`
/// calls skip the `send_telemetry_from_ctx!` branch — the test app does
/// not register a `TelemetryContextProvider` singleton and the macro
/// would otherwise panic.
fn disable_telemetry_path(app: &mut App) {
AISettings::handle(app).update(app, |settings, ctx| {
report_if_error!(settings.show_agent_notifications.set_value(false, ctx));
});
}
/// Pre-populates a `Complete` notification for `conversation_id` so that a
/// subsequent non-terminal status update has something to clear.
fn seed_stale_notification(
notifications: &ModelHandle<AgentNotificationsModel>,
app: &mut App,
conversation_id: AIConversationId,
terminal_view_id: EntityId,
) {
notifications.update(app, |model, ctx| {
model.add_notification(
"Agent task".to_owned(),
"Task completed.".to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::Oz { is_ambient: false },
NotificationOrigin::Conversation(conversation_id),
terminal_view_id,
vec![],
None,
ctx,
);
});
}
#[test]
fn waiting_for_events_clears_stale_notification_and_adds_none() {
App::test((), |mut app| async move {
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
let (history, notifications) = setup_app(&mut app);
disable_telemetry_path(&mut app);
let conversation = AIConversation::new(false, false);
let conversation_id = conversation.id();
let terminal_view_id = EntityId::new();
history.update(&mut app, |model, ctx| {
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
});
seed_stale_notification(&notifications, &mut app, conversation_id, terminal_view_id);
notifications.read(&app, |model, _| {
assert_eq!(
model
.notifications()
.filtered_count(NotificationFilter::All),
1,
"precondition: one stale notification queued"
);
});
history.update(&mut app, |model, ctx| {
let conv = model
.conversation_mut(&conversation_id)
.expect("conversation was just restored");
conv.update_status(ConversationStatus::WaitingForEvents, terminal_view_id, ctx);
});
notifications.read(&app, |model, _| {
assert_eq!(
model
.notifications()
.filtered_count(NotificationFilter::All),
0,
"WaitingForEvents must clear stale notifications and add no new toast"
);
});
});
}
#[test]
fn in_progress_resume_clears_stale_notification_and_adds_none() {
App::test((), |mut app| async move {
let _guard = FeatureFlag::HOANotifications.override_enabled(true);
let (history, notifications) = setup_app(&mut app);
disable_telemetry_path(&mut app);
let conversation = AIConversation::new(false, false);
let conversation_id = conversation.id();
let terminal_view_id = EntityId::new();
history.update(&mut app, |model, ctx| {
model.restore_conversations(terminal_view_id, vec![conversation], ctx);
});
// First move the conversation into WaitingForEvents, then back into
// InProgress. The second transition is the resume signal that
// PRODUCT.md (18) requires not to fire a notification.
history.update(&mut app, |model, ctx| {
let conv = model
.conversation_mut(&conversation_id)
.expect("conversation was just restored");
conv.update_status(ConversationStatus::WaitingForEvents, terminal_view_id, ctx);
});
seed_stale_notification(&notifications, &mut app, conversation_id, terminal_view_id);
notifications.read(&app, |model, _| {
assert_eq!(
model
.notifications()
.filtered_count(NotificationFilter::All),
1,
"precondition: one stale notification queued before the resume transition"
);
});
history.update(&mut app, |model, ctx| {
let conv = model
.conversation_mut(&conversation_id)
.expect("conversation still exists");
conv.update_status(ConversationStatus::InProgress, terminal_view_id, ctx);
});
notifications.read(&app, |model, _| {
assert_eq!(
model
.notifications()
.filtered_count(NotificationFilter::All),
0,
"WaitingForEvents → InProgress resume must not fire a notification \
(covers PRODUCT.md (18))"
);
});
});
}
app/src/ai/agent_management/agent_type_selector.rs
+5 -4
View File
@@ -3,8 +3,8 @@
//! 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 galaxy_core::ui::color::blend::Blend;
use galaxy_core::ui::theme::color::internal_colors;
use galaxyui::elements::{
@@ -17,8 +17,9 @@ use galaxyui::keymap::{FixedBinding, Keystroke};
use galaxyui::platform::Cursor;
use galaxyui::ui_components::components::{Coords, UiComponent, UiComponentStyles};
use galaxyui::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext};
use pathfinder_color::ColorU;
use pathfinder_geometry::vector::vec2f;
use crate::appearance::Appearance;
use crate::ui_components::icons::Icon;
// Modal dimensions based on Figma design.
const MODAL_WIDTH: f32 = 440.;
app/src/ai/agent_management/cloud_setup_guide_view.rs
+23 -20
View File
@@ -1,3 +1,25 @@
use std::collections::HashMap;
use serde::Serialize;
use string_offset::CharCounter;
use warp_completer::signatures::CommandRegistry;
use warp_completer::util::parse_current_commands_and_tokens;
use warp_completer::ParsedTokensSnapshot;
use galaxy_core::report_error;
use galaxy_core::ui::theme::{AnsiColorIdentifier, AnsiColors};
use warpui::clipboard::ClipboardContent;
use warpui::elements::new_scrollable::{ClippedAxisConfiguration, DualAxisConfig, NewScrollable};
use warpui::elements::{
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::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext, ViewHandle};
use crate::ai::agent_management::telemetry::{AgentManagementTelemetryEvent, SetupGuideStep};
use crate::ai::blocklist::code_block::{
render_code_block_plain, CodeBlockOptions, CodeSnippetButtonHandles,
@@ -8,26 +30,6 @@ 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 galaxy_completer::signatures::CommandRegistry;
use galaxy_completer::{util::parse_current_commands_and_tokens, ParsedTokensSnapshot};
use galaxy_core::report_error;
use galaxy_core::ui::theme::{AnsiColorIdentifier, AnsiColors};
use galaxyui::clipboard::ClipboardContent;
use galaxyui::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 galaxyui::fonts::{Properties, Weight};
use galaxyui::prelude::ChildView;
use galaxyui::text_layout::TextStyle;
use galaxyui::ui_components::components::{UiComponent, UiComponentStyles};
use galaxyui::ViewHandle;
use galaxyui::{AppContext, Entity, SingletonEntity, TypedActionView, View, ViewContext};
use serde::Serialize;
use std::collections::HashMap;
use string_offset::CharCounter;
const DOCS_URL: &str = "https://docs.warp.dev/agent-platform/cloud-agents/overview";
const ENV_DOCS_URL: &str =
@@ -412,6 +414,7 @@ impl CloudSetupGuideView {
mouse_handles: Some(handles),
file_path: None,
},
true,
app,
None,
)
app/src/ai/agent_management/details_action_buttons.rs
+5 -9
View File
@@ -4,14 +4,12 @@ use galaxy_core::ui::theme::AnsiColorIdentifier;
use galaxyui::elements::{ChildView, CrossAxisAlignment, Empty, Flex, ParentElement};
use galaxyui::{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::agent_conversations_model::{AgentConversationEntryId, AgentRunDisplayStatus};
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ui_components::icons::Icon;
use crate::view_components::action_button::{ActionButton, ButtonSize, SecondaryTheme};
use crate::view_components::copyable_text_field::COPY_FEEDBACK_DURATION;
use crate::workspace::WorkspaceAction;
const BUTTON_SPACING: f32 = 4.;
@@ -25,7 +23,7 @@ pub struct ActionButtonsConfig {
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>,
pub view_details_item_id: Option<AgentConversationEntryId>,
/// Conversation link URL (either to the transcript or live session) for copy link button.
pub copy_link_url: Option<String>,
}
@@ -87,7 +85,7 @@ pub enum AgentDetailsButtonEvent {
Open,
CancelTask { task_id: AmbientAgentTaskId },
ForkConversation { conversation_id: AIConversationId },
ViewDetails { item_id: ManagementCardItemId },
ViewDetails { item_id: AgentConversationEntryId },
CopyLink { link: String },
}
@@ -260,9 +258,7 @@ impl TypedActionView for ConversationActionButtonsRow {
}
AgentDetailsAction::ViewDetails => {
if let Some(item_id) = &self.config.view_details_item_id {
ctx.emit(AgentDetailsButtonEvent::ViewDetails {
item_id: item_id.clone(),
});
ctx.emit(AgentDetailsButtonEvent::ViewDetails { item_id: *item_id });
}
}
AgentDetailsAction::CopyLink => {
app/src/ai/agent_management/notifications/item.rs
+14 -3
View File
@@ -42,12 +42,23 @@ impl NotificationFilter {
}
}
/// Identifies the agent that produced a notification.
/// Identifies the agent that produced a notification, including whether the run was
/// ambient (cloud) or local. The `is_ambient` flag drives the cloud-lobe rendering in
/// [`render_agent_avatar`].
#[derive(Debug, Clone, Copy)]
#[allow(clippy::upper_case_acronyms)]
pub enum NotificationSourceAgent {
Oz,
CLI(CLIAgent),
Oz { is_ambient: bool },
CLI { agent: CLIAgent, is_ambient: bool },
}
impl NotificationSourceAgent {
pub fn is_ambient(&self) -> bool {
match self {
NotificationSourceAgent::Oz { is_ambient }
| NotificationSourceAgent::CLI { is_ambient, .. } => *is_ambient,
}
}
}
/// Identifies the conversation or session a notification belongs to.
app/src/ai/agent_management/notifications/item_rendering.rs
+14 -21
View File
@@ -1,7 +1,10 @@
use std::sync::Arc;
use pathfinder_color::ColorU;
use galaxy_core::ui::appearance::Appearance as CoreAppearance;
use galaxy_core::ui::icons::Icon;
use galaxy_core::ui::theme::{Fill, GalaxyTheme};
use galaxy_core::ui::theme::color::internal_colors;
use galaxy_core::ui::theme::{Fill, WarpTheme};
use galaxyui::clipboard::ClipboardContent;
use galaxyui::elements::{
ChildView, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, DispatchEventResult,
@@ -11,10 +14,6 @@ use galaxyui::elements::{
use galaxyui::fonts::Weight;
use galaxyui::ui_components::components::{UiComponent, UiComponentStyles};
use galaxyui::{View, ViewContext, ViewHandle};
use pathfinder_color::ColorU;
use galaxy_core::ui::appearance::Appearance as CoreAppearance;
use galaxy_core::ui::theme::color::internal_colors;
use crate::ai::agent::conversation::ConversationStatus;
use crate::ai::agent_management::notifications::item::NotificationSourceAgent;
@@ -25,9 +24,7 @@ use crate::ai::artifacts::{
};
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::ui_components::icon_with_status::{render_icon_with_status, IconWithStatusVariant};
use crate::util::time_format::format_elapsed_since;
use crate::view_components::action_button::ActionButtonTheme;
use crate::workspace::WorkspaceAction;
@@ -399,14 +396,8 @@ fn render_message_text(message: &str, expanded: bool, appearance: &Appearance) -
.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.),
};
/// Total size of the agent avatar component rendered alongside each notification.
const NOTIFICATION_AVATAR_SIZE: f32 = 32.;
fn render_agent_avatar(
agent: NotificationSourceAgent,
@@ -415,18 +406,20 @@ fn render_agent_avatar(
) -> Box<dyn Element> {
let status = notification_category_to_conversation_status(category);
let variant = match agent {
NotificationSourceAgent::Oz => IconWithStatusVariant::OzAgent {
NotificationSourceAgent::Oz { is_ambient } => IconWithStatusVariant::OzAgent {
status: Some(status),
is_ambient: false,
is_ambient,
},
NotificationSourceAgent::CLI(cli) => IconWithStatusVariant::CLIAgent {
agent: cli,
NotificationSourceAgent::CLI { agent, is_ambient } => IconWithStatusVariant::CLIAgent {
agent,
status: Some(status),
is_ambient,
},
};
render_icon_with_status(
variant,
&NOTIFICATION_AVATAR_SIZING,
NOTIFICATION_AVATAR_SIZE,
0.,
theme,
theme.surface_2(),
)
app/src/ai/agent_management/notifications/item_tests.rs
+5 -2
View File
@@ -12,7 +12,7 @@ fn make_conversation_notification(
"test".to_owned(),
"msg".to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::Oz,
NotificationSourceAgent::Oz { is_ambient: false },
NotificationOrigin::Conversation(conversation_id),
false,
terminal_view_id,
@@ -26,7 +26,10 @@ fn make_cli_session_notification(terminal_view_id: EntityId) -> NotificationItem
"cli test".to_owned(),
"cli msg".to_owned(),
NotificationCategory::Complete,
NotificationSourceAgent::CLI(CLIAgent::Claude),
NotificationSourceAgent::CLI {
agent: CLIAgent::Claude,
is_ambient: false,
},
NotificationOrigin::CLISession(terminal_view_id),
false,
terminal_view_id,
app/src/ai/agent_management/telemetry.rs
+20
View File
@@ -108,9 +108,14 @@ pub enum AgentManagementTelemetryEvent {
/// User clicked "Continue locally" in the tombstone
#[cfg(not(target_family = "wasm"))]
TombstoneContinueLocally,
/// User clicked "Continue" in the tombstone to start a cloud follow-up.
TombstoneContinueInCloud { task_id: String },
/// User clicked "Continue locally" in the details panel
#[cfg(not(target_family = "wasm"))]
DetailsPanelContinueLocally,
/// User invoked the /continue-locally slash command
#[cfg(not(target_family = "wasm"))]
SlashCommandContinueLocally,
/// User clicked "Open in Warp" in the tombstone (wasm)
#[cfg(target_family = "wasm")]
TombstoneOpenInWarp,
@@ -190,8 +195,13 @@ impl TelemetryEvent for AgentManagementTelemetryEvent {
}
#[cfg(not(target_family = "wasm"))]
AgentManagementTelemetryEvent::TombstoneContinueLocally => None,
AgentManagementTelemetryEvent::TombstoneContinueInCloud { task_id } => Some(json!({
"task_id": task_id,
})),
#[cfg(not(target_family = "wasm"))]
AgentManagementTelemetryEvent::DetailsPanelContinueLocally => None,
#[cfg(not(target_family = "wasm"))]
AgentManagementTelemetryEvent::SlashCommandContinueLocally => None,
#[cfg(target_family = "wasm")]
AgentManagementTelemetryEvent::TombstoneOpenInWarp => None,
AgentManagementTelemetryEvent::CloudRunCancelled { task_id } => {
@@ -242,8 +252,11 @@ impl TelemetryEventDesc for AgentManagementTelemetryEventDiscriminants {
Self::TombstoneArtifactClicked => "AgentManagement.TombstoneArtifactClicked",
#[cfg(not(target_family = "wasm"))]
Self::TombstoneContinueLocally => "AgentManagement.TombstoneContinueLocally",
Self::TombstoneContinueInCloud => "AgentManagement.TombstoneContinueInCloud",
#[cfg(not(target_family = "wasm"))]
Self::DetailsPanelContinueLocally => "AgentManagement.DetailsPanelContinueLocally",
#[cfg(not(target_family = "wasm"))]
Self::SlashCommandContinueLocally => "AgentManagement.SlashCommandContinueLocally",
#[cfg(target_family = "wasm")]
Self::TombstoneOpenInWarp => "AgentManagement.TombstoneOpenInWarp",
Self::CloudRunCancelled => "AgentManagement.CloudRunCancelled",
@@ -274,10 +287,17 @@ impl TelemetryEventDesc for AgentManagementTelemetryEventDiscriminants {
Self::TombstoneArtifactClicked => "User clicked an artifact in the tombstone view",
#[cfg(not(target_family = "wasm"))]
Self::TombstoneContinueLocally => "User clicked Continue locally in the tombstone",
Self::TombstoneContinueInCloud => {
"User clicked Continue in the tombstone to start a cloud follow-up"
}
#[cfg(not(target_family = "wasm"))]
Self::DetailsPanelContinueLocally => {
"User clicked Continue locally in the details panel"
}
#[cfg(not(target_family = "wasm"))]
Self::SlashCommandContinueLocally => {
"User invoked /continue-locally to fork a cloud conversation locally"
}
#[cfg(target_family = "wasm")]
Self::TombstoneOpenInWarp => "User clicked Open in Warp in the tombstone",
Self::CloudRunCancelled => "User cancelled a cloud run",
app/src/ai/agent_management/view.rs
+305 -308
View File
@@ -8,62 +8,9 @@ use galaxyui::scene::DropShadow;
use galaxyui::ui_components::button::ButtonVariant;
use lazy_static::lazy_static;
use pathfinder_color::ColorU;
use pathfinder_geometry::vector::vec2f;
use settings::Setting;
use siphasher::sip::SipHasher;
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_conversations_model::{
AgentConversationsModel, AgentConversationsModelEvent, AgentManagementFilters, ArtifactFilter,
ConversationOrTask, CreatedOnFilter, CreatorFilter, EnvironmentFilter, HarnessFilter,
OwnerFilter, SessionStatus, SourceFilter, StatusFilter,
};
use crate::ai::agent_management::agent_type_selector::{
AgentType, AgentTypeSelector, AgentTypeSelectorEvent,
};
use crate::ai::agent_management::cloud_setup_guide_view::{
CloudSetupGuideEvent, CloudSetupGuideView,
};
use crate::ai::agent_management::details_action_buttons::{
ActionButtonsConfig, AgentDetailsButtonEvent, ConversationActionButtonsRow,
};
use crate::ai::agent_management::telemetry::{
AgentManagementTelemetryEvent, ArtifactType, FilterType, OpenedFrom,
};
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::ambient_agents::{cancel_task_with_toast, AgentSource};
use crate::ai::artifacts::{Artifact, ArtifactButtonsRow, ArtifactButtonsRowEvent};
use crate::ai::blocklist::history_model::BlocklistAIHistoryModel;
use crate::ai::conversation_details_panel::{
ConversationDetailsData, ConversationDetailsPanel, ConversationDetailsPanelEvent,
};
use crate::ai::conversation_status_ui::render_status_element;
use crate::ai::harness_display;
use crate::app_state::PersistedAgentManagementFilters;
use crate::appearance::Appearance;
use crate::auth::AuthStateProvider;
use crate::editor::{
EditorView, Event as EditorEvent, PropagateAndNoOpNavigationKeys,
PropagateHorizontalNavigationKeys, SingleLineEditorOptions, TextOptions,
};
use crate::menu::{MenuItem, MenuItemFields};
use crate::notebooks::NotebookId;
use crate::settings::ai::AISettings;
use crate::ui_components::avatar::{Avatar, AvatarContent};
use crate::util::time_format::format_approx_duration_from_now_utc;
use crate::view_components::action_button::{
ActionButton, ButtonSize, NakedTheme, PrimaryTheme, SecondaryTheme,
};
use crate::view_components::compactible_action_button::{
CompactibleActionButton, MEDIUM_SIZE_SWITCH_THRESHOLD,
};
use crate::view_components::dropdown::{Dropdown, DropdownAction, DropdownStyle};
use crate::view_components::DismissibleToast;
use crate::view_components::FilterableDropdown;
use crate::workflows::WorkflowType;
use crate::workspace::{ForkedConversationDestination, ToastStack};
use crate::workspace::{RestoreConversationLayout, WorkspaceAction};
use crate::workspaces::user_workspaces::UserWorkspaces;
use crate::{send_telemetry_from_ctx, AgentModeEntrypoint};
use galaxy_cli::agent::Harness;
use galaxy_core::ui::icons::Icon;
use galaxy_core::ui::theme::color::internal_colors;
use galaxy_core::ui::theme::Fill;
@@ -79,15 +26,71 @@ use galaxyui::elements::{
SizeConstraintCondition, SizeConstraintSwitch, Stack, Text, Wrap,
};
use galaxyui::fonts::{Properties, Weight};
use galaxyui::keymap::FixedBinding;
use galaxyui::platform::Cursor;
use galaxyui::ui_components::components::UiComponent;
use galaxyui::ui_components::components::UiComponentStyles;
use galaxyui::ui_components::components::{UiComponent, UiComponentStyles};
use galaxyui::{
keymap::FixedBinding, Action, AppContext, Entity, FocusContext, ModelHandle, SingletonEntity,
TypedActionView, View, ViewContext, ViewHandle, WeakViewHandle,
AppContext, Entity, FocusContext, ModelHandle, SingletonEntity, TypedActionView, View,
ViewContext, ViewHandle, WeakViewHandle,
};
use pathfinder_geometry::vector::vec2f;
use settings::Setting;
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_conversations_model::{
AgentConversationEntry, AgentConversationEntryId, AgentConversationNavigationSubject,
AgentConversationsModel, AgentConversationsModelEvent, AgentManagementFilters, ArtifactFilter,
ConversationUpdateKind, CreatedOnFilter, CreatorFilter, EnvironmentFilter, HarnessFilter,
OwnerFilter, SessionStatus, SourceFilter, StatusFilter,
};
use crate::ai::agent_management::agent_type_selector::{
AgentType, AgentTypeSelector, AgentTypeSelectorEvent,
};
use crate::ai::agent_management::cloud_setup_guide_view::{
CloudSetupGuideEvent, CloudSetupGuideView,
};
use crate::ai::agent_management::details_action_buttons::{
ActionButtonsConfig, AgentDetailsButtonEvent, ConversationActionButtonsRow,
};
use crate::ai::agent_management::telemetry::{
AgentManagementTelemetryEvent, ArtifactType, FilterType, OpenedFrom,
};
use crate::ai::ambient_agents::{cancel_task_with_toast, AgentSource};
use crate::ai::artifacts::{Artifact, ArtifactButtonsRow, ArtifactButtonsRowEvent};
use crate::ai::blocklist::format_credits;
use crate::ai::conversation_details_panel::{
ConversationDetailsData, ConversationDetailsPanel, ConversationDetailsPanelEvent,
};
use crate::ai::harness_availability::HarnessAvailabilityModel;
use crate::ai::harness_display;
use crate::app_state::PersistedAgentManagementFilters;
use crate::appearance::Appearance;
use crate::auth::AuthStateProvider;
use crate::editor::{
EditorView, Event as EditorEvent, PropagateAndNoOpNavigationKeys,
PropagateHorizontalNavigationKeys, SingleLineEditorOptions, TextOptions,
};
use crate::menu::{MenuItem, MenuItemFields};
use crate::notebooks::NotebookId;
use crate::settings::ai::AISettings;
use crate::ui_components::agent_icon::agent_conversation_entry_icon_variant;
use crate::ui_components::avatar::{Avatar, AvatarContent};
use crate::ui_components::icon_with_status::render_icon_with_status;
use crate::util::time_format::format_approx_duration_from_now_utc;
use crate::view_components::action_button::{
ActionButton, ButtonSize, NakedTheme, PrimaryTheme, SecondaryTheme,
};
use crate::view_components::compactible_action_button::{
CompactibleActionButton, MEDIUM_SIZE_SWITCH_THRESHOLD,
};
use crate::view_components::dropdown::{
Dropdown, DropdownAction, DropdownItemAction, DropdownStyle,
};
use crate::view_components::{DismissibleToast, FilterableDropdown};
use crate::workflows::WorkflowType;
use crate::workspace::{
ForkedConversationDestination, RestoreConversationLayout, ToastStack, WorkspaceAction,
};
use crate::workspaces::user_workspaces::UserWorkspaces;
use crate::{send_telemetry_from_ctx, AgentModeEntrypoint};
lazy_static! {
static ref HASHER: SipHasher = SipHasher::new_with_keys(0, 0);
@@ -105,8 +108,9 @@ const CARD_CONTENT_PADDING: f32 = 12.;
const CARD_BORDER_RADIUS: f32 = 4.;
const CARD_MARGIN_BOTTOM: f32 = 8.;
const STATUS_ICON_SIZE: f32 = 12.;
const BUTTON_SIZE: f32 = 20.;
/// Total size of the agent icon-with-status component rendered in each card's header row.
const CARD_AGENT_ICON_SIZE: f32 = 24.;
const CREATOR_AVATAR_FONT_SIZE: f32 = 10.;
const SESSION_EXPIRED_TEXT: &str = "Sessions expire after one week and cannot be opened.";
@@ -125,21 +129,7 @@ fn should_show_artifacts(artifacts: &[Artifact]) -> bool {
!artifacts.is_empty() && FeatureFlag::ConversationArtifacts.is_enabled()
}
/// Identifies a card item - either a task ID or a conversation ID
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ManagementCardItemId {
Task(AmbientAgentTaskId),
Conversation(AIConversationId),
}
impl ManagementCardItemId {
fn as_key(&self) -> String {
match self {
ManagementCardItemId::Task(id) => format!("task_{id}"),
ManagementCardItemId::Conversation(id) => format!("conv_{id}"),
}
}
}
pub type ManagementCardItemId = AgentConversationEntryId;
/// Store state for a given task row
struct CardState {
@@ -222,6 +212,13 @@ impl AgentManagementView {
Self::handle_agent_management_model_event,
);
ctx.subscribe_to_model(
&HarnessAvailabilityModel::handle(ctx),
|me, _, _event, ctx| {
me.update_harness_dropdown(ctx);
},
);
let list_state = Self::construct_fresh_list_state(ctx.handle());
let all_filter_button = ctx.add_typed_action_view(|_ctx| {
@@ -496,7 +493,7 @@ impl AgentManagementView {
let make_status_option =
|label: &str, action: AgentManagementViewAction, icon_data: Option<(Icon, Fill)>| {
let mut fields = MenuItemFields::new(label)
.with_on_select_action(DropdownAction::SelectActionAndClose(action));
.with_on_select_action(DropdownAction::select_action_and_close(action));
if let Some((icon, color)) = icon_data {
fields = fields.with_icon(icon).with_override_icon_color(color);
}
@@ -546,7 +543,7 @@ impl AgentManagementView {
}
/// Build the list of source filter items.
fn build_source_dropdown_items() -> Vec<MenuItem<DropdownAction<AgentManagementViewAction>>> {
fn build_source_dropdown_items() -> Vec<MenuItem<DropdownAction>> {
// Build up the sources list
let mut sources = vec![
AgentSource::WebApp,
@@ -564,14 +561,16 @@ impl AgentManagementView {
}
let mut items = vec![MenuItem::Item(
MenuItemFields::new("All").with_on_select_action(DropdownAction::SelectActionAndClose(
AgentManagementViewAction::SetSourceFilter(SourceFilter::All),
)),
MenuItemFields::new("All").with_on_select_action(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetSourceFilter(SourceFilter::All),
),
),
)];
for source in sources {
items.push(MenuItem::Item(
MenuItemFields::new(source.display_name()).with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetSourceFilter(SourceFilter::Specific(source)),
),
),
@@ -601,22 +600,22 @@ impl AgentManagementView {
let items = vec![
MenuItem::Item(MenuItemFields::new("All").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetCreatedOnFilter(CreatedOnFilter::All),
),
)),
MenuItem::Item(MenuItemFields::new("Last 24 hours").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetCreatedOnFilter(CreatedOnFilter::Last24Hours),
),
)),
MenuItem::Item(MenuItemFields::new("Past 3 days").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetCreatedOnFilter(CreatedOnFilter::Past3Days),
),
)),
MenuItem::Item(MenuItemFields::new("Last week").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetCreatedOnFilter(CreatedOnFilter::LastWeek),
),
)),
@@ -635,29 +634,29 @@ impl AgentManagementView {
let items = vec![
MenuItem::Item(MenuItemFields::new("All").with_on_select_action(
DropdownAction::SelectActionAndClose(AgentManagementViewAction::SetArtifactFilter(
ArtifactFilter::All,
)),
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetArtifactFilter(ArtifactFilter::All),
),
)),
MenuItem::Item(MenuItemFields::new("Pull Request").with_on_select_action(
DropdownAction::SelectActionAndClose(AgentManagementViewAction::SetArtifactFilter(
ArtifactFilter::PullRequest,
)),
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetArtifactFilter(ArtifactFilter::PullRequest),
),
)),
MenuItem::Item(MenuItemFields::new("Plan").with_on_select_action(
DropdownAction::SelectActionAndClose(AgentManagementViewAction::SetArtifactFilter(
ArtifactFilter::Plan,
)),
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetArtifactFilter(ArtifactFilter::Plan),
),
)),
MenuItem::Item(MenuItemFields::new("Screenshot").with_on_select_action(
DropdownAction::SelectActionAndClose(AgentManagementViewAction::SetArtifactFilter(
ArtifactFilter::Screenshot,
)),
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetArtifactFilter(ArtifactFilter::Screenshot),
),
)),
MenuItem::Item(MenuItemFields::new("File").with_on_select_action(
DropdownAction::SelectActionAndClose(AgentManagementViewAction::SetArtifactFilter(
ArtifactFilter::File,
)),
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetArtifactFilter(ArtifactFilter::File),
),
)),
];
@@ -672,17 +671,27 @@ impl AgentManagementView {
let mut dropdown = Dropdown::new(ctx);
Self::setup_filter_menu(&mut dropdown, "Harness", ctx);
// "All" has no leading icon, matching the Status dropdown's "All" row.
let items = Self::build_harness_dropdown_items(ctx);
dropdown.set_rich_items(items, ctx);
dropdown.set_selected_by_index(0, ctx);
dropdown
}
fn build_harness_dropdown_items(app: &AppContext) -> Vec<MenuItem<DropdownAction>> {
let mut items = vec![MenuItem::Item(
MenuItemFields::new("All").with_on_select_action(DropdownAction::SelectActionAndClose(
AgentManagementViewAction::SetHarnessFilter(HarnessFilter::All),
)),
MenuItemFields::new("All").with_on_select_action(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetHarnessFilter(HarnessFilter::All),
),
),
)];
for harness in [Harness::Oz, Harness::Claude, Harness::Gemini] {
let mut fields = MenuItemFields::new(harness_display::display_name(harness))
let availability = HarnessAvailabilityModel::as_ref(app);
for entry in availability.available_harnesses() {
let harness = entry.harness;
let mut fields = MenuItemFields::new(entry.display_name.clone())
.with_icon(harness_display::icon_for(harness))
.with_on_select_action(DropdownAction::SelectActionAndClose(
.with_on_select_action(DropdownAction::select_action_and_close(
AgentManagementViewAction::SetHarnessFilter(HarnessFilter::Specific(harness)),
));
if let Some(color) = harness_display::brand_color(harness) {
@@ -691,9 +700,7 @@ impl AgentManagementView {
items.push(MenuItem::Item(fields));
}
dropdown.set_rich_items(items, ctx);
dropdown.set_selected_by_index(0, ctx);
dropdown
items
}
fn create_environment_dropdown(
@@ -732,7 +739,7 @@ impl AgentManagementView {
}
// Initialize the dropdown menu for the filter dropdowns (status, source)
fn setup_filter_menu<A: Action + Clone>(
fn setup_filter_menu<A: DropdownItemAction>(
dropdown: &mut Dropdown<A>,
label_prefix: &'static str,
ctx: &mut ViewContext<Dropdown<A>>,
@@ -744,7 +751,7 @@ impl AgentManagementView {
}
// Initialize the dropdown menu for the searchable filter dropdowns (creator)
fn setup_searchable_filter_menu<A: Action + Clone>(
fn setup_searchable_filter_menu<A: DropdownItemAction>(
dropdown: &mut FilterableDropdown<A>,
label_prefix: &'static str,
ctx: &mut ViewContext<FilterableDropdown<A>>,
@@ -755,6 +762,13 @@ impl AgentManagementView {
dropdown.set_button_variant(ButtonVariant::Secondary);
}
fn update_harness_dropdown(&mut self, ctx: &mut ViewContext<Self>) {
let items = Self::build_harness_dropdown_items(ctx);
self.harness_dropdown.update(ctx, |dropdown, ctx| {
dropdown.set_rich_items(items, ctx);
});
}
/// Since the valid set of environments depends on what tasks we have loaded in,
/// we use this function to update the available options depending on the most recent
/// set of tasks.
@@ -771,7 +785,7 @@ impl AgentManagementView {
self.environment_dropdown.update(ctx, |dropdown, ctx| {
let mut items = vec![MenuItem::Item(
MenuItemFields::new("All").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetEnvironmentFilter(EnvironmentFilter::All),
),
),
@@ -779,7 +793,7 @@ impl AgentManagementView {
items.push(MenuItem::Item(
MenuItemFields::new("None").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetEnvironmentFilter(
EnvironmentFilter::NoEnvironment,
),
@@ -793,7 +807,7 @@ impl AgentManagementView {
for (environment_id, environment_name) in sorted_envs {
items.push(MenuItem::Item(
MenuItemFields::new(environment_name).with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetEnvironmentFilter(
EnvironmentFilter::Specific(environment_id),
),
@@ -818,7 +832,7 @@ impl AgentManagementView {
self.creator_dropdown.update(ctx, |dropdown, ctx| {
let mut items = vec![MenuItem::Item(
MenuItemFields::new("All").with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetCreatorFilter(CreatorFilter::All),
),
),
@@ -826,7 +840,7 @@ impl AgentManagementView {
for (name, uid) in creators {
items.push(MenuItem::Item(
MenuItemFields::new(&name).with_on_select_action(
DropdownAction::SelectActionAndClose(
DropdownAction::select_action_and_close(
AgentManagementViewAction::SetCreatorFilter(CreatorFilter::Specific {
name,
uid,
@@ -938,46 +952,34 @@ impl AgentManagementView {
let model = AgentConversationsModel::as_ref(ctx);
let search_query = self.search_query.trim().to_lowercase();
let cards: Vec<CardData> = model
.get_tasks_and_conversations(&self.filters, ctx)
.filter(|t| {
.get_entries(&self.filters, ctx)
.into_iter()
.filter(|entry| {
if search_query.is_empty() {
return true;
}
match_indices_case_insensitive(&t.title(ctx), &search_query).is_some()
match_indices_case_insensitive(&entry.display.title, &search_query).is_some()
})
.map(|t| {
let item_id = match t {
ConversationOrTask::Task(task) => ManagementCardItemId::Task(task.task_id),
ConversationOrTask::Conversation(conversation) => {
ManagementCardItemId::Conversation(conversation.nav_data.id)
}
};
let artifacts = t.artifacts(ctx);
let copy_link_url = t.session_or_conversation_link(ctx);
let mut config = match t {
ConversationOrTask::Task(task) => ActionButtonsConfig::for_task(
task.task_id,
&t.display_status(ctx),
None, // Don't show open button in card hover
copy_link_url,
),
ConversationOrTask::Conversation(conversation) => {
ActionButtonsConfig::for_conversation(
conversation.nav_data.id,
None, // Don't show open button in card hover
copy_link_url,
.map(|entry| {
let item_id = entry.id;
let copy_link_url = entry
.capabilities
.can_copy_link
.then(|| {
AgentConversationsModel::resolve_copy_link(
AgentConversationNavigationSubject::Entry(entry.id),
ctx,
)
}
};
// Show info button in card hover for ViewDetails if feature flag enabled
})
.flatten();
let mut config = Self::action_buttons_config_for_entry(&entry, None, copy_link_url);
if FeatureFlag::AgentManagementDetailsView.is_enabled() {
config.view_details_item_id = Some(item_id.clone());
config.view_details_item_id = Some(item_id);
}
CardData {
item_id,
artifacts,
artifacts: entry.display.artifacts,
action_buttons_config: config,
}
})
@@ -1026,11 +1028,8 @@ impl AgentManagementView {
} else {
None
};
let action_buttons_view = self.create_action_buttons_view(
card.item_id.clone(),
card.action_buttons_config,
ctx,
);
let action_buttons_view =
self.create_action_buttons_view(card.item_id, card.action_buttons_config, ctx);
new_items.push(CardState {
hover_state: MouseStateHandle::default(),
@@ -1074,6 +1073,29 @@ impl AgentManagementView {
view
}
fn action_buttons_config_for_entry(
entry: &AgentConversationEntry,
open_action: Option<WorkspaceAction>,
copy_link_url: Option<String>,
) -> ActionButtonsConfig {
if let Some(task_id) = entry.identity.ambient_agent_task_id {
ActionButtonsConfig::for_task(
task_id,
&entry.display.status,
open_action,
copy_link_url,
)
} else if let Some(conversation_id) = entry.identity.local_conversation_id {
ActionButtonsConfig::for_conversation(conversation_id, open_action, copy_link_url)
} else {
ActionButtonsConfig {
open_action,
copy_link_url,
..Default::default()
}
}
}
fn handle_action_buttons_event(
&mut self,
item_id: &ManagementCardItemId,
@@ -1109,6 +1131,7 @@ impl AgentManagementView {
summarize_after_fork: false,
summarization_prompt: None,
initial_prompt: None,
initial_attachments: vec![],
destination: ForkedConversationDestination::NewTab,
});
}
@@ -1122,7 +1145,7 @@ impl AgentManagementView {
);
self.update_details_panel_for_item(item_id, ctx);
self.selected_item_id = Some(item_id.clone());
self.selected_item_id = Some(*item_id);
ctx.notify();
}
AgentDetailsButtonEvent::CopyLink { link } => {
@@ -1136,7 +1159,7 @@ impl AgentManagementView {
ctx
);
}
ManagementCardItemId::Task(task_id) => {
ManagementCardItemId::AmbientRun(task_id) => {
send_telemetry_from_ctx!(
AgentManagementTelemetryEvent::SessionLinkCopied {
task_id: task_id.to_string(),
@@ -1237,13 +1260,9 @@ impl AgentManagementView {
self.refresh_details_panel_if_needed(ctx);
self.get_tasks_from_model(ctx);
}
AgentConversationsModelEvent::ConversationUpdated => {
self.get_tasks_from_model(ctx);
self.refresh_details_panel_if_needed(ctx);
ctx.notify();
AgentConversationsModelEvent::ConversationUpdated { kind } => {
self.handle_conversation_updated(*kind, ctx);
}
// TaskManuallyOpened is handled by the conversation list view, not here.
AgentConversationsModelEvent::TaskManuallyOpened => {}
AgentConversationsModelEvent::ConversationArtifactsUpdated { conversation_id } => {
self.update_artifacts_for_conversation(*conversation_id, ctx);
self.refresh_details_panel_if_needed(ctx);
@@ -1253,11 +1272,48 @@ impl AgentManagementView {
/// Refresh the details panel if it's currently showing an item
fn refresh_details_panel_if_needed(&mut self, ctx: &mut ViewContext<Self>) {
if let Some(item_id) = self.selected_item_id.clone() {
if let Some(item_id) = self.selected_item_id {
self.update_details_panel_for_item(&item_id, ctx);
}
}
/// Decide how much work a `ConversationUpdated` event requires, based on its kind and the
/// active status filter:
/// * `Restored`: the underlying status didn't change, so the visible cards don't change
/// either. Just refresh the details panel.
/// * `MetadataChanged`: rebuild the cards so metadata-derived actions update.
/// * `TitleChanged`: rebuild the cards so filtering and titles update.
/// * `StatusSet` that crosses the active status filter: rebuild the
/// card list via `get_tasks_from_model`.
/// * `StatusSet` that doesn't cross the active filter (or `All` is active):
/// just refresh the details panel re-render so the status icon picks up the new value.
fn handle_conversation_updated(
&mut self,
kind: ConversationUpdateKind,
ctx: &mut ViewContext<Self>,
) {
match kind {
ConversationUpdateKind::Restored => {}
ConversationUpdateKind::MetadataChanged => self.get_tasks_from_model(ctx),
ConversationUpdateKind::TitleChanged => self.get_tasks_from_model(ctx),
ConversationUpdateKind::StatusSet {
prev_filter,
new_filter,
} => {
if self
.filters
.status
.is_membership_crossed(prev_filter, new_filter)
{
self.get_tasks_from_model(ctx);
} else {
ctx.notify();
}
}
}
self.refresh_details_panel_if_needed(ctx);
}
/// Update the details panel with fresh data for the given item.
fn update_details_panel_for_item(
&mut self,
@@ -1265,75 +1321,28 @@ impl AgentManagementView {
ctx: &mut ViewContext<Self>,
) {
let model = AgentConversationsModel::as_ref(ctx);
let data = match item_id {
ManagementCardItemId::Task(task_id) => {
let Some(task_wrapper) = model.get_task(task_id) else {
return;
};
// Agent management view should always open in a new tab
let open_action =
task_wrapper.get_open_action(Some(RestoreConversationLayout::NewTab), ctx);
let copy_link_url = task_wrapper.session_or_conversation_link(ctx);
let Some(task) = model.get_task_data(task_id) else {
return;
};
ConversationDetailsData::from_task(&task, open_action, copy_link_url, ctx)
}
ManagementCardItemId::Conversation(conversation_id) => {
let Some(conversation) = model.get_conversation(conversation_id) else {
return;
};
// Agent management view should always open in a new tab
let open_action =
conversation.get_open_action(Some(RestoreConversationLayout::NewTab), ctx);
let history_model = BlocklistAIHistoryModel::as_ref(ctx);
let ai_conversation = conversation
.navigation_data()
.and_then(|nav| history_model.conversation(&nav.id));
let server_conv_id = ai_conversation
.and_then(|c| c.server_conversation_token())
.map(|t| t.as_str().to_string())
.or_else(|| {
conversation
.navigation_data()
.and_then(|nav| history_model.get_conversation_metadata(&nav.id))
.and_then(|m| m.server_conversation_token.as_ref())
.map(|t| t.as_str().to_string())
});
let artifacts = ai_conversation
.map(|c| c.artifacts().to_vec())
.unwrap_or_default();
let status = Some(conversation.status(ctx));
let navigation_data = conversation.navigation_data();
let copy_link_url = conversation.session_or_conversation_link(ctx);
// Prefer server-reported harness when available; otherwise treat as a pure
// local conversation (always Warp Agent).
let harness = navigation_data
.and_then(|nav| history_model.get_server_conversation_metadata(&nav.id))
.map(|m| Harness::from(m.harness))
.or(Some(Harness::Oz));
ConversationDetailsData::from_conversation_metadata(
*conversation_id,
conversation.title(ctx),
conversation.creator_name(ctx),
conversation.created_at().with_timezone(&chrono::Local),
navigation_data.and_then(|n| n.initial_working_directory.clone()),
conversation.request_usage(ctx),
server_conv_id,
artifacts,
open_action,
status,
navigation_data.and_then(|n| n.initial_query.clone()),
copy_link_url,
harness,
)
}
let Some(entry) = model.get_entry_by_id(item_id, ctx) else {
return;
};
let open_action = AgentConversationsModel::resolve_open_action(
AgentConversationNavigationSubject::Entry(*item_id),
Some(RestoreConversationLayout::NewTab),
ctx,
);
let copy_link_url = AgentConversationsModel::resolve_copy_link(
AgentConversationNavigationSubject::Entry(*item_id),
ctx,
);
let task = entry
.identity
.ambient_agent_task_id
.and_then(|task_id| model.get_task_data(&task_id));
let data = ConversationDetailsData::from_agent_conversation_entry(
&entry,
task.as_ref(),
open_action,
copy_link_url,
);
self.details_panel.update(ctx, |p, ctx| {
p.set_conversation_details(data, ctx);
@@ -1347,19 +1356,14 @@ impl AgentManagementView {
ctx: &mut ViewContext<Self>,
) {
let model = AgentConversationsModel::as_ref(ctx);
let Some(card_data) = model.get_conversation(&conversation_id) else {
return;
};
let artifacts = card_data.artifacts(ctx);
// Find the index of the card for this conversation
let Some(index) = self
.items
.iter()
.position(|card| card.item_id == ManagementCardItemId::Conversation(conversation_id))
else {
let Some((index, entry)) = self.items.iter().enumerate().find_map(|(index, card)| {
let entry = model.get_entry_by_id(&card.item_id, ctx)?;
(entry.identity.local_conversation_id == Some(conversation_id))
.then_some((index, entry))
}) else {
return;
};
let artifacts = entry.display.artifacts;
// Update the artifact buttons for this card
if should_show_artifacts(&artifacts) {
@@ -1497,7 +1501,7 @@ impl AgentManagementView {
fn render_session_status_label(
appearance: &Appearance,
mouse_state: MouseStateHandle,
session_status: SessionStatus,
session_status: &SessionStatus,
) -> Box<dyn Element> {
let theme = appearance.theme();
let font_family = appearance.ui_font_family();
@@ -1604,21 +1608,17 @@ impl AgentManagementView {
};
let model = AgentConversationsModel::as_ref(app);
let card_data = match &card_state.item_id {
ManagementCardItemId::Task(task_id) => model.get_task(task_id),
ManagementCardItemId::Conversation(conv_id) => model.get_conversation(conv_id),
};
let Some(card_data) = card_data else {
let Some(entry) = model.get_entry_by_id(&card_state.item_id, app) else {
return Empty::new().finish();
};
self.render_card(card_state, &card_data, appearance, app)
self.render_card(card_state, &entry, appearance, app)
}
fn render_card(
&self,
card_state: &CardState,
card_data: &ConversationOrTask,
entry: &AgentConversationEntry,
appearance: &Appearance,
app: &AppContext,
) -> Box<dyn Element> {
@@ -1642,10 +1642,8 @@ impl AgentManagementView {
let card_hoverable = Hoverable::new(card_state.hover_state.clone(), move |mouse_state| {
let mut card_content = Flex::column()
.with_spacing(CARD_ROW_SPACING)
.with_child(Self::render_header_row(
card_state, card_data, appearance, app,
))
.with_child(Self::render_metadata_row(card_data, appearance, app));
.with_child(Self::render_header_row(card_state, entry, appearance))
.with_child(Self::render_metadata_row(entry, appearance, app));
// Add artifacts row if there is a buttons view
if let Some(buttons_element) = artifact_buttons_element {
@@ -1706,18 +1704,12 @@ impl AgentManagementView {
})
.with_defer_events_to_children();
// Add click handler to open session if available
let item_id = card_state.item_id.clone();
let card_hoverable = if card_data
.get_open_action(Some(RestoreConversationLayout::NewTab), app)
.is_some()
{
let item_id = card_state.item_id;
let card_hoverable = if entry.capabilities.can_open {
card_hoverable
.with_cursor(Cursor::PointingHand)
.on_click(move |ctx, _, _| {
ctx.dispatch_typed_action(AgentManagementViewAction::OpenSession {
item_id: item_id.clone(),
});
ctx.dispatch_typed_action(AgentManagementViewAction::OpenSession { item_id });
})
} else {
card_hoverable
@@ -1728,29 +1720,30 @@ impl AgentManagementView {
fn render_header_row(
card_state: &CardState,
card_data: &ConversationOrTask,
entry: &AgentConversationEntry,
appearance: &Appearance,
app: &AppContext,
) -> Box<dyn Element> {
let theme = appearance.theme();
let font_family = appearance.ui_font_family();
let font_size = appearance.ui_font_size();
let title = card_data.title(app);
let title_text = Text::new_inline(title, font_family, font_size)
let title_text = Text::new_inline(entry.display.title.clone(), font_family, font_size)
.with_color(theme.active_ui_text_color().into());
let status_icon =
render_status_element(&card_data.display_status(app), STATUS_ICON_SIZE, appearance);
// Build the time and avatar elements
let last_updated = card_data.last_updated();
let time_str = format_approx_duration_from_now_utc(last_updated);
let status_icon = render_icon_with_status(
agent_conversation_entry_icon_variant(entry),
CARD_AGENT_ICON_SIZE,
0.,
theme,
internal_colors::fg_overlay_1(theme),
);
let time_str = format_approx_duration_from_now_utc(entry.display.last_updated);
let time_text = Text::new_inline(time_str, font_family, font_size)
.with_color(theme.nonactive_ui_text_color().into());
let creator_name = card_data
.creator_name(app)
let creator_name = entry
.display
.creator
.name
.clone()
.unwrap_or_else(|| "Unknown".to_string());
let avatar = Self::render_avatar_with_tooltip(
&creator_name,
@@ -1763,12 +1756,11 @@ impl AgentManagementView {
.with_spacing(2.)
.with_child(Container::new(status_icon).with_margin_right(4.).finish())
.with_child(Expanded::new(1., title_text.finish()).finish());
let mut time_and_avatar = Flex::row()
.with_cross_axis_alignment(CrossAxisAlignment::Center)
.with_spacing(4.);
if let Some(session_status) = card_data.get_session_status() {
if let Some(session_status) = &entry.display.session_status {
time_and_avatar.add_child(Self::render_session_status_label(
appearance,
card_state.session_status_hover_state.clone(),
@@ -1792,45 +1784,56 @@ impl AgentManagementView {
}
fn render_metadata_row(
card_data: &ConversationOrTask,
entry: &AgentConversationEntry,
appearance: &Appearance,
app: &AppContext,
) -> Box<dyn Element> {
let theme = appearance.theme();
let font_family = appearance.ui_font_family();
let font_size = appearance.ui_font_size();
// Build metadata parts conditionally
let mut metadata_parts = Vec::new();
if let Some(source) = card_data.source() {
if let Some(source) = &entry.display.source {
metadata_parts.push(format!("Source: {}", source.display_name()));
}
if FeatureFlag::AgentHarness.is_enabled() {
if let Some(harness) = card_data.harness() {
let availability = HarnessAvailabilityModel::as_ref(app);
if availability.should_show_harness_selector() {
if let Some(harness) = entry.display.harness {
metadata_parts.push(format!(
"Harness: {}",
harness_display::display_name(harness)
availability.display_name_for(harness)
));
}
}
if let Some(run_time) = card_data.run_time() {
if let Some(executor) = &entry.display.executor {
let same_as_creator =
executor.uid.is_some() && executor.uid == entry.display.creator.uid;
if !same_as_creator {
if let Some(name) = executor.name.as_deref().or(executor.uid.as_deref()) {
let label = if executor
.principal_type
.is_some_and(|pt| pt.is_service_account())
{
"Agent"
} else {
"Executor"
};
metadata_parts.push(format!("{label}: {name}"));
}
}
}
if let Some(run_time) = &entry.display.run_time {
metadata_parts.push(format!("Run time: {run_time}"));
}
if let Some(usage) = card_data.display_request_usage(app) {
if let Some(usage) = entry.display.request_usage.map(format_credits) {
metadata_parts.push(format!("Credits used: {usage}"));
}
if let Some(tokens) = card_data.display_total_tokens(app) {
metadata_parts.push(format!("Tokens: {tokens}"));
}
let metadata_text = metadata_parts.join(" • ");
Text::new(metadata_text, font_family, font_size)
Text::new(metadata_parts.join(" • "), font_family, font_size)
.with_color(theme.nonactive_ui_text_color().into())
.finish()
}
@@ -1947,7 +1950,7 @@ impl AgentManagementView {
.with_child(ChildView::new(&self.created_on_dropdown).finish())
.with_child(ChildView::new(&self.artifact_dropdown).finish());
if FeatureFlag::AgentHarness.is_enabled() {
if HarnessAvailabilityModel::as_ref(app).should_show_harness_selector() {
filters_wrap.add_child(ChildView::new(&self.harness_dropdown).finish());
}
@@ -2356,17 +2359,11 @@ impl TypedActionView for AgentManagementView {
ctx.notify();
}
AgentManagementViewAction::OpenSession { item_id } => {
let model = AgentConversationsModel::as_ref(ctx);
let card_data = match item_id {
ManagementCardItemId::Task(task_id) => model.get_task(task_id),
ManagementCardItemId::Conversation(conv_id) => model.get_conversation(conv_id),
};
let Some(card_data) = card_data else {
return;
};
let Some(action) =
card_data.get_open_action(Some(RestoreConversationLayout::NewTab), ctx)
else {
let Some(action) = AgentConversationsModel::resolve_open_action(
AgentConversationNavigationSubject::Entry(*item_id),
Some(RestoreConversationLayout::NewTab),
ctx,
) else {
return;
};
@@ -2380,7 +2377,7 @@ impl TypedActionView for AgentManagementView {
ctx
);
}
ManagementCardItemId::Task(task_id) => {
ManagementCardItemId::AmbientRun(task_id) => {
send_telemetry_from_ctx!(
AgentManagementTelemetryEvent::CloudRunOpened {
task_id: task_id.to_string(),
app/src/ai/agent_sdk/admin.rs
+2
View File
@@ -4,6 +4,8 @@ use anyhow::{Context, Result};
use galaxy_cli::agent::OutputFormat;
use galaxyui::{platform::TerminationMode, AppContext, SingletonEntity};
use serde::Serialize;
use galaxyui::platform::TerminationMode;
use galaxyui::{AppContext, SingletonEntity};
use crate::auth::auth_manager::{AuthManager, AuthManagerEvent};
use crate::auth::user::PrincipalType;
app/src/ai/agent_sdk/agent_config.rs
+12 -10
View File
@@ -1,13 +1,15 @@
//! Commands to interact with available agents via the public API.
use warp_cli::agent::ListAgentSkillsArgs;
use warp_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus;
use warp_graphql::queries::user_repo_auth_status::UserRepoAuthStatusEnum;
use warpui::platform::TerminationMode;
use warpui::{AppContext, ModelContext, SingletonEntity};
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::ai::AgentSkillItem;
use crate::server::server_api::ServerApiProvider;
use galaxy_cli::agent::ListAgentConfigsArgs;
use galaxy_graphql::queries::get_oauth_connect_tx_status::OauthConnectTxStatus;
use galaxy_graphql::queries::user_repo_auth_status::UserRepoAuthStatusEnum;
use galaxyui::{platform::TerminationMode, AppContext, ModelContext, SingletonEntity};
const MAX_LINE_WIDTH: usize = 90;
const MAX_AUTH_ATTEMPTS: u32 = 8;
@@ -15,8 +17,8 @@ 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<()> {
/// List all available agent skills.
pub fn list_skills(ctx: &mut AppContext, args: ListAgentSkillsArgs) -> anyhow::Result<()> {
let runner = ctx.add_singleton_model(|_ctx| AgentConfigRunner);
runner.update(ctx, |runner, ctx| runner.list(args.repo.clone(), ctx))
}
@@ -220,7 +222,7 @@ impl AgentConfigRunner {
println!("Fetching agent skills from your Warp environments...");
}
let list_future = async move { ai_client.list_agents(repo).await };
let list_future = async move { ai_client.list_skills(repo).await };
ctx.spawn(list_future, |_, result, ctx| match result {
Ok(agents) => {
@@ -234,9 +236,9 @@ impl AgentConfigRunner {
}
/// Print a list of agents in a card-style format.
fn print_agents_table(agents: &[AgentListItem]) {
fn print_agents_table(agents: &[AgentSkillItem]) {
if agents.is_empty() {
println!("No agents found.");
println!("No skills found.");
return;
}
app/src/ai/agent_sdk/agent_management.rs
+515
View File
@@ -0,0 +1,515 @@
//! Commands to manage named agents via the public API.
use std::cmp::Ordering;
use std::io::Write as _;
use anyhow::anyhow;
use comfy_table::Cell;
use serde::Serialize;
use warp_cli::agent::{
AgentCreateArgs, AgentDeleteArgs, AgentGetArgs, AgentListArgs, AgentSortByArg, AgentUpdateArgs,
OutputFormat,
};
use warp_cli::json_filter::JsonOutput;
use warp_cli::SortOrderArg;
use warpui::platform::TerminationMode;
use warpui::{AppContext, ModelContext, SingletonEntity};
use super::output::TableFormat;
use crate::server::server_api::ai::{
AgentResponse, CreateAgentRequest, SecretRef, UpdateAgentRequest,
};
use crate::server::server_api::ServerApiProvider;
/// Singleton model that runs async work for named-agent CLI commands.
struct AgentManagementRunner;
pub fn list_agents(
ctx: &mut AppContext,
output_format: OutputFormat,
args: AgentListArgs,
) -> anyhow::Result<()> {
let runner = ctx.add_singleton_model(|_ctx| AgentManagementRunner);
runner.update(ctx, |runner, ctx| runner.list(args, output_format, ctx))
}
pub fn get_agent(
ctx: &mut AppContext,
output_format: OutputFormat,
args: AgentGetArgs,
) -> anyhow::Result<()> {
let runner = ctx.add_singleton_model(|_ctx| AgentManagementRunner);
runner.update(ctx, |runner, ctx| runner.get(args, output_format, ctx))
}
pub fn create_agent(
ctx: &mut AppContext,
output_format: OutputFormat,
args: AgentCreateArgs,
) -> anyhow::Result<()> {
let runner = ctx.add_singleton_model(|_ctx| AgentManagementRunner);
runner.update(ctx, |runner, ctx| runner.create(args, output_format, ctx))
}
pub fn update_agent(
ctx: &mut AppContext,
output_format: OutputFormat,
args: AgentUpdateArgs,
) -> anyhow::Result<()> {
let runner = ctx.add_singleton_model(|_ctx| AgentManagementRunner);
runner.update(ctx, |runner, ctx| runner.update(args, output_format, ctx))
}
pub fn delete_agent(
ctx: &mut AppContext,
output_format: OutputFormat,
args: AgentDeleteArgs,
) -> anyhow::Result<()> {
let runner = ctx.add_singleton_model(|_ctx| AgentManagementRunner);
runner.update(ctx, |runner, ctx| runner.delete(args, output_format, ctx))
}
impl AgentManagementRunner {
fn spawn_command(
&self,
future: impl warpui::r#async::Spawnable<Output = anyhow::Result<()>>,
ctx: &mut ModelContext<Self>,
) {
ctx.spawn(future, |_, result, ctx| match result {
Ok(()) => {
ctx.terminate_app(TerminationMode::ForceTerminate, None);
}
Err(err) => {
super::report_fatal_error(err, ctx);
}
});
}
fn list(
&self,
args: AgentListArgs,
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let future = async move {
ensure_json_sort_is_not_requested(output_format, &args.json_output, &args)?;
if matches!(output_format, OutputFormat::Json) || args.json_output.force_json_output() {
let response = ai_client.list_agents_raw().await?;
super::output::print_raw_json(response, &args.json_output)?;
} else {
let mut agents = ai_client.list_agents().await?;
sort_agents(&mut agents, args.sort_by, args.sort_order);
print_agents(&agents, output_format)?;
}
Ok(())
};
self.spawn_command(future, ctx);
Ok(())
}
fn get(
&self,
args: AgentGetArgs,
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let future = async move {
if matches!(output_format, OutputFormat::Json) || args.json_output.force_json_output() {
let response = ai_client.get_agent_raw(&args.uid).await?;
super::output::print_raw_json(response, &args.json_output)?;
} else {
let agent = ai_client.get_agent(&args.uid).await?;
print_single_agent(&agent, output_format)?;
}
Ok(())
};
self.spawn_command(future, ctx);
Ok(())
}
fn create(
&self,
args: AgentCreateArgs,
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let future = async move {
let json_output = args.json_output.clone();
let request = CreateAgentRequest {
name: args.name,
description: args.description,
secrets: secret_refs(args.secrets),
skills: args.skills,
base_model: args.base_model,
environment_id: args.environment,
};
if matches!(output_format, OutputFormat::Json) || json_output.force_json_output() {
let response = ai_client.create_agent_raw(request).await?;
super::output::print_raw_json(response, &json_output)?;
} else {
let agent = ai_client.create_agent(request).await?;
print_single_agent(&agent, output_format)?;
}
Ok(())
};
self.spawn_command(future, ctx);
Ok(())
}
fn update(
&self,
args: AgentUpdateArgs,
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let future = async move {
let uid = args.uid.clone();
let json_output = args.json_output.clone();
let current_agent = if args.add_secrets.is_empty()
&& args.remove_secrets.is_empty()
&& args.add_skills.is_empty()
&& args.remove_skills.is_empty()
{
None
} else {
Some(ai_client.get_agent(&uid).await?)
};
let request = UpdateAgentRequest {
name: args.name,
description: if args.remove_description {
Some(String::new())
} else {
args.description
},
secrets: if args.remove_all_secrets {
Some(vec![])
} else if args.add_secrets.is_empty() && args.remove_secrets.is_empty() {
None
} else {
let current_agent = current_agent
.as_ref()
.expect("current agent is fetched when applying secret deltas");
Some(apply_secret_deltas(
&current_agent.secrets,
args.add_secrets,
args.remove_secrets,
))
},
skills: if args.remove_all_skills {
Some(vec![])
} else if args.add_skills.is_empty() && args.remove_skills.is_empty() {
None
} else {
let current_agent = current_agent
.as_ref()
.expect("current agent is fetched when applying skill deltas");
Some(apply_string_deltas(
&current_agent.skills,
args.add_skills,
args.remove_skills,
))
},
base_model: if args.remove_base_model {
Some(String::new())
} else {
args.base_model
},
environment_id: if args.remove_environment {
Some(String::new())
} else {
args.environment
},
};
if request_is_empty(&request) {
return Err(anyhow!("No updates requested"));
}
if matches!(output_format, OutputFormat::Json) || json_output.force_json_output() {
let response = ai_client.update_agent_raw(&uid, request).await?;
super::output::print_raw_json(response, &json_output)?;
} else {
let agent = ai_client.update_agent(&uid, request).await?;
print_single_agent(&agent, output_format)?;
}
Ok(())
};
self.spawn_command(future, ctx);
Ok(())
}
fn delete(
&self,
args: AgentDeleteArgs,
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let future = async move {
ai_client.delete_agent(&args.uid).await?;
print_delete_result(&args.uid, output_format)?;
Ok(())
};
self.spawn_command(future, ctx);
Ok(())
}
}
fn secret_refs(secrets: Vec<String>) -> Vec<SecretRef> {
secrets.into_iter().map(|name| SecretRef { name }).collect()
}
/// Add and remove the requested secrets, starting with `current` as a baseline.
fn apply_secret_deltas(
current: &[SecretRef],
add_secrets: Vec<String>,
remove_secrets: Vec<String>,
) -> Vec<SecretRef> {
let names = apply_string_deltas(
&current
.iter()
.map(|secret| secret.name.clone())
.collect::<Vec<_>>(),
add_secrets,
remove_secrets,
);
secret_refs(names)
}
/// Add and remove the requested values, starting with `current` as a baseline.
fn apply_string_deltas(
current: &[String],
add_values: Vec<String>,
remove_values: Vec<String>,
) -> Vec<String> {
let mut values = current
.iter()
.filter(|value| !remove_values.contains(value))
.cloned()
.collect::<Vec<_>>();
for value in add_values {
if !values.contains(&value) {
values.push(value);
}
}
values
}
fn request_is_empty(request: &UpdateAgentRequest) -> bool {
request.name.is_none()
&& request.description.is_none()
&& request.secrets.is_none()
&& request.skills.is_none()
&& request.base_model.is_none()
&& request.environment_id.is_none()
}
fn ensure_json_sort_is_not_requested(
output_format: OutputFormat,
json_output: &JsonOutput,
args: &AgentListArgs,
) -> anyhow::Result<()> {
if (matches!(output_format, OutputFormat::Json) || json_output.force_json_output())
&& (args.sort_by.is_some() || args.sort_order.is_some())
{
return Err(anyhow!(
"--sort-by and --sort-order are not supported with JSON output"
));
}
Ok(())
}
fn sort_agents(
agents: &mut [AgentResponse],
sort_by: Option<AgentSortByArg>,
sort_order: Option<SortOrderArg>,
) {
let sort_by = sort_by.unwrap_or(AgentSortByArg::Name);
let default_order = match sort_by {
AgentSortByArg::Name => SortOrderArg::Asc,
AgentSortByArg::CreatedAt => SortOrderArg::Desc,
};
let sort_order = sort_order.unwrap_or(default_order);
agents.sort_by(|left, right| {
let ordering = match sort_by {
AgentSortByArg::Name => left
.name
.to_lowercase()
.cmp(&right.name.to_lowercase())
.then_with(|| left.uid.cmp(&right.uid)),
AgentSortByArg::CreatedAt => left
.created_at
.cmp(&right.created_at)
.then_with(|| left.uid.cmp(&right.uid)),
};
match sort_order {
SortOrderArg::Asc => ordering,
SortOrderArg::Desc => match ordering {
Ordering::Less => Ordering::Greater,
Ordering::Equal => Ordering::Equal,
Ordering::Greater => Ordering::Less,
},
}
});
}
impl TableFormat for AgentResponse {
fn header() -> Vec<Cell> {
vec![
Cell::new("UID"),
Cell::new("Name"),
Cell::new("Created"),
Cell::new("Description"),
Cell::new("Secrets"),
Cell::new("Skills"),
Cell::new("Base model"),
Cell::new("Environment"),
]
}
fn row(&self) -> Vec<Cell> {
vec![
Cell::new(&self.uid),
Cell::new(&self.name),
Cell::new(self.created_at.to_rfc3339()),
Cell::new(display_optional(self.description.as_deref())),
Cell::new(display_list(
self.secrets.iter().map(|secret| secret.name.as_str()),
)),
Cell::new(display_list(self.skills.iter().map(String::as_str))),
Cell::new(display_optional(self.base_model.as_deref())),
Cell::new(display_optional(self.environment_id.as_deref())),
]
}
}
fn print_agents(agents: &[AgentResponse], output_format: OutputFormat) -> anyhow::Result<()> {
match output_format {
OutputFormat::Pretty | OutputFormat::Text => {
let (visible_agents, hidden_count) = visible_agents_and_hidden_count(agents);
match output_format {
OutputFormat::Pretty if visible_agents.is_empty() => {
println!("No agents found.");
print_skills_hint();
}
OutputFormat::Pretty => {
super::output::write_list(visible_agents, output_format, std::io::stdout())?;
print_skills_hint();
}
OutputFormat::Text => {
super::output::write_list(visible_agents, output_format, std::io::stdout())?;
}
OutputFormat::Json | OutputFormat::Ndjson => {
unreachable!("handled by outer match")
}
}
print_disabled_agents_hidden_notice(hidden_count);
}
OutputFormat::Ndjson => {
for agent in agents {
super::output::write_json_line(agent, std::io::stdout())?;
}
}
OutputFormat::Json => unreachable!("JSON output is handled by the raw API path"),
}
Ok(())
}
fn visible_agents_and_hidden_count(agents: &[AgentResponse]) -> (Vec<AgentResponse>, usize) {
let visible_agents = agents
.iter()
.filter(|agent| agent.available)
.cloned()
.collect::<Vec<_>>();
let hidden_count = agents.len() - visible_agents.len();
(visible_agents, hidden_count)
}
fn print_disabled_agents_hidden_notice(hidden_count: usize) {
if hidden_count > 0 {
eprintln!("{hidden_count} disabled agents hidden");
}
}
fn print_single_agent(agent: &AgentResponse, output_format: OutputFormat) -> anyhow::Result<()> {
match output_format {
OutputFormat::Pretty | OutputFormat::Text => {
super::output::write_list([agent.clone()], output_format, std::io::stdout())?;
}
OutputFormat::Ndjson => {
super::output::write_json_line(agent, std::io::stdout())?;
}
OutputFormat::Json => unreachable!("JSON output is handled by the raw API path"),
}
Ok(())
}
fn print_skills_hint() {
let binary_name = warp_cli::binary_name().unwrap_or_else(|| "warp".to_string());
println!("\n\nLooking for your agent skills? Use `{binary_name} agent skills` instead.");
}
#[derive(Serialize)]
struct DeleteAgentResult<'a> {
uid: &'a str,
deleted: bool,
}
fn print_delete_result(uid: &str, output_format: OutputFormat) -> anyhow::Result<()> {
match output_format {
OutputFormat::Pretty => {
println!("Deleted agent {uid}.");
}
OutputFormat::Text => {
let mut stdout = std::io::stdout();
writeln!(stdout, "{uid}")?;
}
OutputFormat::Ndjson => {
super::output::write_json_line(
&DeleteAgentResult { uid, deleted: true },
std::io::stdout(),
)?;
}
OutputFormat::Json => {
super::output::write_json(
&DeleteAgentResult { uid, deleted: true },
std::io::stdout(),
)?;
}
}
Ok(())
}
fn display_optional(value: Option<&str>) -> String {
value
.filter(|value| !value.is_empty())
.unwrap_or("-")
.to_string()
}
fn display_list<'a>(values: impl IntoIterator<Item = &'a str>) -> String {
let values = values
.into_iter()
.filter(|value| !value.is_empty())
.collect::<Vec<_>>();
if values.is_empty() {
"-".to_string()
} else {
values.join(", ")
}
}
impl warpui::Entity for AgentManagementRunner {
type Event = ();
}
impl SingletonEntity for AgentManagementRunner {}
#[cfg(test)]
#[path = "agent_management_tests.rs"]
mod tests;
app/src/ai/agent_sdk/agent_management_tests.rs
+170
View File
@@ -0,0 +1,170 @@
use chrono::{TimeZone as _, Utc};
use super::*;
fn agent(uid: &str, name: &str, created_at_seconds: i64) -> AgentResponse {
agent_with_available(uid, name, created_at_seconds, true)
}
fn agent_with_available(
uid: &str,
name: &str,
created_at_seconds: i64,
available: bool,
) -> AgentResponse {
AgentResponse {
uid: uid.to_string(),
name: name.to_string(),
description: None,
available,
created_at: Utc
.timestamp_opt(created_at_seconds, 0)
.single()
.expect("valid timestamp"),
secrets: vec![],
skills: vec![],
base_model: None,
environment_id: None,
}
}
#[test]
fn table_format_does_not_include_available_column() {
let header = AgentResponse::header()
.into_iter()
.map(|cell| cell.content().to_string())
.collect::<Vec<_>>();
let row = agent("1", "agent", 1).row();
assert_eq!(
header,
[
"UID",
"Name",
"Created",
"Description",
"Secrets",
"Skills",
"Base model",
"Environment",
]
);
assert_eq!(row.len(), header.len());
}
#[test]
fn visible_agents_and_hidden_count_filters_disabled_agents() {
let agents = vec![
agent_with_available("1", "enabled", 1, true),
agent_with_available("2", "disabled", 2, false),
];
let (visible_agents, hidden_count) = visible_agents_and_hidden_count(&agents);
assert_eq!(visible_agents.len(), 1);
assert_eq!(visible_agents[0].name, "enabled");
assert_eq!(hidden_count, 1);
}
#[test]
fn sort_agents_defaults_to_name_ascending() {
let mut agents = vec![agent("2", "zeta", 2), agent("1", "alpha", 1)];
sort_agents(&mut agents, None, None);
assert_eq!(agents[0].name, "alpha");
assert_eq!(agents[1].name, "zeta");
}
#[test]
fn sort_agents_defaults_created_at_to_descending() {
let mut agents = vec![agent("1", "old", 1), agent("2", "new", 2)];
sort_agents(&mut agents, Some(AgentSortByArg::CreatedAt), None);
assert_eq!(agents[0].name, "new");
assert_eq!(agents[1].name, "old");
}
#[test]
fn sort_agents_respects_explicit_sort_order_without_sort_field() {
let mut agents = vec![agent("1", "alpha", 1), agent("2", "zeta", 2)];
sort_agents(&mut agents, None, Some(SortOrderArg::Desc));
assert_eq!(agents[0].name, "zeta");
assert_eq!(agents[1].name, "alpha");
}
#[test]
fn update_request_omits_unset_fields_and_serializes_clears() {
let request = UpdateAgentRequest {
description: Some(String::new()),
secrets: Some(vec![]),
base_model: Some(String::new()),
..Default::default()
};
let json = serde_json::to_value(request).expect("request serializes");
assert_eq!(
json,
serde_json::json!({
"description": "",
"secrets": [],
"base_model": "",
})
);
}
#[test]
fn rejects_sort_for_json_output() {
let args = AgentListArgs {
sort_by: Some(AgentSortByArg::Name),
sort_order: None,
json_output: JsonOutput { filter: None },
};
let err = ensure_json_sort_is_not_requested(OutputFormat::Json, &args.json_output, &args)
.unwrap_err();
assert!(err.to_string().contains("not supported with JSON output"));
}
#[test]
fn apply_string_deltas_removes_and_appends_without_duplicates() {
let values = apply_string_deltas(
&["old".to_string(), "keep".to_string()],
vec!["new".to_string(), "keep".to_string()],
vec!["old".to_string()],
);
assert_eq!(values, ["keep", "new"]);
}
#[test]
fn apply_secret_deltas_uses_secret_names() {
let values = apply_secret_deltas(
&[
SecretRef {
name: "OLD_TOKEN".to_string(),
},
SecretRef {
name: "KEEP_TOKEN".to_string(),
},
],
vec!["NEW_TOKEN".to_string()],
vec!["OLD_TOKEN".to_string()],
);
assert_eq!(
values,
[
SecretRef {
name: "KEEP_TOKEN".to_string()
},
SecretRef {
name: "NEW_TOKEN".to_string()
},
]
);
}
app/src/ai/agent_sdk/ambient.rs
+260 -77
View File
@@ -3,54 +3,50 @@ use std::io::Write as _;
use std::sync::Arc;
use std::time::Duration;
use anyhow::{anyhow, Context as _};
use comfy_table::Cell;
use futures::{future, StreamExt};
use serde::Serialize;
use warp_cli::agent::{Harness, OutputFormat, Prompt, RunCloudArgs};
use warp_cli::json_filter::JsonOutput;
use warp_cli::task::{
ArtifactTypeArg, ExecutionLocationArg, ListTasksArgs, MessageCommand, MessageDeliveredArgs,
MessageListArgs, MessageReadArgs, MessageSendArgs, MessageWatchArgs, RunSortByArg,
RunSourceArg, RunStateArg, TaskGetArgs,
};
use warp_cli::{GlobalOptions, SortOrderArg};
use galaxy_core::channel::ChannelState;
use galaxy_core::features::FeatureFlag;
use warpui::platform::TerminationMode;
use warpui::r#async::{Spawnable, Timer};
use warpui::{AppContext, ModelContext, SingletonEntity};
use super::common::{parse_ambient_task_id, EnvironmentChoice, ResolveConfigurationError};
use crate::ai::agent::{extract_user_query_mode, UserQueryMode};
use crate::ai::agent_sdk::driver::attachments::{
process_attachment, MAX_ATTACHMENT_COUNT_FOR_CLOUD_QUERY,
};
use crate::ai::ambient_agents::spawn::{
spawn_task, AmbientAgentEvent, SessionJoinInfo, TASK_STATUS_POLLING_DURATION,
};
use crate::ai::ambient_agents::task::HarnessConfig;
use crate::ai::ambient_agents::AmbientAgentTaskState;
use crate::ai::ambient_agents::{AgentConfigSnapshot, AmbientAgentTask};
use crate::ai::ambient_agents::{
AgentConfigSnapshot, AmbientAgentTask, AmbientAgentTaskId, AmbientAgentTaskState,
};
use crate::ai::artifacts::Artifact;
use crate::auth::AuthStateProvider;
use crate::cloud_object::model::persistence::CloudModel;
use crate::server::ids::{ServerId, SyncId};
use crate::server::server_api::ai::{
AIClient, AgentMessageHeader, AgentRunEvent, AgentSource, ArtifactType, ExecutionLocation,
ListAgentMessagesRequest, ReadAgentMessageResponse, RunSortBy, RunSortOrder,
SendAgentMessageRequest, SendAgentMessageResponse, SpawnAgentRequest, TaskListFilter,
};
use crate::server::server_api::ServerApi;
use crate::terminal::shared_session;
use crate::util::time_format::format_approx_duration_from_now_utc;
use crate::workspaces::user_workspaces::UserWorkspaces;
use crate::{
terminal::shared_session, util::time_format::format_approx_duration_from_now_utc,
ServerApiProvider,
};
use anyhow::{anyhow, Context as _};
use comfy_table::Cell;
use futures::{future, StreamExt};
use serde::Serialize;
use galaxy_cli::{
agent::{Harness, OutputFormat, Prompt, RunCloudArgs},
json_filter::JsonOutput,
task::{
ArtifactTypeArg, ExecutionLocationArg, ListTasksArgs, MessageCommand, MessageDeliveredArgs,
MessageListArgs, MessageReadArgs, MessageSendArgs, MessageWatchArgs, RunSortByArg,
RunSortOrderArg, RunSourceArg, RunStateArg, TaskGetArgs,
},
GlobalOptions,
};
use galaxy_core::channel::ChannelState;
use galaxy_core::features::FeatureFlag;
use galaxyui::r#async::Timer;
use galaxyui::{
platform::TerminationMode, r#async::Spawnable, AppContext, ModelContext, SingletonEntity,
};
use crate::ai::agent_sdk::driver::attachments::{
process_attachment, MAX_ATTACHMENT_COUNT_FOR_CLOUD_QUERY,
};
use crate::cloud_object::model::persistence::CloudModel;
use crate::server::ids::{ServerId, SyncId};
use super::common::{EnvironmentChoice, ResolveConfigurationError};
use crate::ServerApiProvider;
const MAX_LINE_WIDTH: usize = 90;
const STREAM_RETRY_BACKOFF_STEPS: &[u64] = &[1, 2, 5, 10];
@@ -185,10 +181,10 @@ fn sort_by_from_arg(arg: RunSortByArg) -> RunSortBy {
}
}
fn sort_order_from_arg(arg: RunSortOrderArg) -> RunSortOrder {
fn sort_order_from_arg(arg: SortOrderArg) -> RunSortOrder {
match arg {
RunSortOrderArg::Asc => RunSortOrder::Asc,
RunSortOrderArg::Desc => RunSortOrder::Desc,
SortOrderArg::Asc => RunSortOrder::Asc,
SortOrderArg::Desc => RunSortOrder::Desc,
}
}
@@ -266,11 +262,8 @@ impl AmbientAgentRunner {
);
return;
}
// TODO: Consider making the server's prompt field optional when skill is provided,
// rather than sending an empty string for skill-only invocations.
let prompt_string = match prompt {
Some(Prompt::PlainText(text)) => text,
let prompt = match prompt {
Some(Prompt::PlainText(text)) => Some(text),
Some(Prompt::SavedPrompt(id)) => {
// Resolve the saved prompt to pass along as the ambient agent query.
// We look up the prompt text here, rather than passing along the saved prompt ID,
@@ -294,7 +287,7 @@ impl AmbientAgentRunner {
match workflow {
Some(cloud_workflow) => match cloud_workflow.model().data.prompt() {
Some(prompt_text) => prompt_text.to_string(),
Some(prompt_text) => Some(prompt_text.to_string()),
None => {
super::report_fatal_error(
anyhow::anyhow!("'{id}' is not a saved prompt"),
@@ -312,8 +305,7 @@ impl AmbientAgentRunner {
}
}
}
// Skill-only invocation: use empty prompt, skill provides instructions
None => String::new(),
None => None,
};
let loaded_file = match args.config_file.file.as_deref() {
@@ -327,6 +319,23 @@ impl AmbientAgentRunner {
None => None,
};
// The `--runner` CLI flag is gated in `run_agent`, but a config file
// can also set `runner_id`, which would otherwise bypass the gate.
if loaded_file
.as_ref()
.and_then(|f| f.file.runner_id.as_ref())
.is_some()
&& !FeatureFlag::CloudRunners.is_enabled()
{
super::report_fatal_error(
anyhow::anyhow!(
"`runner_id` is set in the config file but runner support is not enabled"
),
ctx,
);
return;
}
// Validate and process attachments early, before environment selection
// This ensures users don't have to go through env selection if attachment validation fails
if args.attachment_paths.len() > MAX_ATTACHMENT_COUNT_FOR_CLOUD_QUERY {
@@ -416,10 +425,13 @@ impl AmbientAgentRunner {
let harness_override = (args.harness != Harness::Oz).then_some(HarnessConfig {
harness_type: args.harness,
model_id: None,
reasoning_level: None,
});
let harness_auth_secrets = args.claude_auth_secret.clone().map(|name| {
crate::ai::ambient_agents::task::HarnessAuthSecretsConfig {
claude_auth_secret_name: Some(name),
codex_auth_secret_name: None,
}
});
@@ -428,6 +440,7 @@ impl AmbientAgentRunner {
AgentConfigSnapshot {
name: args.name,
environment_id,
runner_id: args.runner,
model_id: args.model.model.clone(),
base_prompt: None,
mcp_servers: cli_mcp_servers,
@@ -471,8 +484,16 @@ impl AmbientAgentRunner {
None
};
let (prompt, mode) = match prompt {
Some(prompt) => {
let (prompt, mode) = extract_user_query_mode(prompt);
(Some(prompt), mode)
}
None => (None, UserQueryMode::Normal),
};
let request = SpawnAgentRequest {
prompt: prompt_string,
prompt,
mode,
config,
title: None,
team: match (args.scope.team, args.scope.personal) {
@@ -480,12 +501,17 @@ impl AmbientAgentRunner {
(_, true) => Some(false),
_ => None,
},
agent_identity_uid: args.agent_uid,
skill,
attachments,
interactive: None,
parent_run_id: None,
runtime_skills: vec![],
referenced_attachments: vec![],
conversation_id: args.conversation,
initial_snapshot_token: None,
snapshot_disabled: None,
orchestration_handoff: None,
};
let should_open = args.open;
@@ -640,17 +666,42 @@ impl AmbientAgentRunner {
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let provider = ServerApiProvider::as_ref(ctx);
let ai_client = provider.get_ai_client();
let server_api = provider.get();
let scoped_task_id = task_id_for_message_send(&args.sender_run_id)?;
let future = async move {
let response = ai_client
.send_agent_message(SendAgentMessageRequest {
to: args.to,
subject: args.subject,
body: args.body,
sender_run_id: args.sender_run_id,
})
.await?;
let request = SendAgentMessageRequest {
to: args.to,
subject: args.subject,
body: args.body,
sender_run_id: args.sender_run_id,
};
let log_context = SendAgentMessageLogContext::new(&request, scoped_task_id.as_ref());
log_context.log_start();
let send_message = async move {
match scoped_task_id {
Some(task_id) => {
server_api
.send_agent_message_for_task(&task_id, request)
.await
}
None => ai_client.send_agent_message(request).await,
}
};
let response = match send_message.await {
Ok(response) => {
log_context.log_success(&response);
response
}
Err(err) => {
let err = err.context(log_context.error_context());
log_context.log_error(&err);
eprintln!("{err:#}");
return Err(err);
}
};
print_send_message_response(&response, output_format)?;
Ok(())
};
@@ -658,25 +709,31 @@ impl AmbientAgentRunner {
Ok(())
}
fn list_messages(
&self,
args: MessageListArgs,
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let provider = ServerApiProvider::as_ref(ctx);
let ai_client = provider.get_ai_client();
let server_api = provider.get();
let future = async move {
let messages = ai_client
.list_agent_messages(
&args.run_id,
ListAgentMessagesRequest {
unread_only: args.unread,
since: args.since,
limit: args.limit,
},
)
.await?;
let request = ListAgentMessagesRequest {
unread_only: args.unread,
since: args.since,
limit: args.limit,
};
let messages = match task_id_from_run_id(&args.run_id) {
Some(task_id) => {
server_api
.list_agent_messages_for_task(&task_id, &args.run_id, request)
.await?
}
None => ai_client.list_agent_messages(&args.run_id, request).await?,
};
super::output::print_list(messages, output_format);
Ok(())
};
@@ -708,10 +765,20 @@ impl AmbientAgentRunner {
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let provider = ServerApiProvider::as_ref(ctx);
let ai_client = provider.get_ai_client();
let server_api = provider.get();
let scoped_task_id = task_id_from_oz_run_id_env()?;
let future = async move {
let message = ai_client.read_agent_message(&args.message_id).await?;
let message = match scoped_task_id {
Some(task_id) => {
server_api
.read_agent_message_for_task(&task_id, &args.message_id)
.await?
}
None => ai_client.read_agent_message(&args.message_id).await?,
};
print_read_message_response(&message, output_format)?;
Ok(())
};
@@ -726,10 +793,20 @@ impl AmbientAgentRunner {
output_format: OutputFormat,
ctx: &mut ModelContext<Self>,
) -> anyhow::Result<()> {
let ai_client = ServerApiProvider::as_ref(ctx).get_ai_client();
let provider = ServerApiProvider::as_ref(ctx);
let ai_client = provider.get_ai_client();
let server_api = provider.get();
let scoped_task_id = task_id_from_oz_run_id_env()?;
let future = async move {
ai_client.mark_message_delivered(&args.message_id).await?;
match scoped_task_id {
Some(task_id) => {
server_api
.mark_message_delivered_for_task(&task_id, &args.message_id)
.await?
}
None => ai_client.mark_message_delivered(&args.message_id).await?,
}
print_mark_message_delivered_result(&args.message_id, output_format)?;
Ok(())
};
@@ -790,6 +867,10 @@ impl AmbientAgentRunner {
table.add_row(vec![title_cell]);
}
if let Some(executor) = task.executor_display_name() {
table.add_row(vec![format!("Executed as: {executor}")]);
}
// Agent config snapshot (if available)
if let Some(config) = task.agent_config_snapshot.as_ref() {
let config_str =
@@ -933,6 +1014,91 @@ fn write_stream_record<T: Serialize>(record: &T) -> anyhow::Result<()> {
stdout.flush().context("unable to flush stdout")?;
Ok(())
}
fn task_id_from_run_id(run_id: &str) -> Option<AmbientAgentTaskId> {
run_id.parse().ok()
}
fn task_id_from_oz_run_id_env() -> anyhow::Result<Option<AmbientAgentTaskId>> {
match std::env::var(warp_cli::OZ_RUN_ID_ENV) {
Ok(run_id) => parse_ambient_task_id(&run_id, "Invalid OZ_RUN_ID").map(Some),
Err(std::env::VarError::NotPresent) => Ok(None),
Err(std::env::VarError::NotUnicode(_)) => Err(anyhow!(
"{} is set but is not valid Unicode",
warp_cli::OZ_RUN_ID_ENV
)),
}
}
fn task_id_for_message_send(sender_run_id: &str) -> anyhow::Result<Option<AmbientAgentTaskId>> {
match task_id_from_run_id(sender_run_id) {
Some(task_id) => Ok(Some(task_id)),
None => task_id_from_oz_run_id_env(),
}
}
#[derive(Debug, Clone)]
struct SendAgentMessageLogContext {
sender_run_id: String,
task_id: Option<String>,
target_agent_ids: Vec<String>,
subject: String,
body_len: usize,
}
impl SendAgentMessageLogContext {
fn new(request: &SendAgentMessageRequest, task_id: Option<&AmbientAgentTaskId>) -> Self {
Self {
sender_run_id: request.sender_run_id.clone(),
task_id: task_id.map(|task_id| task_id.to_string()),
target_agent_ids: request.to.clone(),
subject: request.subject.clone(),
body_len: request.body.chars().count(),
}
}
fn error_context(&self) -> String {
format!(
"Failed to send agent message (sender_run_id={:?}, task_id={:?}, target_agent_ids={:?})",
self.sender_run_id, self.task_id, self.target_agent_ids
)
}
fn log_start(&self) {
log::info!(
"Sending ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={}",
self.sender_run_id,
self.task_id,
self.target_agent_ids,
self.subject,
self.body_len
);
}
fn log_success(&self, response: &SendAgentMessageResponse) {
log::info!(
"Sent ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} message_ids={:?}",
self.sender_run_id,
self.task_id,
self.target_agent_ids,
self.subject,
self.body_len,
response.message_ids
);
}
fn log_error(&self, err: &anyhow::Error) {
log::warn!(
"Failed to send ambient agent message: sender_run_id={:?} task_id={:?} target_agent_ids={:?} subject={:?} body_len={} error={err:#}",
self.sender_run_id,
self.task_id,
self.target_agent_ids,
self.subject,
self.body_len
);
}
}
async fn watch_messages_forever(
server_api: Arc<ServerApi>,
ai_client: Arc<dyn AIClient>,
@@ -940,15 +1106,25 @@ async fn watch_messages_forever(
) -> anyhow::Result<()> {
let run_id = args.run_id;
let watched_run_ids = vec![run_id.clone()];
let scoped_task_id = task_id_from_run_id(&run_id);
let mut last_seen_sequence = args.since_sequence;
let mut initial_connect = true;
let mut failures = 0usize;
loop {
let mut stream = match server_api
.stream_agent_events(&watched_run_ids, last_seen_sequence)
.await
{
let stream_result = match scoped_task_id.as_ref() {
Some(task_id) => {
server_api
.stream_agent_events_for_task(task_id, &watched_run_ids, last_seen_sequence)
.await
}
None => {
server_api
.stream_agent_events(&watched_run_ids, last_seen_sequence)
.await
}
};
let mut stream = match stream_result {
Ok(stream) => {
if !initial_connect {
eprintln!(
@@ -1004,8 +1180,15 @@ async fn watch_messages_forever(
last_seen_sequence = event.sequence;
continue;
};
let message = match ai_client.read_agent_message(&message_id).await {
let message_result = match scoped_task_id.as_ref() {
Some(task_id) => {
server_api
.read_agent_message_for_task(task_id, &message_id)
.await
}
None => ai_client.read_agent_message(&message_id).await,
};
let message = match message_result {
Ok(message) => message,
Err(err) => {
failures += 1;
app/src/ai/agent_sdk/ambient_tests.rs
+53 -7
View File
@@ -1,17 +1,17 @@
//! Unit tests for `filter_from_args`. Verifies the clap enums are faithfully translated into
//! `TaskListFilter` without dropping any fields.
//! Unit tests for ambient agent CLI argument mapping and message helpers.
use chrono::{TimeZone, Utc};
use galaxy_cli::json_filter::JsonOutput;
use galaxy_cli::task::{
ArtifactTypeArg, ExecutionLocationArg, ListTasksArgs, RunSortByArg, RunSortOrderArg,
RunSourceArg, RunStateArg,
ArtifactTypeArg, ExecutionLocationArg, ListTasksArgs, RunSortByArg, RunSourceArg, RunStateArg,
};
use warp_cli::SortOrderArg;
use super::*;
use crate::server::server_api::ai::{ArtifactType, ExecutionLocation, RunSortBy, RunSortOrder};
const TASK_ID: &str = "00000000-0000-0000-0000-000000000001";
const OTHER_TASK_ID: &str = "00000000-0000-0000-0000-000000000002";
/// A `ListTasksArgs` whose fields are all at their defaults.
fn empty_args() -> ListTasksArgs {
ListTasksArgs {
@@ -138,7 +138,7 @@ fn every_field_maps_through() {
updated_after: Some(updated_after),
query: Some("oz run".to_string()),
sort_by: Some(RunSortByArg::CreatedAt),
sort_order: Some(RunSortOrderArg::Asc),
sort_order: Some(SortOrderArg::Asc),
cursor: Some("abcd==".to_string()),
json_output: JsonOutput::default(),
};
@@ -167,3 +167,49 @@ fn every_field_maps_through() {
assert_eq!(filter.sort_order, Some(RunSortOrder::Asc));
assert_eq!(filter.cursor.as_deref(), Some("abcd=="));
}
#[test]
fn task_id_from_run_id_accepts_task_uuid() {
let task_id = task_id_from_run_id(TASK_ID).expect("valid task id");
assert_eq!(task_id.to_string(), TASK_ID);
}
#[test]
fn task_id_from_run_id_ignores_non_task_ids() {
assert!(task_id_from_run_id("local-child-run").is_none());
}
#[test]
#[serial_test::serial]
fn task_id_for_message_send_prefers_sender_run_id() {
std::env::set_var(warp_cli::OZ_RUN_ID_ENV, OTHER_TASK_ID);
let task_id = task_id_for_message_send(TASK_ID)
.expect("valid task id")
.expect("task id");
std::env::remove_var(warp_cli::OZ_RUN_ID_ENV);
assert_eq!(task_id.to_string(), TASK_ID);
}
#[test]
#[serial_test::serial]
fn task_id_for_message_send_falls_back_to_oz_run_id() {
std::env::set_var(warp_cli::OZ_RUN_ID_ENV, TASK_ID);
let task_id = task_id_for_message_send("local-child-run")
.expect("valid env task id")
.expect("task id");
std::env::remove_var(warp_cli::OZ_RUN_ID_ENV);
assert_eq!(task_id.to_string(), TASK_ID);
}
#[test]
#[serial_test::serial]
fn task_id_from_oz_run_id_env_rejects_invalid_value() {
std::env::set_var(warp_cli::OZ_RUN_ID_ENV, "not-a-task-id");
let err = task_id_from_oz_run_id_env().expect_err("invalid task id");
std::env::remove_var(warp_cli::OZ_RUN_ID_ENV);
assert!(err.to_string().contains("Invalid OZ_RUN_ID"));
}
app/src/ai/agent_sdk/api_key.rs
+485
View File
@@ -0,0 +1,485 @@
use std::cmp::Reverse;
use std::fmt;
use std::io::{self, IsTerminal as _};
use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc};
use comfy_table::Cell;
use inquire::{Confirm, InquireError, Select};
use serde::Serialize;
use warp_cli::agent::OutputFormat;
use warp_cli::api_key::{
ApiKeyCommand, ApiKeyExpirationArgs, ApiKeySortByArg, CreateApiKeyArgs, ExpireApiKeyArgs,
ListApiKeysArgs,
};
use warp_cli::{GlobalOptions, SortOrderArg};
use warp_graphql::mutations::expire_api_key::ExpireApiKeyResult;
use warp_graphql::mutations::generate_api_key::GenerateApiKeyResult;
use warp_graphql::queries::api_keys::ApiKeyProperties;
use warp_graphql::scalars::Time;
use warpui::platform::TerminationMode;
use warpui::{AppContext, ModelContext, SingletonEntity};
use super::output::{self, TableFormat};
use crate::server::ids::ApiKeyUid;
use crate::util::time_format::format_approx_duration_from_now_utc;
use crate::ServerApiProvider;
/// Run API key-related commands.
pub fn run(
ctx: &mut AppContext,
global_options: GlobalOptions,
command: ApiKeyCommand,
) -> Result<()> {
let runner = ctx.add_singleton_model(|_ctx| ApiKeyCommandRunner);
match command {
ApiKeyCommand::List(args) => {
runner.update(ctx, |runner, ctx| {
runner.list(global_options.output_format, args, ctx)
});
Ok(())
}
ApiKeyCommand::Create(args) => {
runner.update(ctx, |runner, ctx| {
runner.create(global_options.output_format, args, ctx)
});
Ok(())
}
ApiKeyCommand::Expire(args) => {
runner.update(ctx, |runner, ctx| {
runner.expire(global_options.output_format, args, ctx)
});
Ok(())
}
}
}
struct ApiKeyCommandRunner;
impl ApiKeyCommandRunner {
fn list(
&self,
output_format: OutputFormat,
args: ListApiKeysArgs,
ctx: &mut ModelContext<Self>,
) {
let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client();
ctx.spawn(
async move {
let mut keys: Vec<_> = auth_client
.list_api_keys()
.await?
.into_iter()
.map(ApiKeyInfo::from)
.collect();
sort_api_keys(&mut keys, args.sort_by, args.sort_order);
if args.json_output.force_json_output() {
output::print_raw_json(serde_json::to_value(&keys)?, &args.json_output)?;
} else {
output::print_list(keys, output_format);
}
Ok(())
},
|_, result: Result<()>, ctx| finish_command(result, ctx),
);
}
fn create(
&self,
output_format: OutputFormat,
args: CreateApiKeyArgs,
ctx: &mut ModelContext<Self>,
) {
let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client();
ctx.spawn(
async move {
let json_output = args.json_output;
let expires_at = expires_at_from_args(args.expiration)?;
let agent_uid = args.agent_uid.map(cynic::Id::new);
let result = auth_client
.create_api_key(args.name, None, agent_uid, expires_at)
.await?;
let result = match result {
GenerateApiKeyResult::GenerateApiKeyOutput(output) => CreatedApiKeyInfo {
raw_api_key: output.raw_api_key,
api_key: ApiKeyInfo::from(output.api_key),
},
GenerateApiKeyResult::UserFacingError(e) => {
return Err(anyhow!(
warp_graphql::client::get_user_facing_error_message(e)
));
}
GenerateApiKeyResult::Unknown => {
return Err(anyhow!("failed to create API key"))
}
};
print_created_api_key(result, output_format, json_output)?;
Ok(())
},
|_, result: Result<()>, ctx| finish_command(result, ctx),
);
}
fn expire(
&self,
output_format: OutputFormat,
args: ExpireApiKeyArgs,
ctx: &mut ModelContext<Self>,
) {
let key_identifier = args.key_uid;
let force = args.force;
let json_output = args.json_output;
let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client();
ctx.spawn(
async move {
let keys = auth_client
.list_api_keys()
.await?
.into_iter()
.map(ApiKeyInfo::from)
.collect();
Ok(keys)
},
move |_, result: Result<Vec<ApiKeyInfo>>, ctx| {
let keys = match result {
Ok(keys) => keys,
Err(err) => {
super::report_fatal_error(err, ctx);
return;
}
};
let key = match resolve_api_key_identifier(&keys, &key_identifier) {
Ok(Some(key)) => key,
Ok(None) => {
ctx.terminate_app(TerminationMode::ForceTerminate, None);
return;
}
Err(err) => {
super::report_fatal_error(err, ctx);
return;
}
};
if !force {
if !io::stdin().is_terminal() {
super::report_fatal_error(
anyhow!(
"Refusing to expire API key without confirmation in non-interactive mode (use --force to bypass)"
),
ctx,
);
return;
}
let prompt = format!("Expire API key '{key}'?");
let should_expire = match Confirm::new(&prompt)
.with_default(false)
.with_help_message("This action takes effect immediately")
.prompt()
{
Ok(should_expire) => should_expire,
Err(InquireError::OperationCanceled | InquireError::OperationInterrupted) => {
ctx.terminate_app(TerminationMode::ForceTerminate, None);
return;
}
Err(err) => {
super::report_fatal_error(err.into(), ctx);
return;
}
};
if !should_expire {
println!("Expiration cancelled");
ctx.terminate_app(TerminationMode::ForceTerminate, None);
return;
}
}
let uid = ApiKeyUid::from(key.uid);
let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client();
ctx.spawn(
async move {
let result = auth_client.expire_api_key(&uid).await?;
let expired = match result {
ExpireApiKeyResult::ExpireApiKeyOutput(output) => output.success,
ExpireApiKeyResult::UserFacingError(e) => {
return Err(anyhow!(
warp_graphql::client::get_user_facing_error_message(e)
));
}
ExpireApiKeyResult::Unknown => {
return Err(anyhow!("failed to expire API key"))
}
};
print_expire_api_key_result(
uid.to_string(),
expired,
output_format,
json_output,
)?;
Ok(())
},
|_, result: Result<()>, ctx| finish_command(result, ctx),
);
},
);
}
}
impl warpui::Entity for ApiKeyCommandRunner {
type Event = ();
}
impl SingletonEntity for ApiKeyCommandRunner {}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
struct ApiKeyInfo {
uid: String,
name: String,
key_suffix: String,
scope: String,
created_at: DateTime<Utc>,
last_used_at: Option<DateTime<Utc>>,
expires_at: Option<DateTime<Utc>>,
}
impl From<ApiKeyProperties> for ApiKeyInfo {
fn from(key: ApiKeyProperties) -> Self {
Self {
uid: key.uid.into_inner(),
name: key.name,
key_suffix: key.key_suffix,
scope: key.owner_type.to_string(),
created_at: key.created_at.utc(),
last_used_at: key.last_used_at.map(|t| t.utc()),
expires_at: key.expires_at.map(|t| t.utc()),
}
}
}
impl fmt::Display for ApiKeyInfo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let name = &self.name;
let uid = &self.uid;
let created_at = self.created_at.format("%Y-%m-%d %H:%M:%S UTC");
write!(f, "{name} ({uid}, created {created_at})")
}
}
impl TableFormat for ApiKeyInfo {
fn header() -> Vec<Cell> {
vec![
Cell::new("UID"),
Cell::new("Name"),
Cell::new("Key"),
Cell::new("Scope"),
Cell::new("Created"),
Cell::new("Last Used"),
Cell::new("Expires At"),
]
}
fn row(&self) -> Vec<Cell> {
vec![
Cell::new(&self.uid),
Cell::new(&self.name),
Cell::new(format!("wk-**{}", self.key_suffix)),
Cell::new(&self.scope),
Cell::new(format_approx_duration_from_now_utc(self.created_at)),
Cell::new(
self.last_used_at
.map(format_approx_duration_from_now_utc)
.unwrap_or_else(|| "Never".to_string()),
),
Cell::new(
self.expires_at
.map(|dt| dt.format("%Y-%m-%d %H:%M:%S UTC").to_string())
.unwrap_or_else(|| "Never".to_string()),
),
]
}
}
fn resolve_api_key_identifier(
keys: &[ApiKeyInfo],
key_identifier: &str,
) -> Result<Option<ApiKeyInfo>> {
if let Some(key) = keys.iter().find(|key| key.uid == key_identifier) {
return Ok(Some(key.clone()));
}
let mut matches = keys
.iter()
.filter(|key| key.name == key_identifier)
.cloned()
.collect::<Vec<_>>();
matches.sort_by_key(|key| Reverse(key.created_at));
if matches.is_empty() {
return Err(anyhow!("API key '{key_identifier}' not found"));
} else if matches.len() == 1 {
return Ok(Some(matches[0].clone()));
}
if io::stdin().is_terminal() {
return match Select::new(
&format!("Multiple API keys match '{key_identifier}'. Select a key to expire:"),
matches,
)
.prompt()
{
Ok(key) => Ok(Some(key)),
Err(InquireError::OperationCanceled | InquireError::OperationInterrupted) => Ok(None),
Err(err) => Err(err.into()),
};
}
println!("Multiple API keys match '{key_identifier}':");
for key in matches {
println!(" {key}");
}
Err(anyhow!(
"Multiple API keys match '{key_identifier}'; specify the key by UID"
))
}
#[derive(Debug, Clone, Serialize)]
struct CreatedApiKeyInfo {
raw_api_key: String,
api_key: ApiKeyInfo,
}
#[derive(Debug, Clone, Serialize)]
struct ExpiredApiKeyInfo {
key_uid: String,
expired: bool,
}
fn sort_api_keys(
keys: &mut [ApiKeyInfo],
sort_by: Option<ApiKeySortByArg>,
sort_order: Option<SortOrderArg>,
) {
let Some(sort_by) = sort_by else {
return;
};
let descending = matches!(sort_order, Some(SortOrderArg::Desc));
match sort_by {
ApiKeySortByArg::Name => {
if descending {
keys.sort_by_key(|k| Reverse(k.name.to_lowercase()));
} else {
keys.sort_by_key(|k| k.name.to_lowercase());
}
}
ApiKeySortByArg::CreatedAt => {
if descending {
keys.sort_by_key(|k| Reverse(k.created_at));
} else {
keys.sort_by_key(|k| k.created_at);
}
}
ApiKeySortByArg::LastUsedAt => {
if descending {
keys.sort_by_key(|k| Reverse(k.last_used_at));
} else {
keys.sort_by_key(|k| k.last_used_at);
}
}
ApiKeySortByArg::ExpiresAt => {
if descending {
keys.sort_by_key(|k| Reverse(k.expires_at));
} else {
keys.sort_by_key(|k| k.expires_at);
}
}
ApiKeySortByArg::Scope => {
if descending {
keys.sort_by_key(|k| Reverse(k.scope.clone()));
} else {
keys.sort_by_key(|k| k.scope.clone());
}
}
}
}
fn expires_at_from_args(args: ApiKeyExpirationArgs) -> Result<Option<Time>> {
if args.no_expiration {
return Ok(None);
}
if let Some(expires_at) = args.expires_at {
return Ok(Some(Time::from(expires_at)));
}
if let Some(expires_in) = args.expires_in {
let duration = chrono::Duration::from_std(expires_in.into())
.map_err(|_| anyhow!("expiration duration is too large"))?;
return Ok(Some(Time::from(Utc::now() + duration)));
}
Err(anyhow!("expiration behavior is required"))
}
fn print_created_api_key(
result: CreatedApiKeyInfo,
output_format: OutputFormat,
json_output: warp_cli::json_filter::JsonOutput,
) -> Result<()> {
if json_output.force_json_output() {
output::print_raw_json(serde_json::to_value(&result)?, &json_output)?;
return Ok(());
}
match output_format {
OutputFormat::Json => output::write_json(&result, std::io::stdout())?,
OutputFormat::Ndjson => output::write_json_line(&result, std::io::stdout())?,
OutputFormat::Pretty | OutputFormat::Text => {
println!("API key '{}' created.", result.api_key.name);
println!("UID: {}", result.api_key.uid);
println!("Raw API key: {}", result.raw_api_key);
println!("This secret key is shown only once. Store it securely.");
}
}
Ok(())
}
fn print_expire_api_key_result(
key_uid: String,
expired: bool,
output_format: OutputFormat,
json_output: warp_cli::json_filter::JsonOutput,
) -> Result<()> {
let result = ExpiredApiKeyInfo { key_uid, expired };
if json_output.force_json_output() {
output::print_raw_json(serde_json::to_value(&result)?, &json_output)?;
return Ok(());
}
match output_format {
OutputFormat::Json => output::write_json(&result, std::io::stdout())?,
OutputFormat::Ndjson => output::write_json_line(&result, std::io::stdout())?,
OutputFormat::Pretty | OutputFormat::Text => {
if expired {
println!("API key '{}' expired.", result.key_uid);
} else {
println!("API key '{}' was not expired.", result.key_uid);
}
}
}
Ok(())
}
fn finish_command(result: Result<()>, ctx: &mut ModelContext<ApiKeyCommandRunner>) {
match result {
Ok(()) => ctx.terminate_app(TerminationMode::ForceTerminate, None),
Err(err) => super::report_fatal_error(err, ctx),
}
}
#[cfg(test)]
#[path = "api_key_tests.rs"]
mod tests;
app/src/ai/agent_sdk/api_key_tests.rs
+115
View File
@@ -0,0 +1,115 @@
use super::*;
fn key(name: &str, scope: &str, created_at: DateTime<Utc>) -> ApiKeyInfo {
key_with_uid(name, name, scope, created_at)
}
fn key_with_uid(uid: &str, name: &str, scope: &str, created_at: DateTime<Utc>) -> ApiKeyInfo {
ApiKeyInfo {
uid: uid.to_string(),
name: name.to_string(),
key_suffix: "abcd".to_string(),
scope: scope.to_string(),
created_at,
last_used_at: None,
expires_at: None,
}
}
#[test]
fn sort_api_keys_sorts_by_name_ascending() {
let created_at = Utc::now();
let mut keys = vec![
key("beta", "Team", created_at),
key("alpha", "Personal", created_at),
];
sort_api_keys(
&mut keys,
Some(ApiKeySortByArg::Name),
Some(SortOrderArg::Asc),
);
assert_eq!(keys[0].name, "alpha");
assert_eq!(keys[1].name, "beta");
}
#[test]
fn sort_api_keys_sorts_by_created_at_descending() {
let older = Utc::now() - chrono::Duration::days(1);
let newer = Utc::now();
let mut keys = vec![key("older", "Team", older), key("newer", "Personal", newer)];
sort_api_keys(
&mut keys,
Some(ApiKeySortByArg::CreatedAt),
Some(SortOrderArg::Desc),
);
assert_eq!(keys[0].name, "newer");
assert_eq!(keys[1].name, "older");
}
#[test]
fn resolve_api_key_identifier_prefers_uid_match() {
let created_at = Utc::now();
let keys = vec![
key_with_uid("target", "other-name", "Team", created_at),
key_with_uid("other-uid", "target", "Team", created_at),
];
assert_eq!(
resolve_api_key_identifier(&keys, "target")
.unwrap()
.unwrap(),
keys[0].clone()
);
}
#[test]
fn resolve_api_key_identifier_falls_back_to_name_match() {
let created_at = Utc::now();
let keys = vec![key_with_uid("uid-1", "deploy-key", "Team", created_at)];
assert_eq!(
resolve_api_key_identifier(&keys, "deploy-key")
.unwrap()
.unwrap(),
keys[0].clone()
);
}
#[test]
fn resolve_api_key_identifier_errors_for_ambiguous_name_matches() {
let created_at = Utc::now();
let keys = vec![
key_with_uid("uid-1", "deploy-key", "Team", created_at),
key_with_uid("uid-2", "deploy-key", "Personal", created_at),
];
let err = resolve_api_key_identifier(&keys, "deploy-key").unwrap_err();
assert_eq!(
err.to_string(),
"Multiple API keys match 'deploy-key'; specify the key by UID"
);
}
#[test]
fn resolve_api_key_identifier_errors_when_not_found() {
let created_at = Utc::now();
let keys = vec![key_with_uid("uid-1", "deploy-key", "Team", created_at)];
let err = resolve_api_key_identifier(&keys, "missing-key").unwrap_err();
assert_eq!(err.to_string(), "API key 'missing-key' not found");
}
#[test]
fn api_key_display_includes_creation_date() {
let created_at = "2026-01-02T03:04:05Z".parse().unwrap();
let key = key_with_uid("uid-1", "deploy-key", "Team", created_at);
assert_eq!(
key.to_string(),
"deploy-key (uid-1, created 2026-01-02 03:04:05 UTC)"
);
}
app/src/ai/agent_sdk/artifact.rs
+7 -7
View File
@@ -1,4 +1,5 @@
use std::{path::PathBuf, sync::Arc};
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Context, Result};
use galaxy_cli::agent::OutputFormat;
@@ -6,19 +7,18 @@ use galaxy_cli::artifact::{
ArtifactCommand, DownloadArtifactArgs, GetArtifactArgs, UploadArtifactArgs,
};
use galaxy_cli::GlobalOptions;
use galaxyui::{platform::TerminationMode, AppContext, ModelContext, SingletonEntity};
use serde::Serialize;
use galaxyui::platform::TerminationMode;
use galaxyui::{AppContext, ModelContext, SingletonEntity};
use super::artifact_upload::{
CompletedFileArtifactUpload, FileArtifactUploadRequest, FileArtifactUploader,
};
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,
app/src/ai/agent_sdk/artifact_tests.rs
+1
View File
@@ -1,6 +1,7 @@
use galaxy_cli::agent::OutputFormat;
use std::path::PathBuf;
use super::*;
fn sample_completed_upload() -> CompletedFileArtifactUpload {
app/src/ai/agent_sdk/artifact_upload.rs
+3 -10
View File
@@ -7,7 +7,6 @@ use std::sync::Arc;
use anyhow::{anyhow, bail, Context, Result};
use blocking::unblock;
use galaxy_cli::artifact::UploadArtifactArgs;
use mime_guess::from_path;
use super::common::parse_ambient_task_id;
use crate::ai::agent::api::ServerConversationToken;
@@ -17,10 +16,11 @@ use crate::server::server_api::ai::{
AIClient, CreateFileArtifactUploadRequest, CreateFileArtifactUploadResponse,
FileArtifactRecord, FileArtifactUploadTargetInfo,
};
use crate::server::server_api::harness_support::FileUploadBody;
use crate::server::server_api::presigned_upload::upload_file_to_target;
use crate::server::server_api::ServerApi;
use crate::util::image::{infer_mime_type, MIME_SNIFF_BYTES};
const MIME_SNIFF_BYTES: usize = 8 * 1024;
const OZ_RUN_ID_ENV_VAR: &str = "OZ_RUN_ID";
#[derive(Debug, Clone, Eq, PartialEq)]
@@ -165,8 +165,7 @@ impl FileArtifactUploader {
upload_file_to_target(
self.server_api.http_client(),
target,
&artifact.path,
artifact.file_size,
FileUploadBody::new(artifact.path.clone()),
)
.await
}
@@ -234,12 +233,6 @@ 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()))?;
app/src/ai/agent_sdk/artifact_upload_tests.rs
+4 -6
View File
@@ -1,6 +1,5 @@
use std::env;
use std::fs;
use std::path::PathBuf;
use std::{env, fs};
use chrono::Utc;
use galaxy_cli::artifact::UploadArtifactArgs;
@@ -39,15 +38,14 @@ fn create_conversation_metadata(
was_summarized: false,
context_window_usage: 0.0,
credits_spent: 0.0,
platform_credits_spent: 0.0,
credits_spent_for_last_block: None,
token_usage: vec![],
tool_usage_metadata: Default::default(),
total_cache_read_tokens: 0,
total_cache_write_tokens: 0,
total_cache_miss_tokens: 0,
total_cost_cents: 0.0,
context_window_segments: Vec::new(),
},
metadata: create_mock_server_metadata(),
creator: None,
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()),
app/src/ai/agent_sdk/common.rs
+4 -21
View File
@@ -18,7 +18,7 @@ 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::cloud_object::{CloudObject, CloudObjectLookup as _, Owner};
use crate::server::cloud_objects::update_manager::UpdateManager;
use crate::server::ids::{ServerId, SyncId};
use crate::server::server_api::ai::AIClient;
@@ -37,7 +37,7 @@ pub fn validate_agent_mode_base_model_id(
let llm_id: LLMId = model_id.into();
let valid_ids = llm_prefs
.get_base_llm_choices_for_agent_mode()
.get_base_llm_choices_for_agent_mode(ctx)
.map(|info| info.id.clone())
.collect::<Vec<_>>();
@@ -329,22 +329,5 @@ impl fmt::Display for EnvironmentChoice {
}
#[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'"));
}
}
#[path = "common_tests.rs"]
mod tests;
app/src/ai/agent_sdk/common_tests.rs
+16
View File
@@ -0,0 +1,16 @@
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'"));
}
app/src/ai/agent_sdk/config_file.rs
+5 -1
View File
@@ -20,6 +20,8 @@ pub struct AgentConfigSnapshotFile {
#[serde(default)]
pub environment_id: Option<String>,
#[serde(default)]
pub runner_id: Option<String>,
#[serde(default)]
pub model_id: Option<String>,
#[serde(default)]
pub base_prompt: Option<String>,
@@ -93,7 +95,7 @@ fn parse_yaml(input: &str) -> anyhow::Result<AgentConfigSnapshotFile> {
}
fn supported_keys_context() -> String {
"Supported keys: name, environment_id, model_id, base_prompt, mcp_servers, host, computer_use_enabled".to_string()
"Supported keys: name, environment_id, runner_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.
@@ -148,6 +150,7 @@ pub fn merge_with_precedence(
let name = cli.name.or_else(|| file.name.clone());
let environment_id = cli.environment_id.or_else(|| file.environment_id.clone());
let runner_id = cli.runner_id.or_else(|| file.runner_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());
@@ -158,6 +161,7 @@ pub fn merge_with_precedence(
AgentConfigSnapshot {
name,
environment_id,
runner_id,
model_id,
base_prompt,
mcp_servers,
app/src/ai/agent_sdk/config_file_tests.rs
+2 -1
View File
@@ -3,9 +3,9 @@
use std::io::Write as _;
use serde_json::json;
use warp_cli::mcp::MCPSpec;
use crate::ai::ambient_agents::AgentConfigSnapshot;
use galaxy_cli::mcp::MCPSpec;
fn write_temp(suffix: &str, contents: &str) -> tempfile::NamedTempFile {
let mut file = tempfile::Builder::new().suffix(suffix).tempfile().unwrap();
@@ -96,6 +96,7 @@ fn merge_precedence_cli_over_file_and_merges_mcp() {
let cli = AgentConfigSnapshot {
name: Some("cli-name".to_string()),
environment_id: None,
runner_id: None,
model_id: Some("cli-model".to_string()),
base_prompt: None,
mcp_servers: Some(serde_json::Map::from_iter([(
app/src/ai/agent_sdk/driver.rs
+2155 -622
View File
File diff suppressed because it is too large Load Diff
app/src/ai/agent_sdk/driver/attachments.rs
+16 -2
View File
@@ -2,7 +2,8 @@ use std::path::{Path, PathBuf};
use std::sync::Arc;
use anyhow::Context;
use base64::{engine::general_purpose, Engine};
use base64::engine::general_purpose;
use base64::Engine;
use futures::future::join_all;
use futures::TryStreamExt as _;
use galaxy_core::features::FeatureFlag;
@@ -30,6 +31,7 @@ pub const MAX_ATTACHMENT_COUNT_FOR_CLOUD_QUERY: usize = 25;
///
/// Makes a best-effort attempt to download all attachments.
/// Individual download failures are logged but don't cause the entire function to fail.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true))]
pub(crate) async fn fetch_and_download_attachments(
ai_client: Arc<dyn AIClient>,
http_client: Arc<ServerApi>,
@@ -65,6 +67,7 @@ pub(crate) async fn fetch_and_download_attachments(
/// 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`.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true))]
pub(crate) async fn fetch_and_download_handoff_snapshot_attachments(
ai_client: Arc<dyn AIClient>,
http_client: &http_client::Client,
@@ -219,13 +222,20 @@ async fn download_handoff_entry(
/// 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.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true))]
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 {
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true))]
async fn attempt(
http_client: &http_client::Client,
download_url: &str,
file_path: &Path,
) -> anyhow::Result<()> {
let response = http_client
.get(download_url)
.send()
@@ -254,6 +264,10 @@ async fn download_attachment(
.context("Failed to write file")?;
Ok(())
}
with_bounded_retry(&operation, || async {
attempt(http_client, download_url, file_path).await
})
.await
}
app/src/ai/agent_sdk/driver/bedrock_credentials.rs
+170
View File
@@ -0,0 +1,170 @@
/// Proactive Bedrock OIDC credential refresh for cloud agent sandboxes.
///
/// The `OidcManaged` Bedrock credential path mints an OIDC token at agent startup,
/// exchanges it for 1-hour STS temporary credentials via `AssumeRoleWithWebIdentity`,
/// and stores them in `ApiKeyManager`. Without a proactive refresh, those credentials
/// expire after ~1 hour and all subsequent Bedrock LLM calls fail.
///
/// This module provides a background `refresh_loop` — modelled on `git_credentials`'s
/// loop — that proactively re-mints the OIDC token and re-calls STS every 50 minutes,
/// well ahead of the 1-hour STS expiry. The loop is raced against the run execution
/// future via `futures::select!` and dropped automatically when the run completes.
use std::time::{Duration, SystemTime};
use ai::api_keys::{ApiKeyManager, AwsCredentials, AwsCredentialsState};
use anyhow::{Context as _, Result};
use vec1::vec1;
use warp_managed_secrets::client::IdentityTokenOptions;
use warp_managed_secrets::ManagedSecretManager;
use warpui::{ModelSpawner, SingletonEntity};
use super::AgentDriver;
use crate::ai::aws_credentials::{
aws_role_session_name, sts_client, AWS_BEDROCK_STS_AUDIENCE, BEDROCK_IDENTITY_TOKEN_DURATION,
};
/// How long to wait between Bedrock credential refresh attempts — well ahead of the
/// 1-hour STS temporary credential expiry, matching the approach used for git credentials.
pub(crate) const BEDROCK_CREDENTIALS_REFRESH_INTERVAL: Duration = Duration::from_secs(50 * 60);
/// Perform one Bedrock OIDC credential refresh attempt.
///
/// Returns `Ok(())` on success. Returns `Err` when token minting or the STS call
/// fails — these are transient failures worth retrying.
#[tracing::instrument(
name = "bedrock_credentials::try_refresh",
skip_all,
err,
fields(tags.cloud_agent = true, task_id)
)]
async fn try_refresh(
task_id: &str,
role_arn: &str,
region: &str,
foreground: &ModelSpawner<AgentDriver>,
) -> Result<()> {
// Step 1: Mint a new OIDC identity token via the model context.
let token_future = foreground
.spawn(|_, ctx| {
ManagedSecretManager::handle(ctx)
.as_ref(ctx)
.issue_task_identity_token(IdentityTokenOptions {
audience: AWS_BEDROCK_STS_AUDIENCE.to_string(),
requested_duration: BEDROCK_IDENTITY_TOKEN_DURATION,
subject_template: vec1!["scoped_principal".to_string()],
})
})
.await
.context("Failed to dispatch OIDC token request for Bedrock refresh")?;
let token = token_future
.await
.context("Failed to mint OIDC identity token for Bedrock refresh")?;
// Step 2: Exchange the OIDC token for fresh STS temporary credentials.
let client = sts_client(region).await;
let session_name = aws_role_session_name(task_id);
let sts_creds = client
.assume_role_with_web_identity()
.role_arn(role_arn)
.role_session_name(&session_name)
.web_identity_token(&token.token)
.send()
.await
.map_err(|err| {
log::error!("Bedrock OIDC refresh: STS AssumeRoleWithWebIdentity error: {err:#?}");
let detail = err
.as_service_error()
.map(|e| e.to_string())
.unwrap_or_else(|| err.to_string());
anyhow::anyhow!("STS AssumeRoleWithWebIdentity failed: {detail}")
})?
.credentials
.context("STS response did not include credentials")?;
let aws_creds = AwsCredentials::new(
sts_creds.access_key_id().to_string(),
sts_creds.secret_access_key().to_string(),
Some(sts_creds.session_token().to_string()),
SystemTime::try_from(*sts_creds.expiration()).ok(),
);
// Step 3: Update ApiKeyManager with the fresh credentials.
foreground
.spawn(move |_, ctx| {
ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
manager.set_aws_credentials_state(
AwsCredentialsState::Loaded {
credentials: aws_creds,
loaded_at: SystemTime::now(),
},
ctx,
);
});
})
.await
.context("Failed to dispatch Bedrock credential update to ApiKeyManager")?;
log::info!("Bedrock OIDC: proactive credential refresh succeeded for task {task_id}");
Ok(())
}
/// Infinite async loop that proactively refreshes AWS Bedrock OIDC credentials every
/// [`BEDROCK_CREDENTIALS_REFRESH_INTERVAL`], keeping long-running agents authenticated
/// for their entire duration.
///
/// On each iteration:
/// 1. Issue a new OIDC identity token via warp-server.
/// 2. Call STS `AssumeRoleWithWebIdentity` to get fresh temporary credentials.
/// 3. Update `ApiKeyManager` with the new credentials.
///
/// On transient failure, retries up to three times with exponential backoff
/// (1 min, 2 min, 4 min), keeping all retries well within the buffer before
/// the 1-hour STS credentials expire. If all retries fail, a warning is logged
/// and the next refresh is scheduled after the normal interval.
///
/// This future never resolves — it is designed to be raced with the run execution
/// future via `futures::select!` and dropped automatically when the run completes.
pub(crate) async fn refresh_loop(
task_id: String,
role_arn: String,
region: String,
foreground: &ModelSpawner<AgentDriver>,
) {
loop {
warpui::r#async::Timer::after(BEDROCK_CREDENTIALS_REFRESH_INTERVAL).await;
log::info!("Proactively refreshing AWS Bedrock OIDC credentials for task {task_id}");
let backoff_delays = [
Duration::from_secs(60),
Duration::from_secs(2 * 60),
Duration::from_secs(4 * 60),
];
let mut attempt = 0usize;
loop {
match try_refresh(&task_id, &role_arn, &region, foreground).await {
Ok(()) => break,
Err(e) if attempt < backoff_delays.len() => {
let delay = backoff_delays[attempt];
log::warn!(
"Bedrock credentials refresh failed (attempt {}): {e:#}; \
retrying in {}s",
attempt + 1,
delay.as_secs()
);
warpui::r#async::Timer::after(delay).await;
attempt += 1;
}
Err(e) => {
log::warn!(
"Bedrock credentials refresh failed after {} attempts: {e:#}; \
credentials may expire before next refresh cycle",
attempt + 1
);
break;
}
}
}
}
}
app/src/ai/agent_sdk/driver/cloud_provider.rs
+4 -1
View File
@@ -1,4 +1,7 @@
use std::{collections::HashMap, ffi::OsString, future::Future, pin::Pin};
use std::collections::HashMap;
use std::ffi::OsString;
use std::future::Future;
use std::pin::Pin;
use anyhow::Error;
use galaxyui::ModelSpawner;
app/src/ai/agent_sdk/driver/cloud_provider/aws.rs
+7 -4
View File
@@ -1,4 +1,8 @@
use std::{collections::HashMap, ffi::OsString, future::Future, pin::Pin, time::Duration};
use std::collections::HashMap;
use std::ffi::OsString;
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use anyhow::Context;
use galaxy_core::safe_info;
@@ -7,11 +11,10 @@ use galaxyui::{ModelSpawner, SingletonEntity};
use tempfile::{Builder, NamedTempFile};
use vec1::Vec1;
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};
use crate::ai::aws_credentials::aws_role_session_name;
use crate::ai::cloud_environments::AwsProviderConfig;
/// 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
app/src/ai/agent_sdk/driver/cloud_provider/gcp.rs
+6 -3
View File
@@ -1,11 +1,14 @@
use std::{collections::HashMap, ffi::OsString, future::Future, pin::Pin, time::Duration};
use std::collections::HashMap;
use std::ffi::OsString;
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use anyhow::Context as _;
use galaxy_managed_secrets::{GcpCredentials, GcpFederationConfig};
use crate::ai::cloud_environments::GcpProviderConfig;
use super::{CloudProvider, CloudProviderSetupError, Result};
use crate::ai::cloud_environments::GcpProviderConfig;
/// Token lifetime for GCP executable-sourced credentials. The GCP client
/// libraries handle refreshing automatically, so we keep this short.
app/src/ai/agent_sdk/driver/cloud_provider_tests.rs
+6 -5
View File
@@ -1,11 +1,12 @@
use std::{collections::HashMap, ffi::OsString, path::PathBuf};
use std::collections::HashMap;
use std::ffi::OsString;
use std::path::PathBuf;
use super::aws::AwsCloudProvider;
use super::gcp::GcpCloudProvider;
use super::{collect_env_vars, load_providers, CloudProvider};
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 {
app/src/ai/agent_sdk/driver/environment.rs
+376 -188
View File
@@ -1,25 +1,28 @@
use std::{
collections::HashMap,
future::Future,
path::{Path, PathBuf},
sync::{Arc, Mutex},
time::Duration,
};
use std::collections::HashMap;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::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 galaxy_completer::completer::CommandExitStatus;
use galaxy_core::{command::ExitCode, safe_info, safe_warn};
use galaxyui::{r#async::FutureExt, ModelContext, ModelSpawner, SingletonEntity};
use futures::channel::oneshot;
use futures::future::join_all;
use repo_metadata::repositories::{DetectedRepositories, RepoDetectionSource};
use super::{terminal::TerminalDriver, AgentDriverError};
use galaxy_cli::agent::Harness;
use galaxy_completer::completer::CommandExitStatus;
use galaxy_core::command::ExitCode;
use galaxy_core::{safe_info, safe_warn};
use galaxyui::r#async::FutureExt;
use galaxyui::{ModelContext, ModelSpawner, SingletonEntity};
use super::terminal::TerminalDriver;
use super::AgentDriverError;
use crate::ai::agent_sdk::setup_observability::{SetupClientEventReporter, SetupStep};
use crate::ai::cloud_environments::{AmbientAgentEnvironment, SourceRepo};
use crate::terminal::model::session::command_executor::shell_escape_single_quotes;
use crate::terminal::shell::ShellType;
const CODEBASE_INDEX_SYNC_TIMEOUT: Duration = Duration::from_secs(60);
@@ -54,20 +57,18 @@ pub fn prepare_environment(
working_dir: PathBuf,
is_sandbox: bool,
harness: Harness,
setup_events: SetupClientEventReporter,
ctx: &mut ModelContext<TerminalDriver>,
) -> impl Future<Output = Result<(), PrepareEnvironmentError>> {
let spawner = ctx.spawner();
async move {
let AmbientAgentEnvironment {
github_repos,
setup_commands,
..
} = environment;
let source_repos = environment.effective_repos();
let setup_commands = environment.setup_commands;
// 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 should_subscribe_to_index_updates = should_index_codebase && !source_repos.is_empty();
let repo_channels = Arc::new(Mutex::new(HashMap::<PathBuf, oneshot::Sender<()>>::new()));
if should_subscribe_to_index_updates {
@@ -78,14 +79,15 @@ pub fn prepare_environment(
&spawner,
working_dir.as_path(),
is_sandbox,
&github_repos,
&source_repos,
setup_commands,
should_index_codebase,
Arc::clone(&repo_channels),
setup_events,
)
.await;
if should_subscribe_to_index_updates {
if should_subscribe_to_index_updates && result.is_err() {
let _ = spawner
.spawn(|_, ctx| {
ctx.unsubscribe_from_model(&CodebaseIndexManager::handle(ctx));
@@ -97,14 +99,16 @@ pub fn prepare_environment(
}
}
#[allow(clippy::too_many_arguments)]
async fn prepare_environment_impl(
spawner: &ModelSpawner<TerminalDriver>,
working_dir: &Path,
is_sandbox: bool,
github_repos: &[GithubRepo],
source_repos: &[SourceRepo],
setup_commands: Vec<String>,
should_index_codebase: bool,
repo_channels: Arc<Mutex<HashMap<PathBuf, oneshot::Sender<()>>>>,
setup_events: SetupClientEventReporter,
) -> Result<(), PrepareEnvironmentError> {
let working_dir_string = working_dir.to_string_lossy().to_string();
@@ -120,175 +124,95 @@ async fn prepare_environment_impl(
}
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,
if !source_repos.is_empty() {
setup_events
.record_result(SetupStep::EnvironmentRepoClone, async {
clone_repos(source_repos, working_dir, spawner).await?;
for repo in source_repos {
register_cloned_repo(repo, working_dir, is_sandbox, spawner).await?;
if !is_sandbox && should_index_codebase {
let receiver = index_repo_codebase(
&repo.repo,
working_dir,
Arc::clone(&repo_channels),
spawner,
)
})
})
.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);
.await?;
if let Some(receiver) = receiver {
codebase_context_receivers.push(receiver);
}
}
}
}
Ok::<(), PrepareEnvironmentError>(())
})
.await?;
if should_index_codebase {
record_codebase_indexing(
setup_events.clone(),
spawner.clone(),
codebase_context_receivers,
);
}
}
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?;
setup_events
.record_result(SetupStep::EnvironmentSetupCommands, async {
// 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}")
);
}
// Unset CI after setup commands complete so the agent session
// does not run with CI=true.
execute_command("unset CI".to_string(), spawner).await?;
Ok::<(), PrepareEnvironmentError>(())
})
.await?;
} else if should_index_codebase && source_repos.is_empty() {
let _ = spawner
.spawn(|_, ctx| {
ctx.unsubscribe_from_model(&CodebaseIndexManager::handle(ctx));
})
.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!(
galaxy_isolation_platform::detect(),
Some(
galaxy_isolation_platform::IsolationPlatformType::DockerSandbox
| galaxy_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 should_index_codebase && source_repos.is_empty() {
log::info!("No repositories to index for codebase context");
}
// 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) {
if let Some(repo_name) = single_repo_name(source_repos) {
safe_info!(
safe: ("Changing directory into single repository"),
full: ("Changing directory into single repository: {repo_name}")
@@ -302,6 +226,269 @@ async fn prepare_environment_impl(
Ok(())
}
fn record_codebase_indexing(
setup_events: SetupClientEventReporter,
spawner: ModelSpawner<TerminalDriver>,
codebase_context_receivers: Vec<oneshot::Receiver<()>>,
) {
if codebase_context_receivers.is_empty() {
setup_events.record_value_detached(SetupStep::EnvironmentCodebaseIndexing, async move {
let _ = spawner
.spawn(|_, ctx| {
ctx.unsubscribe_from_model(&CodebaseIndexManager::handle(ctx));
})
.await;
});
return;
}
setup_events.record_value_detached(SetupStep::EnvironmentCodebaseIndexing, async move {
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",
);
}
let _ = spawner
.spawn(|_, ctx| {
ctx.unsubscribe_from_model(&CodebaseIndexManager::handle(ctx));
})
.await;
});
}
fn build_parallel_clone_command(repos: &[SourceRepo], shell_type: ShellType) -> String {
let mut script = String::from(
r#"set +e
failed=0
pids=""
tmp_dir="$(mktemp -d "${TMPDIR:-/tmp}/warp-clone-logs.XXXXXX")"
cleanup_clone_logs() {
rm -rf "$tmp_dir"
}
trap cleanup_clone_logs EXIT
clone_repo() {
repo_name="$1"
repo_url="$2"
target="$3"
if [ -d "$target" ]; then
printf '%s\n' "Repository directory $target already exists, skipping clone..."
return 0
fi
printf '%s\n' "Cloning repository $repo_name..."
git clone --filter=tree:0 "$repo_url" "$target"
}
"#,
);
let mut log_outputs = String::new();
for (index, repo) in repos.iter().enumerate() {
let repo_name = format!("{}/{}", repo.owner, repo.repo);
let repo_url = repo.https_clone_url();
let escaped_repo_name = shell_escape_single_quotes(&repo_name, ShellType::Bash);
let escaped_repo_url = shell_escape_single_quotes(&repo_url, ShellType::Bash);
let escaped_target = shell_escape_single_quotes(&repo.repo, ShellType::Bash);
let log_var = format!("log_file_{index}");
script.push_str(&format!(
"{log_var}=\"$tmp_dir/repo-{index}.log\"\n\
clone_repo '{escaped_repo_name}' '{escaped_repo_url}' '{escaped_target}' >\"${log_var}\" 2>&1 &\n"
));
script.push_str("pids=\"$pids $!\"\n");
log_outputs.push_str(&format!(
"printf '%s\\n' '===== {escaped_repo_name} ====='\n\
if [ -s \"${log_var}\" ]; then\n\
\tcat \"${log_var}\"\n\
else\n\
\tprintf '%s\\n' '(no output)'\n\
fi\n"
));
}
script.push_str(
r#"for pid in $pids; do
if ! wait "$pid"; then
failed=1
fi
done
"#,
);
script.push_str(&log_outputs);
script.push_str(
r#"
exit "$failed"
"#,
);
let escaped_script = shell_escape_single_quotes(&script, shell_type);
format!("sh -c '{escaped_script}'")
}
/// Clone all source repositories to `{working_dir}/{repo.repo}` if they do not already exist.
/// Multiple repositories are cloned in parallel to reduce environment setup time.
pub(super) async fn clone_repos(
repos: &[SourceRepo],
working_dir: &Path,
spawner: &ModelSpawner<TerminalDriver>,
) -> Result<(), PrepareEnvironmentError> {
match repos {
[] => Ok(()),
[repo] => clone_repo(repo, working_dir, spawner).await,
repos => {
let shell_type = spawner
.spawn(|driver, ctx| {
driver
.active_session_shell_type(ctx)
.unwrap_or(ShellType::Bash)
})
.await
.unwrap_or(ShellType::Bash);
let repo_names = repos
.iter()
.map(|repo| format!("{}/{}", repo.owner, repo.repo))
.collect::<Vec<_>>();
safe_info!(
safe: ("Cloning repositories via terminal"),
full: ("Cloning repositories via terminal: {}", repo_names.join(", "))
);
let command = build_parallel_clone_command(repos, shell_type);
let exit_code = execute_command(command, spawner).await?;
if exit_code != 0.into() {
return Err(PrepareEnvironmentError::CloneRepo {
repo_name: repo_names.join(", "),
});
}
safe_info!(
safe: ("Successfully cloned repositories"),
full: ("Successfully cloned repositories: {}", repo_names.join(", "))
);
Ok(())
}
}
}
/// Clone a source repository to `{working_dir}/{repo.repo}` if it does not already exist.
/// This only performs the clone -- it does NOT register the repo with `DetectedRepositories`.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true, repo = %repo))]
pub(super) async fn clone_repo(
repo: &SourceRepo,
working_dir: &Path,
spawner: &ModelSpawner<TerminalDriver>,
) -> Result<(), PrepareEnvironmentError> {
let repo_name = format!("{}/{}", repo.owner, repo.repo);
let repo_url = repo.https_clone_url();
// Get the session's shell type for proper escaping, falling back to Bash
// when the session is not yet bootstrapped or the spawn fails.
let shell_type = spawner
.spawn(|driver, ctx| {
driver
.active_session_shell_type(ctx)
.unwrap_or(ShellType::Bash)
})
.await
.unwrap_or(ShellType::Bash);
let escaped_url = shell_escape_single_quotes(&repo_url, shell_type);
// We do a partial clone here to speed up environment setup time.
let command = format!("git clone --filter=tree:0 '{escaped_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}")
);
}
Ok(())
}
/// Register a cloned source repository with `DetectedRepositories` so that the
/// skill watcher and other repo-aware subsystems can discover it.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true, repo = %repo, is_sandbox = is_sandbox))]
pub(super) async fn register_cloned_repo(
repo: &SourceRepo,
working_dir: &Path,
is_sandbox: bool,
spawner: &ModelSpawner<TerminalDriver>,
) -> Result<(), PrepareEnvironmentError> {
let repo_dir = working_dir.join(&repo.repo);
// 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_local_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_local_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())
);
}
}
Ok(())
}
async fn subscribe_to_codebase_index_events(
spawner: &ModelSpawner<TerminalDriver>,
repo_channels: Arc<Mutex<HashMap<PathBuf, oneshot::Sender<()>>>>,
@@ -309,10 +496,11 @@ async fn subscribe_to_codebase_index_events(
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) {
ctx.subscribe_to_model(&CodebaseIndexManager::handle(ctx), move |_, _, event, ctx| {
if !matches!(
event,
CodebaseIndexManagerEvent::SyncStateUpdated { .. }
) {
return;
}
@@ -348,13 +536,13 @@ async fn subscribe_to_codebase_index_events(
let _ = tx.send(());
}
}
},
);
});
})
.await
.map_err(|_| PrepareEnvironmentError::InvalidRuntimeState)
}
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true, repo = %repo_name))]
async fn index_repo_codebase(
repo_name: &str,
working_dir: &Path,
@@ -454,7 +642,7 @@ async fn cd_in_terminal(
})
}
fn single_repo_name(repos: &[GithubRepo]) -> Option<String> {
fn single_repo_name(repos: &[SourceRepo]) -> Option<String> {
if repos.len() != 1 {
return None;
}
app/src/ai/agent_sdk/driver/environment_tests.rs
+58 -6
View File
@@ -1,9 +1,13 @@
use super::single_repo_name;
use crate::ai::cloud_environments::GithubRepo;
use cloud_object_models::CodeForge;
use super::{build_parallel_clone_command, single_repo_name};
use crate::ai::cloud_environments::SourceRepo;
use crate::terminal::shell::ShellType;
#[test]
fn single_repo_name_returns_repo_when_exactly_one_repo() {
let repos = vec![GithubRepo::new(
let repos = vec![SourceRepo::new(
CodeForge::GitHub,
"warpdotdev".to_string(),
"warp-internal".to_string(),
)];
@@ -13,12 +17,60 @@ fn single_repo_name_returns_repo_when_exactly_one_repo() {
#[test]
fn single_repo_name_returns_none_for_zero_or_many_repos() {
let no_repos = Vec::<GithubRepo>::new();
let no_repos = Vec::<SourceRepo>::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()),
SourceRepo::new(
CodeForge::GitHub,
"warpdotdev".to_string(),
"warp-internal".to_string(),
),
SourceRepo::new(
CodeForge::GitHub,
"warpdotdev".to_string(),
"warp-server".to_string(),
),
];
assert_eq!(single_repo_name(&two_repos), None);
}
#[test]
fn parallel_clone_command_runs_repos_in_background_and_waits() {
let repos = vec![
SourceRepo::new(
CodeForge::GitHub,
"warpdotdev".to_string(),
"warp".to_string(),
),
SourceRepo::new(
CodeForge::GitLab,
"platform/backend".to_string(),
"api".to_string(),
),
];
let command = build_parallel_clone_command(&repos, ShellType::Bash);
assert!(command.starts_with("sh -c '"));
assert!(command.contains("warpdotdev/warp"));
assert!(command.contains("https://github.com/warpdotdev/warp.git"));
assert!(command.contains("platform/backend/api"));
assert!(command.contains("https://gitlab.com/platform/backend/api.git"));
assert_eq!(command.matches("clone_repo").count(), 3);
assert_eq!(command.matches("2>&1 &").count(), 2);
assert!(command.contains("mktemp -d"));
assert!(command.contains("warp-clone-logs"));
assert!(command.contains("trap cleanup_clone_logs EXIT"));
assert!(command.contains("repo-0.log"));
assert!(command.contains("repo-1.log"));
assert!(command.contains(">\"$log_file_0\" 2>&1 &"));
assert!(command.contains(">\"$log_file_1\" 2>&1 &"));
assert!(command.contains("pids=\"$pids $!\""));
assert!(command.contains("wait \"$pid\""));
assert!(command.contains("===== warpdotdev/warp ====="));
assert!(command.contains("cat \"$log_file_0\""));
assert!(command.contains("===== platform/backend/api ====="));
assert!(command.contains("cat \"$log_file_1\""));
assert!(command.contains("exit \"$failed\""));
}
app/src/ai/agent_sdk/driver/error_classification.rs
+72 -10
View File
@@ -1,9 +1,9 @@
use crate::ai::blocklist::task_status_sync_model::classify_renderable_error;
use crate::server::server_api::ai::TaskStatusUpdate;
use galaxy_graphql::ai::{AgentTaskState, PlatformErrorCode};
use super::terminal::ShareSessionError;
use super::AgentDriverError;
use crate::ai::blocklist::local_agent_task_sync_model::classify_renderable_error;
use crate::server::server_api::ai::TaskStatusUpdate;
/// Classify an `AgentDriverError` into a task state and a `TaskStatusUpdate`
/// suitable for reporting via `update_agent_task`.
@@ -17,10 +17,10 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
PlatformErrorCode::InternalError,
),
),
AgentDriverError::BootstrapFailed => (
AgentDriverError::BootstrapFailed { error } => (
AgentTaskState::Error,
TaskStatusUpdate::with_error_code(
"Terminal session failed to start. Please try running your task again.",
format!("Terminal session failed to start: {error}"),
PlatformErrorCode::InternalError,
),
),
@@ -100,13 +100,29 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
PlatformErrorCode::EnvironmentSetupFailed,
),
),
AgentDriverError::MCPStartupFailed => (
AgentDriverError::ManagedMcpResolutionFailed { uid, message } => (
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.",
format!("Managed MCP server {uid} could not be resolved: {message}"),
PlatformErrorCode::EnvironmentSetupFailed,
),
),
AgentDriverError::MCPStartupFailed { details } => {
let server_lines = details
.iter()
.map(|detail| format!("- {detail}"))
.collect::<Vec<_>>()
.join("\n");
(
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
format!(
"One or more MCP servers failed to start:\n\n{server_lines}\n\nCheck that each server's configuration is valid and that it is reachable from the agent's environment."
),
PlatformErrorCode::EnvironmentSetupFailed,
),
)
}
AgentDriverError::MCPJsonParseError(msg) => (
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
@@ -171,7 +187,7 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
// --- 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,
// fallback — LocalAgentTaskSyncModel 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);
@@ -248,7 +264,11 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
PlatformErrorCode::InternalError,
),
),
AgentDriverError::ConversationHarnessMismatch { conversation_id, expected, got } => (
AgentDriverError::ConversationHarnessMismatch {
conversation_id,
expected,
got,
} => (
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
format!(
@@ -258,7 +278,11 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
PlatformErrorCode::EnvironmentSetupFailed,
),
),
AgentDriverError::TaskHarnessMismatch { task_id, expected, got } => (
AgentDriverError::TaskHarnessMismatch {
task_id,
expected,
got,
} => (
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
format!(
@@ -268,7 +292,10 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
PlatformErrorCode::EnvironmentSetupFailed,
),
),
AgentDriverError::ConversationResumeStateMissing { harness, conversation_id } => (
AgentDriverError::ConversationResumeStateMissing {
harness,
conversation_id,
} => (
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
format!(
@@ -299,6 +326,41 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS
PlatformErrorCode::EnvironmentSetupFailed,
),
),
AgentDriverError::HarnessAuthCheckFailed { harness, detail } => {
let message = format!(
"Harness '{harness}' authentication check failed: login credentials \
are invalid or expired. Verify that the authentication secret \
configured for this harness is correct."
);
log::error!("Preflight detail for {harness}: {detail}");
(
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
message,
PlatformErrorCode::AuthenticationRequired,
),
)
}
AgentDriverError::HarnessRuntimeFailureDetected {
harness,
pattern,
excerpt,
} => {
let message = format!(
"Harness '{harness}' could not make a successful API request. \
Matched failure pattern '{pattern}' in harness output: \"{excerpt}\". \
This usually means the API key is invalid, out of credits, or the \
account is misconfigured."
);
log::error!("Runtime failure for {harness}: pattern={pattern}, excerpt={excerpt}");
(
AgentTaskState::Failed,
TaskStatusUpdate::with_error_code(
message,
PlatformErrorCode::AuthenticationRequired,
),
)
}
}
}
app/src/ai/agent_sdk/driver/error_classification_tests.rs
+151 -6
View File
@@ -1,7 +1,7 @@
use galaxy_graphql::ai::{AgentTaskState, PlatformErrorCode};
use super::classify_driver_error;
use crate::ai::agent_sdk::driver::terminal::ShareSessionError;
use crate::ai::agent_sdk::driver::terminal::{BootstrapError, ShareSessionError};
use crate::ai::agent_sdk::driver::AgentDriverError;
fn assert_state_and_code(
@@ -20,14 +20,58 @@ fn assert_state_and_code(
// --- Infrastructure errors → ERROR ---
#[test]
fn bootstrap_failed_is_error_with_internal() {
assert_state_and_code(
AgentDriverError::BootstrapFailed,
AgentTaskState::Error,
Some(PlatformErrorCode::InternalError),
fn bootstrap_pty_spawn_failed_with_reason_includes_reason_in_message() {
let (state, update) = classify_driver_error(&AgentDriverError::BootstrapFailed {
error: BootstrapError::PtySpawnFailed {
reason: Some("Argument list too long (os error 7)".to_string()),
},
});
assert_eq!(state, AgentTaskState::Error);
assert_eq!(update.error_code, Some(PlatformErrorCode::InternalError));
assert!(
update.message.contains("Argument list too long"),
"message should include the specific failure reason: {:?}",
update.message
);
}
#[test]
fn bootstrap_pty_spawn_failed_without_reason_is_generic() {
let (state, update) = classify_driver_error(&AgentDriverError::BootstrapFailed {
error: BootstrapError::PtySpawnFailed { reason: None },
});
assert_eq!(state, AgentTaskState::Error);
assert_eq!(update.error_code, Some(PlatformErrorCode::InternalError));
assert!(
update.message.contains("Shell spawn failed"),
"message should describe the spawn failure: {:?}",
update.message
);
}
#[test]
fn bootstrap_timed_out_is_error_with_internal() {
let (state, update) = classify_driver_error(&AgentDriverError::BootstrapFailed {
error: BootstrapError::TimedOut,
});
assert_eq!(state, AgentTaskState::Error);
assert_eq!(update.error_code, Some(PlatformErrorCode::InternalError));
assert!(
update.message.contains("did not start within"),
"message should describe the timeout: {:?}",
update.message
);
}
#[test]
fn bootstrap_internal_error_is_error_with_internal() {
let (state, update) = classify_driver_error(&AgentDriverError::BootstrapFailed {
error: BootstrapError::InternalError,
});
assert_eq!(state, AgentTaskState::Error);
assert_eq!(update.error_code, Some(PlatformErrorCode::InternalError));
}
#[test]
fn terminal_unavailable_is_error_with_internal() {
assert_state_and_code(
@@ -72,6 +116,40 @@ fn mcp_server_not_found_is_failed_with_env_setup() {
);
}
#[test]
fn managed_mcp_resolution_failed_is_failed_with_env_setup() {
assert_state_and_code(
AgentDriverError::ManagedMcpResolutionFailed {
uid: uuid::Uuid::nil(),
message: "not active".into(),
},
AgentTaskState::Failed,
Some(PlatformErrorCode::EnvironmentSetupFailed),
);
}
#[test]
fn mcp_startup_failed_is_failed_with_env_setup_and_per_server_details() {
let (state, update) = classify_driver_error(&AgentDriverError::MCPStartupFailed {
details: vec![
"'devin' failed to start: connection refused".to_string(),
"'datadog' did not start within 20s".to_string(),
],
});
assert_eq!(state, AgentTaskState::Failed);
assert_eq!(
update.error_code,
Some(PlatformErrorCode::EnvironmentSetupFailed)
);
// Each unavailable server is rendered as its own bullet line.
assert!(update
.message
.contains("- 'devin' failed to start: connection refused"));
assert!(update
.message
.contains("- 'datadog' did not start within 20s"));
}
#[test]
fn environment_setup_failed_is_failed() {
assert_state_and_code(
@@ -99,6 +177,35 @@ fn environment_not_found_is_failed_with_resource_not_found() {
);
}
#[test]
fn conversation_harness_mismatch_is_failed_with_env_setup() {
let (state, update) = classify_driver_error(&AgentDriverError::ConversationHarnessMismatch {
conversation_id: "conv-123".into(),
expected: "claude".into(),
got: "oz".into(),
});
assert_eq!(state, AgentTaskState::Failed);
assert_eq!(
update.error_code,
Some(PlatformErrorCode::EnvironmentSetupFailed)
);
assert!(update.message.contains("conv-123"));
assert!(update.message.contains("--harness claude"));
}
#[test]
fn conversation_resume_state_missing_is_failed_with_resource_not_found() {
let (state, update) =
classify_driver_error(&AgentDriverError::ConversationResumeStateMissing {
harness: "claude".into(),
conversation_id: "conv-123".into(),
});
assert_eq!(state, AgentTaskState::Failed);
assert_eq!(update.error_code, Some(PlatformErrorCode::ResourceNotFound));
assert!(update.message.contains("conv-123"));
assert!(update.message.contains("claude"));
}
// --- ShareSessionFailed variants ---
#[test]
@@ -153,3 +260,41 @@ fn conversation_blocked_is_blocked() {
assert_eq!(state, AgentTaskState::Blocked);
assert!(update.message.contains("rm -rf /"));
}
// --- Harness auth preflight errors ---
#[test]
fn harness_auth_check_failed_is_failed_with_auth_required() {
let (state, update) = classify_driver_error(&AgentDriverError::HarnessAuthCheckFailed {
harness: "claude".into(),
detail: "exit code 1".into(),
});
assert_eq!(state, AgentTaskState::Failed);
assert_eq!(
update.error_code,
Some(PlatformErrorCode::AuthenticationRequired)
);
assert!(update.message.contains("authentication check failed"));
assert!(update.message.contains("claude"));
}
// --- Runtime failure detection ---
#[test]
fn harness_runtime_failure_detected_is_failed_with_auth_required() {
let (state, update) = classify_driver_error(&AgentDriverError::HarnessRuntimeFailureDetected {
harness: "claude".into(),
pattern: "credit balance is too low".into(),
excerpt: "Error: Your credit balance is too low to make this request.".into(),
});
assert_eq!(state, AgentTaskState::Failed);
assert_eq!(
update.error_code,
Some(PlatformErrorCode::AuthenticationRequired)
);
// The user-visible message must surface both the matched pattern and
// the excerpt so on-call/users have actionable context.
assert!(update.message.contains("claude"));
assert!(update.message.contains("credit balance is too low"));
assert!(update.message.contains("Your credit balance is too low"));
}
app/src/ai/agent_sdk/driver/git_credentials.rs
+417
View File
@@ -0,0 +1,417 @@
/// Git credentials management for cloud agent sandboxes.
///
/// This module handles:
/// - Writing provider credentials to `~/.git-credentials`, plus GitHub
/// credentials to `~/.config/gh/hosts.yml`, without requiring environment
/// variables.
/// - One-time git configuration (`credential.helper store`, SSH→HTTPS URL
/// rewrites).
/// - Configuring the git user identity from the server-returned username/email.
/// - An async refresh loop that periodically fetches a fresh token from the
/// server and overwrites the credential files, keeping long-running agents
/// authenticated for their entire duration.
use std::{path::PathBuf, sync::Arc, time::Duration};
use anyhow::{Context, Result};
// Use the project's allowed Command wrapper (not std::process::Command, which is
// disallowed by clippy rules because it flashes a terminal window on Windows).
use command::blocking::Command as BlockingCommand;
use crate::server::server_api::ai::{AIClient, GitCredential};
/// How long to wait between credential refresh attempts (~50 minutes, staying
/// well ahead of the shortest-lived one-hour token expiry).
pub(crate) const GIT_CREDENTIALS_REFRESH_INTERVAL: Duration = Duration::from_secs(50 * 60);
const DEFAULT_GIT_NAME: &str = "Oz";
const DEFAULT_GIT_EMAIL: &str = "oz-agent@warp.dev";
const GITHUB_HOST: &str = "github.com";
const GH_HOSTS_FILENAME: &str = "hosts.yml";
const GLAB_HOST: &str = "gitlab.com";
const GLAB_CONFIG_FILENAME: &str = "config.yml";
fn home_dir() -> Result<PathBuf> {
dirs::home_dir().ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))
}
/// Write `content` to `path` using owner-only (0600) permissions.
///
/// On Unix the file is created with mode 0600 so no other user can read the
/// credential material. On non-Unix platforms the function falls back to the
/// standard write, relying on OS default permissions.
fn write_secret_file(path: &std::path::Path, content: &str) -> Result<()> {
#[cfg(unix)]
{
use std::io::Write as _;
use std::os::unix::fs::{OpenOptionsExt as _, PermissionsExt as _};
let mut file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.mode(0o600)
.open(path)
.with_context(|| format!("Failed to open {} for writing", path.display()))?;
file.set_permissions(std::fs::Permissions::from_mode(0o600))
.with_context(|| format!("Failed to set permissions on {}", path.display()))?;
file.write_all(content.as_bytes())
.with_context(|| format!("Failed to write {}", path.display()))?;
}
#[cfg(not(unix))]
{
std::fs::write(path, content)
.with_context(|| format!("Failed to write {}", path.display()))?;
}
Ok(())
}
fn git_credentials_file_content(credentials: &[GitCredential]) -> String {
let mut content = String::new();
for cred in credentials {
let userinfo = match &cred.username {
Some(username) => format!("{username}:{}", cred.token),
None => format!("x-access-token:{}", cred.token),
};
content.push_str(&format!("https://{}@{}\n", userinfo, cred.host));
}
content
}
/// Write `~/.git-credentials` with the given credentials.
///
/// Each credential entry is formatted as:
/// - `https://{username}:{token}@{host}` when a username is present
/// - `https://x-access-token:{token}@{host}` for service-account tokens
///
/// The write is done atomically: a temporary file is written then renamed.
fn write_git_credentials_file(credentials: &[GitCredential]) -> Result<()> {
if credentials.is_empty() {
return Ok(());
}
let home = home_dir()?;
let path = home.join(".git-credentials");
let tmp_path = home.join(".git-credentials.tmp");
let content = git_credentials_file_content(credentials);
write_secret_file(&tmp_path, &content)?;
std::fs::rename(&tmp_path, &path).with_context(|| {
format!(
"Failed to rename {} to {}",
tmp_path.display(),
path.display()
)
})?;
Ok(())
}
/// Write `~/.config/gh/hosts.yml` so the `gh` CLI is authenticated.
///
/// The YAML format is stable for `gh` v2+:
/// ```yaml
/// github.com:
/// oauth_token: TOKEN
/// git_protocol: https
/// user: USERNAME
/// ```
///
/// The write is atomic: a temporary file is written then renamed.
fn write_gh_hosts_yml(credentials: &[GitCredential], home: &std::path::Path) -> Result<()> {
let github_credentials = credentials
.iter()
.filter(|credential| credential.host == GITHUB_HOST)
.collect::<Vec<_>>();
if github_credentials.is_empty() {
return Ok(());
}
let gh_config_dir = home.join(".config").join("gh");
std::fs::create_dir_all(&gh_config_dir)
.with_context(|| format!("Failed to create {}", gh_config_dir.display()))?;
let path = gh_config_dir.join(GH_HOSTS_FILENAME);
let tmp_path = gh_config_dir.join(format!("{GH_HOSTS_FILENAME}.tmp"));
let mut yaml = String::new();
for cred in github_credentials {
yaml.push_str(&format!("{}:\n", cred.host));
yaml.push_str(&format!(" oauth_token: {}\n", cred.token));
yaml.push_str(" git_protocol: https\n");
if let Some(username) = &cred.username {
yaml.push_str(&format!(" user: {username}\n"));
}
}
write_secret_file(&tmp_path, &yaml)?;
std::fs::rename(&tmp_path, &path).with_context(|| {
format!(
"Failed to rename {} to {}",
tmp_path.display(),
path.display()
)
})?;
Ok(())
}
/// Write `~/.config/glab-cli/config.yml` so the `glab` CLI is authenticated.
///
/// The YAML format for glab is:
/// ```yaml
/// hosts:
/// gitlab.com:
/// token: TOKEN
/// git_protocol: https
/// api_protocol: https
/// ```
///
/// The write is atomic: a temporary file is written then renamed.
fn write_glab_config(credentials: &[GitCredential], home: &std::path::Path) -> Result<()> {
let gitlab_credentials = credentials
.iter()
.filter(|credential| credential.host == GLAB_HOST)
.collect::<Vec<_>>();
if gitlab_credentials.is_empty() {
return Ok(());
}
let glab_config_dir = home.join(".config").join("glab-cli");
std::fs::create_dir_all(&glab_config_dir)
.with_context(|| format!("Failed to create {}", glab_config_dir.display()))?;
let path = glab_config_dir.join(GLAB_CONFIG_FILENAME);
let tmp_path = glab_config_dir.join(format!("{GLAB_CONFIG_FILENAME}.tmp"));
let mut yaml = String::new();
yaml.push_str("hosts:\n");
for cred in gitlab_credentials {
yaml.push_str(&format!(" {}:\n", cred.host));
yaml.push_str(&format!(" token: {}\n", cred.token));
yaml.push_str(" git_protocol: https\n");
yaml.push_str(" api_protocol: https\n");
}
write_secret_file(&tmp_path, &yaml)?;
std::fs::rename(&tmp_path, &path).with_context(|| {
format!(
"Failed to rename {} to {}",
tmp_path.display(),
path.display()
)
})?;
Ok(())
}
/// Formats non-sensitive metadata for verifying local credential injection.
pub(crate) fn credential_diagnostics(credentials: &[GitCredential]) -> String {
credentials
.iter()
.map(|credential| {
format!(
"{}(token_present={}, username_present={})",
credential.host,
!credential.token.is_empty(),
credential.username.is_some()
)
})
.collect::<Vec<_>>()
.join(", ")
}
pub(crate) fn write_git_credentials(credentials: &[GitCredential]) -> Result<()> {
if credentials.is_empty() {
return Ok(());
}
write_git_credentials_file(credentials)?;
let home = home_dir()?;
write_gh_hosts_yml(credentials, &home)?;
write_glab_config(credentials, &home)?;
log::info!(
"Wrote {} git credential(s) to the local credential store: {}",
credentials.len(),
credential_diagnostics(credentials)
);
Ok(())
}
pub(crate) fn configure_git_credentials(credentials: &[GitCredential]) -> Result<()> {
if credentials.is_empty() {
return Ok(());
}
setup_git_config(credentials);
configure_git_identity(credentials);
write_git_credentials(credentials)
}
/// Run a git config command, logging a warning on failure rather than
/// propagating the error (git may not be installed in all sandboxes).
fn run_git_config(key: &str, value: &str) {
match BlockingCommand::new("git")
.args(["config", "--global", key, value])
.output()
{
Ok(output) if output.status.success() => {}
Ok(output) => {
log::warn!(
"git config --global {key} failed: {}",
String::from_utf8_lossy(&output.stderr)
);
}
Err(e) => {
log::warn!("Failed to run git config --global {key}: {e}");
}
}
}
/// Like [`run_git_config`] but passes `--add` so the new value is appended to
/// any existing values for `key` rather than replacing them.
fn run_git_config_add(key: &str, value: &str) {
match BlockingCommand::new("git")
.args(["config", "--global", "--add", key, value])
.output()
{
Ok(output) if output.status.success() => {}
Ok(output) => {
log::warn!(
"git config --global --add {key} failed: {}",
String::from_utf8_lossy(&output.stderr)
);
}
Err(e) => {
log::warn!("Failed to run git config --global --add {key}: {e}");
}
}
}
/// Run one-time git configuration that is set at startup and never needs to
/// be refreshed:
/// - `credential.helper store` so git reads `~/.git-credentials`
/// - SSH→HTTPS URL rewrites for each credential host, covering both the
/// scp-style (`git@{host}:`) and explicit-protocol (`ssh://git@{host}/`)
/// URL forms, so operations on either form use HTTPS credentials instead
/// of looking for an SSH key.
pub(crate) fn setup_git_config(credentials: &[GitCredential]) {
run_git_config("credential.helper", "store");
// Use --add for both forms per host so all values coexist as a
// multi-value key rather than each entry overwriting the previous one.
for cred in credentials {
let host = &cred.host;
run_git_config_add(
&format!("url.https://{host}/.insteadOf"),
&format!("ssh://git@{host}/"),
);
run_git_config_add(
&format!("url.https://{host}/.insteadOf"),
&format!("git@{host}:"),
);
}
}
/// Configure the git user identity from the server-returned credential.
///
/// Uses the first credential's `username`/`email` fields, falling back to the
/// Oz defaults when either is absent (e.g. service-account principals).
pub(crate) fn configure_git_identity(credentials: &[GitCredential]) {
let (name, email) = credentials
.first()
.map(|c| {
(
c.username.as_deref().unwrap_or(DEFAULT_GIT_NAME),
c.email.as_deref().unwrap_or(DEFAULT_GIT_EMAIL),
)
})
.unwrap_or((DEFAULT_GIT_NAME, DEFAULT_GIT_EMAIL));
run_git_config("user.name", name);
run_git_config("user.email", email);
}
/// Perform one git credentials refresh attempt.
///
/// Returns `Ok(())` on success (including when the server returns no
/// credentials). Returns `Err` when the workload-token issuance or the server
/// API call fails — these are transient failures worth retrying.
#[tracing::instrument(name = "git_credentials::try_refresh", skip_all, err, fields(
tags.cloud_agent = true,
task_id,
))]
async fn try_refresh(task_id: &str, ai_client: &Arc<dyn AIClient>) -> Result<()> {
let workload_token =
warp_isolation_platform::issue_workload_token(Some(Duration::from_secs(5 * 60)))
.await
.context("Failed to issue workload token for git credentials refresh")?
.token;
let credentials = ai_client
.get_task_git_credentials(task_id.to_string(), workload_token)
.await
.context("Failed to fetch git credentials from server")?;
if credentials.is_empty() {
log::debug!("No git credentials returned during refresh; skipping file write");
return Ok(());
}
if let Err(e) = write_git_credentials(&credentials) {
log::warn!("Failed to write refreshed git credentials: {e:#}");
} else {
log::info!("Git credentials refreshed successfully");
}
Ok(())
}
/// Infinite async loop that refreshes git credentials every
/// [`GIT_CREDENTIALS_REFRESH_INTERVAL`].
///
/// On each iteration:
/// 1. Issue a short-lived workload token.
/// 2. Call `taskGitCredentials` to get a fresh token from the server.
/// 3. Overwrite `~/.git-credentials` and refresh GitHub credentials in
/// `~/.config/gh/hosts.yml`.
///
/// On transient failure, the refresh is retried up to three times with
/// exponential backoff (1 min, 2 min, 4 min), keeping all retries within the
/// ~10-minute buffer before the one-hour token expires. If all retries fail,
/// a warning is logged and the next refresh is scheduled after the normal
/// interval.
///
/// This future never resolves — it is designed to be raced with the harness
/// execution future via `futures::select!` and dropped when the harness
/// completes.
pub(crate) async fn refresh_loop(task_id: String, ai_client: Arc<dyn AIClient>) {
loop {
warpui::r#async::Timer::after(GIT_CREDENTIALS_REFRESH_INTERVAL).await;
log::info!("Refreshing git credentials for task {task_id}");
let backoff_delays = [
Duration::from_secs(60),
Duration::from_secs(2 * 60),
Duration::from_secs(4 * 60),
];
let mut attempt = 0usize;
loop {
match try_refresh(&task_id, &ai_client).await {
Ok(()) => break,
Err(e) if attempt < backoff_delays.len() => {
let delay = backoff_delays[attempt];
log::warn!(
"Git credentials refresh failed (attempt {}): {e:#}; retrying in {}s",
attempt + 1,
delay.as_secs()
);
warpui::r#async::Timer::after(delay).await;
attempt += 1;
}
Err(e) => {
log::warn!(
"Git credentials refresh failed after {} attempts: {e:#}; \
credentials may expire before next refresh cycle",
attempt + 1
);
break;
}
}
}
}
}
#[cfg(test)]
#[path = "git_credentials_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/git_credentials_tests.rs
+204
View File
@@ -0,0 +1,204 @@
use super::*;
#[test]
fn write_gh_hosts_yml_uses_gh_cli_filename() -> Result<()> {
let temp_dir = tempfile::tempdir()?;
let gh_config_dir = temp_dir.path().join(".config").join("gh");
write_gh_hosts_yml(
&[GitCredential {
token: "token".to_string(),
username: Some("octocat".to_string()),
email: Some("octocat@example.com".to_string()),
host: "github.com".to_string(),
}],
temp_dir.path(),
)?;
let hosts_path = gh_config_dir.join(GH_HOSTS_FILENAME);
assert!(hosts_path.exists());
assert!(!gh_config_dir
.join(format!("{GH_HOSTS_FILENAME}.tmp"))
.exists());
let hosts = std::fs::read_to_string(hosts_path)?;
assert!(hosts.contains("github.com:"));
assert!(hosts.contains(" oauth_token: token"));
assert!(hosts.contains(" git_protocol: https"));
assert!(hosts.contains(" user: octocat"));
Ok(())
}
#[test]
fn write_gh_hosts_yml_excludes_gitlab_credentials() -> Result<()> {
let temp_dir = tempfile::tempdir()?;
let gh_config_dir = temp_dir.path().join(".config").join("gh");
write_gh_hosts_yml(
&[
GitCredential {
token: "github-token".to_string(),
username: Some("octocat".to_string()),
email: None,
host: "github.com".to_string(),
},
GitCredential {
token: "gitlab-token".to_string(),
username: Some("oauth2".to_string()),
email: None,
host: "gitlab.com".to_string(),
},
],
temp_dir.path(),
)?;
let hosts = std::fs::read_to_string(gh_config_dir.join(GH_HOSTS_FILENAME))?;
assert!(hosts.contains("github.com:"));
assert!(!hosts.contains("gitlab.com:"));
assert!(!hosts.contains("gitlab-token"));
Ok(())
}
#[test]
fn write_gh_hosts_yml_skips_gitlab_only_credentials() -> Result<()> {
let temp_dir = tempfile::tempdir()?;
write_gh_hosts_yml(
&[GitCredential {
token: "gitlab-token".to_string(),
username: Some("oauth2".to_string()),
email: None,
host: "gitlab.com".to_string(),
}],
temp_dir.path(),
)?;
assert!(!temp_dir.path().join(".config").join("gh").exists());
Ok(())
}
#[test]
fn git_credentials_file_content_includes_each_provider_host() {
let content = git_credentials_file_content(&[
GitCredential {
token: "github-token".to_string(),
username: None,
email: None,
host: "github.com".to_string(),
},
GitCredential {
token: "gitlab-token".to_string(),
username: Some("oauth2".to_string()),
email: None,
host: "gitlab.com".to_string(),
},
]);
assert_eq!(
content,
"https://x-access-token:github-token@github.com\n\
https://oauth2:gitlab-token@gitlab.com\n"
);
}
#[test]
fn credential_diagnostics_reports_presence_without_values() {
let diagnostics = credential_diagnostics(&[GitCredential {
token: "secret-token".to_string(),
username: Some("oauth2".to_string()),
email: Some("user@example.com".to_string()),
host: "gitlab.com".to_string(),
}]);
assert_eq!(
diagnostics,
"gitlab.com(token_present=true, username_present=true)"
);
assert!(!diagnostics.contains("secret-token"));
assert!(!diagnostics.contains("oauth2"));
assert!(!diagnostics.contains("user@example.com"));
}
#[test]
fn write_glab_config_uses_glab_cli_filename() -> Result<()> {
let temp_dir = tempfile::tempdir()?;
let glab_config_dir = temp_dir.path().join(".config").join("glab-cli");
write_glab_config(
&[GitCredential {
token: "gitlab-token".to_string(),
username: Some("oauth2".to_string()),
email: Some("user@example.com".to_string()),
host: "gitlab.com".to_string(),
}],
temp_dir.path(),
)?;
let config_path = glab_config_dir.join(GLAB_CONFIG_FILENAME);
assert!(config_path.exists());
assert!(!glab_config_dir
.join(format!("{GLAB_CONFIG_FILENAME}.tmp"))
.exists());
let config = std::fs::read_to_string(config_path)?;
assert!(config.contains("hosts:"));
assert!(config.contains(" gitlab.com:"));
assert!(config.contains(" token: gitlab-token"));
assert!(config.contains(" git_protocol: https"));
assert!(config.contains(" api_protocol: https"));
Ok(())
}
#[test]
fn write_glab_config_excludes_github_credentials() -> Result<()> {
let temp_dir = tempfile::tempdir()?;
let glab_config_dir = temp_dir.path().join(".config").join("glab-cli");
write_glab_config(
&[
GitCredential {
token: "github-token".to_string(),
username: Some("octocat".to_string()),
email: None,
host: "github.com".to_string(),
},
GitCredential {
token: "gitlab-token".to_string(),
username: Some("oauth2".to_string()),
email: None,
host: "gitlab.com".to_string(),
},
],
temp_dir.path(),
)?;
let config = std::fs::read_to_string(glab_config_dir.join(GLAB_CONFIG_FILENAME))?;
assert!(config.contains("gitlab.com:"));
assert!(!config.contains("github.com:"));
assert!(!config.contains("github-token"));
Ok(())
}
#[test]
fn write_glab_config_skips_github_only_credentials() -> Result<()> {
let temp_dir = tempfile::tempdir()?;
write_glab_config(
&[GitCredential {
token: "github-token".to_string(),
username: Some("octocat".to_string()),
email: None,
host: "github.com".to_string(),
}],
temp_dir.path(),
)?;
assert!(!temp_dir.path().join(".config").join("glab-cli").exists());
Ok(())
}
app/src/ai/agent_sdk/driver/harness/claude_code.rs
+281 -127
View File
@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::ffi::{OsStr, OsString};
use std::fmt::Write as _;
use std::path::{Path, PathBuf};
use std::sync::Arc;
@@ -12,42 +13,55 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use tempfile::NamedTempFile;
use uuid::Uuid;
use galaxy_managed_secrets::ManagedSecretValue;
use super::super::terminal::{CommandHandle, TerminalDriver};
use super::super::{AgentDriver, AgentDriverError};
use super::claude_transcript::{
claude_config_dir, home_dir_for_claude_config, read_envelope, rehydrate_claude_transcript,
ClaudeResumeInfo, ClaudeTranscriptEnvelope,
};
use super::json_utils::{read_json_file_or_default, write_json_file};
use super::{
cli_agent_session_status, write_temp_file, HarnessCleanupDisposition, HarnessRunner,
JSONMCPServer, ResumePayload, SavePoint, ThirdPartyHarness,
};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_sdk::setup_observability::{
OzRunTimelineEvent, SetupClientEventReporter, SetupStep,
};
use crate::ai::ambient_agents::task::HarnessModelConfig;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::mcp::JSONTransportType;
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;
mod wake_driver;
#[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,
parent_bridge_char_count, parent_bridge_event_cursor_file, 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,
prime_parent_bridge_staged_for_self_managed_wake, read_parent_bridge_event_cursor,
render_parent_bridge_message_block, stage_parent_bridge_message,
write_parent_bridge_event_cursor, MessageBridgeHookOutput, MessageBridgeMessageRecord,
MESSAGE_BRIDGE_CONTEXT_PREAMBLE,
};
use parent_bridge::{MessageBridge, MessageBridgeCleanupDisposition};
#[cfg(test)]
use shell_words::quote as shell_quote;
#[cfg(test)]
use wake_driver::{ClaudeWakeRemoteContext, CLAUDE_WAKE_PROMPT_FILE_NAME};
#[cfg(test)]
use super::super::OZ_MESSAGE_LISTENER_STATE_ROOT_ENV;
pub(crate) struct ClaudeHarness;
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl ThirdPartyHarness for ClaudeHarness {
@@ -63,18 +77,33 @@ impl ThirdPartyHarness for ClaudeHarness {
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,
}
})
fn auth_check_command(&self) -> Option<String> {
let cli = self.cli_agent().command_prefix();
Some(format!("{cli} auth status --json"))
}
fn runtime_error_patterns(&self) -> &'static [&'static str] {
&[
// Out-of-credits / billing.
"Credit balance too low",
// Plan/usage limits emitted as `You've hit your <kind> limit`.
// We match on the common prefix so the variants (session,
// weekly, Opus, etc.) all hit.
"You've hit your",
// Invalid or malformed API key.
"Invalid API key",
"This organization has been disabled",
"belongs to a disabled organization",
// OAuth / login state.
"Not logged in",
"OAuth token revoked",
"OAuth token has expired",
// Routines disabled by org policy.
"Routines are disabled by your organization's policy",
// Generic upstream API failures Claude Code surfaces verbatim.
"API Error: Request rejected (429)",
"authentication_error",
]
}
/// Fetch the Claude Code transcript for the current task's conversation and wrap it
@@ -86,28 +115,9 @@ impl ThirdPartyHarness for ClaudeHarness {
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 envelope: ClaudeTranscriptEnvelope =
super::fetch_transcript_envelope("claude", conversation_id, harness_support_client)
.await?;
let session_id = envelope.uuid;
Ok(Some(ResumePayload::Claude(ClaudeResumeInfo {
conversation_id: *conversation_id,
@@ -121,24 +131,43 @@ impl ThirdPartyHarness for ClaudeHarness {
prompt: &str,
system_prompt: Option<&str>,
resumption_prompt: Option<&str>,
context: Option<&str>,
working_dir: &Path,
task_id: Option<AmbientAgentTaskId>,
server_api: Arc<ServerApi>,
terminal_driver: ModelHandle<TerminalDriver>,
resume: Option<ResumePayload>,
resolved_env_vars: &HashMap<OsString, OsString>,
_resolved_secrets: &HashMap<String, ManagedSecretValue>,
resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
_third_party_harness_model_config: Option<&HarnessModelConfig>,
) -> 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,
});
// Prepare the environment config files.
prepare_claude_environment_config(working_dir, resolved_env_vars).map_err(|error| {
AgentDriverError::HarnessConfigSetupFailed {
harness: self.cli_agent().command_prefix().to_owned(),
error,
}
})?;
// The ResumePayload shouldn't contain non-Claude information, error if it does.
let claude_resume = resume.map(ClaudeResumeInfo::try_from).transpose()?;
// 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(),
};
// and server context are most reliable when prepended directly to the prompt that gets
// piped into the CLI. Order: resumption_prompt → context → prompt
let mut parts: Vec<&str> = Vec::new();
if let Some(preamble) = resumption_prompt {
if !preamble.is_empty() {
parts.push(preamble);
}
}
if let Some(ctx) = context {
if !ctx.is_empty() {
parts.push(ctx);
}
}
parts.push(prompt);
let owned_prompt = parts.join("\n\n");
Ok(Box::new(ClaudeHarnessRunner::new(
self.cli_agent().command_prefix(),
&owned_prompt,
@@ -148,8 +177,13 @@ impl ThirdPartyHarness for ClaudeHarness {
server_api,
terminal_driver,
claude_resume,
resolved_mcp_servers,
)?))
}
fn requires_verified_platform_plugin(&self) -> bool {
true
}
}
/// Format slug sent to the server when creating a Claude Code conversation.
@@ -169,6 +203,7 @@ fn claude_command(
session_id: &Uuid,
prompt_path: &str,
system_prompt_path: Option<&str>,
mcp_config_path: Option<&str>,
resuming: bool,
) -> String {
let flag = if resuming { "--resume" } else { "--session-id" };
@@ -176,6 +211,9 @@ fn claude_command(
if let Some(sp_path) = system_prompt_path {
let _ = write!(cmd, " --append-system-prompt-file '{sp_path}'");
}
if let Some(mcp_path) = mcp_config_path {
let _ = write!(cmd, " --mcp-config '{mcp_path}'");
}
format!("{cmd} < '{prompt_path}'")
}
@@ -198,6 +236,8 @@ struct ClaudeHarnessRunner {
_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>,
/// Held so the MCP config temp file lives until the CLI exits.
_temp_mcp_config_file: Option<NamedTempFile>,
client: Arc<dyn HarnessSupportClient>,
server_api: Arc<ServerApi>,
terminal_driver: ModelHandle<TerminalDriver>,
@@ -224,49 +264,44 @@ impl ClaudeHarnessRunner {
server_api: Arc<ServerApi>,
terminal_driver: ModelHandle<TerminalDriver>,
resume: Option<ClaudeResumeInfo>,
resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
) -> 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 temp_file = write_temp_file("oz_prompt_", prompt, ".txt")?;
let prompt_path = temp_file.path().display().to_string();
let (session_id, preexisting_conversation_id, resuming) = match resume {
let (session_id, preexisting_conversation_id) = 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)
rehydrate_claude_transcript(&mut envelope, working_dir)
.map_err(AgentDriverError::ConfigBuildFailed)?;
(session_id, Some(conversation_id))
}
None => (Uuid::new_v4(), None, false),
None => (Uuid::new_v4(), None),
};
let temp_system_prompt_file = system_prompt
.map(|sp| write_temp_file("oz_system_prompt_", sp))
.map(|sp| write_temp_file("oz_system_prompt_", sp, ".txt"))
.transpose()?;
let system_prompt_path = temp_system_prompt_file
.as_ref()
.map(|f| f.path().display().to_string());
let temp_mcp_config_file = (!resolved_mcp_servers.is_empty())
.then(|| {
let mcp_json = serialize_claude_mcp_config(resolved_mcp_servers)
.map_err(AgentDriverError::ConfigBuildFailed)?;
write_temp_file("oz_mcp_config_", &mcp_json, ".json")
})
.transpose()?;
let mcp_config_path = temp_mcp_config_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()
@@ -279,11 +314,13 @@ impl ClaudeHarnessRunner {
&session_id,
&prompt_path,
system_prompt_path.as_deref(),
resuming,
mcp_config_path.as_deref(),
preexisting_conversation_id.is_some(),
),
cli_name: cli_command.to_string(),
_temp_prompt_file: temp_file,
_temp_system_prompt_file: temp_system_prompt_file,
_temp_mcp_config_file: temp_mcp_config_file,
client,
server_api,
terminal_driver,
@@ -362,9 +399,33 @@ impl ClaudeHarnessRunner {
.await
}
fn cleanup_parent_bridge(&self) -> Result<()> {
async fn should_preserve_parent_bridge(
&self,
cleanup_disposition: HarnessCleanupDisposition,
foreground: &ModelSpawner<AgentDriver>,
) -> bool {
if !matches!(
cleanup_disposition,
HarnessCleanupDisposition::PreserveResumptionStateIfSupported
) {
return false;
}
!matches!(
cli_agent_session_status(&self.terminal_driver, foreground).await,
Some(crate::terminal::cli_agent_sessions::CLIAgentSessionStatus::Blocked { .. })
| Some(crate::terminal::cli_agent_sessions::CLIAgentSessionStatus::InProgress)
)
}
fn cleanup_parent_bridge(&self, preserve_state: bool) -> Result<()> {
if let Some(parent_bridge) = self.parent_bridge.as_ref() {
parent_bridge.cleanup()?;
let cleanup_disposition = if preserve_state {
MessageBridgeCleanupDisposition::PreserveState
} else {
MessageBridgeCleanupDisposition::RemoveState
};
parent_bridge.cleanup(cleanup_disposition)?;
}
Ok(())
}
@@ -373,9 +434,14 @@ impl ClaudeHarnessRunner {
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl HarnessRunner for ClaudeHarnessRunner {
fn harness_name(&self) -> &str {
&self.cli_name
}
async fn start(
&self,
foreground: &ModelSpawner<AgentDriver>,
setup_events: &SetupClientEventReporter,
) -> 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.
@@ -387,14 +453,17 @@ impl HarnessRunner for ClaudeHarnessRunner {
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)
})?;
let id = setup_events
.record_result(SetupStep::ThirdPartyHarnessExternalConversation, async {
self.client
.create_external_conversation(CLAUDE_CODE_FORMAT)
.await
.map_err(|e| {
log::error!("Failed to create external conversation: {e}");
AgentDriverError::ConfigBuildFailed(e)
})
})
.await?;
log::info!("Created external conversation {id}");
id
}
@@ -414,7 +483,7 @@ impl HarnessRunner for ClaudeHarnessRunner {
{
Ok(command_handle) => command_handle,
Err(err) => {
self.cleanup_parent_bridge()
self.cleanup_parent_bridge(false)
.map_err(AgentDriverError::ConfigBuildFailed)?;
return Err(err);
}
@@ -426,6 +495,10 @@ impl HarnessRunner for ClaudeHarnessRunner {
block_id: command_handle.block_id().clone(),
};
setup_events
.post_timeline_event(OzRunTimelineEvent::AgentStarted)
.await;
Ok(command_handle)
}
@@ -494,9 +567,16 @@ impl HarnessRunner for ClaudeHarnessRunner {
Ok(())
}
async fn cleanup(&self, _foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
async fn cleanup(
&self,
cleanup_disposition: HarnessCleanupDisposition,
foreground: &ModelSpawner<AgentDriver>,
) -> Result<()> {
self.flush_parent_bridge_acks().await?;
self.cleanup_parent_bridge()
let preserve_state = self
.should_preserve_parent_bridge(cleanup_disposition, foreground)
.await;
self.cleanup_parent_bridge(preserve_state)
}
}
@@ -526,21 +606,31 @@ async fn upload_transcript(
.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(
pub(crate) fn prepare_claude_environment_config(
working_dir: &Path,
secrets: &HashMap<String, ManagedSecretValue>,
resolved_env_vars: &HashMap<OsString, OsString>,
) -> 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_json_path = claude_global_config_path()?;
let claude_settings_path = claude_config_dir()?.join(CLAUDE_SETTINGS_FILE_NAME);
let api_key_suffix = resolve_anthropic_api_key_suffix(secrets);
let api_key_suffix = resolve_anthropic_api_key_suffix(resolved_env_vars);
prepare_claude_config(&claude_json_path, working_dir, api_key_suffix.as_deref())?;
prepare_claude_settings(&claude_settings_path)?;
Ok(())
}
// This function is used specifically for determining where to land `.claude.json`.
fn claude_global_config_path() -> Result<PathBuf> {
if let Ok(dir) = std::env::var("CLAUDE_CONFIG_DIR") {
if !dir.is_empty() {
return Ok(PathBuf::from(dir).join(CLAUDE_JSON_FILE_NAME));
}
}
home_dir_for_claude_config()
.map(|home| home.join(CLAUDE_JSON_FILE_NAME))
.ok_or_else(|| anyhow::anyhow!("could not determine home directory"))
}
fn prepare_claude_config(
claude_json_path: &Path,
working_dir: &Path,
@@ -628,27 +718,23 @@ struct ClaudeSettings {
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.
/// Try to get the last 20 chars of the ANTHROPIC_API_KEY, where 20 chars is the
/// suffix length that Claude Code truncates keys to.
fn resolve_anthropic_api_key_suffix(
secrets: &HashMap<String, ManagedSecretValue>,
resolved_env_vars: &HashMap<OsString, OsString>,
) -> 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);
// Worker-injected process env wins.
if let Ok(key) = std::env::var(ANTHROPIC_API_KEY_ENV) {
if !key.is_empty() {
return suffix_of(&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))
// Otherwise use the resolved value from the secrets map.
resolved_env_vars
.get(OsStr::new(ANTHROPIC_API_KEY_ENV))
.and_then(|v| v.to_str())
.and_then(suffix_of)
.map(str::to_owned)
}
fn suffix_of(key: &str) -> Option<&str> {
@@ -659,6 +745,74 @@ fn suffix_of(key: &str) -> Option<&str> {
}
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ClaudeMcpConfig {
mcp_servers: HashMap<String, ClaudeMcpServerEntry>,
}
#[derive(Serialize)]
#[serde(tag = "type")]
enum ClaudeMcpServerEntry {
#[serde(rename = "stdio")]
Stdio {
command: String,
args: Vec<String>,
#[serde(skip_serializing_if = "HashMap::is_empty")]
env: HashMap<String, String>,
#[serde(skip_serializing_if = "Option::is_none")]
cwd: Option<String>,
},
#[serde(rename = "http")]
Http {
url: String,
#[serde(skip_serializing_if = "HashMap::is_empty")]
headers: HashMap<String, String>,
},
}
impl ClaudeMcpServerEntry {
fn from_json_mcp_server(server: &JSONMCPServer) -> Self {
match &server.transport_type {
JSONTransportType::CLIServer {
command,
args,
env,
working_directory,
} => Self::Stdio {
command: command.clone(),
args: args.clone(),
env: env.clone(),
cwd: working_directory.clone(),
},
JSONTransportType::SSEServer { url, headers } => Self::Http {
url: url.clone(),
headers: headers.clone(),
},
}
}
}
/// Serialize resolved MCP servers into Claude Code's `--mcp-config` JSON format.
///
/// Produces `{ "mcpServers": { "name": { "type": "stdio"|"http", ... }, ... } }`.
pub(crate) fn serialize_claude_mcp_config(
servers: &HashMap<String, JSONMCPServer>,
) -> Result<String> {
let config = ClaudeMcpConfig {
mcp_servers: servers
.iter()
.map(|(name, server)| {
(
name.clone(),
ClaudeMcpServerEntry::from_json_mcp_server(server),
)
})
.collect(),
};
serde_json::to_string_pretty(&config).context("Failed to serialize Claude MCP config")
}
#[cfg(test)]
#[path = "claude_code_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs
+155 -16
View File
@@ -25,22 +25,25 @@ use uuid::Uuid;
use crate::ai::agent_events::{
run_agent_event_driver, AgentEventConsumer, AgentEventConsumerControlFlow,
AgentEventDriverConfig, MessageHydrator, ServerApiAgentEventSource,
AgentEventDriverConfig, AgentMessageEventMetadata, MessageHydrator, ServerApiAgentEventSource,
};
use crate::ai::agent_sdk::driver::{AgentDriver, OZ_MESSAGE_LISTENER_STATE_ROOT_ENV};
use crate::server::server_api::ai::AgentRunEvent;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::server::server_api::ai::{AIClient, 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_EVENT_CURSOR_FILE_NAME: &str = "event-cursor.json";
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";
pub(super) const MESSAGE_BRIDGE_CONTEXT_PREAMBLE: &str =
"Oz mailbox update for this child run.\nSource: lead agent\nContext type: user-level coordination messages\n";
const PARENT_BRIDGE_REMAINING_MESSAGES_NOTE: &str =
"\n\nMore lead-agent messages are still staged and will be surfaced on a later turn.";
"\n\nAdditional lead agent mailbox messages remain queued and may be surfaced later.";
pub(super) struct MessageBridge {
run_id: String,
@@ -48,6 +51,16 @@ pub(super) struct MessageBridge {
runtime: Mutex<Option<MessageBridgeRuntime>>,
state_lock: AsyncMutex<()>,
}
pub(super) enum MessageBridgeCleanupDisposition {
RemoveState,
PreserveState,
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct MessageBridgeEventCursor {
since_sequence: i64,
}
struct MessageBridgeRuntime {
task: SpawnedFutureHandle,
}
@@ -55,6 +68,7 @@ struct MessageBridgeRuntime {
struct MessageBridgeEventConsumer {
run_id: String,
state_dir: PathBuf,
server_api: Arc<ServerApi>,
}
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
@@ -91,6 +105,13 @@ impl AgentEventConsumer for MessageBridgeEventConsumer {
Ok(AgentEventConsumerControlFlow::Continue)
}
async fn persist_cursor(&mut self, sequence: i64) -> anyhow::Result<()> {
write_parent_bridge_event_cursor(&self.state_dir, sequence)?;
self.server_api
.update_event_sequence_on_server(&self.run_id, sequence)
.await
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
@@ -130,6 +151,12 @@ struct SelectedMessageBridgeMessages {
}
impl MessageBridge {
fn hydrator(&self, server_api: Arc<ServerApi>) -> MessageHydrator {
match self.run_id.parse::<AmbientAgentTaskId>() {
Ok(task_id) => MessageHydrator::for_task(server_api, task_id),
Err(_) => MessageHydrator::new(server_api),
}
}
pub(super) fn new(run_id: String, session_id: Uuid) -> Result<Self> {
Ok(Self {
run_id,
@@ -177,8 +204,7 @@ impl MessageBridge {
if !self.state_dir.exists() {
return Ok(());
}
let hydrator = MessageHydrator::new(server_api);
let hydrator = self.hydrator(server_api);
let _guard = self.state_lock.lock().await;
acknowledge_parent_bridge_hook_output(&hydrator, &self.state_dir).await?;
prepare_parent_bridge_hook_output(
@@ -193,16 +219,18 @@ impl MessageBridge {
if !self.state_dir.exists() {
return Ok(());
}
let hydrator = MessageHydrator::new(server_api);
let hydrator = self.hydrator(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<()> {
pub(super) fn cleanup(&self, disposition: MessageBridgeCleanupDisposition) -> Result<()> {
if let Some(runtime) = self.runtime.lock().take() {
runtime.task.abort();
}
if matches!(disposition, MessageBridgeCleanupDisposition::PreserveState) {
return Ok(());
}
match fs::remove_dir_all(&self.state_dir) {
Ok(()) => {}
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
@@ -249,6 +277,10 @@ pub(super) fn parent_bridge_hook_output_ack_file(state_dir: &Path) -> PathBuf {
state_dir.join(PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME)
}
pub(super) fn parent_bridge_event_cursor_file(state_dir: &Path) -> PathBuf {
state_dir.join(PARENT_BRIDGE_EVENT_CURSOR_FILE_NAME)
}
fn parent_bridge_message_path(dir: &Path, sequence: i64, message_id: &str) -> PathBuf {
dir.join(format!("{sequence:020}-{message_id}.json"))
}
@@ -277,6 +309,28 @@ pub(super) fn ensure_parent_bridge_state_dir(state_dir: &Path) -> Result<()> {
Ok(())
}
pub(super) fn read_parent_bridge_event_cursor(state_dir: &Path) -> Result<i64> {
let path = parent_bridge_event_cursor_file(state_dir);
if !path.exists() {
return Ok(0);
}
let cursor = serde_json::from_slice::<MessageBridgeEventCursor>(
&fs::read(&path).with_context(|| format!("Failed to read {}", path.display()))?,
)
.with_context(|| format!("Failed to parse {}", path.display()))?;
Ok(cursor.since_sequence)
}
pub(super) fn write_parent_bridge_event_cursor(state_dir: &Path, sequence: i64) -> Result<()> {
write_parent_bridge_json_atomically(
&parent_bridge_event_cursor_file(state_dir),
&MessageBridgeEventCursor {
since_sequence: sequence,
},
)
}
pub(super) fn stage_parent_bridge_message(
state_dir: &Path,
record: &MessageBridgeMessageRecord,
@@ -288,7 +342,57 @@ pub(super) fn stage_parent_bridge_message(
Ok(())
}
fn parent_bridge_max_context_chars() -> usize {
pub(super) async fn prime_parent_bridge_staged_for_self_managed_wake(
hydrator: &MessageHydrator,
state_dir: &Path,
wake_message: Option<&AgentMessageEventMetadata>,
) -> Result<()> {
remove_file_if_exists(&parent_bridge_hook_output_file(state_dir))?;
remove_file_if_exists(&parent_bridge_hook_output_ack_file(state_dir))?;
move_parent_bridge_surfaced_messages_to_staged(state_dir)?;
let Some(wake_message) = wake_message else {
return Ok(());
};
let record = hydrate_parent_bridge_message_record(
hydrator,
&MessageBridgeMessageRecord {
sequence: wake_message.sequence,
message_id: wake_message.message_id.clone(),
sender_run_id: String::new(),
subject: String::new(),
body: String::new(),
occurred_at: wake_message.occurred_at.clone(),
},
)
.await?;
stage_parent_bridge_message(state_dir, &record)?;
write_parent_bridge_event_cursor(state_dir, wake_message.sequence)
}
fn move_parent_bridge_surfaced_messages_to_staged(state_dir: &Path) -> Result<()> {
let surfaced_records = parent_bridge_message_records(&parent_bridge_surfaced_dir(state_dir))?;
for (path, record) in surfaced_records {
let target =
parent_bridge_staged_message_path(state_dir, record.sequence, &record.message_id);
if target.exists() {
write_parent_bridge_json_atomically(&target, &record)?;
remove_file_if_exists(&path)?;
} else {
fs::rename(&path, &target).with_context(|| {
format!(
"Failed to move message bridge record {} back to {}",
path.display(),
target.display()
)
})?;
}
}
Ok(())
}
pub(super) 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())
@@ -338,7 +442,7 @@ pub(super) fn render_parent_bridge_message_block(record: &MessageBridgeMessageRe
record.subject.as_str()
};
let mut block = String::from("---\nLead-agent message");
let mut block = String::from("---\nLead agent mailbox message");
if record.sequence != 0 {
let _ = write!(block, " #{}", record.sequence);
}
@@ -579,15 +683,50 @@ async fn run_parent_bridge_forever(
state_dir: PathBuf,
) -> Result<()> {
ensure_parent_bridge_state_dir(&state_dir)?;
let since_sequence =
read_parent_bridge_resume_cursor(server_api.as_ref(), &run_id, &state_dir).await?;
// 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 };
// loop and we also persist it inside the session state dir so dormant runs
// can resume without replaying already handled events.
let config =
AgentEventDriverConfig::retry_forever_run_ids(vec![run_id.clone()], since_sequence);
let source = ServerApiAgentEventSource::new(server_api.clone());
let mut consumer = MessageBridgeEventConsumer {
run_id,
state_dir,
server_api,
};
run_agent_event_driver(source, config, &mut consumer).await
}
async fn read_parent_bridge_resume_cursor(
server_api: &ServerApi,
run_id: &str,
state_dir: &Path,
) -> Result<i64> {
// The server cursor is the durable cross-client source of truth, but the
// bridge also keeps a local cursor for same-machine recovery. If Warp or
// Claude restarts after the bridge has staged events locally but before the
// server cursor update is visible, the local cursor prevents replaying
// messages already handed to this Claude session.
let local_sequence = read_parent_bridge_event_cursor(state_dir)?;
let Ok(task_id) = run_id.parse() else {
return Ok(local_sequence);
};
let server_sequence = match server_api.get_ambient_agent_task(&task_id).await {
Ok(task) => task.last_event_sequence.unwrap_or(0),
Err(err) => {
log::warn!(
"Failed to read server-backed event cursor for Claude message bridge run {run_id}: {err:#}"
);
0
}
};
Ok(local_sequence.max(server_sequence))
}
fn write_parent_bridge_json_atomically<T: Serialize>(path: &Path, value: &T) -> Result<()> {
write_parent_bridge_bytes_atomically(path, &serde_json::to_vec(value)?)
}
app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver.rs
+296
View File
@@ -0,0 +1,296 @@
use std::collections::HashMap;
use std::ffi::OsString;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Context, Result};
use shell_words::quote as shell_quote;
use uuid::Uuid;
use warp_cli::agent::Harness;
use warp_graphql::ai::AgentTaskState;
use super::super::claude_transcript::{
claude_config_dir, write_envelope, write_session_index_entry, ClaudeTranscriptEnvelope,
};
use super::super::{remove_claude_externally_managed_listener_env_vars, task_env_vars};
use super::parent_bridge::{
ensure_parent_bridge_state_dir, parent_bridge_root,
prime_parent_bridge_staged_for_self_managed_wake,
};
use super::{claude_command, prepare_claude_environment_config, ClaudeHarness};
use crate::ai::agent::conversation::{AIConversation, ConversationStatus};
use crate::ai::agent_events::{AgentMessageEventMetadata, MessageHydrator};
use crate::ai::ambient_agents::{AmbientAgentTaskId, AmbientAgentTaskState};
use crate::server::server_api::ai::AIClient;
use crate::server::server_api::harness_support::ResolvePromptRequest;
use crate::server::server_api::ServerApi;
use crate::terminal::CLIAgent;
const CLAUDE_WAKE_PROMPT: &str =
"A lead agent mailbox message is available for this child run. Review the mailbox context and continue the task.";
pub(super) const CLAUDE_WAKE_PROMPT_FILE_NAME: &str = "wake-turn-prompt.txt";
#[derive(Debug)]
pub(super) struct ClaudeWakeRemoteContext {
pub(super) session_id: Uuid,
pub(super) envelope: ClaudeTranscriptEnvelope,
pub(super) wake_prompt: String,
}
struct ClaudeWakeCandidate {
task_id: AmbientAgentTaskId,
parent_run_id: Option<String>,
working_dir: Option<PathBuf>,
}
impl ClaudeHarness {
pub(crate) async fn wake_dormant_session(
server_api: Arc<ServerApi>,
conversation: AIConversation,
parent_conversation: Option<AIConversation>,
working_dir: Option<PathBuf>,
wake_message: Option<AgentMessageEventMetadata>,
) -> Result<Option<String>> {
let Some(candidate) =
Self::local_wake_candidate(&conversation, parent_conversation.as_ref(), working_dir)
else {
return Ok(None);
};
let ClaudeWakeCandidate {
task_id,
parent_run_id,
working_dir,
} = candidate;
let task = server_api.get_ambient_agent_task(&task_id).await?;
let harness = task
.agent_config_snapshot
.as_ref()
.and_then(|snapshot| snapshot.harness.as_ref())
.map(|config| config.harness_type);
log::info!(
"Evaluating dormant Claude wake: task_id={task_id} server_task_state={:?} harness={harness:?}",
task.state
);
if !is_local_wake_task_state_ready(task.state.clone()) || harness != Some(Harness::Claude) {
log::info!(
"Skipping dormant Claude wake: task_id={task_id} server_task_state={:?} harness={harness:?}",
task.state
);
return Ok(None);
}
let remote = Self::fetch_local_wake_remote_context(task_id, server_api.clone()).await?;
let command = Self::prepare_local_wake_command(
server_api.clone(),
task_id,
parent_run_id,
working_dir,
remote,
wake_message,
)
.await?;
log::info!("Reopening dormant Claude task before wake command: task_id={task_id}");
server_api
.update_agent_task(task_id, Some(AgentTaskState::InProgress), None, None, None)
.await
.map_err(|err| {
anyhow::anyhow!(
"Failed to reopen dormant Claude task {task_id} before wake: {err:#}"
)
})?;
log::info!("Reopened dormant Claude task before wake command: task_id={task_id}");
Ok(Some(command))
}
fn local_wake_candidate(
conversation: &AIConversation,
parent_conversation: Option<&AIConversation>,
working_dir: Option<PathBuf>,
) -> Option<ClaudeWakeCandidate> {
let conversation_id = conversation.id();
if !matches!(conversation.status(), ConversationStatus::Success) {
log::info!(
"Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_success status={:?}",
conversation.status()
);
return None;
}
if !conversation.is_child_agent_conversation() || conversation.is_remote_child() {
log::info!(
"Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=not_local_child is_child_agent_conversation={} is_remote_child={}",
conversation.is_child_agent_conversation(),
conversation.is_remote_child()
);
return None;
}
let Some(task_id) = conversation.task_id() else {
log::info!(
"Skipping dormant Claude wake candidate: conversation_id={conversation_id:?} reason=missing_task_id"
);
return None;
};
let parent_run_id = conversation
.parent_agent_id()
.map(str::to_owned)
.or_else(|| parent_conversation.and_then(AIConversation::run_id));
Some(ClaudeWakeCandidate {
task_id,
parent_run_id,
working_dir,
})
}
async fn fetch_local_wake_remote_context(
task_id: AmbientAgentTaskId,
server_api: Arc<ServerApi>,
) -> Result<ClaudeWakeRemoteContext> {
let resolved = server_api
.resolve_prompt_for_task(
&task_id,
ResolvePromptRequest {
skill: None,
attachments_dir: None,
},
)
.await
.with_context(|| format!("Failed to resolve Claude wake prompt for task {task_id}"))?;
let bytes = server_api
.fetch_transcript_for_task(&task_id)
.await
.with_context(|| format!("Failed to fetch Claude transcript for task {task_id}"))?;
let envelope: ClaudeTranscriptEnvelope =
serde_json::from_slice(&bytes).with_context(|| {
format!("Failed to deserialize Claude transcript for wake task {task_id}")
})?;
let wake_prompt = match resolved.resumption_prompt {
Some(resumption_prompt) if !resumption_prompt.is_empty() => {
format!(
"{resumption_prompt}
{CLAUDE_WAKE_PROMPT}"
)
}
_ => CLAUDE_WAKE_PROMPT.to_string(),
};
Ok(ClaudeWakeRemoteContext {
session_id: envelope.uuid,
envelope,
wake_prompt,
})
}
pub(super) async fn prepare_local_wake_command(
server_api: Arc<ServerApi>,
task_id: AmbientAgentTaskId,
parent_run_id: Option<String>,
working_dir: Option<PathBuf>,
mut remote: ClaudeWakeRemoteContext,
wake_message: Option<AgentMessageEventMetadata>,
) -> Result<String> {
let working_dir = working_dir.unwrap_or_else(|| remote.envelope.cwd.clone());
prepare_claude_environment_config(&working_dir, &HashMap::new())
.context("Failed to prepare Claude environment for wake")?;
remote.envelope.cwd = working_dir.clone();
let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?;
write_envelope(&remote.envelope, &config_root)
.context("Failed to rehydrate Claude transcript for wake")?;
if let Err(error) = write_session_index_entry(remote.session_id, &working_dir, &config_root)
{
log::warn!("Failed to update Claude sessions-index.json for wake: {error:#}");
}
let state_dir = parent_bridge_root()?.join(remote.session_id.to_string());
ensure_parent_bridge_state_dir(&state_dir)?;
let hydrator = MessageHydrator::for_task(server_api, task_id);
prime_parent_bridge_staged_for_self_managed_wake(
&hydrator,
&state_dir,
wake_message.as_ref(),
)
.await?;
let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME);
std::fs::write(&prompt_path, remote.wake_prompt.as_bytes())
.with_context(|| format!("Failed to write {}", prompt_path.display()))?;
let command = claude_command(
CLIAgent::Claude.command_prefix(),
&remote.session_id,
&prompt_path.display().to_string(),
None,
None,
true,
);
let env_vars = local_wake_task_env_vars(Some(&task_id), parent_run_id.as_deref());
Ok(prefix_command_with_env_vars(command, env_vars))
}
}
fn local_wake_task_env_vars(
task_id: Option<&AmbientAgentTaskId>,
parent_run_id: Option<&str>,
) -> HashMap<OsString, OsString> {
let mut env_vars = task_env_vars(task_id, parent_run_id, Harness::Claude);
// The local wake command is executed directly in the existing child
// terminal, not through `AgentDriver::run_harness`, so Warp does not start
// `MessageBridge` for this resumed Claude process. Leave the listener in
// the Claude plugin's self-managed mode; otherwise the hook waits for
// state files that no managed bridge is producing and the wake message is
// never surfaced to Claude.
remove_claude_externally_managed_listener_env_vars(&mut env_vars);
env_vars
}
fn is_local_wake_task_state_ready(state: AmbientAgentTaskState) -> bool {
match state {
AmbientAgentTaskState::Succeeded => true,
// The local conversation status is already gated on `Success` before
// this function is called. The server task update is fire-and-forget,
// so it can still report `InProgress` for a short window after the
// local Claude process has actually stopped. Treat that stale server
// state as wakeable for local children.
AmbientAgentTaskState::InProgress => true,
AmbientAgentTaskState::Queued
| AmbientAgentTaskState::Pending
| AmbientAgentTaskState::Claimed
| AmbientAgentTaskState::Failed
| AmbientAgentTaskState::Error
| AmbientAgentTaskState::Blocked
| AmbientAgentTaskState::Cancelled
| AmbientAgentTaskState::Unknown => false,
}
}
fn prefix_command_with_env_vars(command: String, env_vars: HashMap<OsString, OsString>) -> String {
if env_vars.is_empty() {
return command;
}
let mut env_pairs = env_vars
.into_iter()
.map(|(key, value)| {
(
key.to_string_lossy().into_owned(),
value.to_string_lossy().into_owned(),
)
})
.collect::<Vec<_>>();
env_pairs.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
let assignments = env_pairs
.into_iter()
.map(|(key, value)| format!("{key}={}", shell_quote(&value)))
.collect::<Vec<_>>()
.join(" ");
format!("env {assignments} {command}")
}
#[cfg(test)]
#[path = "wake_driver_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/harness/claude_code/wake_driver_tests.rs
+24
View File
@@ -0,0 +1,24 @@
use super::*;
#[test]
fn local_wake_task_state_ready_allows_success_and_stale_in_progress() {
assert!(is_local_wake_task_state_ready(
AmbientAgentTaskState::Succeeded
));
assert!(is_local_wake_task_state_ready(
AmbientAgentTaskState::InProgress
));
for state in [
AmbientAgentTaskState::Queued,
AmbientAgentTaskState::Pending,
AmbientAgentTaskState::Claimed,
AmbientAgentTaskState::Failed,
AmbientAgentTaskState::Error,
AmbientAgentTaskState::Blocked,
AmbientAgentTaskState::Cancelled,
AmbientAgentTaskState::Unknown,
] {
assert!(!is_local_wake_task_state_ready(state));
}
}
app/src/ai/agent_sdk/driver/harness/claude_code_tests.rs
+441 -65
View File
@@ -1,15 +1,21 @@
use mockall::predicate::eq;
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::sync::Arc;
use mockall::predicate::eq;
use tempfile::TempDir;
use uuid::Uuid;
use warp_cli::{OZ_HARNESS_ENV, OZ_PARENT_RUN_ID_ENV, OZ_RUN_ID_ENV};
use super::*;
use crate::ai::agent_events::MessageHydrator;
use crate::server::server_api::ai::{MockAIClient, ReadAgentMessageResponse};
use crate::ai::agent_events::{AgentMessageEventMetadata, MessageHydrator};
use crate::ai::agent_sdk::driver::harness::claude_transcript::{
encode_cwd, write_session_index_entry,
};
use crate::ai::agent_sdk::driver::OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV;
use crate::server::server_api::ai::{AIClient, MockAIClient, ReadAgentMessageResponse};
use crate::server::server_api::ServerApiProvider;
fn sample_parent_bridge_message(
sequence: i64,
@@ -52,7 +58,7 @@ fn write_surfaced_parent_bridge_message(state_dir: &Path, record: &MessageBridge
#[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);
let cmd = claude_command("claude", &uuid, "/tmp/prompt.txt", None, None, false);
assert!(
cmd.contains(&format!("--session-id {uuid}")),
"expected --session-id flag in non-resume command, got: {cmd}"
@@ -66,7 +72,7 @@ fn claude_command_uses_session_id_when_not_resuming() {
#[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);
let cmd = claude_command("claude", &uuid, "/tmp/prompt.txt", None, None, true);
assert!(
cmd.contains(&format!("--resume {uuid}")),
"expected --resume flag in resume command, got: {cmd}"
@@ -80,7 +86,14 @@ fn claude_command_uses_resume_flag_when_resuming() {
#[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);
let cmd = claude_command(
"claude",
&uuid,
"/tmp/prompt with spaces.txt",
None,
None,
true,
);
assert!(
cmd.contains("< '/tmp/prompt with spaces.txt'"),
"expected single-quoted stdin redirect of the prompt path, got: {cmd}"
@@ -91,6 +104,111 @@ fn claude_command_pipes_prompt_path() {
);
}
#[test]
fn write_session_index_entry_creates_expected_entry() {
let tmp = TempDir::new().unwrap();
let cwd = Path::new("/my/project");
let session_id = Uuid::new_v4();
write_session_index_entry(session_id, cwd, tmp.path()).unwrap();
let index_path = tmp.path().join("sessions-index.json");
let index: Value = serde_json::from_slice(&fs::read(index_path).unwrap()).unwrap();
let session_key = session_id.to_string();
let entry = &index[&session_key];
let encoded = encode_cwd(cwd);
assert_eq!(entry["sessionId"], Value::String(session_key.clone()));
assert_eq!(
entry["cwd"],
Value::String(cwd.to_string_lossy().into_owned())
);
assert_eq!(entry["projectPath"], Value::String(encoded.clone()));
assert_eq!(
entry["transcriptPath"],
Value::String(format!("projects/{encoded}/{session_id}.jsonl"))
);
}
#[test]
fn serialize_claude_mcp_config_cli_server() {
let servers = HashMap::from([(
"test-server".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::CLIServer {
command: "node".to_string(),
args: vec!["server.js".to_string()],
env: HashMap::from([("API_KEY".to_string(), "secret".to_string())]),
working_directory: None,
},
},
)]);
let json = serialize_claude_mcp_config(&servers).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
let server = &parsed["mcpServers"]["test-server"];
assert_eq!(server["type"], "stdio");
assert_eq!(server["command"], "node");
assert_eq!(server["args"][0], "server.js");
assert_eq!(server["env"]["API_KEY"], "secret");
}
#[test]
fn serialize_claude_mcp_config_cli_server_with_cwd() {
let servers = HashMap::from([(
"test-server".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::CLIServer {
command: "node".to_string(),
args: vec!["server.js".to_string()],
env: HashMap::new(),
working_directory: Some("/opt/mcp".to_string()),
},
},
)]);
let json = serialize_claude_mcp_config(&servers).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
let server = &parsed["mcpServers"]["test-server"];
assert_eq!(server["cwd"], "/opt/mcp");
}
#[test]
fn serialize_claude_mcp_config_cli_server_omits_cwd_when_none() {
let servers = HashMap::from([(
"test-server".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::CLIServer {
command: "node".to_string(),
args: vec![],
env: HashMap::new(),
working_directory: None,
},
},
)]);
let json = serialize_claude_mcp_config(&servers).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
let server = &parsed["mcpServers"]["test-server"];
assert!(server.get("cwd").is_none());
}
#[test]
fn serialize_claude_mcp_config_sse_server() {
let servers = HashMap::from([(
"remote".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::SSEServer {
url: "https://mcp.example.com".to_string(),
headers: HashMap::from([("Authorization".to_string(), "Bearer tok".to_string())]),
},
},
)]);
let json = serialize_claude_mcp_config(&servers).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
let server = &parsed["mcpServers"]["remote"];
assert_eq!(server["type"], "http");
assert_eq!(server["url"], "https://mcp.example.com");
assert_eq!(server["headers"]["Authorization"], "Bearer tok");
}
#[test]
#[serial_test::serial]
fn parent_bridge_root_prefers_environment_override() {
@@ -119,6 +237,77 @@ fn stage_parent_bridge_message_writes_message_record() {
assert!(staged_record.sender_run_id.is_empty());
}
#[tokio::test]
async fn parent_bridge_event_cursor_defaults_to_zero_when_missing() {
let tmp = TempDir::new().unwrap();
let state_dir = tmp.path().join("session-123");
ensure_parent_bridge_state_dir(&state_dir).unwrap();
assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 0);
assert!(!parent_bridge_event_cursor_file(&state_dir).exists());
}
#[tokio::test]
async fn parent_bridge_event_cursor_round_trips() {
let tmp = TempDir::new().unwrap();
let state_dir = tmp.path().join("session-123");
ensure_parent_bridge_state_dir(&state_dir).unwrap();
write_parent_bridge_event_cursor(&state_dir, 42).unwrap();
assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42);
assert!(parent_bridge_event_cursor_file(&state_dir).exists());
}
#[test]
#[serial_test::serial]
fn message_bridge_cleanup_preserves_state_for_wakeable_runs() {
let tmp = TempDir::new().unwrap();
std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path());
let session_id = Uuid::new_v4();
let bridge = MessageBridge::new("run-123".to_string(), session_id).unwrap();
let state_dir = tmp.path().join(session_id.to_string());
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();
write_parent_bridge_event_cursor(&state_dir, 42).unwrap();
bridge
.cleanup(MessageBridgeCleanupDisposition::PreserveState)
.unwrap();
std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV);
assert!(state_dir.exists());
assert!(parent_bridge_staged_message_path(&state_dir, 42, "msg-123").exists());
assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42);
}
#[test]
#[serial_test::serial]
fn message_bridge_cleanup_removes_state_for_non_wakeable_runs() {
let tmp = TempDir::new().unwrap();
std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, tmp.path());
let session_id = Uuid::new_v4();
let bridge = MessageBridge::new("run-123".to_string(), session_id).unwrap();
let state_dir = tmp.path().join(session_id.to_string());
ensure_parent_bridge_state_dir(&state_dir).unwrap();
stage_parent_bridge_message(
&state_dir,
&sample_staged_parent_bridge_message(42, "msg-123"),
)
.unwrap();
write_parent_bridge_event_cursor(&state_dir, 42).unwrap();
bridge
.cleanup(MessageBridgeCleanupDisposition::RemoveState)
.unwrap();
std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV);
assert!(!state_dir.exists());
}
#[tokio::test]
async fn prepare_parent_bridge_hook_output_moves_selected_messages_to_surfaced_dir() {
let tmp = TempDir::new().unwrap();
@@ -481,75 +670,262 @@ fn prepare_claude_config_none_suffix_preserves_existing_responses() {
}
#[test]
fn resolve_suffix_from_raw_value_secret() {
#[serial_test::serial]
fn prepare_claude_environment_config_without_config_dir_uses_home_global_config() {
let home_dir = TempDir::new().unwrap();
let old_home = std::env::var_os("HOME");
let old_config_dir = std::env::var_os("CLAUDE_CONFIG_DIR");
std::env::set_var("HOME", home_dir.path());
std::env::remove_var("CLAUDE_CONFIG_DIR");
let working_dir = home_dir.path().join("workspace/project");
prepare_claude_environment_config(&working_dir, &HashMap::new()).unwrap();
assert!(home_dir.path().join(CLAUDE_JSON_FILE_NAME).exists());
assert!(home_dir
.path()
.join(".claude")
.join(CLAUDE_SETTINGS_FILE_NAME)
.exists());
assert!(!home_dir
.path()
.join(".claude")
.join(CLAUDE_JSON_FILE_NAME)
.exists());
match old_home {
Some(home) => std::env::set_var("HOME", home),
None => std::env::remove_var("HOME"),
}
match old_config_dir {
Some(dir) => std::env::set_var("CLAUDE_CONFIG_DIR", dir),
None => std::env::remove_var("CLAUDE_CONFIG_DIR"),
}
}
#[test]
#[serial_test::serial]
fn prepare_claude_environment_config_with_config_dir_uses_dir_global_config() {
let home_dir = TempDir::new().unwrap();
let claude_config_dir = TempDir::new().unwrap();
let old_home = std::env::var_os("HOME");
let old_config_dir = std::env::var_os("CLAUDE_CONFIG_DIR");
std::env::set_var("HOME", home_dir.path());
std::env::set_var("CLAUDE_CONFIG_DIR", claude_config_dir.path());
let working_dir = home_dir.path().join("workspace/project");
prepare_claude_environment_config(&working_dir, &HashMap::new()).unwrap();
assert!(claude_config_dir
.path()
.join(CLAUDE_JSON_FILE_NAME)
.exists());
assert!(claude_config_dir
.path()
.join(CLAUDE_SETTINGS_FILE_NAME)
.exists());
assert!(!home_dir.path().join(CLAUDE_JSON_FILE_NAME).exists());
match old_home {
Some(home) => std::env::set_var("HOME", home),
None => std::env::remove_var("HOME"),
}
match old_config_dir {
Some(dir) => std::env::set_var("CLAUDE_CONFIG_DIR", dir),
None => std::env::remove_var("CLAUDE_CONFIG_DIR"),
}
}
#[test]
#[serial_test::serial]
fn resolve_suffix_from_resolved_env_vars() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
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);
let resolved = HashMap::from([(OsString::from("ANTHROPIC_API_KEY"), OsString::from(key))]);
let suffix = resolve_anthropic_api_key_suffix(&resolved);
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]
#[serial_test::serial]
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);
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
let resolved = HashMap::from([(OsString::from("ANTHROPIC_API_KEY"), OsString::from("short"))]);
assert_eq!(resolve_anthropic_api_key_suffix(&resolved), 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"),
#[serial_test::serial]
fn resolve_suffix_returns_none_when_empty() {
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
assert_eq!(resolve_anthropic_api_key_suffix(&HashMap::new()), None);
}
#[test]
#[serial_test::serial]
fn prepare_local_wake_command_rehydrates_transcript_with_self_managed_listener() {
let home_dir = TempDir::new().unwrap();
let claude_config_dir = TempDir::new().unwrap();
let bridge_state_root = TempDir::new().unwrap();
let working_dir = home_dir.path().join("workspace/project");
fs::create_dir_all(&working_dir).unwrap();
std::env::set_var("HOME", home_dir.path());
std::env::set_var("CLAUDE_CONFIG_DIR", claude_config_dir.path());
std::env::set_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV, bridge_state_root.path());
let session_id = Uuid::new_v4();
let remote = ClaudeWakeRemoteContext {
session_id,
envelope: ClaudeTranscriptEnvelope {
cwd: Path::new("/stale/cwd").to_path_buf(),
uuid: session_id,
claude_version: None,
entries: vec![serde_json::json!({"type": "assistant", "text": "done"})],
subagents: HashMap::new(),
todos: HashMap::new(),
},
wake_prompt: "resume prompt\n\nwake prompt".to_string(),
};
let task_id: AmbientAgentTaskId = "550e8400-e29b-41d4-a716-446655440010".parse().unwrap();
let parent_run_id = "parent-run-456".to_string();
let command = futures::executor::block_on(ClaudeHarness::prepare_local_wake_command(
ServerApiProvider::new_for_test().get(),
task_id,
Some(parent_run_id.clone()),
Some(working_dir.clone()),
remote,
None,
))
.unwrap();
let state_dir = bridge_state_root.path().join(session_id.to_string());
let prompt_path = state_dir.join(CLAUDE_WAKE_PROMPT_FILE_NAME);
assert!(command.contains("--resume"));
assert!(command.starts_with("env "));
assert!(command.contains(&session_id.to_string()));
assert!(command.contains(CLAUDE_WAKE_PROMPT_FILE_NAME));
assert!(command.contains(&format!(
"{OZ_RUN_ID_ENV}={}",
shell_quote(&task_id.to_string())
)));
assert!(command.contains(&format!(
"{OZ_PARENT_RUN_ID_ENV}={}",
shell_quote(&parent_run_id)
)));
assert!(command.contains(&format!("{OZ_HARNESS_ENV}={}", shell_quote("claude"))));
assert!(!command.contains(OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV));
assert!(!command.contains("OZ_PARENT_LISTENER_MANAGED_EXTERNALLY"));
assert_eq!(
fs::read_to_string(&prompt_path).unwrap(),
"resume prompt\n\nwake prompt"
);
assert!(!parent_bridge_hook_output_file(&state_dir).exists());
let restored_envelope =
read_envelope(session_id, &working_dir, claude_config_dir.path()).unwrap();
assert_eq!(restored_envelope.cwd, working_dir);
assert_eq!(
restored_envelope.entries,
vec![serde_json::json!({"type": "assistant", "text": "done"})]
);
assert!(claude_config_dir
.path()
.join(CLAUDE_JSON_FILE_NAME)
.exists());
assert!(claude_config_dir
.path()
.join(CLAUDE_SETTINGS_FILE_NAME)
.exists());
std::env::remove_var("HOME");
std::env::remove_var("CLAUDE_CONFIG_DIR");
std::env::remove_var(OZ_MESSAGE_LISTENER_STATE_ROOT_ENV);
}
#[tokio::test]
async fn prime_parent_bridge_staged_for_self_managed_wake_keeps_message_in_staged() {
let tmp = TempDir::new().unwrap();
let state_dir = tmp.path().join("session-123");
ensure_parent_bridge_state_dir(&state_dir).unwrap();
let stale = sample_parent_bridge_message(
41,
"stale-msg",
"Old direction",
"This message should be returned to staged.",
);
write_surfaced_parent_bridge_message(&state_dir, &stale);
fs::write(parent_bridge_hook_output_file(&state_dir), "stale context").unwrap();
fs::write(parent_bridge_hook_output_ack_file(&state_dir), "").unwrap();
let wake_message = AgentMessageEventMetadata {
sequence: 42,
message_id: "msg-123".to_string(),
occurred_at: "2026-04-17T15:47:00Z".to_string(),
};
let expected = sample_parent_bridge_message(
42,
"msg-123",
"Please pivot",
"Inspect the failing tests first.",
);
let mut ai_client = MockAIClient::new();
let expected_message = expected.clone();
ai_client
.expect_read_agent_message()
.with(eq("msg-123"))
.times(1)
.returning(move |_| {
Ok(ReadAgentMessageResponse {
message_id: expected_message.message_id.clone(),
sender_run_id: expected_message.sender_run_id.clone(),
subject: expected_message.subject.clone(),
body: expected_message.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 AIClient>);
prime_parent_bridge_staged_for_self_managed_wake(&hydrator, &state_dir, Some(&wake_message))
.await
.unwrap();
assert_eq!(read_parent_bridge_event_cursor(&state_dir).unwrap(), 42);
assert!(!parent_bridge_hook_output_file(&state_dir).exists());
assert!(!parent_bridge_hook_output_ack_file(&state_dir).exists());
assert!(parent_bridge_staged_message_path(&state_dir, 41, "stale-msg").exists());
assert!(!parent_bridge_surfaced_message_path(&state_dir, 41, "stale-msg").exists());
let staged_path = parent_bridge_staged_message_path(&state_dir, 42, "msg-123");
assert!(staged_path.exists());
assert!(!parent_bridge_surfaced_message_path(&state_dir, 42, "msg-123").exists());
let staged_record: MessageBridgeMessageRecord =
serde_json::from_slice(&fs::read(&staged_path).unwrap()).unwrap();
assert_eq!(staged_record.subject, expected.subject);
assert_eq!(staged_record.body, expected.body);
assert_eq!(staged_record.occurred_at, wake_message.occurred_at);
}
#[test]
#[serial_test::serial]
fn suffix_uses_worker_injected_env_when_present() {
let worker_key = "sk-ant-api03-WWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWWW-worker-suffix!";
std::env::set_var(ANTHROPIC_API_KEY_ENV, worker_key);
// Even when the resolved map has a different value, the worker env wins.
let resolved = HashMap::from([(
OsString::from("ANTHROPIC_API_KEY"),
OsString::from(
"sk-ant-api03-RRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRRR-resolved-val!",
),
)]);
assert_eq!(resolve_anthropic_api_key_suffix(&secrets), None);
let suffix = resolve_anthropic_api_key_suffix(&resolved);
let expected = &worker_key[worker_key.len() - 20..];
assert_eq!(suffix.as_deref(), Some(expected));
std::env::remove_var(ANTHROPIC_API_KEY_ENV);
}
#[test]
app/src/ai/agent_sdk/driver/harness/claude_transcript.rs
+129 -23
View File
@@ -15,7 +15,8 @@
//! 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::fs::{create_dir_all, write};
use std::io::{BufRead, BufReader, Read};
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
@@ -24,6 +25,7 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use super::json_utils::entries_to_jsonl;
use crate::ai::agent::conversation::AIConversationId;
/// JSON envelope sent to the server representing a complete Claude Code session.
@@ -66,6 +68,11 @@ pub(crate) struct ClaudeResumeInfo {
pub(crate) envelope: ClaudeTranscriptEnvelope,
}
#[derive(Debug)]
pub(crate) struct ClaudeLocalContinuation {
pub(crate) command: String,
}
/// Encode a filesystem path as a Claude config directory name, matching the
/// Claude CLI convention of replacing every `/` with `-`.
///
@@ -83,11 +90,23 @@ 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()
home_dir_for_claude_config()
.map(|h| h.join(".claude"))
.ok_or_else(|| anyhow::anyhow!("could not determine home directory"))
}
/// In tests on Windows, `dirs::home_dir()` ignores `HOME`, so we check it
/// manually so that tests can override the home directory.
pub(super) fn home_dir_for_claude_config() -> Option<PathBuf> {
#[cfg(test)]
if let Some(home) = std::env::var_os("HOME") {
if !home.is_empty() {
return Some(PathBuf::from(home));
}
}
dirs::home_dir()
}
/// Assemble a [`ClaudeTranscriptEnvelope`] from the Claude config directory.
///
/// Reads:
@@ -184,12 +203,12 @@ pub(crate) fn write_envelope(
) -> Result<()> {
let encoded = encode_cwd(&envelope.cwd);
let projects_dir = config_root.join("projects").join(&encoded);
std::fs::create_dir_all(&projects_dir)
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)?)
write(&session_file, entries_to_jsonl(&envelope.entries)?)
.with_context(|| format!("Failed to write {}", session_file.display()))?;
// Subagent JSONLs.
@@ -197,11 +216,11 @@ pub(crate) fn write_envelope(
let subagents_dir = projects_dir
.join(envelope.uuid.to_string())
.join("subagents");
std::fs::create_dir_all(&subagents_dir)
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)?)
write(&path, entries_to_jsonl(entries)?)
.with_context(|| format!("Failed to write {}", path.display()))?;
}
}
@@ -209,11 +228,11 @@ pub(crate) fn write_envelope(
// Per-agent todo lists.
if !envelope.todos.is_empty() {
let todos_dir = config_root.join("todos");
std::fs::create_dir_all(&todos_dir)
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)?)
write(&path, serde_json::to_vec(value)?)
.with_context(|| format!("Failed to write {}", path.display()))?;
}
}
@@ -221,10 +240,108 @@ pub(crate) fn write_envelope(
Ok(())
}
pub(crate) fn rehydrate_claude_transcript(
envelope: &mut ClaudeTranscriptEnvelope,
local_cwd: &Path,
) -> Result<ClaudeLocalContinuation> {
envelope.cwd = local_cwd.to_path_buf();
let session_id = envelope.uuid;
let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?;
write_envelope(envelope, &config_root).context("Failed to rehydrate Claude transcript")?;
if let Err(e) = write_session_index_entry(session_id, local_cwd, &config_root) {
log::warn!("Failed to update Claude sessions-index.json: {e:#}");
}
Ok(ClaudeLocalContinuation {
command: format!("claude --resume {session_id}"),
})
}
/// Write a [`ClaudeTranscriptEnvelope`] to a project directory derived from `storage_cwd`,
/// without mutating `envelope.cwd`.
///
/// Used by the local continuation path so the transcript's recorded working directory
/// (the original cloud cwd) is preserved as-is while the file is placed under
/// `~/.claude/projects/<encoded(storage_cwd)>/` where Claude's per-project lookup can find it.
/// Cloud resume uses [`write_envelope`] instead, which derives the path from `envelope.cwd`.
///
/// Creates:
/// - `<config_root>/projects/<encoded(storage_cwd)>/<uuid>.jsonl` — main transcript
/// - `<config_root>/projects/<encoded(storage_cwd)>/<uuid>/subagents/<stem>.jsonl` — subagents
/// - `<config_root>/todos/<stem>.json` — per-agent todo lists (same location as cloud resume)
pub(crate) fn write_envelope_for_local_continuation(
envelope: &ClaudeTranscriptEnvelope,
storage_cwd: &Path,
config_root: &Path,
) -> Result<()> {
let projects_dir = config_root.join("projects").join(encode_cwd(storage_cwd));
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));
write(&session_file, entries_to_jsonl(&envelope.entries)?)
.with_context(|| format!("Failed to write {}", session_file.display()))?;
// Subagent JSONLs — same relative layout as write_envelope.
if !envelope.subagents.is_empty() {
let subagents_dir = projects_dir
.join(envelope.uuid.to_string())
.join("subagents");
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"));
write(&path, entries_to_jsonl(entries)?)
.with_context(|| format!("Failed to write {}", path.display()))?;
}
}
// Per-agent todo lists are written to the same global location as cloud resume.
if !envelope.todos.is_empty() {
let todos_dir = config_root.join("todos");
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"));
write(&path, serde_json::to_vec(value)?)
.with_context(|| format!("Failed to write {}", path.display()))?;
}
}
Ok(())
}
/// Rehydrate a Claude transcript downloaded from a remote cloud run for local continuation.
///
/// Unlike [`rehydrate_claude_transcript`] (used by the cloud resume harness runner), this
/// function does **not** mutate the envelope's `cwd` field — the remote session's original
/// working directory is preserved as-is in the transcript. The session file is stored under
/// `~/.claude/projects/<encoded(home_dir)>/` so Claude's per-project session lookup finds it
/// when the user runs `claude --resume <uuid>` from their home directory.
pub(crate) fn rehydrate_claude_transcript_from_reader(
reader: impl Read,
) -> Result<ClaudeLocalContinuation> {
let envelope: ClaudeTranscriptEnvelope =
serde_json::from_reader(reader).context("Failed to parse Claude transcript envelope")?;
let session_id = envelope.uuid;
let config_root = claude_config_dir().context("Failed to resolve Claude config dir")?;
let home_dir = home_dir_for_claude_config()
.ok_or_else(|| anyhow::anyhow!("could not determine home directory"))?;
write_envelope_for_local_continuation(&envelope, &home_dir, &config_root)
.context("Failed to rehydrate Claude transcript for local continuation")?;
if let Err(e) = write_session_index_entry(session_id, &home_dir, &config_root) {
log::warn!("Failed to update Claude sessions-index.json: {e:#}");
}
Ok(ClaudeLocalContinuation {
command: format!("claude --resume {session_id}"),
})
}
/// 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
/// 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
@@ -243,7 +360,7 @@ pub(crate) fn write_session_index_entry(
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.
// we'd rather clobber an unparsable 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,
@@ -281,10 +398,9 @@ pub(crate) fn write_session_index_entry(
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()))?;
create_dir_all(parent).with_context(|| format!("Failed to create {}", parent.display()))?;
}
std::fs::write(
write(
&index_path,
serde_json::to_vec_pretty(&Value::Object(index))
.context("Failed to serialize sessions-index.json")?,
@@ -293,16 +409,6 @@ pub(crate) fn write_session_index_entry(
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
app/src/ai/agent_sdk/driver/harness/codex.rs
+939
View File
@@ -0,0 +1,939 @@
use std::collections::HashMap;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};
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 galaxy_core::features::FeatureFlag;
use warp_managed_secrets::ManagedSecretValue;
use warpui::{ModelHandle, ModelSpawner, SingletonEntity};
use super::super::terminal::{CommandHandle, TerminalDriver};
use super::super::{AgentDriver, AgentDriverError};
use super::claude_transcript::read_jsonl;
use super::codex_transcript::{
codex_sessions_root, find_session_file, parse_session_meta, rehydrate_codex_transcript,
CodexResumeInfo, CodexTranscriptEnvelope,
};
use super::json_utils::read_json_file_or_default;
use super::{
write_temp_file, HarnessRunner, JSONMCPServer, ResumePayload, SavePoint, ThirdPartyHarness,
};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_sdk::setup_observability::{
OzRunTimelineEvent, SetupClientEventReporter, SetupStep,
};
use crate::ai::ambient_agents::task::HarnessModelConfig;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::mcp::JSONTransportType;
use crate::server::server_api::harness_support::{upload_to_target, HarnessSupportClient};
use crate::server::server_api::ServerApi;
use crate::terminal::cli_agent_sessions::CLIAgentSessionsModel;
use crate::terminal::model::block::BlockId;
use crate::terminal::CLIAgent;
pub(crate) struct CodexHarness;
/// Format slug sent to the server when creating a Codex conversation.
const CODEX_CLI_FORMAT: &str = "codex_cli";
/// Slash command Codex's TUI recognises as a graceful shutdown.
const CODEX_EXIT_COMMAND: &str = "/exit";
/// Allow the Warp-installed Codex plugin hooks to run in vetted driver sessions
/// without requiring an unattended `/hooks` review step.
const CODEX_BYPASS_HOOK_TRUST_FLAG: &str = "--dangerously-bypass-hook-trust";
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl ThirdPartyHarness for CodexHarness {
fn harness(&self) -> Harness {
Harness::Codex
}
fn cli_agent(&self) -> CLIAgent {
CLIAgent::Codex
}
fn install_docs_url(&self) -> Option<&'static str> {
Some("https://developers.openai.com/codex/cli")
}
fn auth_check_command(&self) -> Option<String> {
let cli = self.cli_agent().command_prefix();
Some(format!("{cli} login status"))
}
fn runtime_error_patterns(&self) -> &'static [&'static str] {
&[
// Quota / billing.
"Quota exceeded. Check your plan and billing details.",
"You've hit your usage limit",
// Upstream HTTP failures Codex surfaces verbatim. The 401 form
// matches invalid-API-key and wrong-endpoint variants.
"unexpected status 401",
"Incorrect API key provided",
"invalid API key",
// Region/endpoint block (Anthropic-style global vs US-only
// routing surfaced through Codex's upstream client).
"Access blocked by Cloudflare",
// OAuth refresh failures — all five Codex variants share this
// substring (see upstream session/token messages).
"could not be refreshed",
]
}
fn requires_verified_platform_plugin(&self) -> bool {
FeatureFlag::CodexPlugin.is_enabled()
}
/// Fetch the codex transcript for the current task's conversation and wrap it into a
/// [`ResumePayload::Codex`].
async fn fetch_resume_payload(
&self,
conversation_id: &AIConversationId,
harness_support_client: Arc<dyn HarnessSupportClient>,
) -> Result<Option<ResumePayload>, AgentDriverError> {
let envelope: CodexTranscriptEnvelope =
super::fetch_transcript_envelope("codex", conversation_id, harness_support_client)
.await?;
let session_id = envelope.session_id;
Ok(Some(ResumePayload::Codex(CodexResumeInfo {
conversation_id: *conversation_id,
session_id,
envelope,
})))
}
fn build_runner(
&self,
prompt: &str,
system_prompt: Option<&str>,
resumption_prompt: Option<&str>,
context: Option<&str>,
working_dir: &Path,
_task_id: Option<AmbientAgentTaskId>,
server_api: Arc<ServerApi>,
terminal_driver: ModelHandle<TerminalDriver>,
resume: Option<ResumePayload>,
resolved_env_vars: &HashMap<OsString, OsString>,
resolved_secrets: &HashMap<String, ManagedSecretValue>,
resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
third_party_harness_model_config: Option<&HarnessModelConfig>,
) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
// Prepare the environment config files.
prepare_codex_environment_config(
working_dir,
system_prompt,
resolved_env_vars,
resolved_secrets,
resolved_mcp_servers,
third_party_harness_model_config,
)
.map_err(|error| AgentDriverError::HarnessConfigSetupFailed {
harness: self.cli_agent().command_prefix().to_owned(),
error,
})?;
// The ResumePayload shouldn't contain non-Codex information, error if it does.
let codex_resume = resume.map(CodexResumeInfo::try_from).transpose()?;
// Mirror Claude harness behavior: prepend the resumption preamble and server context
// to the user-turn prompt so codex treats it as immediate intent.
// Order: resumption_prompt → context → prompt
let mut parts: Vec<&str> = Vec::new();
if let Some(preamble) = resumption_prompt {
if !preamble.is_empty() {
parts.push(preamble);
}
}
if let Some(ctx) = context {
if !ctx.is_empty() {
parts.push(ctx);
}
}
parts.push(prompt);
let owned_prompt = parts.join("\n\n");
let client: Arc<dyn HarnessSupportClient> = server_api;
Ok(Box::new(CodexHarnessRunner::new(
self.cli_agent().command_prefix(),
&owned_prompt,
system_prompt,
working_dir,
client,
terminal_driver,
codex_resume,
)?))
}
}
/// Build the shell command that launches the Codex TUI.
///
/// `--dangerously-bypass-approvals-and-sandbox` disables both the sandbox and approval
/// prompts so the agent can run autonomously.
/// `--dangerously-bypass-hook-trust` allows the orchestration plugin hooks installed by
/// Warp to run without a manual hook review in unattended driver sessions. Driver setup
/// verifies the Codex platform plugin before launching commands with this flag.
/// `Some(session_id)` indicates that we want to resume that prior session. Unlike claude,
/// codex does not support assigning a session_id to a new conversation.
fn codex_command(cli_name: &str, session_id: Option<&Uuid>, prompt_path: &str) -> String {
match session_id {
Some(session_id) => format!(
"{cli_name} resume --dangerously-bypass-approvals-and-sandbox {CODEX_BYPASS_HOOK_TRUST_FLAG} {session_id} \
\"$(cat '{prompt_path}')\""
),
None => {
format!(
"{cli_name} --dangerously-bypass-approvals-and-sandbox {CODEX_BYPASS_HOOK_TRUST_FLAG} \"$(cat '{prompt_path}')\""
)
}
}
}
enum CodexRunnerState {
Preexec,
Running {
conversation_id: AIConversationId,
block_id: BlockId,
},
}
struct CodexHarnessRunner {
command: String,
cli_name: 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<CodexRunnerState>,
/// Codex session UUID. Populated lazily by [`HarnessRunner::handle_session_update`]
/// once the codex hooks emit `SessionStart`. Set once (using `OnceLock`).
session_id: OnceLock<Uuid>,
/// Path to the codex session rollout JSONL file. Populated by the first
/// successful [`find_session_file`] walk so that subsequent saves skip the YYYY/MM/DD
/// directory walk and read the JSONL file directly.
transcript_path: OnceLock<PathBuf>,
/// Optionally supply an existing conversation ID.
preexisting_conversation_id: Option<AIConversationId>,
}
impl CodexHarnessRunner {
#[allow(clippy::too_many_arguments)]
fn new(
cli_command: &str,
prompt: &str,
_system_prompt: Option<&str>,
_working_dir: &Path,
client: Arc<dyn HarnessSupportClient>,
terminal_driver: ModelHandle<TerminalDriver>,
resume: Option<CodexResumeInfo>,
) -> Result<Self, AgentDriverError> {
let temp_file = write_temp_file("oz_prompt_", prompt, ".txt")?;
let prompt_path = temp_file.path().display().to_string();
let (session_id, preexisting_conversation_id, transcript_path) = match resume {
Some(CodexResumeInfo {
conversation_id,
session_id,
mut envelope,
}) => {
let continuation = rehydrate_codex_transcript(&mut envelope, _working_dir)
.map_err(AgentDriverError::ConfigBuildFailed)?;
(
Some(session_id),
Some(conversation_id),
Some(continuation.transcript_path),
)
}
None => (None, None, None),
};
let command = codex_command(cli_command, session_id.as_ref(), &prompt_path);
let session_id_cell: OnceLock<Uuid> = OnceLock::new();
if let Some(id) = session_id {
let _ = session_id_cell.set(id);
}
let transcript_path_cell: OnceLock<PathBuf> = OnceLock::new();
if let Some(p) = transcript_path {
let _ = transcript_path_cell.set(p);
}
Ok(Self {
command,
cli_name: cli_command.to_string(),
_temp_prompt_file: temp_file,
client,
terminal_driver,
state: Mutex::new(CodexRunnerState::Preexec),
session_id: session_id_cell,
transcript_path: transcript_path_cell,
preexisting_conversation_id,
})
}
/// Return the filepath for the session transcript, walking the codex sessions tree to find it on the
/// first save call.
async fn resolve_transcript_path(&self) -> Option<PathBuf> {
if let Some(cached) = self.transcript_path.get() {
return Some(cached.clone());
}
let session_id = self.session_id.get().copied()?;
let resolved = tokio::task::spawn_blocking(move || -> Option<PathBuf> {
let root = codex_sessions_root().ok()?;
find_session_file(&root, session_id)
})
.await
.ok()
.flatten()?;
let _ = self.transcript_path.set(resolved.clone());
Some(resolved)
}
}
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl HarnessRunner for CodexHarnessRunner {
fn harness_name(&self) -> &str {
&self.cli_name
}
async fn start(
&self,
foreground: &ModelSpawner<AgentDriver>,
setup_events: &SetupClientEventReporter,
) -> Result<CommandHandle, AgentDriverError> {
// Resume runs reuse the prior server conversation id; fresh runs mint a new one.
let conversation_id = match self.preexisting_conversation_id {
Some(id) => {
log::info!("Resuming external conversation {id}");
id
}
None => {
let id = setup_events
.record_result(SetupStep::ThirdPartyHarnessExternalConversation, async {
self.client
.create_external_conversation(CODEX_CLI_FORMAT)
.await
.map_err(|e| {
log::error!("Failed to create external conversation: {e}");
AgentDriverError::ConfigBuildFailed(e)
})
})
.await?;
log::info!("Created external conversation {id}");
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?;
*self.state.lock() = CodexRunnerState::Running {
conversation_id,
block_id: command_handle.block_id().clone(),
};
setup_events
.post_timeline_event(OzRunTimelineEvent::AgentStarted)
.await;
Ok(command_handle)
}
async fn exit(&self, foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
log::info!("Sending /exit to Codex CLI");
let terminal_driver = self.terminal_driver.clone();
foreground
.spawn(move |_, ctx| {
terminal_driver.update(ctx, |driver, ctx| {
driver.send_text_to_cli(CODEX_EXIT_COMMAND.to_string(), ctx);
});
})
.await
.map_err(|_| anyhow::anyhow!("Agent driver dropped while sending /exit"))
}
/// Capture the codex session ID from the `SessionStart` event picked up by the `CLIAgentSessionsModel`.
///
/// Relies on codex hooks being set up to emit this event correctly.
async fn handle_session_update(&self, foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
if self.session_id.get().is_some() {
return Ok(());
}
let terminal_driver = self.terminal_driver.clone();
let session_id_str = foreground
.spawn(move |_, ctx| {
let terminal_view_id = terminal_driver.as_ref(ctx).terminal_view().id();
CLIAgentSessionsModel::handle(ctx)
.as_ref(ctx)
.session(terminal_view_id)
.and_then(|s| s.session_context.session_id.clone())
})
.await
.ok()
.flatten();
let Some(session_id_str) = session_id_str else {
return Ok(());
};
match Uuid::parse_str(&session_id_str) {
Ok(uuid) => {
log::info!("Captured codex session id {uuid}");
let _ = self.session_id.set(uuid);
}
Err(e) => log::warn!("Failed to parse codex session id '{session_id_str}': {e}"),
}
Ok(())
}
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, Codex not in progress");
return Ok(());
}
let (conversation_id, block_id) = match &*self.state.lock() {
CodexRunnerState::Preexec => {
log::warn!("save_conversation called before start");
return Ok(());
}
CodexRunnerState::Running {
conversation_id,
block_id,
} => (*conversation_id, block_id.clone()),
};
let session_id = self.session_id.get().copied();
let rollout_path = self.resolve_transcript_path().await;
let client = self.client.as_ref();
let is_final = matches!(save_point, SavePoint::Final);
futures::try_join!(
super::upload_current_block_snapshot(
foreground,
&self.terminal_driver,
client,
conversation_id,
block_id,
),
upload_transcript(client, conversation_id, session_id, rollout_path, is_final),
)?;
Ok(())
}
}
/// Upload the codex session transcript to the server. No-ops if the session UUID hasn't
/// been captured yet or no rollout file is on disk yet.
async fn upload_transcript(
client: &dyn HarnessSupportClient,
conversation_id: AIConversationId,
session_id: Option<Uuid>,
transcript_path: Option<PathBuf>,
is_final: bool,
) -> Result<()> {
let Some(session_id) = session_id else {
if is_final {
log::warn!(
"Codex session id still unknown at final save; transcript was never uploaded"
);
} else {
log::debug!("Codex session id not yet known; skipping transcript upload");
}
return Ok(());
};
let Some(transcript_path) = transcript_path else {
if is_final {
log::warn!(
"No codex rollout file found at final save for session {session_id}; transcript was never uploaded"
);
} else {
log::debug!("No codex rollout file yet for session {session_id}");
}
return Ok(());
};
log::info!("Uploading codex transcript to conversation {conversation_id}");
let body = tokio::task::spawn_blocking(move || -> Result<Vec<u8>> {
let entries = read_jsonl(&transcript_path)?;
let metadata = parse_session_meta(entries.first()).unwrap_or_default();
let envelope = CodexTranscriptEnvelope::new(session_id, metadata, entries);
serde_json::to_vec(&envelope).context("Failed to serialize codex transcript")
})
.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?;
Ok(())
}
const CODEX_CONFIG_DIR: &str = ".codex";
const CODEX_HOME_ENV: &str = "CODEX_HOME";
const CODEX_AGENTS_OVERRIDE_FILE_NAME: &str = "AGENTS.override.md";
const CODEX_AUTH_FILE_NAME: &str = "auth.json";
const CODEX_CONFIG_TOML_FILE_NAME: &str = "config.toml";
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const CODEX_AUTH_MODE_API_KEY: &str = "apikey";
/// Lowercase string Codex's `TrustLevel` enum serializes to (codex
/// `protocol/src/config_types.rs::TrustLevel`).
const CODEX_TRUST_LEVEL_TRUSTED: &str = "trusted";
/// Top-level config key codex reads to override the built-in `openai` provider's base URL
/// (codex `core/src/config/mod.rs`).
const CODEX_OPENAI_BASE_URL_KEY: &str = "openai_base_url";
const CODEX_CHECK_FOR_UPDATE_ON_STARTUP_KEY: &str = "check_for_update_on_startup";
const CODEX_MODEL_KEY: &str = "model";
const CODEX_MODEL_REASONING_EFFORT_KEY: &str = "model_reasoning_effort";
/// Target model for the `[notice.model_migrations]` table that suppresses Codex's
/// "choose a newer model" upgrade prompt at session launch. We stamp this for any
/// pinned model id (even when it already matches the target) so the unattended
/// cloud run never blocks on the prompt.
///
/// TODO: Ideally, we would make this server-driven so we don't depend on a client
/// release to change this.
const CODEX_MODEL_MIGRATIONS_TARGET: &str = "gpt-5.4";
fn prepare_codex_environment_config(
working_dir: &Path,
system_prompt: Option<&str>,
resolved_env_vars: &HashMap<OsString, OsString>,
resolved_secrets: &HashMap<String, ManagedSecretValue>,
resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
third_party_harness_model_config: Option<&HarnessModelConfig>,
) -> Result<()> {
let codex_dir = codex_config_dir()?;
if let Some(prompt) = system_prompt {
write_codex_agents_override(&codex_dir, prompt)?;
}
match resolve_openai_api_key(resolved_env_vars) {
Some(api_key) => prepare_codex_auth(&codex_dir.join(CODEX_AUTH_FILE_NAME), &api_key)?,
None => log::info!("No OPENAI_API_KEY available; skipping Codex auth.json seed"),
}
// Resolve the base URL directly from the typed OpenAI secret. This avoids
// leaking base_url into the child process environment and ensures we only
// apply it when the typed secret is the active API key source.
let openai_base_url = resolve_openai_base_url_from_secret(resolved_secrets, resolved_env_vars);
prepare_codex_config_toml(
&codex_dir.join(CODEX_CONFIG_TOML_FILE_NAME),
working_dir,
resolved_mcp_servers,
third_party_harness_model_config,
openai_base_url.as_deref(),
)?;
Ok(())
}
fn codex_config_dir() -> Result<PathBuf> {
if let Ok(dir) = std::env::var(CODEX_HOME_ENV) {
if !dir.is_empty() {
return Ok(PathBuf::from(dir));
}
}
dirs::home_dir()
.map(|home| home.join(CODEX_CONFIG_DIR))
.ok_or_else(|| anyhow::anyhow!("could not determine home directory"))
}
fn write_codex_agents_override(codex_dir: &Path, system_prompt: &str) -> Result<()> {
fs::create_dir_all(codex_dir).with_context(|| {
format!(
"Failed to create Codex config dir at {}",
codex_dir.display()
)
})?;
// Note: this currently works because we are only doing this for cloud agents; if we enable
// this for local runs we'll want to make sure we don't clobber any existing file overrides.
let prompt_path = codex_dir.join(CODEX_AGENTS_OVERRIDE_FILE_NAME);
fs::write(&prompt_path, system_prompt).with_context(|| {
format!(
"Failed to write Codex system prompt to {}",
prompt_path.display()
)
})
}
/// Mirrors the subset of Codex's `AuthDotJson` (codex `login/src/auth/storage.rs`) that we
/// need to seed. Unknown fields (`tokens`, `last_refresh`, `agent_identity`, ...) are
/// preserved via `extra` so we don't clobber an existing login.
#[derive(Default, Deserialize, Serialize, Debug)]
struct CodexAuthDotJson {
#[serde(default, skip_serializing_if = "Option::is_none")]
auth_mode: Option<String>,
#[serde(
rename = "OPENAI_API_KEY",
default,
skip_serializing_if = "Option::is_none"
)]
openai_api_key: Option<String>,
#[serde(flatten)]
extra: Map<String, Value>,
}
fn prepare_codex_auth(auth_path: &Path, api_key: &str) -> Result<()> {
let mut auth: CodexAuthDotJson = read_json_file_or_default(auth_path)?;
auth.openai_api_key = Some(api_key.to_owned());
if auth.auth_mode.is_none() {
auth.auth_mode = Some(CODEX_AUTH_MODE_API_KEY.to_owned());
}
write_codex_auth_json(auth_path, &auth)
}
/// Write Codex's `auth.json` with restrictive (0o600) permissions, mirroring how
/// codex sets up this file itself.
fn write_codex_auth_json(path: &Path, auth: &CodexAuthDotJson) -> Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.with_context(|| format!("Failed to create {}", parent.display()))?;
}
let bytes = serde_json::to_vec_pretty(auth).context("Failed to serialize Codex auth.json")?;
#[cfg(unix)]
{
use std::io::Write as _;
use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
let mut file = fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.mode(0o600)
.open(path)
.with_context(|| format!("Failed to open {} for writing", path.display()))?;
file.set_permissions(fs::Permissions::from_mode(0o600))
.with_context(|| format!("Failed to set permissions on {}", path.display()))?;
file.write_all(&bytes)
.with_context(|| format!("Failed to write {}", path.display()))?;
}
#[cfg(not(unix))]
fs::write(path, &bytes).with_context(|| format!("Failed to write {}", path.display()))?;
Ok(())
}
/// Returns the OpenAI API key for Codex auth.
///
/// Checks the worker-injected process env first (not in the resolved map since
/// `build_secret_env_vars` skips env vars already present in the process env),
/// then falls back to the resolved secret env vars map.
fn resolve_openai_api_key(resolved_env_vars: &HashMap<OsString, OsString>) -> Option<String> {
// Worker-injected process env wins.
if let Ok(value) = std::env::var(OPENAI_API_KEY_ENV) {
let trimmed = value.trim();
if !trimmed.is_empty() {
return Some(trimmed.to_owned());
}
}
// Otherwise use the resolved value from the secrets map.
resolved_env_vars
.get(OsStr::new(OPENAI_API_KEY_ENV))
.and_then(|v| v.to_str())
.map(|s| s.trim().to_owned())
.filter(|s| !s.is_empty())
}
/// Returns the OpenAI base URL from the typed secret, if applicable.
///
/// The base URL is only used when the typed `OpenaiApiKey` secret is the active
/// source of `OPENAI_API_KEY`. If a worker-injected process env already provides
/// the API key, the typed-secret base URL is not applied (the worker controls
/// both the key and endpoint).
fn resolve_openai_base_url_from_secret(
secrets: &HashMap<String, ManagedSecretValue>,
resolved_env_vars: &HashMap<OsString, OsString>,
) -> Option<String> {
// If the worker already injected an API key, the typed secret lost
// precedence — do not apply its base URL.
if std::env::var(OPENAI_API_KEY_ENV)
.ok()
.is_some_and(|v| !v.trim().is_empty())
{
return None;
}
// Only apply when the resolved env vars actually contain OPENAI_API_KEY
// from the typed secret (i.e. the secret was not skipped).
resolved_env_vars.get(OsStr::new(OPENAI_API_KEY_ENV))?;
secrets.values().find_map(|secret| match secret {
ManagedSecretValue::OpenaiApiKey { base_url, .. } => base_url
.as_ref()
.map(|s| s.trim().to_owned())
.filter(|s| !s.is_empty()),
_ => None,
})
}
/// Edit `~/.codex/config.toml` via `toml_edit` to seed the harness defaults
/// while preserving anything that might already exist there. We handle:
/// - project trust: for a working dir and all of its git repo subdirectories,
/// set the projects to `trusted`.
/// - base URL: when `openai_base_url` is provided (from the secret's `base_url`
/// field), write it to config.toml. When absent, skip the key entirely so
/// Codex uses the provider's default global endpoint.
/// - update checks: disable Codex's startup update prompt for unattended runs.
/// - model override: when a non-default harness model config is
/// supplied, write the top-level `model` key so Codex pins the chosen model
/// for new sessions.
fn prepare_codex_config_toml(
config_toml_path: &Path,
working_dir: &Path,
resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
third_party_harness_model_config: Option<&HarnessModelConfig>,
openai_base_url: Option<&str>,
) -> Result<()> {
let existing = match fs::read_to_string(config_toml_path) {
Ok(content) => content,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => String::new(),
Err(e) => {
return Err(anyhow::Error::from(e).context(format!(
"Failed to read Codex config.toml at {}",
config_toml_path.display()
)));
}
};
let mut doc: toml_edit::DocumentMut = existing.parse().with_context(|| {
format!(
"Failed to parse Codex config.toml at {}",
config_toml_path.display()
)
})?;
// Only write openai_base_url when the secret specifies one.
if let Some(url) = openai_base_url {
set_codex_openai_base_url(&mut doc, url);
}
set_codex_check_for_update_on_startup(&mut doc, false);
set_codex_model(&mut doc, third_party_harness_model_config);
set_codex_model_reasoning_effort(&mut doc, third_party_harness_model_config);
let canonical = working_dir.canonicalize().with_context(|| {
format!(
"Failed to canonicalize Codex working dir at {}",
working_dir.display()
)
})?;
let project_key = canonical.to_string_lossy().into_owned();
set_codex_project_trust_level(&mut doc, &project_key, CODEX_TRUST_LEVEL_TRUSTED);
// Codex's trust check is not recursive (see openai/codex#19426) -- since we
// clone the git repos into workspace/ for cloud agents, we usually have git
// repo children that we also want to trust.
for child_repo in find_child_git_repos(&canonical) {
let key = child_repo.to_string_lossy().into_owned();
set_codex_project_trust_level(&mut doc, &key, CODEX_TRUST_LEVEL_TRUSTED);
}
write_codex_mcp_servers(&mut doc, resolved_mcp_servers);
if let Some(parent) = config_toml_path.parent() {
fs::create_dir_all(parent).with_context(|| {
format!("Failed to create Codex config dir at {}", parent.display())
})?;
}
fs::write(config_toml_path, doc.to_string()).with_context(|| {
format!(
"Failed to write Codex config.toml at {}",
config_toml_path.display()
)
})
}
/// Set the top-level `openai_base_url` key, overwriting any existing value.
fn set_codex_openai_base_url(doc: &mut toml_edit::DocumentMut, base_url: &str) {
doc[CODEX_OPENAI_BASE_URL_KEY] = toml_edit::value(base_url);
}
fn set_codex_check_for_update_on_startup(doc: &mut toml_edit::DocumentMut, enabled: bool) {
doc[CODEX_CHECK_FOR_UPDATE_ON_STARTUP_KEY] = toml_edit::value(enabled);
}
fn set_codex_model_reasoning_effort(
doc: &mut toml_edit::DocumentMut,
third_party_harness_model_config: Option<&HarnessModelConfig>,
) {
let Some(reasoning_level) = third_party_harness_model_config
.and_then(|config| config.reasoning_level.as_deref())
.filter(|level| !level.is_empty())
else {
doc.remove(CODEX_MODEL_REASONING_EFFORT_KEY);
return;
};
doc[CODEX_MODEL_REASONING_EFFORT_KEY] = toml_edit::value(reasoning_level);
}
fn set_codex_model(
doc: &mut toml_edit::DocumentMut,
third_party_harness_model_config: Option<&HarnessModelConfig>,
) {
let Some(model_id) = third_party_harness_model_config
.map(|config| config.model_id.as_str())
.filter(|id| !id.is_empty() && *id != "default")
else {
// No model specified or "default" selected — remove any pre-existing
// key so Codex uses its own default.
doc.remove(CODEX_MODEL_KEY);
return;
};
doc[CODEX_MODEL_KEY] = toml_edit::value(model_id);
// Codex's TUI prompts the user to upgrade older models on session launch even when
// a `model` key has been pinned. Stamping a migration entry keyed on the chosen
// model id suppresses that prompt for the unattended cloud run. We do this
// unconditionally rather than enumerating a list of "old" models on the client:
// mapping the migration target to itself (e.g. `gpt-5.4 = "gpt-5.4"`) is a no-op
// for Codex, and keeping the client free of model-version knowledge means we
// don't have to ship a client update every time Anthropic/OpenAI ages out a model.
set_codex_model_migration(doc, model_id, CODEX_MODEL_MIGRATIONS_TARGET);
}
fn set_codex_model_migration(
doc: &mut toml_edit::DocumentMut,
from_model_id: &str,
to_model_id: &str,
) {
if !doc.contains_table("notice") {
let mut notice_tbl = toml_edit::Table::new();
notice_tbl.set_implicit(true);
doc.insert("notice", toml_edit::Item::Table(notice_tbl));
}
let migrations_tbl = doc["notice"]
.as_table_mut()
.expect("notice table inserted above")
.entry("model_migrations")
.or_insert_with(toml_edit::table)
.as_table_mut()
.expect("model_migrations entry is a table");
migrations_tbl.set_implicit(false);
migrations_tbl[from_model_id] = toml_edit::value(to_model_id);
}
/// Return immediate subdirectories of `dir` that contain a `.git`.
fn find_child_git_repos(dir: &Path) -> Vec<std::path::PathBuf> {
let Ok(entries) = fs::read_dir(dir) else {
return Vec::new();
};
entries
.flatten()
.filter_map(|entry| {
let path = entry.path();
(path.is_dir() && path.join(".git").exists()).then_some(path)
})
.collect()
}
/// Insert/update `[projects."<project_key>"] trust_level = <trust_level>`.
///
/// Codex itself always writes `projects` as an explicit table, so we don't
/// handle the inline-table form here.
fn set_codex_project_trust_level(
doc: &mut toml_edit::DocumentMut,
project_key: &str,
trust_level: &str,
) {
if !doc.contains_table("projects") {
let mut projects_tbl = toml_edit::Table::new();
projects_tbl.set_implicit(true);
doc.insert("projects", toml_edit::Item::Table(projects_tbl));
}
let proj_tbl = doc["projects"]
.as_table_mut()
.expect("projects table inserted above")
.entry(project_key)
.or_insert_with(toml_edit::table)
.as_table_mut()
.expect("project entry is a table");
proj_tbl.set_implicit(false);
proj_tbl["trust_level"] = toml_edit::value(trust_level);
}
/// Write resolved MCP servers into `[mcp_servers.<name>]` sections in the Codex config.
fn write_codex_mcp_servers(
doc: &mut toml_edit::DocumentMut,
servers: &HashMap<String, JSONMCPServer>,
) {
if servers.is_empty() {
return;
}
if !doc.contains_table("mcp_servers") {
let mut tbl = toml_edit::Table::new();
tbl.set_implicit(true);
doc.insert("mcp_servers", toml_edit::Item::Table(tbl));
}
let mcp_tbl = doc["mcp_servers"]
.as_table_mut()
.expect("mcp_servers table inserted above");
for (name, server) in servers {
let entry = mcp_tbl
.entry(name)
.or_insert_with(toml_edit::table)
.as_table_mut()
.expect("mcp_servers entry is a table");
entry.set_implicit(false);
match &server.transport_type {
JSONTransportType::CLIServer {
command,
args,
env,
working_directory,
} => {
entry["command"] = toml_edit::value(command.as_str());
if !args.is_empty() {
let mut arr = toml_edit::Array::new();
for arg in args {
arr.push(arg.as_str());
}
entry["args"] = toml_edit::value(arr);
}
if !env.is_empty() {
let mut env_tbl = toml_edit::InlineTable::new();
for (k, v) in env {
env_tbl.insert(k, v.as_str().into());
}
entry["env"] = toml_edit::value(env_tbl);
}
if let Some(cwd) = working_directory {
entry["cwd"] = toml_edit::value(cwd.as_str());
}
}
JSONTransportType::SSEServer { url, headers } => {
entry["url"] = toml_edit::value(url.as_str());
if !headers.is_empty() {
let mut hdrs_tbl = toml_edit::InlineTable::new();
for (k, v) in headers {
hdrs_tbl.insert(k, v.as_str().into());
}
entry["http_headers"] = toml_edit::value(hdrs_tbl);
}
}
}
}
}
#[cfg(test)]
#[path = "codex_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/harness/codex_tests.rs
+849
View File
@@ -0,0 +1,849 @@
use std::collections::HashMap;
use std::ffi::OsString;
use std::fs;
use std::sync::Arc;
use serde_json::Value;
use tempfile::TempDir;
use uuid::Uuid;
use super::super::codex_transcript::CodexTranscriptEnvelope;
use super::*;
use crate::ai::agent::conversation::AIConversationId;
use crate::server::server_api::harness_support::MockHarnessSupportClient;
#[test]
fn prepare_codex_auth_writes_fresh_file_with_api_key_mode() {
let tmp = TempDir::new().unwrap();
let auth_path = tmp.path().join(".codex/auth.json");
prepare_codex_auth(&auth_path, "sk-test-key").unwrap();
let auth: Value = serde_json::from_slice(&fs::read(&auth_path).unwrap()).unwrap();
assert_eq!(auth["OPENAI_API_KEY"], "sk-test-key");
assert_eq!(auth["auth_mode"], "apikey");
}
#[test]
fn prepare_codex_auth_preserves_unrelated_fields() {
let tmp = TempDir::new().unwrap();
let auth_path = tmp.path().join("auth.json");
fs::write(
&auth_path,
r#"{"tokens":{"access_token":"tok"},"last_refresh":"2026-01-01T00:00:00Z"}"#,
)
.unwrap();
prepare_codex_auth(&auth_path, "sk-new-key").unwrap();
let auth: Value = serde_json::from_slice(&fs::read(&auth_path).unwrap()).unwrap();
assert_eq!(auth["OPENAI_API_KEY"], "sk-new-key");
assert_eq!(auth["auth_mode"], "apikey");
assert_eq!(auth["tokens"]["access_token"], "tok");
assert_eq!(auth["last_refresh"], "2026-01-01T00:00:00Z");
}
#[test]
fn prepare_codex_auth_does_not_overwrite_existing_auth_mode() {
let tmp = TempDir::new().unwrap();
let auth_path = tmp.path().join("auth.json");
fs::write(&auth_path, r#"{"auth_mode":"Chatgpt"}"#).unwrap();
prepare_codex_auth(&auth_path, "sk-new-key").unwrap();
let auth: Value = serde_json::from_slice(&fs::read(&auth_path).unwrap()).unwrap();
assert_eq!(auth["auth_mode"], "Chatgpt");
assert_eq!(auth["OPENAI_API_KEY"], "sk-new-key");
}
#[test]
fn prepare_codex_auth_overwrites_stale_openai_api_key() {
let tmp = TempDir::new().unwrap();
let auth_path = tmp.path().join("auth.json");
fs::write(
&auth_path,
r#"{"auth_mode":"apikey","OPENAI_API_KEY":"sk-old"}"#,
)
.unwrap();
prepare_codex_auth(&auth_path, "sk-new").unwrap();
let auth: Value = serde_json::from_slice(&fs::read(&auth_path).unwrap()).unwrap();
assert_eq!(auth["OPENAI_API_KEY"], "sk-new");
}
#[cfg(unix)]
#[test]
fn prepare_codex_auth_writes_with_0600_perms() {
use std::os::unix::fs::PermissionsExt;
let tmp = TempDir::new().unwrap();
let auth_path = tmp.path().join(".codex/auth.json");
prepare_codex_auth(&auth_path, "sk-test-key").unwrap();
let mode = fs::metadata(&auth_path).unwrap().permissions().mode() & 0o777;
assert_eq!(mode, 0o600);
}
#[test]
fn resolve_openai_api_key_returns_value_from_resolved_map() {
let resolved = HashMap::from([(
OsString::from("OPENAI_API_KEY"),
OsString::from("sk-from-secret"),
)]);
assert_eq!(
resolve_openai_api_key(&resolved).as_deref(),
Some("sk-from-secret")
);
}
#[test]
#[serial_test::serial]
fn resolve_openai_api_key_falls_back_to_env_var() {
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::set_var(OPENAI_API_KEY_ENV, "sk-from-env");
let result = resolve_openai_api_key(&HashMap::new());
match prev {
Some(v) => std::env::set_var(OPENAI_API_KEY_ENV, v),
None => std::env::remove_var(OPENAI_API_KEY_ENV),
}
assert_eq!(result.as_deref(), Some("sk-from-env"));
}
#[test]
#[serial_test::serial]
fn resolve_openai_api_key_returns_none_when_map_and_env_empty() {
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::remove_var(OPENAI_API_KEY_ENV);
let result = resolve_openai_api_key(&HashMap::new());
if let Some(v) = prev {
std::env::set_var(OPENAI_API_KEY_ENV, v);
}
assert_eq!(result, None);
}
#[test]
#[serial_test::serial]
fn resolve_openai_api_key_prefers_env_over_resolved_map() {
// Worker-injected env var wins over the resolved secret map because
// build_secret_env_vars skips secrets that collide with process env.
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::set_var(OPENAI_API_KEY_ENV, "sk-from-env");
let resolved = HashMap::from([(
OsString::from("OPENAI_API_KEY"),
OsString::from("sk-from-secret"),
)]);
let result = resolve_openai_api_key(&resolved);
match prev {
Some(v) => std::env::set_var(OPENAI_API_KEY_ENV, v),
None => std::env::remove_var(OPENAI_API_KEY_ENV),
}
assert_eq!(result.as_deref(), Some("sk-from-env"));
}
#[test]
#[serial_test::serial]
fn resolve_openai_api_key_uses_resolved_map_when_env_empty() {
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::set_var(OPENAI_API_KEY_ENV, " ");
let resolved = HashMap::from([(
OsString::from("OPENAI_API_KEY"),
OsString::from("sk-from-secret"),
)]);
let result = resolve_openai_api_key(&resolved);
match prev {
Some(v) => std::env::set_var(OPENAI_API_KEY_ENV, v),
None => std::env::remove_var(OPENAI_API_KEY_ENV),
}
assert_eq!(result.as_deref(), Some("sk-from-secret"));
}
#[test]
#[serial_test::serial]
fn prepare_codex_environment_config_honors_codex_home() {
let tmp = TempDir::new().unwrap();
let codex_home = tmp.path().join("codex-home");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
let prev_codex_home = std::env::var(CODEX_HOME_ENV).ok();
let prev_openai_api_key = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::set_var(CODEX_HOME_ENV, &codex_home);
std::env::remove_var(OPENAI_API_KEY_ENV);
let resolved = HashMap::from([(
OsString::from(OPENAI_API_KEY_ENV),
OsString::from("sk-from-secret"),
)]);
let model_config = harness_model_config("gpt-5.5", None);
let result = prepare_codex_environment_config(
&working_dir,
Some("system prompt"),
&resolved,
&HashMap::new(),
&HashMap::new(),
Some(&model_config),
);
match prev_codex_home {
Some(v) => std::env::set_var(CODEX_HOME_ENV, v),
None => std::env::remove_var(CODEX_HOME_ENV),
}
match prev_openai_api_key {
Some(v) => std::env::set_var(OPENAI_API_KEY_ENV, v),
None => std::env::remove_var(OPENAI_API_KEY_ENV),
}
result.unwrap();
assert_eq!(
fs::read_to_string(codex_home.join(CODEX_AGENTS_OVERRIDE_FILE_NAME)).unwrap(),
"system prompt"
);
let auth: Value =
serde_json::from_slice(&fs::read(codex_home.join(CODEX_AUTH_FILE_NAME)).unwrap()).unwrap();
assert_eq!(auth["OPENAI_API_KEY"], "sk-from-secret");
let cfg = read_codex_config(&codex_home.join(CODEX_CONFIG_TOML_FILE_NAME));
assert_eq!(cfg["model"].as_str(), Some("gpt-5.5"));
assert!(!cfg.contains_key("openai_base_url"));
assert!(!tmp.path().join(CODEX_CONFIG_DIR).exists());
}
fn read_codex_config(path: &std::path::Path) -> toml::Table {
let content = fs::read_to_string(path).unwrap();
toml::from_str(&content).unwrap()
}
fn harness_model_config(model_id: &str, reasoning_level: Option<&str>) -> HarnessModelConfig {
HarnessModelConfig {
model_id: model_id.to_string(),
reasoning_level: reasoning_level.map(str::to_string),
}
}
#[test]
fn prepare_codex_config_toml_writes_fresh_config() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join(".codex/config.toml");
let working_dir = tmp.path().join("workspace/proj");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let canonical = working_dir.canonicalize().unwrap();
let key = canonical.to_string_lossy().into_owned();
let cfg = read_codex_config(&config_path);
assert_eq!(cfg["check_for_update_on_startup"].as_bool(), Some(false));
assert_eq!(
cfg["projects"][&key]["trust_level"].as_str(),
Some("trusted")
);
}
#[test]
fn prepare_codex_config_toml_preserves_unrelated_keys() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
fs::write(
&config_path,
"model = \"gpt-5\"\n\n[projects.\"/other/path\"]\ntrust_level = \"trusted\"\n",
)
.unwrap();
// Pass `None` — the `model` key is intentionally removed (managed
// key), but unrelated keys like existing project entries are kept.
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let canonical = working_dir.canonicalize().unwrap();
let key = canonical.to_string_lossy().into_owned();
let cfg = read_codex_config(&config_path);
// `model` is a managed key — removed when no override is provided.
assert!(!cfg.contains_key("model"));
assert_eq!(
cfg["projects"]["/other/path"]["trust_level"].as_str(),
Some("trusted")
);
assert_eq!(
cfg["projects"][&key]["trust_level"].as_str(),
Some("trusted")
);
}
#[test]
fn prepare_codex_config_toml_is_idempotent() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let after_first = fs::read_to_string(&config_path).unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let after_second = fs::read_to_string(&config_path).unwrap();
assert_eq!(after_first, after_second);
let canonical = working_dir.canonicalize().unwrap();
let key = canonical.to_string_lossy().into_owned();
let cfg: toml::Table = toml::from_str(&after_second).unwrap();
let projects = cfg["projects"].as_table().unwrap();
assert_eq!(projects.len(), 1);
assert_eq!(projects[&key]["trust_level"].as_str(), Some("trusted"));
}
#[test]
fn prepare_codex_config_toml_upgrades_untrusted_entry() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
let canonical = working_dir.canonicalize().unwrap();
let key = canonical.to_string_lossy().into_owned();
// Use a TOML literal-string key ('...') so Windows backslashes in `key`
// (e.g. `\\?\C:\...`) are not interpreted as escape sequences.
fs::write(
&config_path,
format!("[projects.'{key}']\ntrust_level = \"untrusted\"\n"),
)
.unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let cfg = read_codex_config(&config_path);
assert_eq!(
cfg["projects"][&key]["trust_level"].as_str(),
Some("trusted")
);
}
#[test]
fn prepare_codex_config_toml_trusts_multiple_child_repos() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
let repo_a = working_dir.join("a");
let repo_b = working_dir.join("b");
fs::create_dir_all(repo_a.join(".git")).unwrap();
fs::create_dir_all(repo_b.join(".git")).unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let cfg = read_codex_config(&config_path);
let projects = cfg["projects"].as_table().unwrap();
let canonical_a = repo_a.canonicalize().unwrap();
let canonical_b = repo_b.canonicalize().unwrap();
assert_eq!(
projects[canonical_a.to_str().unwrap()]["trust_level"].as_str(),
Some("trusted")
);
assert_eq!(
projects[canonical_b.to_str().unwrap()]["trust_level"].as_str(),
Some("trusted")
);
}
#[test]
fn prepare_codex_config_toml_overwrites_stale_openai_base_url() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
fs::write(
&config_path,
"openai_base_url = \"https://api.openai.com/v1\"\n",
)
.unwrap();
prepare_codex_config_toml(
&config_path,
&working_dir,
&HashMap::new(),
None,
Some("https://custom.api.openai.com/v1"),
)
.unwrap();
let cfg = read_codex_config(&config_path);
assert_eq!(
cfg["openai_base_url"].as_str(),
Some("https://custom.api.openai.com/v1")
);
}
#[test]
fn write_codex_mcp_servers_cli_server() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
let servers = HashMap::from([(
"my-mcp".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::CLIServer {
command: "npx".to_string(),
args: vec!["-y".to_string(), "@some/mcp".to_string()],
env: HashMap::from([("TOKEN".to_string(), "abc".to_string())]),
working_directory: None,
},
},
)]);
prepare_codex_config_toml(&config_path, &working_dir, &servers, None, None).unwrap();
let cfg = read_codex_config(&config_path);
let mcp = &cfg["mcp_servers"]["my-mcp"];
assert_eq!(mcp["command"].as_str(), Some("npx"));
let args: Vec<&str> = mcp["args"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_str().unwrap())
.collect();
assert_eq!(args, vec!["-y", "@some/mcp"]);
assert_eq!(mcp["env"]["TOKEN"].as_str(), Some("abc"));
}
#[test]
fn write_codex_mcp_servers_sse_server() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
let servers = HashMap::from([(
"remote-mcp".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::SSEServer {
url: "https://mcp.example.com/sse".to_string(),
headers: HashMap::from([("X-Key".to_string(), "val".to_string())]),
},
},
)]);
prepare_codex_config_toml(&config_path, &working_dir, &servers, None, None).unwrap();
let cfg = read_codex_config(&config_path);
let mcp = &cfg["mcp_servers"]["remote-mcp"];
assert_eq!(mcp["url"].as_str(), Some("https://mcp.example.com/sse"));
assert_eq!(mcp["http_headers"]["X-Key"].as_str(), Some("val"));
}
#[test]
fn write_codex_mcp_servers_cli_server_with_cwd() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
let servers = HashMap::from([(
"my-mcp".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::CLIServer {
command: "node".to_string(),
args: vec!["server.js".to_string()],
env: HashMap::new(),
working_directory: Some("/opt/mcp-server".to_string()),
},
},
)]);
prepare_codex_config_toml(&config_path, &working_dir, &servers, None, None).unwrap();
let cfg = read_codex_config(&config_path);
let mcp = &cfg["mcp_servers"]["my-mcp"];
assert_eq!(mcp["command"].as_str(), Some("node"));
assert_eq!(mcp["cwd"].as_str(), Some("/opt/mcp-server"));
}
#[test]
fn write_codex_mcp_servers_cli_server_without_cwd_omits_key() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
let servers = HashMap::from([(
"my-mcp".to_string(),
JSONMCPServer {
transport_type: JSONTransportType::CLIServer {
command: "npx".to_string(),
args: vec![],
env: HashMap::new(),
working_directory: None,
},
},
)]);
prepare_codex_config_toml(&config_path, &working_dir, &servers, None, None).unwrap();
let cfg = read_codex_config(&config_path);
let mcp = &cfg["mcp_servers"]["my-mcp"];
assert!(mcp.get("cwd").is_none());
}
#[test]
fn prepare_codex_config_toml_writes_model_when_specified() {
// A non-default model id is written to the top-level `model` key so Codex pins it
// for new sessions launched from this `~/.codex/config.toml`. Even for the
// current target model, we stamp a self-referential migration entry so the
// upgrade prompt is suppressed regardless of what the user selected.
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(
&config_path,
&working_dir,
&HashMap::new(),
Some(&harness_model_config("gpt-5.5", None)),
None,
)
.unwrap();
let cfg = read_codex_config(&config_path);
assert_eq!(cfg["model"].as_str(), Some("gpt-5.5"));
assert_eq!(
cfg["notice"]["model_migrations"]["gpt-5.5"].as_str(),
Some(CODEX_MODEL_MIGRATIONS_TARGET),
);
}
#[test]
fn prepare_codex_config_toml_writes_model_migration_for_older_model() {
// For an older model id, the migration entry maps it to the current target
// so Codex's "choose a newer model" prompt is suppressed at session launch.
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(
&config_path,
&working_dir,
&HashMap::new(),
Some(&harness_model_config("gpt-5.2", None)),
None,
)
.unwrap();
let cfg = read_codex_config(&config_path);
assert_eq!(cfg["model"].as_str(), Some("gpt-5.2"));
assert_eq!(
cfg["notice"]["model_migrations"]["gpt-5.2"].as_str(),
Some(CODEX_MODEL_MIGRATIONS_TARGET),
);
}
#[test]
fn prepare_codex_config_toml_skips_model_for_default_sentinel() {
// The literal "default" sentinel means "let Codex pick its own default model";
// we should NOT write a `model` key (or a migration entry) in that case.
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(
&config_path,
&working_dir,
&HashMap::new(),
Some(&harness_model_config("default", None)),
None,
)
.unwrap();
let cfg = read_codex_config(&config_path);
assert!(
cfg.get("model").is_none(),
"`model` should not be written for the default sentinel"
);
assert!(
cfg.get("notice").is_none(),
"`[notice]` table should not be written without a pinned model id"
);
}
#[test]
fn prepare_codex_config_toml_skips_model_when_none() {
// No model id supplied means the user didn't pick one; we should not write a
// `model` key or any `[notice.model_migrations]` entries.
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let cfg = read_codex_config(&config_path);
assert!(
cfg.get("model").is_none(),
"`model` should not be written when no override is supplied"
);
assert!(
cfg.get("notice").is_none(),
"`[notice]` table should not be written without a pinned model id"
);
}
#[test]
fn prepare_codex_config_toml_writes_model_reasoning_effort_when_specified() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
prepare_codex_config_toml(
&config_path,
&working_dir,
&HashMap::new(),
Some(&harness_model_config("gpt-5.5", Some("medium"))),
None,
)
.unwrap();
let cfg = read_codex_config(&config_path);
assert_eq!(cfg["model"].as_str(), Some("gpt-5.5"));
assert_eq!(cfg["model_reasoning_effort"].as_str(), Some("medium"));
}
#[test]
fn prepare_codex_config_toml_removes_stale_model_reasoning_effort_when_none() {
let tmp = TempDir::new().unwrap();
let config_path = tmp.path().join("config.toml");
let working_dir = tmp.path().join("workspace");
fs::create_dir_all(&working_dir).unwrap();
fs::write(&config_path, "model_reasoning_effort = \"high\"\n").unwrap();
prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None, None).unwrap();
let cfg = read_codex_config(&config_path);
assert!(cfg.get("model_reasoning_effort").is_none());
}
#[test]
fn find_child_git_repos_returns_only_repo_children() {
let tmp = TempDir::new().unwrap();
let workspace = tmp.path().join("workspace");
let repo = workspace.join("repo");
let other = workspace.join("other");
fs::create_dir_all(repo.join(".git")).unwrap();
fs::create_dir_all(&other).unwrap();
let found = find_child_git_repos(&workspace);
let canonical_repo = repo.canonicalize().unwrap();
assert_eq!(found.len(), 1);
assert_eq!(found[0].canonicalize().unwrap(), canonical_repo);
}
#[test]
fn find_child_git_repos_returns_empty_when_dir_missing() {
let tmp = TempDir::new().unwrap();
let missing = tmp.path().join("does-not-exist");
assert!(find_child_git_repos(&missing).is_empty());
}
#[test]
fn codex_command_with_session_id_invokes_resume_subcommand() {
let uuid = Uuid::new_v4();
let cmd = codex_command("codex", Some(&uuid), "/tmp/prompt.txt");
assert!(
cmd.contains(&format!(
"resume --dangerously-bypass-approvals-and-sandbox --dangerously-bypass-hook-trust {uuid}"
)),
"resume command should pass UUID to `resume`: {cmd}"
);
assert!(
cmd.contains("\"$(cat '/tmp/prompt.txt')\""),
"resume command should pipe prompt: {cmd}"
);
}
#[test]
fn codex_command_without_session_id_bypasses_hook_trust() {
let cmd = codex_command("codex", None, "/tmp/prompt.txt");
assert!(
cmd.contains("--dangerously-bypass-approvals-and-sandbox"),
"command should bypass approvals and sandbox: {cmd}"
);
assert!(
cmd.contains("--dangerously-bypass-hook-trust"),
"command should bypass hook trust for driver-installed hooks: {cmd}"
);
assert!(
cmd.contains("\"$(cat '/tmp/prompt.txt')\""),
"command should pipe prompt: {cmd}"
);
}
#[tokio::test]
async fn fetch_resume_payload_maps_404_to_resume_state_missing() {
let mut mock = MockHarnessSupportClient::new();
mock.expect_fetch_transcript()
.returning(|| Err(anyhow::anyhow!("upstream returned status 404")));
let conversation_id = AIConversationId::new();
let result = CodexHarness
.fetch_resume_payload(&conversation_id, Arc::new(mock))
.await;
match result {
Err(AgentDriverError::ConversationResumeStateMissing { harness, .. }) => {
assert_eq!(harness, "codex");
}
other => panic!("expected ConversationResumeStateMissing, got {other:?}"),
}
}
#[tokio::test]
async fn fetch_resume_payload_maps_other_errors_to_load_failed() {
let mut mock = MockHarnessSupportClient::new();
mock.expect_fetch_transcript()
.returning(|| Err(anyhow::anyhow!("connection reset")));
let conversation_id = AIConversationId::new();
let result = CodexHarness
.fetch_resume_payload(&conversation_id, Arc::new(mock))
.await;
assert!(
matches!(result, Err(AgentDriverError::ConversationLoadFailed(_))),
"expected ConversationLoadFailed, got {result:?}"
);
}
#[test]
#[serial_test::serial]
fn resolve_openai_base_url_from_secret_returns_base_url_when_typed_secret_active() {
// When the typed OpenAI secret is the active API key source, the base URL
// should be extracted from the structured secret.
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::remove_var(OPENAI_API_KEY_ENV);
let secrets = HashMap::from([(
"openai-key".to_string(),
ManagedSecretValue::openai_api_key(
"sk-test",
Some("https://us.api.openai.com/v1".to_string()),
),
)]);
let resolved_env =
HashMap::from([(OsString::from("OPENAI_API_KEY"), OsString::from("sk-test"))]);
let result = resolve_openai_base_url_from_secret(&secrets, &resolved_env);
if let Some(v) = prev {
std::env::set_var(OPENAI_API_KEY_ENV, v);
}
assert_eq!(result.as_deref(), Some("https://us.api.openai.com/v1"));
}
#[test]
#[serial_test::serial]
fn resolve_openai_base_url_from_secret_returns_none_when_worker_env_wins() {
// When a worker-injected OPENAI_API_KEY already exists in process env,
// the typed-secret base_url should NOT be applied.
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::set_var(OPENAI_API_KEY_ENV, "sk-worker-key");
let secrets = HashMap::from([(
"openai-key".to_string(),
ManagedSecretValue::openai_api_key(
"sk-secret",
Some("https://us.api.openai.com/v1".to_string()),
),
)]);
let resolved_env = HashMap::new();
let result = resolve_openai_base_url_from_secret(&secrets, &resolved_env);
match prev {
Some(v) => std::env::set_var(OPENAI_API_KEY_ENV, v),
None => std::env::remove_var(OPENAI_API_KEY_ENV),
}
assert_eq!(result, None);
}
#[test]
#[serial_test::serial]
fn resolve_openai_base_url_from_secret_returns_none_when_no_base_url() {
// When the typed OpenAI secret has no base_url, None is returned.
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::remove_var(OPENAI_API_KEY_ENV);
let secrets = HashMap::from([(
"openai-key".to_string(),
ManagedSecretValue::openai_api_key("sk-test", None),
)]);
let resolved_env =
HashMap::from([(OsString::from("OPENAI_API_KEY"), OsString::from("sk-test"))]);
let result = resolve_openai_base_url_from_secret(&secrets, &resolved_env);
if let Some(v) = prev {
std::env::set_var(OPENAI_API_KEY_ENV, v);
}
assert_eq!(result, None);
}
#[test]
#[serial_test::serial]
fn resolve_openai_base_url_from_secret_returns_none_when_api_key_not_in_resolved() {
// When OPENAI_API_KEY is not in the resolved env vars (e.g. the secret was
// skipped due to collision), the base URL should not be applied.
let prev = std::env::var(OPENAI_API_KEY_ENV).ok();
std::env::remove_var(OPENAI_API_KEY_ENV);
let secrets = HashMap::from([(
"openai-key".to_string(),
ManagedSecretValue::openai_api_key(
"sk-test",
Some("https://us.api.openai.com/v1".to_string()),
),
)]);
let resolved_env = HashMap::new();
let result = resolve_openai_base_url_from_secret(&secrets, &resolved_env);
if let Some(v) = prev {
std::env::set_var(OPENAI_API_KEY_ENV, v);
}
assert_eq!(result, None);
}
#[tokio::test]
async fn fetch_resume_payload_returns_codex_variant_on_success() {
let uuid = Uuid::new_v4();
let envelope = CodexTranscriptEnvelope {
cwd: "/cloud/work".into(),
session_id: uuid,
codex_version: Some("0.55.0".to_string()),
session_start_timestamp: None,
entries: vec![serde_json::json!({"type": "event_msg"})],
};
let bytes = serde_json::to_vec(&envelope).unwrap();
let mut mock = MockHarnessSupportClient::new();
mock.expect_fetch_transcript()
.returning(move || Ok(bytes::Bytes::from(bytes.clone())));
let conversation_id = AIConversationId::new();
let payload = CodexHarness
.fetch_resume_payload(&conversation_id, Arc::new(mock))
.await
.unwrap()
.unwrap();
match payload {
ResumePayload::Codex(info) => {
assert_eq!(info.session_id, uuid);
assert_eq!(info.conversation_id, conversation_id);
assert_eq!(info.envelope.codex_version.as_deref(), Some("0.55.0"));
}
other => panic!("expected ResumePayload::Codex, got {other:?}"),
}
}
app/src/ai/agent_sdk/driver/harness/codex_transcript.rs
+248
View File
@@ -0,0 +1,248 @@
//! Codex session transcript envelope + rehydration helpers.
//!
//! Owns:
//! - [`CodexTranscriptEnvelope`] — the on-wire/on-GCS shape of a saved Codex rollout
//! (parsed JSONL entries plus session-level metadata). Reader/writer functions
//! interoperate with Codex's own `~/.codex/sessions/YYYY/MM/DD/rollout-<ts>-<uuid>.jsonl`
//! layout (codex `rollout/src/recorder.rs`).
//! - [`CodexResumeInfo`] — everything the harness runner needs to resume an existing
//! Codex conversation: the Warp server conversation id to reuse, the codex session
//! uuid (`ThreadId`) to pass to `codex resume`, and the decoded envelope to rehydrate
//! onto disk.
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use chrono::{DateTime, Datelike, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use super::json_utils::entries_to_jsonl;
use crate::ai::agent::conversation::AIConversationId;
/// Env var codex honors to override `~/.codex` (see codex `core/src/config/mod.rs`).
const CODEX_HOME_ENV: &str = "CODEX_HOME";
const CODEX_HOME_DIRNAME: &str = ".codex";
/// Subdirectory under `$CODEX_HOME` where rollouts live.
const CODEX_SESSIONS_SUBDIR: &str = "sessions";
/// JSON envelope sent to the server representing a complete Codex session.
///
/// The transcript is the parsed JSONL content of the rollout file; codex's resume
/// path re-reads this JSONL line by line.
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub(crate) struct CodexTranscriptEnvelope {
/// The directory the codex session started in (recovered from the `SessionMeta` line).
pub(crate) cwd: PathBuf,
/// Codex session/thread UUID. Matches the trailing `-<uuid>` in the rollout filename.
pub(crate) session_id: Uuid,
/// `cli_version` from `SessionMeta`, surfaced separately for the server.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) codex_version: Option<String>,
/// Timestamp from the `SessionMeta` line, used to derive the YYYY/MM/DD directory
/// path when writing the rollout file back to disk.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) session_start_timestamp: Option<DateTime<Utc>>,
/// Parsed JSONL entries.
pub(crate) entries: Vec<Value>,
}
impl CodexTranscriptEnvelope {
pub(crate) fn new(session_id: Uuid, meta: CodexSessionMetadata, entries: Vec<Value>) -> Self {
Self {
cwd: meta.cwd,
session_id,
codex_version: meta.codex_version,
session_start_timestamp: meta.session_start_timestamp,
entries,
}
}
}
/// Session-level metadata pulled from the rollout's `SessionMeta` line.
#[derive(Clone, Debug, Default, PartialEq)]
pub(crate) struct CodexSessionMetadata {
pub(crate) cwd: PathBuf,
pub(crate) codex_version: Option<String>,
pub(crate) session_start_timestamp: Option<DateTime<Utc>>,
}
/// Everything needed to resume an existing Codex conversation.
///
/// Built from a `--conversation` id after the client fetches the stored envelope from
/// the server. Passed into `CodexHarnessRunner::new` so the runner reuses the existing
/// session and server conversation ids instead of minting fresh ones.
#[derive(Debug)]
pub(crate) struct CodexResumeInfo {
/// Warp server-side conversation id. Reused so subsequent transcript/block-snapshot
/// uploads overwrite the same GCS objects.
pub(crate) conversation_id: AIConversationId,
/// Codex session uuid passed to `codex resume <session_id>`. Matches `envelope.session_id`.
pub(crate) session_id: Uuid,
/// Envelope fetched from the server, written back to disk before launching codex.
pub(crate) envelope: CodexTranscriptEnvelope,
}
#[derive(Debug)]
pub(crate) struct CodexLocalContinuation {
pub(crate) command: String,
pub(crate) transcript_path: PathBuf,
}
/// Resolve the codex sessions root, honoring `$CODEX_HOME` then falling back to `~/.codex`.
pub(crate) fn codex_sessions_root() -> anyhow::Result<PathBuf> {
let home = if let Ok(dir) = std::env::var(CODEX_HOME_ENV) {
PathBuf::from(dir)
} else {
dirs::home_dir()
.ok_or_else(|| anyhow::anyhow!("could not determine home directory"))?
.join(CODEX_HOME_DIRNAME)
};
Ok(home.join(CODEX_SESSIONS_SUBDIR))
}
/// Walk `<sessions_root>/YYYY/MM/DD/` looking for a `rollout-*-<session_id>.jsonl`.
///
/// Returns `None` if `sessions_root` doesn't exist yet or no matching file is found.
pub(crate) fn find_session_file(sessions_root: &Path, session_id: Uuid) -> Option<PathBuf> {
if !sessions_root.exists() {
return None;
}
let suffix = format!("-{session_id}.jsonl");
for year_dir in read_subdirs(sessions_root) {
for month_dir in read_subdirs(&year_dir) {
for day_dir in read_subdirs(&month_dir) {
let entries = match fs::read_dir(&day_dir) {
Ok(e) => e,
Err(_) => continue,
};
for entry in entries.flatten() {
let path = entry.path();
let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
continue;
};
if name.starts_with("rollout-") && name.ends_with(&suffix) {
return Some(path);
}
}
}
}
}
None
}
fn read_subdirs(parent: &Path) -> impl Iterator<Item = PathBuf> {
fs::read_dir(parent)
.into_iter()
.flatten()
.filter_map(|entry| {
let entry = entry.ok()?;
entry.file_type().ok()?.is_dir().then(|| entry.path())
})
}
/// Pull `cwd` and `cli_version` out of the first JSONL line if it's a `SessionMeta`.
pub(crate) fn parse_session_meta(first: Option<&Value>) -> Option<CodexSessionMetadata> {
let entry = first?;
if entry.get("type").and_then(|v| v.as_str()) != Some("session_meta") {
return None;
}
let payload = entry.get("payload")?;
let cwd = PathBuf::from(payload.get("cwd").and_then(|v| v.as_str())?);
let codex_version = payload
.get("cli_version")
.and_then(|v| v.as_str())
.map(str::to_owned);
let session_start_timestamp = payload
.get("timestamp")
.and_then(|v| v.as_str())
.and_then(|s| DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc));
Some(CodexSessionMetadata {
cwd,
codex_version,
session_start_timestamp,
})
}
/// Write `envelope` back under `<sessions_root>/YYYY/MM/DD/rollout-<ts>-<uuid>.jsonl`.
///
/// YYYY/MM/DD and `<ts>` come from `envelope.session_start_timestamp`. Falls back to
/// today's UTC date if absent — codex's lookup is by UUID so the precise path doesn't
/// matter for resume to work.
pub(crate) fn write_envelope(
envelope: &CodexTranscriptEnvelope,
sessions_root: &Path,
) -> Result<PathBuf> {
let timestamp = envelope.session_start_timestamp.unwrap_or_else(Utc::now);
let day_dir = sessions_root
.join(format!("{:04}", timestamp.year()))
.join(format!("{:02}", timestamp.month()))
.join(format!("{:02}", timestamp.day()));
fs::create_dir_all(&day_dir)
.with_context(|| format!("Failed to create {}", day_dir.display()))?;
// Codex's filename format: `[year]-[month]-[day]T[hour]-[minute]-[second]`
// (codex `rollout/src/recorder.rs::precompute_log_file_info`).
let date_str = timestamp.format("%Y-%m-%dT%H-%M-%S").to_string();
let file_path = day_dir.join(format!(
"rollout-{date_str}-{session_id}.jsonl",
session_id = envelope.session_id
));
fs::write(&file_path, entries_to_jsonl(&envelope.entries)?)
.with_context(|| format!("Failed to write {}", file_path.display()))?;
Ok(file_path)
}
pub(crate) fn rehydrate_codex_transcript(
envelope: &mut CodexTranscriptEnvelope,
local_cwd: &Path,
) -> Result<CodexLocalContinuation> {
envelope.cwd = local_cwd.to_path_buf();
if let Some(Value::Object(entry)) = envelope.entries.first_mut() {
if entry.get("type").and_then(|value| value.as_str()) == Some("session_meta") {
if let Some(Value::Object(payload)) = entry.get_mut("payload") {
payload.insert(
"cwd".to_string(),
Value::String(local_cwd.to_string_lossy().to_string()),
);
}
}
}
let session_id = envelope.session_id;
let sessions_root = codex_sessions_root().context("Failed to resolve codex sessions root")?;
let transcript_path =
write_envelope(envelope, &sessions_root).context("Failed to rehydrate codex transcript")?;
Ok(CodexLocalContinuation {
command: format!("codex resume {session_id}"),
transcript_path,
})
}
/// Rehydrate a Codex transcript downloaded from a remote cloud run for local continuation.
///
/// Unlike [`rehydrate_codex_transcript`] (used by the cloud resume harness runner), this
/// function does **not** mutate the envelope's `cwd` field or patch the `session_meta` payload
/// — the remote session's working directory is preserved as-is in the transcript.
pub(crate) fn rehydrate_codex_transcript_from_reader(
reader: impl Read,
) -> Result<CodexLocalContinuation> {
let envelope: CodexTranscriptEnvelope =
serde_json::from_reader(reader).context("Failed to parse codex transcript envelope")?;
let session_id = envelope.session_id;
let sessions_root = codex_sessions_root().context("Failed to resolve codex sessions root")?;
// Write as-is: no cwd mutation, no session_meta patch.
let transcript_path = write_envelope(&envelope, &sessions_root)
.context("Failed to rehydrate codex transcript")?;
Ok(CodexLocalContinuation {
command: format!("codex resume {session_id}"),
transcript_path,
})
}
#[cfg(test)]
#[path = "codex_transcript_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/harness/codex_transcript_tests.rs
+205
View File
@@ -0,0 +1,205 @@
use std::fs;
use std::path::Path;
use anyhow::Result;
use tempfile::TempDir;
use uuid::Uuid;
use super::super::claude_transcript::read_jsonl;
use super::*;
/// Walk `sessions_root` for `session_id`'s rollout and assemble an envelope.
fn read_envelope(
session_id: Uuid,
sessions_root: &Path,
) -> Result<Option<CodexTranscriptEnvelope>> {
let Some(path) = find_session_file(sessions_root, session_id) else {
return Ok(None);
};
let entries = read_jsonl(&path)?;
let meta = parse_session_meta(entries.first()).unwrap_or_default();
Ok(Some(CodexTranscriptEnvelope::new(
session_id, meta, entries,
)))
}
/// Minimal SessionMeta line in the same shape codex writes (codex
/// `protocol/src/protocol.rs::RolloutItem`): `{type, payload}`.
fn session_meta_line(uuid: Uuid, cwd: &str, timestamp: &str, cli_version: &str) -> String {
serde_json::json!({
"type": "session_meta",
"payload": {
"id": uuid.to_string(),
"timestamp": timestamp,
"cwd": cwd,
"originator": "test",
"cli_version": cli_version,
},
})
.to_string()
}
#[test]
#[serial_test::serial]
fn codex_sessions_root_honors_codex_home_env() {
let tmp = TempDir::new().unwrap();
let prev = std::env::var(CODEX_HOME_ENV).ok();
std::env::set_var(CODEX_HOME_ENV, tmp.path());
let root = codex_sessions_root().unwrap();
match prev {
Some(v) => std::env::set_var(CODEX_HOME_ENV, v),
None => std::env::remove_var(CODEX_HOME_ENV),
}
assert_eq!(root, tmp.path().join(CODEX_SESSIONS_SUBDIR));
}
#[test]
fn find_session_file_walks_yyyy_mm_dd_tree() {
let tmp = TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let day = tmp.path().join("2026").join("04").join("30");
fs::create_dir_all(&day).unwrap();
let file = day.join(format!("rollout-ignored-ts-{uuid}.jsonl"));
fs::write(&file, "").unwrap();
let found = find_session_file(tmp.path(), uuid);
assert_eq!(found, Some(file));
}
#[test]
fn find_session_file_returns_none_when_no_match() {
let tmp = TempDir::new().unwrap();
let day = tmp.path().join("2026").join("04").join("30");
fs::create_dir_all(&day).unwrap();
fs::write(
day.join(format!("rollout-ignored-ts-{}.jsonl", Uuid::new_v4())),
"",
)
.unwrap();
assert!(find_session_file(tmp.path(), Uuid::new_v4()).is_none());
}
#[test]
fn find_session_file_returns_none_when_root_missing() {
let tmp = TempDir::new().unwrap();
assert!(find_session_file(&tmp.path().join("missing"), Uuid::new_v4()).is_none());
}
#[test]
fn read_envelope_recovers_cwd_and_version_from_session_meta() {
let tmp = TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let day = tmp.path().join("2026").join("04").join("30");
fs::create_dir_all(&day).unwrap();
let meta = session_meta_line(uuid, "/work/proj", "2026-04-30T01:54:20.000Z", "0.55.0");
let body = format!("{meta}\n{{\"type\":\"event_msg\",\"payload\":{{\"x\":1}}}}\n");
fs::write(day.join(format!("rollout-ignored-ts-{uuid}.jsonl")), body).unwrap();
let envelope = read_envelope(uuid, tmp.path()).unwrap().unwrap();
assert_eq!(envelope.session_id, uuid);
assert_eq!(envelope.cwd, std::path::PathBuf::from("/work/proj"));
assert_eq!(envelope.codex_version.as_deref(), Some("0.55.0"));
assert_eq!(envelope.entries.len(), 2);
}
#[test]
fn read_envelope_returns_none_when_missing() {
let tmp = TempDir::new().unwrap();
assert!(read_envelope(Uuid::new_v4(), tmp.path()).unwrap().is_none());
}
#[test]
fn write_envelope_uses_session_meta_timestamp_for_path() {
let tmp = TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let envelope = CodexTranscriptEnvelope {
cwd: "/work".into(),
session_id: uuid,
codex_version: Some("0.55.0".to_string()),
session_start_timestamp: Some(
chrono::DateTime::parse_from_rfc3339("2026-04-30T01:54:20.000Z")
.unwrap()
.with_timezone(&chrono::Utc),
),
entries: vec![
serde_json::from_str::<serde_json::Value>(&session_meta_line(
uuid,
"/work",
"2026-04-30T01:54:20.000Z",
"0.55.0",
))
.unwrap(),
],
};
let path = write_envelope(&envelope, tmp.path()).unwrap();
let expected = tmp
.path()
.join("2026")
.join("04")
.join("30")
.join(format!("rollout-2026-04-30T01-54-20-{uuid}.jsonl"));
assert_eq!(path, expected);
assert!(path.exists());
}
#[test]
fn write_envelope_round_trip_preserves_entries() {
let tmp = TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let entries = vec![
serde_json::from_str::<serde_json::Value>(&session_meta_line(
uuid,
"/work",
"2026-04-30T01:54:20.000Z",
"0.55.0",
))
.unwrap(),
serde_json::json!({"type": "event_msg", "payload": {"x": 1}}),
];
let original = CodexTranscriptEnvelope {
cwd: "/work".into(),
session_id: uuid,
codex_version: Some("0.55.0".to_string()),
session_start_timestamp: Some(
chrono::DateTime::parse_from_rfc3339("2026-04-30T01:54:20.000Z")
.unwrap()
.with_timezone(&chrono::Utc),
),
entries: entries.clone(),
};
write_envelope(&original, tmp.path()).unwrap();
let decoded = read_envelope(uuid, tmp.path()).unwrap().unwrap();
assert_eq!(decoded.session_id, uuid);
assert_eq!(decoded.entries, entries);
// `cwd` and `codex_version` are recovered from the SessionMeta line.
assert_eq!(decoded.cwd, std::path::PathBuf::from("/work"));
assert_eq!(decoded.codex_version.as_deref(), Some("0.55.0"));
}
#[test]
fn write_envelope_falls_back_to_today_when_timestamp_missing() {
let tmp = TempDir::new().unwrap();
let uuid = Uuid::new_v4();
// No SessionMeta first line — just one event entry.
let envelope = CodexTranscriptEnvelope {
cwd: "/work".into(),
session_id: uuid,
codex_version: None,
session_start_timestamp: None,
entries: vec![serde_json::json!({"type": "event_msg"})],
};
let path = write_envelope(&envelope, tmp.path()).unwrap();
// Path should still be under sessions_root and findable by uuid.
assert!(path.starts_with(tmp.path()));
let found = find_session_file(tmp.path(), uuid);
assert_eq!(found, Some(path));
}
app/src/ai/agent_sdk/driver/harness/gemini.rs
+63 -29
View File
@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::ffi::OsString;
use std::path::Path;
use std::sync::Arc;
@@ -12,18 +13,24 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use tempfile::NamedTempFile;
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, HarnessCleanupDisposition, HarnessRunner, JSONMCPServer, ResumePayload,
SavePoint, ThirdPartyHarness,
};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_sdk::setup_observability::{
OzRunTimelineEvent, SetupClientEventReporter, SetupStep,
};
use crate::ai::ambient_agents::task::HarnessModelConfig;
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.
@@ -46,38 +53,42 @@ impl ThirdPartyHarness for GeminiHarness {
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>,
context: Option<&str>,
working_dir: &Path,
_task_id: Option<AmbientAgentTaskId>,
server_api: Arc<ServerApi>,
terminal_driver: ModelHandle<TerminalDriver>,
_resume: Option<ResumePayload>,
_resolved_env_vars: &HashMap<OsString, OsString>,
_resolved_secrets: &HashMap<String, ManagedSecretValue>,
_resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
_third_party_harness_model_config: Option<&HarnessModelConfig>,
) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
// Prepare the environment config files.
prepare_gemini_environment_config(working_dir, system_prompt).map_err(|error| {
AgentDriverError::HarnessConfigSetupFailed {
harness: self.cli_agent().command_prefix().to_owned(),
error,
}
})?;
// 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.
// Prepend server context to the prompt if available.
let effective_prompt = match context {
Some(ctx) if !ctx.is_empty() => format!("{ctx}\n\n{prompt}"),
_ => prompt.to_string(),
};
let client: Arc<dyn HarnessSupportClient> = server_api;
Ok(Box::new(GeminiHarnessRunner::new(
self.cli_agent().command_prefix(),
prompt,
&effective_prompt,
system_prompt,
working_dir,
client,
@@ -104,6 +115,8 @@ enum GeminiRunnerState {
struct GeminiHarnessRunner {
command: String,
/// The CLI name used to invoke Gemini.
cli_name: String,
/// Held so the temp file is cleaned up when the runner is dropped.
_temp_prompt_file: NamedTempFile,
client: Arc<dyn HarnessSupportClient>,
@@ -120,11 +133,12 @@ impl GeminiHarnessRunner {
client: Arc<dyn HarnessSupportClient>,
terminal_driver: ModelHandle<TerminalDriver>,
) -> Result<Self, AgentDriverError> {
let temp_file = write_temp_file("oz_prompt_", prompt)?;
let temp_file = write_temp_file("oz_prompt_", prompt, ".txt")?;
let prompt_path = temp_file.path().display().to_string();
Ok(Self {
command: gemini_command(cli_command, &prompt_path),
cli_name: cli_command.to_string(),
_temp_prompt_file: temp_file,
client,
terminal_driver,
@@ -136,19 +150,27 @@ impl GeminiHarnessRunner {
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
impl HarnessRunner for GeminiHarnessRunner {
fn harness_name(&self) -> &str {
&self.cli_name
}
async fn start(
&self,
foreground: &ModelSpawner<AgentDriver>,
setup_events: &SetupClientEventReporter,
) -> 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)
})?;
let conversation_id = setup_events
.record_result(SetupStep::ThirdPartyHarnessExternalConversation, async {
self.client
.create_external_conversation(GEMINI_CLI_FORMAT)
.await
.map_err(|e| {
log::error!("Failed to create external conversation: {e}");
AgentDriverError::ConfigBuildFailed(e)
})
})
.await?;
log::info!("Created external conversation {conversation_id}");
let command = self.command.clone();
@@ -166,6 +188,10 @@ impl HarnessRunner for GeminiHarnessRunner {
block_id: command_handle.block_id().clone(),
};
setup_events
.post_timeline_event(OzRunTimelineEvent::AgentStarted)
.await;
Ok(command_handle)
}
@@ -215,6 +241,14 @@ impl HarnessRunner for GeminiHarnessRunner {
)
.await
}
async fn cleanup(
&self,
_cleanup_disposition: HarnessCleanupDisposition,
_foreground: &ModelSpawner<AgentDriver>,
) -> Result<()> {
Ok(())
}
}
fn prepare_gemini_environment_config(
app/src/ai/agent_sdk/driver/harness/json_utils.rs
+11
View File
@@ -10,6 +10,7 @@ use std::path::Path;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// Read a JSON file as `T`, or return `T::default()` if the file does not exist.
///
@@ -54,3 +55,13 @@ where
)
.with_context(|| format!("Failed to write {}", path.display()))
}
/// Serialize a slice of JSON values as a JSONL byte string (one value per line).
pub(super) 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)
}
app/src/ai/agent_sdk/driver/harness/mod.rs
+208 -36
View File
@@ -1,6 +1,6 @@
use std::borrow::Cow;
use std::collections::HashMap;
use std::ffi::OsString;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::io::Write;
use std::path::Path;
@@ -12,15 +12,6 @@ use galaxy_cli::agent::Harness;
use galaxy_managed_secrets::ManagedSecretValue;
use galaxyui::{ModelHandle, ModelSpawner, SingletonEntity};
use tempfile::NamedTempFile;
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 galaxy_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,
@@ -33,24 +24,97 @@ use super::{
LEGACY_OZ_PARENT_STATE_ROOT_ENV, OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV,
OZ_MESSAGE_LISTENER_STATE_ROOT_ENV,
};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_sdk::setup_observability::SetupClientEventReporter;
use crate::ai::ambient_agents::task::HarnessModelConfig;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::mcp::JSONMCPServer;
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;
mod claude_code;
pub(crate) mod claude_code;
pub(crate) mod claude_transcript;
mod codex;
pub(crate) mod codex_transcript;
mod gemini;
mod json_utils;
mod telemetry;
pub(crate) use claude_code::ClaudeHarness;
use claude_transcript::ClaudeResumeInfo;
use codex::CodexHarness;
use codex_transcript::CodexResumeInfo;
use gemini::GeminiHarness;
pub(crate) use telemetry::ThirdPartyHarnessTelemetryEvent;
/// 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`].
#[derive(Debug)]
pub(crate) enum ResumePayload {
/// Claude Code session state fetched from the server's transcript endpoint.
Claude(ClaudeResumeInfo),
/// Codex session state fetched from the server's transcript endpoint.
Codex(CodexResumeInfo),
}
impl TryFrom<ResumePayload> for ClaudeResumeInfo {
type Error = AgentDriverError;
fn try_from(payload: ResumePayload) -> Result<Self, Self::Error> {
match payload {
ResumePayload::Claude(info) => Ok(info),
_ => {
log::error!("ClaudeHarness given non-Claude ResumePayload variant");
Err(AgentDriverError::InvalidRuntimeState)
}
}
}
}
impl TryFrom<ResumePayload> for CodexResumeInfo {
type Error = AgentDriverError;
fn try_from(payload: ResumePayload) -> Result<Self, Self::Error> {
match payload {
ResumePayload::Codex(info) => Ok(info),
_ => {
log::error!("CodexHarness given non-Codex ResumePayload variant");
Err(AgentDriverError::InvalidRuntimeState)
}
}
}
}
/// Fetch the harness transcript for `conversation_id` and deserialize it into `E`.
pub(super) async fn fetch_transcript_envelope<E: serde::de::DeserializeOwned>(
harness_label: &str,
conversation_id: &AIConversationId,
client: Arc<dyn HarnessSupportClient>,
) -> Result<E, AgentDriverError> {
let bytes = 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: harness_label.to_string(),
conversation_id: conversation_id.to_string(),
}
} else {
AgentDriverError::ConversationLoadFailed(format!("{err:#}"))
}
})?;
serde_json::from_slice(&bytes).map_err(|err| {
AgentDriverError::ConversationLoadFailed(format!(
"Failed to deserialize {harness_label} transcript for {conversation_id}: {err:#}"
))
})
}
/// Trait for third-party agent harnesses that execute prompts via their own CLIs.
@@ -77,14 +141,27 @@ pub(crate) trait ThirdPartyHarness: Send + Sync {
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(())
/// Shell command to verify authentication credentials are valid.
/// Exit code 0 = pass; non-zero = fail.
fn auth_check_command(&self) -> Option<String> {
None
}
/// Substrings to scan for in the running harness block's output. A hit
/// indicates the harness can't make a successful API request (e.g.
/// invalid key, no billing, quota exhausted). The driver matches
/// case-insensitively against the block's plaintext via the same DFA
/// machinery used by the find feature.
fn runtime_error_patterns(&self) -> &'static [&'static str] {
&[]
}
/// Whether this harness must verify its Oz platform plugin before launch.
/// Codex opts into this because its unattended launch command bypasses hook
/// trust globally, so we should fail setup instead of running without the
/// Warp-installed orchestration hooks at the required version.
fn requires_verified_platform_plugin(&self) -> bool {
false
}
/// Fetch the harness-specific resume payload for an existing conversation.
@@ -107,25 +184,34 @@ pub(crate) trait ThirdPartyHarness: Send + Sync {
/// 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.
/// Responsible for all harness-specific setup: writing config files (auth,
/// trust, system prompt, MCP, etc.) and constructing the runner that will
/// execute the CLI command.
///
/// `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.
/// `resolved_env_vars` contains already-resolved secret env vars (worker
/// env > typed secrets > raw values precedence already applied).
///
/// `resolved_secrets` provides the raw typed managed secrets so harnesses
/// can read structured fields (e.g. `base_url`) without relying on env vars.
///
/// If `resume` is `Some`, the harness matches on its own [`ResumePayload`]
/// variant and reuses stored session/conversation ids.
#[allow(clippy::too_many_arguments)]
fn build_runner(
&self,
prompt: &str,
system_prompt: Option<&str>,
resumption_prompt: Option<&str>,
context: Option<&str>,
working_dir: &Path,
task_id: Option<AmbientAgentTaskId>,
server_api: Arc<ServerApi>,
terminal_driver: ModelHandle<TerminalDriver>,
resume: Option<ResumePayload>,
resolved_env_vars: &HashMap<OsString, OsString>,
resolved_secrets: &HashMap<String, ManagedSecretValue>,
resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
third_party_harness_model_config: Option<&HarnessModelConfig>,
) -> Result<Box<dyn HarnessRunner>, AgentDriverError>;
}
@@ -165,12 +251,30 @@ pub(crate) fn harness_kind(harness: Harness) -> Result<HarnessKind, AgentDriverE
match harness {
Harness::Oz => Ok(HarnessKind::Oz),
Harness::Claude => Ok(HarnessKind::ThirdParty(Box::new(ClaudeHarness))),
Harness::Codex => Ok(HarnessKind::ThirdParty(Box::new(CodexHarness))),
Harness::OpenCode => Ok(HarnessKind::Unsupported(Harness::OpenCode)),
Harness::Gemini => Ok(HarnessKind::ThirdParty(Box::new(GeminiHarness))),
Harness::Unknown => Err(AgentDriverError::InvalidRuntimeState),
}
}
/// Returns the harness's auth-check preflight command, if any.
///
/// The viewer uses this to recognize preflight blocks via exact string
/// equality (so they stay grouped under "Set up environment commands"
/// rather than being mistaken for the main harness invocation, which
/// shares the same CLI prefix).
///
/// Returns `None` for [`Harness::Oz`], for unsupported harnesses, and
/// for any third-party harness whose `auth_check_command` returns `None`
/// (e.g. Gemini today).
pub(crate) fn auth_check_command_for(harness: Harness) -> Option<String> {
let HarnessKind::ThirdParty(third_party) = harness_kind(harness).ok()? else {
return None;
};
third_party.auth_check_command()
}
/// 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(
@@ -302,6 +406,17 @@ fn task_env_vars_for_harness_name(
env_vars
}
pub(crate) fn remove_claude_externally_managed_listener_env_vars(
env_vars: &mut HashMap<OsString, OsString>,
) {
for env_name in [
OZ_MESSAGE_LISTENER_MANAGED_EXTERNALLY_ENV,
LEGACY_OZ_PARENT_LISTENER_MANAGED_EXTERNALLY_ENV,
] {
env_vars.remove(OsStr::new(env_name));
}
}
pub(crate) fn task_env_vars(
task_id: Option<&AmbientAgentTaskId>,
parent_run_id: Option<&str>,
@@ -310,6 +425,34 @@ pub(crate) fn task_env_vars(
task_env_vars_for_harness_name(task_id, parent_run_id, selected_harness)
}
/// Returns environment variables that configure the model for a third-party harness.
/// Returns an empty map for Oz or when no model is specified.
///
/// We use the `ANTHROPIC_MODEL` env var rather than the `--model` CLI flag because
/// the env var is the most reliable mechanism and avoids precedence conflicts with
/// Claude Code's `settings.json`.
pub(crate) fn harness_model_env_vars(
selected_harness: Harness,
third_party_harness_model_config: Option<&HarnessModelConfig>,
) -> HashMap<OsString, OsString> {
let mut env_vars = HashMap::new();
let Some(model_id) = third_party_harness_model_config
.map(|config| config.model_id.as_str())
.filter(|id| !id.is_empty())
else {
return env_vars;
};
match selected_harness {
Harness::Claude => {
env_vars.insert(OsString::from("ANTHROPIC_MODEL"), OsString::from(model_id));
}
Harness::Oz | Harness::OpenCode | Harness::Gemini | Harness::Codex | Harness::Unknown => {}
}
env_vars
}
/// Indicates when the harness conversation is being saved.
/// Implementations may use this to customize the saved data, such as
/// recording additional metadata on completion.
@@ -322,6 +465,18 @@ pub(crate) enum SavePoint {
PostTurn,
}
/// Controls how much harness-owned state should survive cleanup after the CLI
/// exits.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum HarnessCleanupDisposition {
/// Tear down all harness-owned resume and wake state.
DropResumptionState,
/// The harness exited cleanly and its final save completed, so wake/resume
/// state may be preserved if the harness-specific runtime also considers
/// the run complete.
PreserveResumptionStateIfSupported,
}
/// Stateful per-run representation of an external harness produced
/// by [`ThirdPartyHarness::build_runner`].
///
@@ -333,6 +488,8 @@ pub(crate) enum SavePoint {
#[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
pub(crate) trait HarnessRunner: Send + Sync {
fn harness_name(&self) -> &str;
/// Create the external conversation on the server and start the harness
/// command in the terminal.
///
@@ -342,6 +499,7 @@ pub(crate) trait HarnessRunner: Send + Sync {
async fn start(
&self,
foreground: &ModelSpawner<AgentDriver>,
setup_events: &SetupClientEventReporter,
) -> Result<CommandHandle, AgentDriverError>;
/// Save the current conversation state (transcript upload, etc.).
@@ -359,7 +517,11 @@ pub(crate) trait HarnessRunner: Send + Sync {
}
/// Clean up any harness-owned background state after the harness exits.
async fn cleanup(&self, _foreground: &ModelSpawner<AgentDriver>) -> Result<()> {
async fn cleanup(
&self,
_cleanup_disposition: HarnessCleanupDisposition,
_foreground: &ModelSpawner<AgentDriver>,
) -> Result<()> {
Ok(())
}
}
@@ -370,20 +532,29 @@ pub(crate) async fn has_running_cli_agent(
terminal_driver: &ModelHandle<TerminalDriver>,
foreground: &ModelSpawner<AgentDriver>,
) -> bool {
matches!(
cli_agent_session_status(terminal_driver, foreground).await,
Some(CLIAgentSessionStatus::InProgress)
)
}
/// Returns the tracked CLI agent session status for the terminal, if any.
pub(crate) async fn cli_agent_session_status(
terminal_driver: &ModelHandle<TerminalDriver>,
foreground: &ModelSpawner<AgentDriver>,
) -> Option<CLIAgentSessionStatus> {
let driver = terminal_driver.clone();
let Ok(running) = foreground
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)
.map(|session| session.status.clone())
})
.await
else {
return false;
};
running
.ok()
.flatten()
}
/// Create a [`NamedTempFile`] with the given prefix and write `content` into it.
@@ -393,10 +564,11 @@ pub(crate) async fn has_running_cli_agent(
pub(super) fn write_temp_file(
prefix: &str,
content: &str,
suffix: &str,
) -> Result<NamedTempFile, AgentDriverError> {
let mut file = tempfile::Builder::new()
.prefix(prefix)
.suffix(".txt")
.suffix(suffix)
.tempfile()
.map_err(|e| {
AgentDriverError::ConfigBuildFailed(anyhow::anyhow!(
@@ -457,5 +629,5 @@ pub(super) async fn upload_current_block_snapshot(
}
#[cfg(test)]
#[path = "mod_test.rs"]
#[path = "mod_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/harness/mod_test.rs
-33
View File
@@ -1,33 +0,0 @@
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"));
}
app/src/ai/agent_sdk/driver/harness/mod_tests.rs
+82
View File
@@ -0,0 +1,82 @@
use warp_cli::agent::Harness;
use super::{auth_check_command_for, 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"));
}
// --- Runtime error pattern tests ---
#[test]
fn claude_runtime_error_patterns_returns_slice() {
use super::claude_code::ClaudeHarness;
use super::ThirdPartyHarness;
// Patterns are initially empty until validated needles are filled in.
// The trait method must still be callable.
let _: &[&str] = ClaudeHarness.runtime_error_patterns();
}
#[test]
fn codex_runtime_error_patterns_returns_slice() {
use super::codex::CodexHarness;
let _: &[&str] = CodexHarness.runtime_error_patterns();
}
#[test]
fn gemini_runtime_error_patterns_is_empty_by_default() {
use super::gemini::GeminiHarness;
assert!(GeminiHarness.runtime_error_patterns().is_empty());
}
#[test]
fn auth_check_command_for_gemini_is_none() {
assert!(auth_check_command_for(Harness::Gemini).is_none());
}
#[test]
fn auth_check_command_for_oz_is_none() {
assert!(auth_check_command_for(Harness::Oz).is_none());
}
#[test]
fn auth_check_command_for_unsupported_is_none() {
// OpenCode is mapped to HarnessKind::Unsupported and therefore has no
// auth check command of its own.
assert!(auth_check_command_for(Harness::OpenCode).is_none());
}
#[test]
fn auth_check_command_for_unknown_is_none() {
// Harness::Unknown causes harness_kind to return Err; the helper still
// returns None instead of panicking.
assert!(auth_check_command_for(Harness::Unknown).is_none());
}
app/src/ai/agent_sdk/driver/harness/telemetry.rs
+84
View File
@@ -0,0 +1,84 @@
use serde_json::{json, Value};
use strum_macros::{EnumDiscriminants, EnumIter};
use galaxy_core::telemetry::{EnablementState, TelemetryEvent, TelemetryEventDesc};
/// Telemetry events emitted by the third-party harness runtime layer.
#[derive(Debug, EnumDiscriminants)]
#[strum_discriminants(derive(EnumIter))]
pub(crate) enum ThirdPartyHarnessTelemetryEvent {
/// The runtime output scanner observed one of the harness's known
/// failure substrings. Fires once per detection, before any suppression
/// logic, so dashboards can compare raw trigger volume vs. detections
/// that actually fail the run.
RuntimeErrorDetected {
/// CLI command prefix for the harness whose block was scanned
/// (e.g. `"claude"`, `"codex"`).
harness: String,
/// The originating needle from `runtime_error_patterns` that hit.
pattern: String,
},
}
impl TelemetryEvent for ThirdPartyHarnessTelemetryEvent {
fn name(&self) -> &'static str {
ThirdPartyHarnessTelemetryEventDiscriminants::from(self).name()
}
fn payload(&self) -> Option<Value> {
match self {
ThirdPartyHarnessTelemetryEvent::RuntimeErrorDetected { harness, pattern } => {
Some(json!({
"harness": harness,
"pattern": pattern,
}))
}
}
}
fn description(&self) -> &'static str {
ThirdPartyHarnessTelemetryEventDiscriminants::from(self).description()
}
fn enablement_state(&self) -> EnablementState {
ThirdPartyHarnessTelemetryEventDiscriminants::from(self).enablement_state()
}
fn contains_ugc(&self) -> bool {
match self {
ThirdPartyHarnessTelemetryEvent::RuntimeErrorDetected { .. } => false,
}
}
fn event_descs() -> impl Iterator<Item = Box<dyn TelemetryEventDesc>> {
galaxy_core::telemetry::enum_events::<Self>()
}
}
impl TelemetryEventDesc for ThirdPartyHarnessTelemetryEventDiscriminants {
fn name(&self) -> &'static str {
match self {
ThirdPartyHarnessTelemetryEventDiscriminants::RuntimeErrorDetected => {
"AmbientAgents.ThirdPartyHarness.RuntimeError.Detected"
}
}
}
fn description(&self) -> &'static str {
match self {
ThirdPartyHarnessTelemetryEventDiscriminants::RuntimeErrorDetected => {
"Runtime output scanner detected a known failure substring in a third-party \
harness block."
}
}
}
fn enablement_state(&self) -> EnablementState {
match self {
ThirdPartyHarnessTelemetryEventDiscriminants::RuntimeErrorDetected => {
EnablementState::Always
}
}
}
}
galaxy_core::register_telemetry_event!(ThirdPartyHarnessTelemetryEvent);
app/src/ai/agent_sdk/driver/harness_output_monitor.rs
+217
View File
@@ -0,0 +1,217 @@
//! Background monitor that scans the running harness block for known runtime
//! failure substrings (e.g. invalid API key, exhausted credits) and reports
//! the first hit so the driver can fail the task fast instead of letting the
//! harness hang.
use std::sync::Arc;
use std::time::Duration;
use regex::escape;
use warpui::ModelSpawner;
use super::terminal::BlockOutputMatch;
use super::AgentDriver;
use crate::terminal::cli_agent_sessions::CLIAgentSessionStatus;
use crate::terminal::model::block::BlockId;
use crate::terminal::model::find::RegexDFAs;
const SCAN_INTERVALS: &[Duration] = &[
Duration::from_secs(5),
Duration::from_secs(5),
Duration::from_secs(5),
Duration::from_secs(5),
Duration::from_secs(5),
Duration::from_secs(5),
Duration::from_secs(15),
Duration::from_secs(15),
Duration::from_secs(15),
Duration::from_secs(15),
];
const STALL_POLL_INTERVAL: Duration = Duration::from_secs(10);
const STALL_CONFIRMATION_BUDGET: Duration = Duration::from_secs(60);
#[derive(Debug, Clone)]
pub(crate) struct DetectedHarnessError {
pub pattern: String,
pub excerpt: String,
}
/// Build a combined case-insensitive DFA from the harness's static patterns.
pub(crate) fn build_dfas(patterns: &[&'static str]) -> Option<RegexDFAs> {
if patterns.is_empty() {
return None;
}
let escaped: Vec<String> = patterns.iter().map(|p| escape(p)).collect();
let refs: Vec<&str> = escaped.iter().map(String::as_str).collect();
match RegexDFAs::new_many(
&refs, false, // enable_unicode_word_boundary
false, // case_sensitive
) {
Ok(dfas) => Some(dfas),
Err(err) => {
log::warn!("Failed to build harness output DFAs: {err}");
None
}
}
}
/// Map a `matched_text` produced by the combined DFA back to the originating
/// needle. The DFA matched `(p1|p2|…)`, so the matched substring is exactly
/// one of `patterns` up to case — we lowercase-compare to identify it.
pub(crate) fn pattern_for_match(
matched_text: &str,
patterns: &[&'static str],
) -> Option<&'static str> {
let matched_lower = matched_text.to_lowercase();
patterns
.iter()
.copied()
.find(|p| p.to_lowercase() == matched_lower)
}
fn outputs_stalled(before: Option<&str>, after: Option<&str>) -> bool {
matches!((before, after), (Some(a), Some(b)) if a == b)
}
async fn find_match_once(
block_id: &BlockId,
dfas: &Arc<RegexDFAs>,
foreground: &ModelSpawner<AgentDriver>,
) -> Option<BlockOutputMatch> {
let block_id_for_tick = block_id.clone();
let dfas_for_tick = Arc::clone(dfas);
foreground
.spawn(move |me, ctx| {
me.terminal_driver
.as_ref(ctx)
.find_first_match_in_block_output(&block_id_for_tick, &dfas_for_tick, ctx)
})
.await
.ok()
.flatten()
}
async fn fetch_plaintext(
block_id: &BlockId,
foreground: &ModelSpawner<AgentDriver>,
) -> Option<String> {
let block_id_for_tick = block_id.clone();
foreground
.spawn(move |me, ctx| {
me.terminal_driver
.as_ref(ctx)
.block_output_plaintext(&block_id_for_tick, ctx)
})
.await
.ok()
.flatten()
}
/// Run the stall-confirmation loop after a pattern hit.
async fn confirm_stall(
block_id: &BlockId,
dfas: &Arc<RegexDFAs>,
foreground: &ModelSpawner<AgentDriver>,
) -> (Option<BlockOutputMatch>, Duration) {
let mut previous = fetch_plaintext(block_id, foreground).await;
let mut elapsed = Duration::ZERO;
while elapsed < STALL_CONFIRMATION_BUDGET {
warpui::r#async::Timer::after(STALL_POLL_INTERVAL).await;
elapsed += STALL_POLL_INTERVAL;
let current = fetch_plaintext(block_id, foreground).await;
if outputs_stalled(previous.as_deref(), current.as_deref()) {
// Output settled. Re-run the DFA so we report the post-dwell
// match in case the row positions shifted while we waited.
return (find_match_once(block_id, dfas, foreground).await, elapsed);
}
previous = current;
}
(None, elapsed)
}
/// Watch the given block for harness output errors on the
/// [`SCAN_INTERVALS`] cadence.
///
/// On every pattern hit, run a stall-confirmation loop (up to
/// [`STALL_CONFIRMATION_BUDGET`]) and only resolve with
/// `Some(DetectedHarnessError)` when the harness output stabilizes with
/// the pattern still present. If the harness keeps producing output
/// (e.g. spinner frames during an automatic retry), the detection is
/// dropped and the scanner resumes normal polling.
///
/// Returns `None` when the schedule completes without a confirmed hit
/// (or when `patterns` is empty / DFA construction fails).
/// Cancellation-safe: dropping the future stops both loops.
pub(crate) async fn watch_block_for_errors(
block_id: BlockId,
patterns: &'static [&'static str],
foreground: &ModelSpawner<AgentDriver>,
) -> Option<DetectedHarnessError> {
if patterns.is_empty() {
return None;
}
let dfas = Arc::new(build_dfas(patterns)?);
// Total observation budget = sum of all scan intervals. Used as an
// early-exit guard after stall confirmation in case a flaky harness
// has burned most of the window into confirmation loops.
let total_budget: Duration = SCAN_INTERVALS.iter().copied().sum();
let mut elapsed = Duration::ZERO;
for &interval in SCAN_INTERVALS {
warpui::r#async::Timer::after(interval).await;
elapsed += interval;
if find_match_once(&block_id, &dfas, foreground)
.await
.is_none()
{
continue;
}
// Candidate match observed. Confirm via the stall loop so we
// don't false-positive while the harness is mid-retry.
let (confirmed, confirmation_elapsed) = confirm_stall(&block_id, &dfas, foreground).await;
elapsed += confirmation_elapsed;
let Some(hit) = confirmed else {
if elapsed >= total_budget {
break;
}
continue;
};
// The DFA matched one of our needles, so `pattern_for_match`
// should always resolve. Fall back to the matched substring
// verbatim if it somehow doesn't.
let pattern = pattern_for_match(&hit.matched_text, patterns)
.map(str::to_owned)
.unwrap_or_else(|| hit.matched_text.clone());
return Some(DetectedHarnessError {
pattern,
excerpt: cap_excerpt(&hit.excerpt),
});
}
None
}
pub(crate) fn should_suppress_runtime_failure(status: Option<&CLIAgentSessionStatus>) -> bool {
matches!(status, Some(CLIAgentSessionStatus::Success))
}
/// Cap excerpt length so we don't blow up status messages or logs with
/// terminal-width rows.
const EXCERPT_MAX_LEN: usize = 240;
fn cap_excerpt(excerpt: &str) -> String {
if excerpt.chars().count() <= EXCERPT_MAX_LEN {
return excerpt.to_owned();
}
let mut out: String = excerpt.chars().take(EXCERPT_MAX_LEN).collect();
out.push('…');
out
}
#[cfg(test)]
#[path = "harness_output_monitor_tests.rs"]
mod tests;
app/src/ai/agent_sdk/driver/harness_output_monitor_tests.rs
+79
View File
@@ -0,0 +1,79 @@
use std::time::Duration;
use super::{outputs_stalled, pattern_for_match, STALL_CONFIRMATION_BUDGET, STALL_POLL_INTERVAL};
#[test]
fn pattern_for_match_returns_originating_needle() {
let patterns: &[&'static str] = &["credit balance is too low", "invalid_api_key"];
// The DFA gives us the matched substring (e.g. lifted from the grid as
// mixed case); we must map back to the original `'static` needle.
let resolved = pattern_for_match("Credit Balance Is Too Low", patterns);
assert_eq!(resolved, Some("credit balance is too low"));
}
#[test]
fn pattern_for_match_is_case_insensitive() {
let patterns: &[&'static str] = &["INVALID_API_KEY"];
let resolved = pattern_for_match("invalid_api_key", patterns);
assert_eq!(resolved, Some("INVALID_API_KEY"));
}
#[test]
fn pattern_for_match_returns_none_when_no_match() {
let patterns: &[&'static str] = &["needle a", "needle b"];
assert!(pattern_for_match("entirely different text", patterns).is_none());
}
#[test]
fn pattern_for_match_picks_first_matching_needle() {
// When two needles lowercase-equal the same matched text (degenerate
// case), we deterministically return the first one in the slice.
let patterns: &[&'static str] = &["Foo", "foo"];
let resolved = pattern_for_match("FOO", patterns);
assert_eq!(resolved, Some("Foo"));
}
// --- outputs_stalled / stall confirmation ---
#[test]
fn outputs_stalled_returns_true_when_inputs_equal() {
let snapshot = "Error: credit balance is too low\n";
assert!(outputs_stalled(Some(snapshot), Some(snapshot)));
}
#[test]
fn outputs_stalled_returns_false_when_inputs_differ_by_any_byte() {
// Spinner case: a single character differs between frames. The
// confirmation loop must treat this as "still moving" and not
// mistakenly declare the harness stalled.
let before = "Retrying \u{2807}";
let after = "Retrying \u{2826}";
assert!(!outputs_stalled(Some(before), Some(after)));
}
#[test]
fn outputs_stalled_returns_false_when_either_input_is_none() {
// A failed snapshot fetch must default to "not confirmed" rather than
// killing the harness on a transient lookup error.
assert!(!outputs_stalled(None, Some("data")));
assert!(!outputs_stalled(Some("data"), None));
assert!(!outputs_stalled(None, None));
}
#[test]
fn stall_confirmation_budget_matches_six_poll_intervals() {
// The loop guarantees up to BUDGET/INTERVAL iterations. Pin the ratio
// so a careless tweak to either constant can't accidentally turn the
// confirmation into a single comparison or extend it indefinitely.
assert_eq!(
STALL_CONFIRMATION_BUDGET.as_secs() / STALL_POLL_INTERVAL.as_secs(),
6
);
// And that they're cleanly divisible (no leftover sub-interval window).
assert_eq!(
STALL_CONFIRMATION_BUDGET.as_secs() % STALL_POLL_INTERVAL.as_secs(),
0
);
// Sanity: STALL_POLL_INTERVAL must be non-zero or the loop would spin.
assert!(STALL_POLL_INTERVAL > Duration::ZERO);
}
app/src/ai/agent_sdk/driver/output.rs
+52 -42
View File
@@ -1,26 +1,22 @@
pub mod text {
use std::{
collections::HashSet,
fmt,
io::{self, Write},
};
use std::collections::HashSet;
use std::fmt;
use std::io::{self, Write};
const CANCELLED_MESSAGE: &str = "<cancelled>";
use ai::agent::action_result::{FetchConversationResult, ReadSkillResult, UseComputerResult};
use itertools::Itertools;
use crate::{
ai::agent::{
AIAgentActionType, AIAgentInput, AIAgentOutput, AIAgentOutputMessageType, AIAgentTodo,
ArtifactCreatedData, CallMCPToolResult, FileGlobResult, FileGlobV2Result, GrepResult,
ReadFilesResult, ReadMCPResourceResult, RequestCommandOutputResult,
RequestFileEditsResult, SearchCodebaseResult, SuggestNewConversationResult,
SuggestPromptResult, TodoOperation, UploadArtifactResult, WebFetchStatus,
WebSearchStatus, WriteToLongRunningShellCommandResult,
},
AIAgentActionResultType,
use crate::ai::agent::{
AIAgentActionType, AIAgentInput, AIAgentOutput, AIAgentOutputMessageType, AIAgentTodo,
ArtifactCreatedData, CallMCPToolResult, FileGlobResult, FileGlobV2Result, GrepResult,
ReadFilesResult, ReadMCPResourceResult, RequestCommandOutputResult, RequestFileEditsResult,
SearchCodebaseResult, SuggestNewConversationResult, SuggestPromptResult, TodoOperation,
UploadArtifactResult, WebFetchStatus, WebSearchStatus,
WriteToLongRunningShellCommandResult,
};
use crate::AIAgentActionResultType;
/// Format an agent input as a human-readable string. For action results, it's assumed that
/// the action is shown immediately before this result.
@@ -42,7 +38,8 @@ pub mod text {
| AIAgentInput::StartFromAmbientRunPrompt { .. }
| AIAgentInput::MessagesReceivedFromAgents { .. }
| AIAgentInput::PassiveSuggestionResult { .. }
| AIAgentInput::EventsFromAgents { .. } => {
| AIAgentInput::EventsFromAgents { .. }
| AIAgentInput::OrchestrationConfigUpdate { .. } => {
// Do not include the user query, since it's already provided as input to the agent.
Ok(())
}
@@ -303,6 +300,10 @@ pub mod text {
// SendMessageToAgent is a client-side orchestration action, not used in SDK
AIAgentActionResultType::SendMessageToAgent(_) => Ok(()),
AIAgentActionResultType::AskUserQuestion(_) => Ok(()),
// RunAgents is a desktop-client-only action; not used in the SDK.
AIAgentActionResultType::RunAgents(_) => Ok(()),
// No user-visible payload to emit.
AIAgentActionResultType::WaitForEvents(_) => Ok(()),
},
}
}
@@ -427,6 +428,9 @@ pub mod text {
)?;
}
AIAgentActionType::AskUserQuestion { .. } => (),
// RunAgents is desktop-client-only; SDK driver renders nothing.
AIAgentActionType::RunAgents(_) => (),
AIAgentActionType::WaitForEvents { .. } => (),
},
AIAgentOutputMessageType::TodoOperation(operation) => match operation {
TodoOperation::UpdateTodos { todos } => {
@@ -559,26 +563,20 @@ pub mod text {
}
pub mod json {
use crate::{
ai::agent::{
AIAgentActionType, AIAgentInput, AIAgentOutput, AIAgentOutputMessage,
AIAgentOutputMessageType, AIAgentTodo, ArtifactCreatedData, CallMCPToolResult,
FileContext, FileGlobResult, FileGlobV2Result, GrepResult, ReadFilesResult,
ReadMCPResourceResult, RequestCommandOutputResult, RequestFileEditsResult,
SearchCodebaseResult, SubagentCall, TodoOperation, UploadArtifactResult,
WriteToLongRunningShellCommandResult,
},
AIAgentActionResultType,
};
use std::borrow::Cow;
use std::ops::Range;
use serde::Serialize;
use crate::ai::agent::comment::ReviewComment;
use serde::Serialize;
use std::path::Path;
use std::{
borrow::Cow,
io::{self, Write},
ops::Range,
use crate::ai::agent::{
AIAgentActionType, AIAgentInput, AIAgentOutput, AIAgentOutputMessage,
AIAgentOutputMessageType, AIAgentTodo, ArtifactCreatedData, CallMCPToolResult, FileContext,
FileGlobResult, FileGlobV2Result, GrepResult, ReadFilesResult, ReadMCPResourceResult,
RequestCommandOutputResult, RequestFileEditsResult, SearchCodebaseResult, SubagentCall,
TodoOperation, UploadArtifactResult, WriteToLongRunningShellCommandResult,
};
use crate::code::buffer_location::LocalOrRemotePath;
/// JSON representation of messages in an agent conversation. This is intentionally not 1:1 with our internal `AIAgent*` types - it's
/// a stable interface for callers.
@@ -742,7 +740,7 @@ pub mod json {
#[derive(Serialize)]
struct JsonComment<'a> {
comment_text: &'a str,
file_path: Option<&'a Path>,
file_path: Option<String>,
line_number: Option<usize>,
head_title: Option<&'a str>,
}
@@ -791,7 +789,8 @@ pub mod json {
| AIAgentInput::StartFromAmbientRunPrompt { .. }
| AIAgentInput::MessagesReceivedFromAgents { .. }
| AIAgentInput::EventsFromAgents { .. }
| AIAgentInput::PassiveSuggestionResult { .. } => None,
| AIAgentInput::PassiveSuggestionResult { .. }
| AIAgentInput::OrchestrationConfigUpdate { .. } => None,
// These input types should not occur in a SDK-run agent.
AIAgentInput::ResumeConversation { .. }
| AIAgentInput::TriggerPassiveSuggestion { .. } => None,
@@ -1108,6 +1107,10 @@ pub mod json {
| AIAgentActionType::SendMessageToAgent { .. }
| AIAgentActionType::TransferShellCommandControlToUser { .. } => None,
AIAgentActionType::AskUserQuestion { .. } => None,
// RunAgents is desktop-client-only; SDK has no JSON
// representation for it.
AIAgentActionType::RunAgents(_) => None,
AIAgentActionType::WaitForEvents { .. } => None,
},
AIAgentOutputMessageType::TodoOperation(operation) => match operation {
TodoOperation::UpdateTodos { todos } => Some(JsonMessage::UpdateTodos {
@@ -1169,7 +1172,11 @@ pub mod json {
fn from(review_comment: &'a ReviewComment) -> Self {
Self {
comment_text: review_comment.content.as_str(),
file_path: review_comment.diff.file_path.as_deref(),
file_path: review_comment
.diff
.file_path
.as_ref()
.map(LocalOrRemotePath::display_path),
line_number: review_comment.diff.line_number,
head_title: review_comment.head_title.as_deref(),
}
@@ -1286,10 +1293,12 @@ pub mod json {
}
}
use std::io::{self, BufWriter, Write};
use galaxy_core::channel::ChannelState;
use crate::ai::agent::{AIAgentText, AIAgentTextSection};
use crate::code::editor_management::CodeSource;
use galaxy_core::channel::ChannelState;
use std::io::{self, BufWriter, Write};
/// Constructs the Oz dashboard URL for a given run ID.
fn run_url(run_id: &str) -> String {
@@ -1328,8 +1337,8 @@ fn format_agent_text<W: Write>(text: &AIAgentText, w: &mut W) -> io::Result<()>
}
match source {
Some(CodeSource::ProjectRules { path }) => {
writeln!(w, " rules_path={}", path.display())?;
Some(CodeSource::ProjectRules { location }) => {
writeln!(w, " rules_path={}", location.display_path())?;
}
Some(CodeSource::Link {
path,
@@ -1348,12 +1357,13 @@ fn format_agent_text<W: Write>(text: &AIAgentText, w: &mut W) -> io::Result<()>
writeln!(w)?;
}
Some(CodeSource::Skill { path, .. }) => {
writeln!(w, " skill_path={}", path.display())?;
Some(CodeSource::Skill { location, .. }) => {
writeln!(w, " skill_path={}", location.display_path())?;
}
Some(CodeSource::AIAction { .. })
| Some(CodeSource::New { .. })
| Some(CodeSource::FileTree { .. })
| Some(CodeSource::CommandPalette { .. })
| Some(CodeSource::Finder { .. })
| None => {}
}
app/src/ai/agent_sdk/driver/snapshot.rs
+454 -75
View File
@@ -30,15 +30,23 @@ use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use anyhow::Result;
use anyhow::{Context as _, Result};
use command::r#async::Command;
use command::Stdio;
use futures::future::join_all;
use tokio::fs::{self as tokio_fs, OpenOptions};
use tokio::io::AsyncWriteExt as _;
use tokio::sync::{mpsc, oneshot};
use galaxy_core::report_error;
use galaxyui::r#async::executor::Background;
use galaxyui::r#async::FutureExt as _;
use crate::ai::agent_sdk::retry::with_bounded_retry;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::server::server_api::ai::{
AIClient, InitialSnapshotToken, SnapshotUploadFileInfo as AiSnapshotUploadFileInfo,
UploadLocalHandoffSnapshotRequest,
};
use crate::server::server_api::harness_support::{
upload_to_target, HarnessSupportClient, SnapshotFileInfo, SnapshotUploadRequest, UploadTarget,
};
@@ -99,6 +107,15 @@ struct DeclarationLine {
path: String,
}
/// Serialize-only sibling of [`DeclarationLine`] used by the writer task to emit `file`
/// entries with a fixed `version` and `kind`.
#[derive(serde::Serialize)]
struct FileDeclaration<'a> {
version: u32,
kind: &'a str,
path: &'a str,
}
/// Invoke `snapshot-declarations.sh` to (re)generate the declarations file consumed by the
/// rest of the upload pipeline.
///
@@ -107,10 +124,11 @@ struct DeclarationLine {
/// from `task_id`. The script appends to the file if it already exists, so
/// repeated invocations within a single run accumulate repos instead of clobbering.
///
/// A missing env var, a missing script, a non-zero exit status, a spawn failure, or a runtime
/// exceeding `script_timeout` are each logged at `log::error!` and returned without aborting the
/// caller — if a previous invocation already produced a declarations file on disk it remains
/// usable; otherwise the upload pipeline becomes a no-op.
/// A missing env var is expected in some paths and is logged at `log::info!`. A missing script,
/// a non-zero exit status, a spawn failure, or a runtime exceeding `script_timeout` are each
/// logged at `log::error!` and returned without aborting the caller — if a previous invocation
/// already produced a declarations file on disk it remains usable; otherwise the upload pipeline
/// becomes a no-op.
///
/// Exposed as a standalone helper so future call sites can trigger declarations generation at
/// other points in the run lifecycle (e.g. periodic mid-run snapshots).
@@ -120,7 +138,7 @@ pub(super) async fn run_declarations_script(
script_timeout: Duration,
) {
let Some(script_path) = std::env::var_os(DECLARATIONS_SCRIPT_PATH_ENV_VAR) else {
log::error!(
log::info!(
"{DECLARATIONS_SCRIPT_PATH_ENV_VAR} is not set; skipping snapshot declarations script (task {task_id})"
);
return;
@@ -199,7 +217,7 @@ fn resolve_declarations_path(task_id: Option<&AmbientAgentTaskId>) -> PathBuf {
/// 3. `{DEFAULT_DECLARATIONS_DIR}/{DEFAULT_DECLARATIONS_FILENAME}` as a final fallback.
fn resolve_declarations_path_with_override(
task_id: Option<&AmbientAgentTaskId>,
override_path: Option<std::ffi::OsString>,
override_path: Option<OsString>,
) -> PathBuf {
if let Some(override_path) = override_path {
return PathBuf::from(override_path);
@@ -216,31 +234,28 @@ fn resolve_declarations_path_with_override(
///
/// Returns `None` when the file is missing, unreadable, or yields no valid entries; logs a
/// WARN describing why in each case. A returned `Some(entries)` is guaranteed non-empty.
fn read_and_parse_declarations(
path: &Path,
task_id: &AmbientAgentTaskId,
) -> Option<Vec<DeclarationEntry>> {
fn read_and_parse_declarations(path: &Path) -> Option<Vec<DeclarationEntry>> {
let contents = match std::fs::read_to_string(path) {
Ok(c) => c,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
log::warn!(
"Snapshot declarations file not found at '{}'; skipping upload (task {task_id})",
"Snapshot declarations file not found at '{}'; skipping upload",
path.display()
);
return None;
}
Err(e) => {
log::warn!(
"Failed to read snapshot declarations file '{}': {e:#}; skipping upload (task {task_id})",
"Failed to read snapshot declarations file '{}': {e:#}; skipping upload",
path.display()
);
return None;
}
};
let entries = parse_declarations(&contents, task_id);
let entries = parse_declarations(&contents);
if entries.is_empty() {
log::warn!(
"Snapshot declarations file '{}' has no valid entries; skipping upload (task {task_id})",
"Snapshot declarations file '{}' has no valid entries; skipping upload",
path.display()
);
return None;
@@ -254,7 +269,7 @@ fn read_and_parse_declarations(
/// and `path` (absolute path). Blank lines are ignored. Malformed lines (invalid JSON, missing
/// fields, unsupported versions, unknown kind, non-absolute path) are logged at WARN and skipped;
/// they never abort parsing.
fn parse_declarations(contents: &str, task_id: &AmbientAgentTaskId) -> Vec<DeclarationEntry> {
fn parse_declarations(contents: &str) -> Vec<DeclarationEntry> {
let mut entries = Vec::new();
let mut seen = HashSet::new();
for (index, raw) in contents.lines().enumerate() {
@@ -267,26 +282,26 @@ fn parse_declarations(contents: &str, task_id: &AmbientAgentTaskId) -> Vec<Decla
Ok(declaration) => declaration,
Err(e) => {
log::warn!(
"Malformed snapshot declarations JSONL line {line_number}: {e:#}: {raw:?} (task {task_id})"
"Malformed snapshot declarations JSONL line {line_number}: {e:#}: {raw:?}"
);
continue;
}
};
if declaration.version != Some(DECLARATION_VERSION) {
log::warn!(
"Malformed snapshot declarations line {line_number} (missing or unsupported version): {raw:?} (task {task_id})"
"Malformed snapshot declarations line {line_number} (missing or unsupported version): {raw:?}"
);
continue;
}
if declaration.path.is_empty() {
log::warn!(
"Malformed snapshot declarations line {line_number} (missing path): {raw:?} (task {task_id})"
"Malformed snapshot declarations line {line_number} (missing path): {raw:?}"
);
continue;
}
if !Path::new(&declaration.path).is_absolute() {
log::warn!(
"Malformed snapshot declarations line {line_number} (non-absolute path): {raw:?} (task {task_id})"
"Malformed snapshot declarations line {line_number} (non-absolute path): {raw:?}"
);
continue;
}
@@ -295,7 +310,7 @@ fn parse_declarations(contents: &str, task_id: &AmbientAgentTaskId) -> Vec<Decla
"file" => EntryKind::File,
other => {
log::warn!(
"Malformed snapshot declarations line {line_number} (unknown kind '{other}'): {raw:?} (task {task_id})"
"Malformed snapshot declarations line {line_number} (unknown kind '{other}'): {raw:?}"
);
continue;
}
@@ -311,6 +326,243 @@ fn parse_declarations(contents: &str, task_id: &AmbientAgentTaskId) -> Vec<Decla
entries
}
/// Drop `file` declarations whose path is already covered by a declared `repo` path so the
/// gather step does not double-upload files the repo patch already carries.
fn drop_files_covered_by_repos(entries: Vec<DeclarationEntry>) -> Vec<DeclarationEntry> {
let repo_paths: Vec<PathBuf> = entries
.iter()
.filter(|entry| entry.kind == EntryKind::Repo)
.map(|entry| PathBuf::from(&entry.path))
.collect();
if repo_paths.is_empty() {
return entries;
}
entries
.into_iter()
.filter(|entry| {
if entry.kind != EntryKind::File {
return true;
}
let file_path = Path::new(&entry.path);
for repo in &repo_paths {
if file_path.starts_with(repo) {
log::info!(
"Dropping file declaration '{}' covered by repo '{}'",
entry.path,
repo.display()
);
return false;
}
}
true
})
.collect()
}
// --- Declarations writer: SDK driver → declarations file ---
/// Commands accepted by the async declarations writer task.
enum WriterCommand {
/// Append `file` entries for the given paths to the declarations file.
Append(Vec<String>),
/// Acknowledge once every previously-queued command has finished its fs writes.
Flush(oneshot::Sender<()>),
}
/// Handle used by the SDK driver to enqueue `file` declaration appends from the subscription
/// thread without ever touching the filesystem inline.
///
/// The handle owns an unbounded `mpsc` sender into a dedicated writer task spawned by
/// [`DeclarationsWriterHandle::new`]. The writer task owns the `seen: HashSet<String>` and
/// the resolved declarations path, and processes commands sequentially, which serializes
/// writes within the process. Handles are cheaply cloneable because the underlying sender is;
/// dropping every handle closes the channel and lets the writer task exit cleanly.
#[derive(Clone)]
pub(super) struct DeclarationsWriterHandle {
tx: mpsc::UnboundedSender<WriterCommand>,
}
impl DeclarationsWriterHandle {
/// Spawn the writer task on `background` and return a fire-and-forget handle.
pub(super) fn new(
task_id: AmbientAgentTaskId,
working_dir: PathBuf,
background: &Background,
) -> Self {
let (tx, rx) = mpsc::unbounded_channel();
let declarations_path = resolve_declarations_path(Some(&task_id));
background
.spawn(writer_task(rx, declarations_path, working_dir, task_id))
.detach();
Self { tx }
}
/// Test-facing constructor that bypasses env-var-dependent path resolution and uses
/// `tokio::spawn` directly so tests can run without standing up a `Background`.
#[cfg(all(test, not(windows)))]
pub(super) fn new_at_path(
declarations_path: PathBuf,
working_dir: PathBuf,
task_id: AmbientAgentTaskId,
) -> Self {
let (tx, rx) = mpsc::unbounded_channel();
tokio::spawn(writer_task(rx, declarations_path, working_dir, task_id));
Self { tx }
}
/// Enqueue `paths` for appending as `file` entries.
///
/// Non-blocking; the subscription handler can call this from a sync context. Empty
/// input is a no-op.
pub(super) fn append(&self, paths: Vec<String>) {
if paths.is_empty() {
return;
}
if let Err(e) = self.tx.send(WriterCommand::Append(paths)) {
log::warn!("Declarations writer channel closed; dropping append: {e}");
}
}
/// Awaits until every previously-queued `append` has finished its fs writes.
///
/// Called once from `AgentDriver::run_snapshot_upload` immediately before
/// `snapshot::run_declarations_script`, so no driver-side write is in flight when the
/// bash script starts its own appends.
pub(super) async fn flush(&self) {
let (ack_tx, ack_rx) = oneshot::channel();
if self.tx.send(WriterCommand::Flush(ack_tx)).is_err() {
// Writer task has already exited; nothing is queued, nothing to drain.
return;
}
if ack_rx.await.is_err() {
log::warn!("Declarations writer flush oneshot dropped without ack");
}
}
}
/// Writer task loop: owns the `seen` set, lazily opens the file per write, and services
/// `Append` and `Flush` commands in order.
async fn writer_task(
mut rx: mpsc::UnboundedReceiver<WriterCommand>,
declarations_path: PathBuf,
working_dir: PathBuf,
task_id: AmbientAgentTaskId,
) {
let mut seen: HashSet<String> = HashSet::new();
while let Some(cmd) = rx.recv().await {
match cmd {
WriterCommand::Append(paths) => {
for path in paths {
process_append_path(
path,
&declarations_path,
&working_dir,
&task_id,
&mut seen,
)
.await;
}
}
WriterCommand::Flush(ack) => {
let _ = ack.send(());
}
}
}
}
/// Normalize, preempt against existing repos, and write one JSONL line for `raw_path`.
/// All failures log at WARN and return without advancing `seen`.
async fn process_append_path(
raw_path: String,
declarations_path: &Path,
working_dir: &Path,
task_id: &AmbientAgentTaskId,
seen: &mut HashSet<String>,
) {
let candidate = Path::new(&raw_path);
let absolute = if candidate.is_absolute() {
candidate.to_path_buf()
} else {
working_dir.join(candidate)
};
if !absolute.is_absolute() {
log::warn!(
"Skipping non-absolute file-edit path {absolute:?} for declarations (task {task_id})"
);
return;
}
let Some(absolute_str) = absolute.to_str().map(str::to_owned) else {
log::warn!(
"Skipping non-UTF-8 file-edit path {absolute:?} for declarations (task {task_id})"
);
return;
};
if seen.contains(&absolute_str) {
return;
}
if path_is_under_existing_repo(&absolute).await {
log::debug!(
"Skipping file declaration for '{absolute_str}': already inside an existing git repo (task {task_id})"
);
seen.insert(absolute_str);
return;
}
match append_declaration_line(declarations_path, &absolute_str).await {
Ok(()) => {
seen.insert(absolute_str);
}
Err(e) => {
log::warn!(
"Failed to append file declaration for '{absolute_str}': {e:#} (task {task_id})"
);
}
}
}
/// Walk ancestors of `path` and return `true` if any of them already contains a `.git`
/// directory. Cheap enough to run per path: one `stat(2)` per ancestor up to `/`.
async fn path_is_under_existing_repo(path: &Path) -> bool {
let mut current = path.parent();
while let Some(dir) = current {
let git_dir = dir.join(".git");
if tokio_fs::try_exists(&git_dir).await.unwrap_or(false) {
return true;
}
current = dir.parent();
}
false
}
/// Open the declarations file in append-create mode and write one JSONL line for `path`.
/// The serialized shape matches the schema the parser expects.
async fn append_declaration_line(declarations_path: &Path, path: &str) -> Result<()> {
if let Some(parent) = declarations_path.parent() {
tokio_fs::create_dir_all(parent)
.await
.with_context(|| format!("create_dir_all {}", parent.display()))?;
}
let mut line = serde_json::to_string(&FileDeclaration {
version: DECLARATION_VERSION,
kind: "file",
path,
})
.context("serialize file declaration")?;
line.push('\n');
let mut file = OpenOptions::new()
.append(true)
.create(true)
.open(declarations_path)
.await
.with_context(|| format!("open declarations file {}", declarations_path.display()))?;
file.write_all(line.as_bytes())
.await
.with_context(|| format!("write declarations file {}", declarations_path.display()))?;
file.flush()
.await
.with_context(|| format!("flush declarations file {}", declarations_path.display()))?;
Ok(())
}
// --- Gather phase: upload blobs and per-entry results ---
struct SnapshotUploadFile {
@@ -447,7 +699,7 @@ pub(super) async fn upload_snapshot_from_declarations(
task_id: &AmbientAgentTaskId,
) {
let declarations_path = resolve_declarations_path(Some(task_id));
let _ = upload_snapshot_from_declarations_file(&declarations_path, client, task_id).await;
let _ = upload_snapshot_from_declarations_file(&declarations_path, client).await;
}
/// Internal entry that reads from an explicit path and returns the structured outcome so tests
@@ -456,13 +708,10 @@ pub(super) async fn upload_snapshot_from_declarations(
async fn upload_snapshot_from_declarations_file(
path: &Path,
client: Arc<dyn HarnessSupportClient>,
task_id: &AmbientAgentTaskId,
) -> Option<SnapshotOutcome> {
log::info!(
"Snapshot upload starting from {} (task {task_id})",
path.display()
);
let declarations = read_and_parse_declarations(path, task_id)?;
log::info!("Snapshot upload starting from {}", path.display());
let declarations = read_and_parse_declarations(path)?;
let declarations = drop_files_covered_by_repos(declarations);
let (repo_count, file_count) = declarations
.iter()
.fold((0usize, 0usize), |(r, f), e| match e.kind {
@@ -470,14 +719,149 @@ async fn upload_snapshot_from_declarations_file(
EntryKind::File => (r, f + 1),
});
log::info!(
"Snapshot declarations: {} entries ({repo_count} repo, {file_count} file) (task {task_id})",
"Snapshot declarations: {} entries ({repo_count} repo, {file_count} file)",
declarations.len(),
);
let outcome = run_pipeline(declarations, client, task_id).await?;
log_snapshot_outcome(&outcome, task_id);
let outcome = run_pipeline(declarations, client).await?;
log_snapshot_outcome(&outcome);
Some(outcome)
}
/// Build the snapshot for a local-to-cloud handoff: gather repo patches and orphan file
/// contents, allocate an initial snapshot token plus presigned upload URLs via
/// `AIClient::upload_local_handoff_snapshot`, and upload the artifacts.
///
/// Returns:
/// - `Ok(Some(initial_snapshot_token))` when a token was minted **and the manifest landed in GCS**.
/// Individual blob uploads may still have failed; the manifest catalogues their status so the
/// cloud agent rehydrates against whatever did land, matching the cloud→cloud best-effort
/// posture.
/// - `Ok(None)` when the workspace was empty (no repos, no orphan files) **or** when the
/// manifest itself failed to upload. Without the manifest the snapshot is unusable, so
/// callers should spawn the cloud agent without an initial snapshot token instead of pointing
/// it at an incomplete prefix. Manifest-upload failures are also routed through
/// `report_error!` so on-call alerting catches the silent regression.
/// - `Err(_)` only for hard failures of `upload_local_handoff_snapshot` itself (auth, etc.).
pub(crate) async fn upload_snapshot_for_handoff(
repo_paths: Vec<PathBuf>,
orphan_file_paths: Vec<PathBuf>,
client: Arc<dyn AIClient>,
http: &http_client::Client,
) -> Result<Option<InitialSnapshotToken>> {
if repo_paths.is_empty() && orphan_file_paths.is_empty() {
log::info!("Handoff snapshot has no declarations; skipping upload");
return Ok(None);
}
let declarations: Vec<DeclarationEntry> = repo_paths
.into_iter()
.map(|path| DeclarationEntry {
kind: EntryKind::Repo,
path: path.display().to_string(),
})
.chain(orphan_file_paths.into_iter().map(|path| DeclarationEntry {
kind: EntryKind::File,
path: path.display().to_string(),
}))
.collect();
let GatheredSnapshot {
manifest_filename,
mut upload_files,
mut repos,
mut files,
mut pre_upload_entries,
} = gather_snapshot_entries(declarations).await;
apply_per_run_cap(
&mut upload_files,
&mut repos,
&mut files,
&mut pre_upload_entries,
);
let mut file_infos: Vec<SnapshotFileInfo> = upload_files
.iter()
.map(|file| SnapshotFileInfo {
filename: file.filename.clone(),
mime_type: file.mime_type.clone(),
})
.collect();
file_infos.push(SnapshotFileInfo {
filename: manifest_filename.clone(),
mime_type: "application/json".to_string(),
});
let upload_request = UploadLocalHandoffSnapshotRequest {
files: file_infos
.iter()
.map(|file| AiSnapshotUploadFileInfo {
filename: file.filename.clone(),
mime_type: file.mime_type.clone(),
})
.collect(),
};
let response = client
.upload_local_handoff_snapshot(upload_request)
.await
.context("failed to allocate initial snapshot token")?;
log::info!(
"Initial snapshot token allocated; expires_at={}, uploads={}",
response.expires_at,
response.uploads.len(),
);
let initial_snapshot_token = response.initial_snapshot_token;
// Server returns `uploads` aligned by index with the request `files` array (and does
// not echo per-entry filenames), so we zip them positionally into a filename-keyed map.
// Any request file the server omits lands in `upload_entry` with no target and is
// marked `skipped` downstream.
if response.uploads.len() != file_infos.len() {
log::warn!(
"Handoff snapshot upload-target response length {} does not match request length {}; \
extras will be marked skipped",
response.uploads.len(),
file_infos.len(),
);
}
let mut target_map: HashMap<String, UploadTarget> = HashMap::new();
for (file, target) in file_infos.iter().zip(response.uploads.into_iter()) {
target_map.insert(file.filename.clone(), target);
}
let Some(outcome) = upload_prepared_snapshot_files(
http,
manifest_filename,
upload_files,
repos,
files,
pre_upload_entries,
target_map,
)
.await
else {
// Manifest serialization failed (already reported via `report_error!` inside
// the helper). Without a manifest the snapshot is unusable, so refuse the token.
return Ok(None);
};
let summary = SnapshotSummary::from_entries(&outcome.entries, outcome.manifest_uploaded);
log_snapshot_outcome(&outcome);
if !summary.manifest_uploaded {
// Without the manifest the cloud agent has no catalogue to rehydrate from, even
// when individual blobs landed. Alert on-call and refuse the token so we don't
// silently spawn a cloud agent with no recoverable state.
report_error!(anyhow::anyhow!(
"Handoff snapshot manifest failed to upload (blobs: {}/{}); cloud agent will start with no rehydration content",
summary.uploaded,
summary.total,
));
return Ok(None);
}
Ok(Some(initial_snapshot_token))
}
/// Core upload pipeline.
///
/// Gather/read/upload failures are captured in [`SnapshotOutcome::entries`] and never abort
@@ -486,7 +870,6 @@ async fn upload_snapshot_from_declarations_file(
async fn run_pipeline(
declarations: Vec<DeclarationEntry>,
client: Arc<dyn HarnessSupportClient>,
task_id: &AmbientAgentTaskId,
) -> Option<SnapshotOutcome> {
let GatheredSnapshot {
manifest_filename,
@@ -494,11 +877,10 @@ async fn run_pipeline(
repos,
files,
pre_upload_entries,
} = gather_snapshot_entries(declarations, task_id).await;
} = gather_snapshot_entries(declarations).await;
upload_gathered_snapshot(
client,
task_id,
manifest_filename,
upload_files,
repos,
@@ -516,10 +898,7 @@ struct GatheredSnapshot {
pre_upload_entries: Vec<EntryResult>,
}
async fn gather_snapshot_entries(
declarations: Vec<DeclarationEntry>,
task_id: &AmbientAgentTaskId,
) -> GatheredSnapshot {
async fn gather_snapshot_entries(declarations: Vec<DeclarationEntry>) -> GatheredSnapshot {
let mut used_filenames = HashSet::new();
let manifest_filename = unique_filename("snapshot_state.json", &mut used_filenames);
@@ -542,7 +921,6 @@ async fn gather_snapshot_entries(
&mut upload_files,
&mut repos,
&mut pre_upload_entries,
task_id,
)
.await;
}
@@ -553,7 +931,6 @@ async fn gather_snapshot_entries(
&mut upload_files,
&mut files,
&mut pre_upload_entries,
task_id,
)
.await;
}
@@ -571,7 +948,6 @@ async fn gather_snapshot_entries(
async fn upload_gathered_snapshot(
client: Arc<dyn HarnessSupportClient>,
task_id: &AmbientAgentTaskId,
manifest_filename: String,
mut upload_files: Vec<SnapshotUploadFile>,
mut repos: Vec<RepoManifestEntry>,
@@ -587,7 +963,6 @@ async fn upload_gathered_snapshot(
&mut repos,
&mut files,
&mut pre_upload_entries,
task_id,
);
// Ask the server for presigned URLs for every filename we intend to upload —
@@ -619,16 +994,14 @@ async fn upload_gathered_snapshot(
Err(e) => {
// Pipeline-abort: route through report_error! so Sentry captures the structured
// error chain and on-call alerting can fire.
report_error!(e.context(format!(
"Failed to get snapshot upload targets; skipping upload (task {task_id})"
)));
report_error!(e.context("Failed to get snapshot upload targets; skipping upload"));
return None;
}
};
if targets.len() != chunk.len() {
log::warn!(
"Snapshot upload-target response length {} does not match request length {}; \
extras will be marked skipped (task {task_id})",
extras will be marked skipped",
targets.len(),
chunk.len(),
);
@@ -637,14 +1010,33 @@ async fn upload_gathered_snapshot(
target_map.insert(file.filename.clone(), target);
}
}
upload_prepared_snapshot_files(
client.http_client(),
manifest_filename,
upload_files,
repos,
files,
pre_upload_entries,
target_map,
)
.await
}
async fn upload_prepared_snapshot_files(
http: &http_client::Client,
manifest_filename: String,
upload_files: Vec<SnapshotUploadFile>,
mut repos: Vec<RepoManifestEntry>,
mut files: Vec<FileManifestEntry>,
pre_upload_entries: Vec<EntryResult>,
target_map: HashMap<String, UploadTarget>,
) -> Option<SnapshotOutcome> {
// Upload non-manifest blobs concurrently, each with bounded retries on transient errors.
let http = client.http_client();
let upload_futures = upload_files
.iter()
.map(|file| upload_entry(http, file, &target_map, task_id));
.map(|file| upload_entry(http, file, &target_map));
let upload_entries: Vec<EntryResult> = join_all(upload_futures).await;
fold_upload_results(&mut repos, &mut files, &upload_entries, task_id);
fold_upload_results(&mut repos, &mut files, &upload_entries);
// Build and upload the manifest last, with the real outcomes baked in.
let manifest = SnapshotManifest {
@@ -656,9 +1048,8 @@ async fn upload_gathered_snapshot(
Ok(b) => b,
Err(e) => {
// Pipeline-abort: route through report_error! so Sentry captures it.
report_error!(anyhow::Error::from(e).context(format!(
"Failed to serialize snapshot manifest; skipping upload (task {task_id})"
)));
report_error!(anyhow::Error::from(e)
.context("Failed to serialize snapshot manifest; skipping upload"));
return None;
}
};
@@ -671,9 +1062,7 @@ async fn upload_gathered_snapshot(
Err(e) => {
// Capture the full chain for the manifest's `error` field, then surface it
// to Sentry via report_error!.
let e = e.context(format!(
"Failed to upload manifest '{manifest_filename}' (task {task_id})"
));
let e = e.context(format!("Failed to upload manifest '{manifest_filename}'"));
let msg = format!("{e:#}");
report_error!(e);
(false, Some(msg))
@@ -714,7 +1103,6 @@ async fn gather_repo(
upload_files: &mut Vec<SnapshotUploadFile>,
repos: &mut Vec<RepoManifestEntry>,
pre_upload_entries: &mut Vec<EntryResult>,
task_id: &AmbientAgentTaskId,
) {
let repo = Path::new(repo_path);
let metadata = repo_metadata(repo).await;
@@ -756,7 +1144,7 @@ async fn gather_repo(
}
Err(e) => {
let err_str = format!("{e:#}");
log::warn!("Failed to snapshot repo '{repo_path}': {err_str} (task {task_id})");
log::warn!("Failed to snapshot repo '{repo_path}': {err_str}");
repos.push(RepoManifestEntry {
path: repo_path.to_string(),
repo_name: metadata.repo_name,
@@ -783,7 +1171,6 @@ async fn gather_file(
upload_files: &mut Vec<SnapshotUploadFile>,
files: &mut Vec<FileManifestEntry>,
pre_upload_entries: &mut Vec<EntryResult>,
task_id: &AmbientAgentTaskId,
) {
let path = Path::new(file_path);
match tokio::fs::read(path).await {
@@ -812,7 +1199,7 @@ async fn gather_file(
}
Err(e) => {
let err_str = format!("Failed to read file '{file_path}': {e:#}");
log::warn!("{err_str} (task {task_id})");
log::warn!("{err_str}");
files.push(FileManifestEntry {
path: file_path.to_string(),
snapshot_file: None,
@@ -836,13 +1223,9 @@ async fn upload_entry(
http: &http_client::Client,
file: &SnapshotUploadFile,
target_map: &HashMap<String, UploadTarget>,
task_id: &AmbientAgentTaskId,
) -> EntryResult {
let Some(target) = target_map.get(&file.filename) else {
log::warn!(
"No upload target for file '{}', skipping (task {task_id})",
file.filename
);
log::warn!("No upload target for file '{}', skipping", file.filename);
return EntryResult {
label: file.filename.clone(),
status: EntryStatus::Skipped,
@@ -860,10 +1243,7 @@ async fn upload_entry(
},
Err(e) => {
let msg = format!("{e:#}");
log::warn!(
"Failed to upload '{}': {msg} (task {task_id})",
file.filename
);
log::warn!("Failed to upload '{}': {msg}", file.filename);
EntryResult {
label: file.filename.clone(),
status: EntryStatus::Failed,
@@ -879,7 +1259,6 @@ fn fold_upload_results(
repos: &mut [RepoManifestEntry],
files: &mut [FileManifestEntry],
upload_entries: &[EntryResult],
task_id: &AmbientAgentTaskId,
) {
for entry in upload_entries {
if let Some(repo_entry) = repos
@@ -903,7 +1282,7 @@ fn fold_upload_results(
}
EntryStatus::GatherFailed | EntryStatus::ReadFailed => {
log::error!(
"fold_upload_results: unexpected pre-upload status {:?} for repo patch '{}' (task {task_id})",
"fold_upload_results: unexpected pre-upload status {:?} for repo patch '{}'",
entry.status,
entry.label
);
@@ -930,7 +1309,7 @@ fn fold_upload_results(
}
EntryStatus::GatherFailed | EntryStatus::ReadFailed => {
log::error!(
"fold_upload_results: unexpected pre-upload status {:?} for file '{}' (task {task_id})",
"fold_upload_results: unexpected pre-upload status {:?} for file '{}'",
entry.status,
entry.label
);
@@ -949,7 +1328,6 @@ fn apply_per_run_cap(
repos: &mut [RepoManifestEntry],
files: &mut [FileManifestEntry],
pre_upload_entries: &mut Vec<EntryResult>,
task_id: &AmbientAgentTaskId,
) {
let blob_limit = MAX_SNAPSHOT_FILES_PER_RUN.saturating_sub(1);
if upload_files.len() <= blob_limit {
@@ -958,7 +1336,7 @@ fn apply_per_run_cap(
let total_including_manifest = upload_files.len() + 1;
let dropped = upload_files.split_off(blob_limit);
log::warn!(
"Snapshot exceeds per-run cap of {MAX_SNAPSHOT_FILES_PER_RUN} files ({total_including_manifest} declared); dropping {} blob(s) from upload (task {task_id})",
"Snapshot exceeds per-run cap of {MAX_SNAPSHOT_FILES_PER_RUN} files ({total_including_manifest} declared); dropping {} blob(s) from upload",
dropped.len(),
);
let err_msg = format!("exceeded per-run snapshot cap of {MAX_SNAPSHOT_FILES_PER_RUN} files");
@@ -1011,13 +1389,14 @@ fn merge_content_type(target: &UploadTarget, mime_type: &str) -> UploadTarget {
url: target.url.clone(),
method: target.method.clone(),
headers,
fields: target.fields.clone(),
}
}
/// Log the final outcome at INFO when everything uploaded, WARN otherwise. The log line
/// includes per-entry statuses so operators can diagnose partial state without parsing any
/// downstream logs.
fn log_snapshot_outcome(outcome: &SnapshotOutcome, task_id: &AmbientAgentTaskId) {
fn log_snapshot_outcome(outcome: &SnapshotOutcome) {
let summary = SnapshotSummary::from_entries(&outcome.entries, outcome.manifest_uploaded);
let manifest_bit = if summary.manifest_uploaded {
"manifest: uploaded"
@@ -1025,7 +1404,7 @@ fn log_snapshot_outcome(outcome: &SnapshotOutcome, task_id: &AmbientAgentTaskId)
"manifest: failed"
};
let header = format!(
"Snapshot upload: {}/{} uploaded (failed: {}, skipped: {}, gather_failed: {}, read_failed: {}; {manifest_bit}) (task {task_id})",
"Snapshot upload: {}/{} uploaded (failed: {}, skipped: {}, gather_failed: {}, read_failed: {}; {manifest_bit})",
summary.uploaded,
summary.total,
summary.failed,
app/src/ai/agent_sdk/driver/snapshot_tests.rs
+311 -30
View File
@@ -86,6 +86,10 @@ impl HarnessSupportClient for TestClient {
unimplemented!("not used by upload_snapshot_from_declarations_file")
}
async fn fetch_transcript(&self) -> Result<bytes::Bytes> {
unimplemented!("not used by upload_snapshot_from_declarations_file")
}
async fn get_block_snapshot_upload_target(
&self,
_conversation_id: &AIConversationId,
@@ -112,6 +116,18 @@ impl HarnessSupportClient for TestClient {
unimplemented!("not used by upload_snapshot_from_declarations_file")
}
async fn report_clean_shutdown(&self) -> Result<()> {
unimplemented!("not used by upload_snapshot_from_declarations_file")
}
async fn report_error_shutdown(
&self,
_error_category: String,
_error_message: String,
) -> Result<()> {
unimplemented!("not used by upload_snapshot_from_declarations_file")
}
async fn get_snapshot_upload_targets(
&self,
request: &SnapshotUploadRequest,
@@ -131,6 +147,7 @@ impl HarnessSupportClient for TestClient {
url: format!("{}/upload/{}", self.server_base_url, f.filename),
method: "PUT".to_string(),
headers: HashMap::new(),
fields: Vec::new(),
})
.collect();
let keep = targets.len().saturating_sub(self.drop_trailing_targets);
@@ -138,10 +155,6 @@ impl HarnessSupportClient for TestClient {
Ok(targets)
}
async fn fetch_transcript(&self) -> Result<bytes::Bytes> {
unimplemented!("not used by upload_snapshot_from_declarations_file")
}
fn http_client(&self) -> &http_client::Client {
&self.http
}
@@ -248,7 +261,6 @@ fn run(
.block_on(upload_snapshot_from_declarations_file(
&declarations_path,
client,
&fake_task_id(),
))
.expect("pipeline returned None");
let summary = SnapshotSummary::from_entries(&outcome.entries, outcome.manifest_uploaded);
@@ -276,7 +288,7 @@ fn parse_declarations_ignores_blank_lines() {
"{\"version\":1,\"kind\":\"file\",\"path\":\"/abs/file.txt\"}\n",
"\n",
);
let entries = parse_declarations(contents, &fake_task_id());
let entries = parse_declarations(contents);
assert_eq!(
entries,
vec![
@@ -302,7 +314,7 @@ fn parse_declarations_skips_malformed_lines_without_aborting() {
"{\"version\":1,\"kind\":\"file\"}\n",
"{\"version\":1,\"kind\":\"file\",\"path\":\"/abs/also-good\",\"extra\":true}\n",
);
let entries = parse_declarations(contents, &fake_task_id());
let entries = parse_declarations(contents);
assert_eq!(
entries,
vec![
@@ -325,7 +337,7 @@ fn parse_declarations_skips_missing_or_unsupported_versions() {
"{\"version\":2,\"kind\":\"repo\",\"path\":\"/abs/unsupported-version\"}\n",
"{\"version\":1,\"kind\":\"file\",\"path\":\"/abs/good\"}\n",
);
let entries = parse_declarations(contents, &fake_task_id());
let entries = parse_declarations(contents);
assert_eq!(
entries,
vec![DeclarationEntry {
@@ -342,7 +354,7 @@ fn parse_declarations_tolerates_crlf_line_endings() {
"{\"version\":1,\"kind\":\"repo\",\"path\":\"/abs/good\"}\r\n",
"{\"version\":1,\"kind\":\"file\",\"path\":\"/abs/also-good\"}\r\n",
);
let entries = parse_declarations(contents, &fake_task_id());
let entries = parse_declarations(contents);
assert_eq!(
entries,
vec![
@@ -365,7 +377,7 @@ fn parse_declarations_skips_lines_with_empty_path() {
"{\"version\":1,\"kind\":\"file\",\"path\":\" \"}\n",
"{\"version\":1,\"kind\":\"repo\",\"path\":\"/abs/still-good\"}\n",
);
let entries = parse_declarations(contents, &fake_task_id());
let entries = parse_declarations(contents);
assert_eq!(
entries,
vec![DeclarationEntry {
@@ -382,7 +394,7 @@ fn parse_declarations_deduplicates_kind_path_pairs() {
"{\"version\":1,\"kind\":\"repo\",\"path\":\"/abs/repo\"}\n",
"{\"version\":1,\"kind\":\"file\",\"path\":\"/abs/repo\"}\n",
);
let entries = parse_declarations(contents, &fake_task_id());
let entries = parse_declarations(contents);
assert_eq!(
entries,
vec![
@@ -406,11 +418,7 @@ fn upload_skipped_when_declarations_file_missing() {
let client = TestClient::new(server.url());
let outcome = Runtime::new()
.unwrap()
.block_on(upload_snapshot_from_declarations_file(
&missing,
client,
&fake_task_id(),
));
.block_on(upload_snapshot_from_declarations_file(&missing, client));
assert!(
outcome.is_none(),
"missing declarations file should skip the upload"
@@ -426,11 +434,7 @@ fn upload_skipped_when_declarations_file_empty() {
let client = TestClient::new(server.url());
let outcome = Runtime::new()
.unwrap()
.block_on(upload_snapshot_from_declarations_file(
&decl,
client,
&fake_task_id(),
));
.block_on(upload_snapshot_from_declarations_file(&decl, client));
assert!(outcome.is_none(), "empty declarations file should skip");
}
@@ -447,11 +451,7 @@ fn upload_skipped_when_declarations_file_has_no_valid_jsonl_entries() {
let client = TestClient::new(server.url());
let outcome = Runtime::new()
.unwrap()
.block_on(upload_snapshot_from_declarations_file(
&decl,
client,
&fake_task_id(),
));
.block_on(upload_snapshot_from_declarations_file(&decl, client));
assert!(
outcome.is_none(),
"declarations file with no valid entries should skip"
@@ -631,10 +631,12 @@ fn e2e_gather_failed_entry_captured_in_manifest() {
#[test]
fn e2e_read_failed_for_missing_file_continues_pipeline() {
// Point a `file` entry at a path that doesn't exist → read_failed, with a clean repo also
// included so we verify the pipeline didn't abort after the read failure.
// included so we verify the pipeline didn't abort after the read failure. Keep the missing
// file outside the repo so the repo-overlap filter does not strip it before gather.
let tempdir = snaptest_tempdir();
init_git_repo(tempdir.path(), false);
let missing_file = tempdir.path().join("does-not-exist.txt");
let missing_dir = snaptest_tempdir();
let missing_file = missing_dir.path().join("does-not-exist.txt");
let decl_dir = snaptest_tempdir();
let mut server = Server::new();
@@ -978,7 +980,6 @@ fn e2e_get_snapshot_upload_targets_failure_returns_none() {
.block_on(upload_snapshot_from_declarations_file(
&declarations_path,
client,
&fake_task_id(),
));
assert!(
outcome.is_none(),
@@ -1022,7 +1023,6 @@ fn e2e_short_response_leaves_trailing_file_without_target() {
.block_on(upload_snapshot_from_declarations_file(
&declarations_path,
client,
&fake_task_id(),
))
.expect("pipeline returned None");
let summary = SnapshotSummary::from_entries(&outcome.entries, outcome.manifest_uploaded);
@@ -1182,3 +1182,284 @@ fn e2e_per_run_cap_drops_excess_blobs_as_skipped() {
}
upload_mock.assert();
}
// ------------------------------------------------------------------------------------------------
// REMOTE-1465: repo-overlap dedup + DeclarationsWriterHandle.
// ------------------------------------------------------------------------------------------------
/// Build a `DeclarationEntry` without exposing the private type to call sites.
fn repo_entry(path: &str) -> DeclarationEntry {
DeclarationEntry {
kind: EntryKind::Repo,
path: path.to_string(),
}
}
fn file_entry(path: &str) -> DeclarationEntry {
DeclarationEntry {
kind: EntryKind::File,
path: path.to_string(),
}
}
#[test]
fn drop_files_covered_by_repos_keeps_everything_when_no_repos_declared() {
let entries = vec![
file_entry("/abs/outside.txt"),
file_entry("/other/also-outside.txt"),
];
let after = drop_files_covered_by_repos(entries.clone());
assert_eq!(after, entries);
}
#[test]
fn drop_files_covered_by_repos_drops_file_inside_repo_keeps_file_outside() {
let entries = vec![
repo_entry("/workspace/my-repo"),
file_entry("/workspace/my-repo/src/foo.rs"),
file_entry("/tmp/outside.txt"),
];
let after = drop_files_covered_by_repos(entries);
assert_eq!(
after,
vec![
repo_entry("/workspace/my-repo"),
file_entry("/tmp/outside.txt"),
]
);
}
#[test]
fn drop_files_covered_by_repos_handles_nested_repo_paths() {
// A file under /a/b/sub should be filtered by either /a or /a/b/sub.
let entries = vec![
repo_entry("/a"),
repo_entry("/a/b/sub"),
file_entry("/a/b/sub/file.txt"),
file_entry("/a/top.txt"),
file_entry("/unrelated.txt"),
];
let after = drop_files_covered_by_repos(entries);
assert_eq!(
after,
vec![
repo_entry("/a"),
repo_entry("/a/b/sub"),
file_entry("/unrelated.txt"),
]
);
}
/// Parse the declarations file written by `DeclarationsWriterHandle` into the paths we care
/// about for assertions, ignoring any lines the helper tests weren't asked to produce.
fn parsed_file_paths(path: &Path) -> Vec<String> {
let contents = fs::read_to_string(path).unwrap_or_default();
let entries = parse_declarations(&contents);
entries
.into_iter()
.filter(|e| e.kind == EntryKind::File)
.map(|e| e.path)
.collect()
}
#[test]
fn declarations_writer_appends_unique_absolute_paths_once() {
let tempdir = snaptest_tempdir();
let workspace = tempdir.path().join("workspace");
fs::create_dir_all(&workspace).unwrap();
let decl_path = tempdir.path().join("declarations.jsonl");
let task_id = fake_task_id();
let rt = Runtime::new().unwrap();
rt.block_on(async {
let handle =
DeclarationsWriterHandle::new_at_path(decl_path.clone(), workspace.clone(), task_id);
let path_one = workspace.join("one.txt");
let path_two = workspace.join("two.txt");
handle.append(vec![
path_one.to_string_lossy().into_owned(),
path_two.to_string_lossy().into_owned(),
// Duplicate: should still only produce one entry per unique path.
path_one.to_string_lossy().into_owned(),
]);
handle.flush().await;
// A second batch that re-declares path_one should also be a no-op.
handle.append(vec![path_one.to_string_lossy().into_owned()]);
handle.flush().await;
});
let paths = parsed_file_paths(&decl_path);
assert_eq!(
paths,
vec![
workspace.join("one.txt").to_string_lossy().into_owned(),
workspace.join("two.txt").to_string_lossy().into_owned(),
]
);
}
#[test]
fn declarations_writer_resolves_relative_paths_against_working_dir() {
let tempdir = snaptest_tempdir();
let workspace = tempdir.path().join("workspace");
fs::create_dir_all(&workspace).unwrap();
let decl_path = tempdir.path().join("declarations.jsonl");
let task_id = fake_task_id();
let rt = Runtime::new().unwrap();
rt.block_on(async {
let handle =
DeclarationsWriterHandle::new_at_path(decl_path.clone(), workspace.clone(), task_id);
handle.append(vec!["notes/relative.txt".to_string()]);
handle.flush().await;
});
let paths = parsed_file_paths(&decl_path);
assert_eq!(
paths,
vec![workspace
.join("notes/relative.txt")
.to_string_lossy()
.into_owned()]
);
}
#[test]
fn declarations_writer_continues_after_per_path_write_failures() {
// Pre-create a directory at the declarations file path so the first append's open call
// fails. Once we remove the directory, a subsequent append must succeed, proving the
// writer task absorbed the failure and kept servicing commands.
let tempdir = snaptest_tempdir();
let workspace = tempdir.path().join("workspace");
fs::create_dir_all(&workspace).unwrap();
let decl_path = tempdir.path().join("declarations.jsonl");
fs::create_dir(&decl_path).unwrap();
let task_id = fake_task_id();
let rt = Runtime::new().unwrap();
rt.block_on(async {
let handle =
DeclarationsWriterHandle::new_at_path(decl_path.clone(), workspace.clone(), task_id);
handle.append(vec![workspace
.join("first.txt")
.to_string_lossy()
.into_owned()]);
handle.flush().await;
// Replace the staged directory so the next append's open call can succeed.
fs::remove_dir(&decl_path).unwrap();
handle.append(vec![workspace
.join("second.txt")
.to_string_lossy()
.into_owned()]);
handle.flush().await;
});
let paths = parsed_file_paths(&decl_path);
assert_eq!(
paths,
vec![workspace.join("second.txt").to_string_lossy().into_owned()],
"writer task should absorb the first failure and process the second append"
);
}
#[test]
fn declarations_writer_preempts_paths_inside_existing_repo() {
let tempdir = snaptest_tempdir();
// Simulate an existing repo by creating the `.git` directory the ancestor walker checks.
let repo = tempdir.path().join("existing-repo");
fs::create_dir_all(repo.join(".git")).unwrap();
let inside = repo.join("inside.txt");
let outside = tempdir.path().join("outside.txt");
let decl_path = tempdir.path().join("declarations.jsonl");
let task_id = fake_task_id();
let rt = Runtime::new().unwrap();
rt.block_on(async {
let handle = DeclarationsWriterHandle::new_at_path(
decl_path.clone(),
tempdir.path().to_path_buf(),
task_id,
);
handle.append(vec![
inside.to_string_lossy().into_owned(),
outside.to_string_lossy().into_owned(),
]);
handle.flush().await;
});
let paths = parsed_file_paths(&decl_path);
assert_eq!(paths, vec![outside.to_string_lossy().into_owned()]);
}
#[test]
fn e2e_repo_plus_inside_and_outside_files_filters_overlap() {
// The writer-written declarations file feeds straight into the upload pipeline. Pair one
// `repo` with two `file` entries (one inside the repo, one outside). The gather-time
// overlap filter should drop the inside-repo file entry before upload so only the repo's
// patch + the outside-repo file + the manifest land on the server.
let repo_dir = snaptest_tempdir();
init_git_repo(repo_dir.path(), true);
let inside_file = repo_dir.path().join("new-untracked.txt");
fs::write(&inside_file, b"tracked-or-not, handled by the patch\n").unwrap();
let outside_dir = snaptest_tempdir();
let outside_file = outside_dir.path().join("standalone_log.txt");
fs::write(&outside_file, b"agent-produced log\n").unwrap();
let decl_dir = snaptest_tempdir();
let decl_path = decl_dir.path().join("snapshot-declarations.jsonl");
let contents = format!(
concat!(
"{{\"version\":1,\"kind\":\"repo\",\"path\":{repo:?}}}\n",
"{{\"version\":1,\"kind\":\"file\",\"path\":{inside:?}}}\n",
"{{\"version\":1,\"kind\":\"file\",\"path\":{outside:?}}}\n",
),
repo = repo_dir.path().to_string_lossy(),
inside = inside_file.to_string_lossy(),
outside = outside_file.to_string_lossy(),
);
fs::write(&decl_path, contents).unwrap();
let mut server = Server::new();
let patch_mock = server
.mock("PUT", upload_path(r".+\.patch"))
.with_status(200)
.expect(1)
.create();
let file_mock = server
.mock("PUT", upload_path("standalone_log\\.txt"))
.with_status(200)
.expect(1)
.create();
let manifest_mock = server
.mock("PUT", upload_path("snapshot_state\\.json"))
.match_body(Matcher::PartialJson(serde_json::json!({
"files": [
{
"path": outside_file.to_string_lossy(),
"status": "uploaded",
"uploaded": true,
}
],
})))
.with_status(200)
.expect(1)
.create();
let client = TestClient::new(server.url());
let outcome = Runtime::new()
.unwrap()
.block_on(upload_snapshot_from_declarations_file(&decl_path, client))
.expect("pipeline returned None");
let summary = SnapshotSummary::from_entries(&outcome.entries, outcome.manifest_uploaded);
assert!(summary.all_uploaded(), "expected all uploads to succeed");
// repo patch + outside file + manifest = 3 uploaded entries total; the inside-repo file
// entry was filtered before gather so it never hits the entries list.
assert_eq!(summary.uploaded, 3);
assert_eq!(summary.total, 3);
assert!(outcome.manifest_uploaded);
patch_mock.assert();
file_mock.assert();
manifest_mock.assert();
}

Some files were not shown because too many files have changed in this diff Show More