diff --git a/app/src/ai/acp/transport.rs b/app/src/ai/acp/transport.rs index 37992bc3..04db119a 100644 --- a/app/src/ai/acp/transport.rs +++ b/app/src/ai/acp/transport.rs @@ -172,7 +172,7 @@ pub(crate) async fn acp_output_stream( match translator.translate_steered_user_content(content) { Ok(response_events) => { for response_event in response_events { - yield Ok(response_event); + yield Ok(api::StreamEvent::Response(response_event)); } } Err(message) => { @@ -198,20 +198,20 @@ pub(crate) async fn acp_output_stream( } Ok(AcpSteeringOutcome::StartedNewTurn) => { for response_event in translator.steering_started_new_turn() { - yield Ok(response_event); + yield Ok(api::StreamEvent::Response(response_event)); } } Ok(AcpSteeringOutcome::Failed) => { for response_event in translator.steering_failed( "the ACP agent could not inject it into the active turn", ) { - yield Ok(response_event); + yield Ok(api::StreamEvent::Response(response_event)); } } Err(error) => { log::warn!("ACP live steering failed: {error}"); for response_event in translator.steering_failed(&error.to_string()) { - yield Ok(response_event); + yield Ok(api::StreamEvent::Response(response_event)); } } } @@ -250,7 +250,7 @@ pub(crate) async fn acp_output_stream( match translator.translate(event) { Ok(response_events) => { for response_event in response_events { - yield Ok(response_event); + yield Ok(api::StreamEvent::Response(response_event)); } } Err(message) => { @@ -363,7 +363,11 @@ fn translated_startup_error_stream( message: &str, ) -> api::ResponseStream { let events = translator.startup_error(message); - Box::pin(futures::stream::iter(events.into_iter().map(Ok))) + Box::pin(futures::stream::iter( + events + .into_iter() + .map(|event| Ok(api::StreamEvent::Response(event))), + )) } #[cfg(test)] diff --git a/app/src/ai/agent/api.rs b/app/src/ai/agent/api.rs index 7cf4f601..28295f4e 100644 --- a/app/src/ai/agent/api.rs +++ b/app/src/ai/agent/api.rs @@ -24,7 +24,7 @@ use mcp::TemplatableMCPServerInfo; pub use r#impl::generate_multi_agent_output; use serde::Serialize; -use super::{AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions}; +use super::{AIAgentAction, AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions}; use crate::ai::agent::conversation::AIConversationId; use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput, SessionContext}; @@ -162,17 +162,35 @@ pub struct RequestParams { pub global_rules: Vec<(String, String)>, } -pub type Event = Result>; +/// Provider/runtime events consumed by the local conversation controller. +/// +/// The legacy response envelope remains at the UI boundary while the local Rig runtime emits +/// executable tool proposals directly as Galaxy domain actions. This avoids translating Rig tool +/// calls into protobuf only to immediately translate them back before execution. +#[derive(Debug)] +pub enum StreamEvent { + Response(warp_multi_agent_api::ResponseEvent), + ToolProposed(AIAgentAction), +} + +pub type Event = Result>; +pub type LegacyEvent = Result>; #[cfg(not(target_family = "wasm"))] pub type ResponseStream = Pin + Send + 'static>>; +#[cfg(not(target_family = "wasm"))] +pub type LegacyResponseStream = Pin + Send + 'static>>; + // The WASM version of this type has no bound on `Send`, which is an unnecessary bound when // targeting wasm because the browser is single-threaded (and we don't leverage WebWorkers for async // execution in WoW). #[cfg(target_family = "wasm")] pub type ResponseStream = Pin>>; +#[cfg(target_family = "wasm")] +pub type LegacyResponseStream = Pin>>; + #[derive(Debug, Clone)] pub struct ConversationData { pub id: AIConversationId, diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index 5ded002d..1b2529c4 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -6,7 +6,7 @@ use galaxy_core::features::FeatureFlag; use warp_multi_agent_api as api; use super::convert_to::convert_input; -use super::{ConvertToAPITypeError, RequestParams, ResponseStream}; +use super::{ConvertToAPITypeError, RequestParams, ResponseStream, StreamEvent}; use crate::ai::agent::redaction; use crate::ai::openai::translator as openai_translator; use crate::ai::provider::ProviderConfig; @@ -164,7 +164,9 @@ pub async fn generate_multi_agent_output( match openai_translator::execute(translator_request, &mut request).await { Ok(stream) => { - let output_stream = stream.take_until(cancellation_rx); + let output_stream = stream + .map(|event| event.map(StreamEvent::Response)) + .take_until(cancellation_rx); Ok(Box::pin(output_stream)) } Err(e) => { @@ -196,7 +198,9 @@ pub async fn generate_multi_agent_output( match crate::ai::bedrock::translator::execute(translator_request, &mut request).await { Ok(stream) => { - let output_stream = stream.take_until(cancellation_rx); + let output_stream = stream + .map(|event| event.map(StreamEvent::Response)) + .take_until(cancellation_rx); Ok(Box::pin(output_stream)) } Err(e) => { diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index e04e1ec0..5e134cc4 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -2227,6 +2227,47 @@ impl AIConversation { Ok(()) } + pub fn apply_domain_tool_proposal( + &mut self, + stream_id: &ResponseStreamId, + terminal_surface_id: EntityId, + action: AIAgentAction, + ctx: &mut ModelContext, + ) -> Result<(), UpdateConversationError> { + let added_exchanges = self + .added_exchanges_by_response + .get(stream_id) + .ok_or(UpdateConversationError::NoPendingRequest)?; + let exchange_id = added_exchanges + .iter() + .find(|added| added.task_id == action.task_id) + .map(|added| added.exchange_id) + .ok_or(UpdateConversationError::TaskNotFound)?; + let message_id = MessageId::new(action.id.to_string()); + let exchange = self.get_exchange_to_update(exchange_id)?; + match &exchange.output_status { + AIAgentOutputStatus::Streaming { + output: Some(output), + } => output + .get_mut() + .messages + .push(AIAgentOutputMessage::action(message_id, action)), + AIAgentOutputStatus::Streaming { output: None } => { + return Err(UpdateConversationError::OutputNeverInitialized); + } + AIAgentOutputStatus::Finished { .. } => { + return Err(UpdateConversationError::OutputAlreadyFinished); + } + } + ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange { + exchange_id, + terminal_surface_id, + conversation_id: self.id, + is_hidden: self.hidden_exchanges.contains(&exchange_id), + }); + Ok(()) + } + pub fn update_cost_and_usage_for_request( &mut self, request_cost: Option, diff --git a/app/src/ai/bedrock/client.rs b/app/src/ai/bedrock/client.rs index 6cf06950..fd133056 100644 --- a/app/src/ai/bedrock/client.rs +++ b/app/src/ai/bedrock/client.rs @@ -11,7 +11,7 @@ use super::diagnostic::BedrockDiagnosticLogger; use super::external_config::ExternalBedrockConfig; use super::models::apply_cross_region_prefix; use super::response_translator::bedrock_stream_to_response_events; -use crate::ai::agent::api::ResponseStream; +use crate::ai::agent::api::LegacyResponseStream; use crate::settings::ai::BedrockAuthMethod; fn strip_context_marker(model_id: &str) -> String { @@ -185,7 +185,7 @@ impl BedrockClient { diagnostic_logger: Option>, messages_sent: Arc>>, tool_result_archive: Vec, - ) -> Result { + ) -> Result { let base_model_id = strip_context_marker(model_id); let effective_model_id = if cross_region_inference { apply_cross_region_prefix(&base_model_id, &self.region) diff --git a/app/src/ai/bedrock/response_translator.rs b/app/src/ai/bedrock/response_translator.rs index c468425a..0b6b11f1 100644 --- a/app/src/ai/bedrock/response_translator.rs +++ b/app/src/ai/bedrock/response_translator.rs @@ -15,7 +15,7 @@ use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; use super::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole}; use super::diagnostic::BedrockDiagnosticLogger; -use crate::ai::agent::api::Event; +use crate::ai::agent::api::LegacyEvent; use crate::server::server_api::AIApiError; fn json_to_prost_struct(value: &serde_json::Value) -> prost_types::Struct { @@ -69,7 +69,7 @@ pub fn bedrock_stream_to_response_events( messages_sent: Arc>>, model_id: String, tool_result_archive: Vec, -) -> BoxStream<'static, Event> { +) -> BoxStream<'static, LegacyEvent> { let request_id = Uuid::new_v4().to_string(); let conversation_id = Uuid::new_v4().to_string(); diff --git a/app/src/ai/bedrock/translator.rs b/app/src/ai/bedrock/translator.rs index c60d9ccf..53817b74 100644 --- a/app/src/ai/bedrock/translator.rs +++ b/app/src/ai/bedrock/translator.rs @@ -4,7 +4,7 @@ use std::sync::{Arc, Mutex}; use warp_multi_agent_api as api; -use crate::ai::agent::api::ResponseStream; +use crate::ai::agent::api::LegacyResponseStream; use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig, BedrockError}; use crate::ai::bedrock::convert::ConversationMessage; use crate::ai::bedrock::diagnostic::BedrockDiagnosticLogger; @@ -25,7 +25,7 @@ pub struct TranslatorRequest { pub async fn execute( params: TranslatorRequest, request: &mut api::Request, -) -> Result { +) -> Result { let config = params.config.with_external_fallbacks(); let cross_region_inference = config.cross_region_inference; let bedrock = BedrockClient::from_config(config).await?; diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index ffcd80ff..8f4c2fcc 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -267,9 +267,15 @@ fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bo } else { ToolResultStatus::Success }; + let content = action_result.result.model_content(); + let content = if permission_denied { + format!("Permission denied by the user. {content}") + } else { + content + }; ToolResult { call_id: action_result.id.to_string(), - content: action_result.result.model_content(), + content, status, } } diff --git a/app/src/ai/blocklist/action_model_tests.rs b/app/src/ai/blocklist/action_model_tests.rs index 3b19c4d0..e8a8346e 100644 --- a/app/src/ai/blocklist/action_model_tests.rs +++ b/app/src/ai/blocklist/action_model_tests.rs @@ -139,6 +139,8 @@ fn domain_tool_results_preserve_success_failure_cancellation_and_denial() { assert_eq!(failure.call_id, "failure"); assert_eq!(cancelled.call_id, "cancelled"); assert_eq!(denied.call_id, "cancelled"); + assert!(!cancelled.content.contains("Permission denied")); + assert!(denied.content.contains("Permission denied by the user")); } #[test] diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index f83ce0cb..425b66d7 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -3423,7 +3423,23 @@ impl BlocklistAIController { }; let history_model = BlocklistAIHistoryModel::handle(ctx); match event { - Ok(event) => { + Ok(api::StreamEvent::ToolProposed(action)) => { + let apply_result = history_model.update(ctx, |history_model, ctx| { + history_model.apply_domain_tool_proposal( + &stream_id, + conversation_id, + self.terminal_surface_id, + action, + ctx, + ) + }); + if let Err(error) = apply_result { + log::error!( + "Failed to apply Rig tool proposal to conversation: {error:?}" + ); + } + } + Ok(api::StreamEvent::Response(event)) => { // If this controller is part of a shared session, forward the entire response event to viewers first. if FeatureFlag::AgentSharedSessions.is_enabled() && !response_stream.as_ref(ctx).is_acp() diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index f9296c01..fade9c19 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -186,7 +186,7 @@ impl ResponseStream { ctx: &mut ModelContext, ) { ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(Ok( - event, + api::StreamEvent::Response(event), )))); } #[cfg(test)] @@ -718,7 +718,16 @@ impl ResponseStream { self.time_to_latest_event = Local::now().signed_duration_since(self.start_time); match &event { - Ok(response_event) => { + Ok(api::StreamEvent::ToolProposed(action)) => { + self.has_received_client_actions = true; + log::debug!( + "Rig proposed domain tool action {} for task {}", + action.id, + action.task_id + ); + ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(event))); + } + Ok(api::StreamEvent::Response(response_event)) => { let event_type_name = match &response_event.r#type { Some(warp_multi_agent_api::response_event::Type::Init(_)) => "Init", Some(warp_multi_agent_api::response_event::Type::ClientActions(a)) => { diff --git a/app/src/ai/blocklist/history_model.rs b/app/src/ai/blocklist/history_model.rs index 36b43257..7f11636d 100644 --- a/app/src/ai/blocklist/history_model.rs +++ b/app/src/ai/blocklist/history_model.rs @@ -33,9 +33,9 @@ use crate::ai::agent::conversation::{ use crate::ai::agent::task::helper::{MessageExt, ToolCallExt}; use crate::ai::agent::task::TaskId; use crate::ai::agent::{ - AIAgentActionId, AIAgentExchange, AIAgentExchangeId, AIAgentInput, AIAgentOutputStatus, - CancellationReason, FinishedAIAgentOutput, MessageId, RenderableAIError, RequestCost, - Suggestions, + AIAgentAction, AIAgentActionId, AIAgentExchange, AIAgentExchangeId, AIAgentInput, + AIAgentOutputStatus, CancellationReason, FinishedAIAgentOutput, MessageId, RenderableAIError, + RequestCost, Suggestions, }; use crate::ai::artifacts::Artifact; use crate::ai::document::ai_document_model::AIDocumentModel; @@ -1852,6 +1852,21 @@ impl BlocklistAIHistoryModel { Ok(()) } + pub fn apply_domain_tool_proposal( + &mut self, + response_stream_id: &ResponseStreamId, + conversation_id: AIConversationId, + terminal_surface_id: EntityId, + action: AIAgentAction, + ctx: &mut ModelContext, + ) -> Result<(), UpdateHistoryError> { + self.conversations_by_id + .get_mut(&conversation_id) + .ok_or(UpdateHistoryError::ConversationNotFound(conversation_id))? + .apply_domain_tool_proposal(response_stream_id, terminal_surface_id, action, ctx)?; + Ok(()) + } + pub fn update_conversation_cost_and_usage_for_request( &mut self, conversation_id: AIConversationId, diff --git a/app/src/ai/blocklist/passive_suggestions/maa.rs b/app/src/ai/blocklist/passive_suggestions/maa.rs index 71d5b974..4622cf1d 100644 --- a/app/src/ai/blocklist/passive_suggestions/maa.rs +++ b/app/src/ai/blocklist/passive_suggestions/maa.rs @@ -486,7 +486,7 @@ async fn extract_suggestion_from_stream( let mut client_actions: Vec = Vec::new(); let mut server_request_token: Option = None; while let Some(event) = stream.next().await { - let Ok(response_event) = event else { + let Ok(crate::ai::agent::api::StreamEvent::Response(response_event)) = event else { continue; }; match response_event.r#type { diff --git a/app/src/ai/openai/response_translator.rs b/app/src/ai/openai/response_translator.rs index f6024399..5565f07b 100644 --- a/app/src/ai/openai/response_translator.rs +++ b/app/src/ai/openai/response_translator.rs @@ -9,7 +9,7 @@ use uuid::Uuid; use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; -use crate::ai::agent::api::Event; +use crate::ai::agent::api::LegacyEvent; use crate::ai::bedrock::response_translator::{ build_create_task, build_stream_init, context_window_for_model, }; @@ -47,7 +47,7 @@ pub(crate) struct StreamUsage { pub fn openai_stream_to_response_events( byte_stream: impl Stream> + Send + 'static, context: OpenAIStreamContext, -) -> BoxStream<'static, Event> { +) -> BoxStream<'static, LegacyEvent> { use futures::StreamExt; let OpenAIStreamContext { diff --git a/app/src/ai/openai/translator.rs b/app/src/ai/openai/translator.rs index 03bcbe9a..e11d9db5 100644 --- a/app/src/ai/openai/translator.rs +++ b/app/src/ai/openai/translator.rs @@ -6,7 +6,7 @@ use super::client::{OpenAIClient, OpenAIClientConfig, OpenAIError}; use super::convert::build_openai_request; use super::request_translator::sanitize_messages_for_openai; use super::response_translator::{openai_stream_to_response_events, OpenAIStreamContext}; -use crate::ai::agent::api::ResponseStream; +use crate::ai::agent::api::LegacyResponseStream; use crate::ai::bedrock::request_translator; use crate::ai::provider::types::{ConversationMessage, MessageContent, MessageRole}; @@ -108,7 +108,7 @@ pub(crate) fn prepare_turn(params: &TranslatorRequest, request: &mut api::Reques pub async fn execute( params: TranslatorRequest, request: &mut api::Request, -) -> Result { +) -> Result { let client = OpenAIClient::from_config(params.config.clone()); let PreparedTurn { task_id, diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index 8c894658..8a6797b7 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -1,6 +1,7 @@ mod provider; mod rig; mod rig_request; +mod rig_tool; pub(crate) use provider::ProviderRuntime; pub(crate) use rig::rig_openai_response_stream; diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index b430253e..d48eae5e 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -12,7 +12,9 @@ use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent, ToolType}; use super::rig_request::{prepare_rig_turn, PreparedRigTurn}; -use crate::ai::agent::api::{Event, RequestParams, ResponseStream}; +use super::rig_tool::action_from_tool_call; +use crate::ai::agent::api::{Event, RequestParams, ResponseStream, StreamEvent}; +use crate::ai::agent::AIAgentAction; use crate::ai::bedrock::response_translator::{ build_add_agent_output_message, build_append_text, build_create_task, build_stream_init, build_user_query_message, @@ -29,6 +31,7 @@ pub(crate) fn rig_openai_response_stream( supported_cli_agent_tools: Vec, cancellation_rx: oneshot::Receiver<()>, ) -> ResponseStream { + let skill_path_origin = params.session_context.skill_path_origin(); let PreparedRigTurn { task_id, needs_create_task, @@ -111,30 +114,30 @@ pub(crate) fn rig_openai_response_stream( match event { AgentEvent::TurnStarted { .. } => { initialized = true; - yield Ok(build_stream_init(&request_id, &conversation_id)); + yield Ok(StreamEvent::Response(build_stream_init(&request_id, &conversation_id))); if needs_create_task { - yield Ok(build_create_task(&task_id)); + yield Ok(StreamEvent::Response(build_create_task(&task_id))); } if let Some(user_query) = &user_query { - yield Ok(build_user_query_message(&task_id, user_query)); + yield Ok(StreamEvent::Response(build_user_query_message(&task_id, user_query))); } } AgentEvent::TextDelta { text } => { full_text.push_str(&text); if let Some(message_id) = ¤t_text_message_id { - yield Ok(build_append_text(&task_id, message_id, &text)); + yield Ok(StreamEvent::Response(build_append_text(&task_id, message_id, &text))); } else { let message_id = Uuid::new_v4().to_string(); - yield Ok(build_add_agent_output_message(&task_id, &message_id, &text)); + yield Ok(StreamEvent::Response(build_add_agent_output_message(&task_id, &message_id, &text))); current_text_message_id = Some(message_id); } } AgentEvent::ReasoningDelta { text } => { if let Some(message_id) = ¤t_reasoning_message_id { - yield Ok(build_append_reasoning(&task_id, message_id, &text)); + yield Ok(StreamEvent::Response(build_append_reasoning(&task_id, message_id, &text))); } else { let message_id = Uuid::new_v4().to_string(); - yield Ok(build_add_reasoning(&task_id, &message_id, &text)); + yield Ok(StreamEvent::Response(build_add_reasoning(&task_id, &message_id, &text))); current_reasoning_message_id = Some(message_id); } } @@ -155,7 +158,16 @@ pub(crate) fn rig_openai_response_stream( .unwrap_or_default(); match tool_policy.decide(&call, &history, &tool_result_archive) { ToolCallDecision::Execute => { - yield Ok(build_tool_proposed(&task_id, &call)); + match build_tool_proposed(&task_id, &call, &skill_path_origin) { + Ok(action) => yield Ok(StreamEvent::ToolProposed(action)), + Err(message) => { + yield Err(agent_error(AgentError::new( + galaxy_agent_core::AgentErrorKind::Protocol, + message, + ))); + return; + } + } } ToolCallDecision::Inline(result) => { append_tool_result(&messages_sent, result); @@ -172,17 +184,17 @@ pub(crate) fn rig_openai_response_stream( ); append_tool_result(&messages_sent, result); let message_id = Uuid::new_v4().to_string(); - yield Ok(build_add_agent_output_message( + yield Ok(StreamEvent::Response(build_add_agent_output_message( &task_id, &message_id, &error_display, - )); + ))); } } } AgentEvent::TurnStopped { reason } => { if !initialized { - yield Ok(build_stream_init(&request_id, &conversation_id)); + yield Ok(StreamEvent::Response(build_stream_init(&request_id, &conversation_id))); } sync_assistant_turn( &messages_sent, @@ -190,7 +202,7 @@ pub(crate) fn rig_openai_response_stream( &proposed_tools, &mut assistant_history_index, ); - yield Ok(build_stream_finished( + yield Ok(StreamEvent::Response(build_stream_finished( map_stop_reason(reason), StreamUsage { input_tokens: saturating_i32(usage.input_tokens), @@ -203,7 +215,7 @@ pub(crate) fn rig_openai_response_stream( model_id, max_context_tokens, }, - )); + ))); return; } AgentEvent::Tool { .. } => { @@ -304,11 +316,12 @@ fn sync_assistant_turn( sent.push(message); } -fn build_tool_proposed(task_id: &str, call: &ToolCall) -> ResponseEvent { - let arguments = serde_json::to_string(&call.arguments).unwrap_or_else(|_| "{}".to_string()); - crate::ai::bedrock::response_translator::build_tool_call_message( - task_id, &call.id, &call.name, &arguments, - ) +fn build_tool_proposed( + task_id: &str, + call: &ToolCall, + skill_path_origin: &ai::skills::SkillPathOrigin, +) -> Result { + action_from_tool_call(task_id, call, skill_path_origin) } fn build_add_reasoning(task_id: &str, message_id: &str, text: &str) -> ResponseEvent { diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index f6277ad9..b0b82fe5 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -395,7 +395,7 @@ fn tool_definitions( .collect::>(); for server in &mcp_context.servers { for tool in &server.tools { - let name = format!("mcp__{}__{}", server.name, tool.name); + let name = format!("mcp__{}__{}", server.id, tool.name); if seen.insert(name.clone()) { tools.push(ToolDefinition { name, diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index 1bc9e6d1..af228c83 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -4,9 +4,11 @@ use std::sync::Arc; use galaxy_agent_core::{ContentPart, MessageContent, MessageRole, ToolResult, ToolResultStatus}; use warp_multi_agent_api::ToolType; -use super::{input_messages, prepare_rig_turn}; +use super::{input_messages, prepare_rig_turn, tool_definitions}; use crate::ai::agent::api::RequestParams; -use crate::ai::agent::{AIAgentContext, AIAgentInput, AnyFileContent, FileContext, UserQueryMode}; +use crate::ai::agent::{ + AIAgentContext, AIAgentInput, AnyFileContent, FileContext, MCPContext, MCPServer, UserQueryMode, +}; use crate::ai::llms::LLMId; use crate::ai::openai::client::OpenAIClientConfig; @@ -112,6 +114,40 @@ fn builds_a_rig_turn_directly_from_galaxy_request_state() { )); } +#[test] +#[allow(deprecated)] +fn grouped_mcp_tool_names_use_the_installation_id_not_the_display_name() { + let tool = serde_json::from_value(serde_json::json!({ + "name": "echo", + "description": "Echo input", + "inputSchema": { + "type": "object", + "properties": {"message": {"type": "string"}} + } + })) + .unwrap(); + let context = MCPContext { + resources: Vec::new(), + tools: Vec::new(), + servers: vec![MCPServer { + id: "11111111-1111-4111-8111-111111111111".to_string(), + name: "Friendly Server".to_string(), + description: String::new(), + resources: Vec::new(), + tools: vec![tool], + }], + }; + + let tools = tool_definitions(&[ToolType::CallMcpTool], Some(&context)); + + assert!(tools + .iter() + .any(|tool| { tool.name == "mcp__11111111-1111-4111-8111-111111111111__echo" })); + assert!(!tools + .iter() + .any(|tool| tool.name == "mcp__Friendly Server__echo")); +} + #[test] fn normalized_tool_outcomes_are_the_only_action_results_sent_to_rig() { let statuses = [ diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index 322aa777..22f0278f 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use ai::skills::SkillPathOrigin; use galaxy_agent_core::{ MessageContent, MessageRole, StopReason, ToolCall, ToolResult, ToolResultStatus, }; @@ -66,8 +67,8 @@ fn reasoning_events_match_the_existing_ui_message_contract() { } #[test] -fn tool_proposal_matches_the_existing_permission_ui_contract() { - let event = build_tool_proposed( +fn tool_proposal_matches_the_domain_permission_contract() { + let action = build_tool_proposed( "task", &ToolCall { id: "call-1".to_string(), @@ -77,67 +78,43 @@ fn tool_proposal_matches_the_existing_permission_ui_contract() { "is_read_only": true }), }, - ); + &SkillPathOrigin::Local, + ) + .unwrap(); - let Some(warp_multi_agent_api::response_event::Type::ClientActions(actions)) = event.r#type - else { - panic!("expected client actions"); - }; - let Some(warp_multi_agent_api::client_action::Action::AddMessagesToTask(add)) = - &actions.actions[0].action - else { - panic!("expected add-message action"); - }; - let Some(warp_multi_agent_api::message::Message::ToolCall(tool_call)) = - &add.messages[0].message - else { - panic!("expected tool-call message"); - }; - let Some(warp_multi_agent_api::message::tool_call::Tool::RunShellCommand(command)) = - &tool_call.tool - else { - panic!("expected run-shell-command payload"); - }; - - assert_eq!(tool_call.tool_call_id, "call-1"); - assert_eq!(command.command, "cargo test"); - assert!(command.is_read_only); + assert_eq!(action.id.to_string(), "call-1"); + assert!(matches!( + action.action, + crate::ai::agent::AIAgentActionType::RequestCommandOutput { + command, + is_read_only: Some(true), + .. + } if command == "cargo test" + )); } #[test] -fn mcp_tool_proposal_routes_through_the_existing_mcp_executor_contract() { - let event = build_tool_proposed( +fn mcp_tool_proposal_routes_directly_to_the_mcp_executor_contract() { + let action = build_tool_proposed( "task", &ToolCall { id: "call-mcp".to_string(), - name: "mcp__filesystem__read_file".to_string(), + name: "mcp__11111111-1111-4111-8111-111111111111__read_file".to_string(), arguments: serde_json::json!({"path": "Cargo.toml"}), }, - ); + &SkillPathOrigin::Local, + ) + .unwrap(); - let Some(warp_multi_agent_api::response_event::Type::ClientActions(actions)) = event.r#type - else { - panic!("expected client actions"); - }; - let Some(warp_multi_agent_api::client_action::Action::AddMessagesToTask(add)) = - &actions.actions[0].action - else { - panic!("expected add-message action"); - }; - let Some(warp_multi_agent_api::message::Message::ToolCall(tool_call)) = - &add.messages[0].message - else { - panic!("expected tool-call message"); - }; - let Some(warp_multi_agent_api::message::tool_call::Tool::CallMcpTool(call)) = &tool_call.tool - else { - panic!("expected MCP tool payload"); - }; - - assert_eq!(tool_call.tool_call_id, "call-mcp"); - assert_eq!(call.server_id, "filesystem"); - assert_eq!(call.name, "read_file"); - assert!(call.args.is_some()); + assert!(matches!( + action.action, + crate::ai::agent::AIAgentActionType::CallMCPTool { + server_id: Some(server_id), + name, + .. + } if server_id.to_string() == "11111111-1111-4111-8111-111111111111" + && name == "read_file" + )); } #[test] diff --git a/app/src/ai/runtime/rig_tool.rs b/app/src/ai/runtime/rig_tool.rs new file mode 100644 index 00000000..f36613e3 --- /dev/null +++ b/app/src/ai/runtime/rig_tool.rs @@ -0,0 +1,319 @@ +use std::time::Duration; + +use ai::diff_validation::ParsedDiff; +use ai::skills::{SkillPathOrigin, SkillReference}; +use galaxy_agent_core::ToolCall; +use uuid::Uuid; + +use crate::ai::agent::task::TaskId; +use crate::ai::agent::{ + AIAgentAction, AIAgentActionType, AIAgentPtyWriteMode, AskUserQuestionItem, + AskUserQuestionOption, AskUserQuestionType, CreateDocumentsRequest, DocumentDiff, + DocumentToCreate, EditDocumentsRequest, FileEdit, FileLocations, ReadDocumentsRequest, + ReadFilesRequest, ReadSkillRequest, SearchCodebaseRequest, ShellCommandDelay, + StartAgentExecutionMode, StartAgentVersion, +}; +use crate::ai::document::ai_document_model::AIDocumentId; + +pub(super) fn action_from_tool_call( + task_id: &str, + call: &ToolCall, + skill_path_origin: &SkillPathOrigin, +) -> Result { + let input = &call.arguments; + let action = match call.name.as_str() { + "run_shell_command" => AIAgentActionType::RequestCommandOutput { + command: string(input, "command"), + is_read_only: Some(boolean(input, "is_read_only")), + is_risky: Some(boolean(input, "is_risky")), + wait_until_completion: boolean(input, "wait_until_complete"), + uses_pager: Some(boolean(input, "uses_pager")), + rationale: None, + citations: Vec::new(), + }, + "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest { + locations: input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(file_location) + .collect(), + }), + "apply_file_diffs" => AIAgentActionType::RequestFileEdits { + file_edits: file_edits(input), + title: nonempty_string(input, "summary"), + }, + "grep" => AIAgentActionType::Grep { + queries: strings(input, "queries"), + path: string(input, "path"), + }, + "file_glob" => AIAgentActionType::FileGlob { + patterns: strings(input, "patterns"), + path: nonempty_string(input, "path"), + }, + "search_codebase" => AIAgentActionType::SearchCodebase(SearchCodebaseRequest { + query: string(input, "query"), + partial_paths: nonempty_strings(input, "path_filters"), + codebase_path: nonempty_string(input, "path"), + }), + "write_to_long_running_shell_command" => { + AIAgentActionType::WriteToLongRunningShellCommand { + block_id: string(input, "command_id").into(), + input: string(input, "input").into_bytes().into(), + mode: match input.get("mode").and_then(serde_json::Value::as_str) { + Some("line") => AIAgentPtyWriteMode::Line, + Some("block") => AIAgentPtyWriteMode::Block, + Some("raw") | Some(_) | None => AIAgentPtyWriteMode::Raw, + }, + } + } + "interrupt_shell_command" => AIAgentActionType::WriteToLongRunningShellCommand { + block_id: string(input, "command_id").into(), + input: vec![galaxy_terminal::model::escape_sequences::C0::ETX].into(), + mode: AIAgentPtyWriteMode::Raw, + }, + "read_shell_command_output" => AIAgentActionType::ReadShellCommandOutput { + block_id: string(input, "command_id").into(), + delay: Some(ShellCommandDelay::Duration(Duration::from_secs( + input + .get("wait_seconds") + .and_then(serde_json::Value::as_u64) + .unwrap_or(2) + .min(crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS), + ))), + }, + "read_mcp_resource" => AIAgentActionType::ReadMCPResource { + server_id: uuid(input, "server_id"), + name: String::new(), + uri: nonempty_string(input, "uri"), + }, + "read_plan" | "read_documents" | "read_notebook" => { + AIAgentActionType::ReadDocuments(ReadDocumentsRequest { + document_ids: strings(input, "document_ids") + .into_iter() + .filter_map(|id| AIDocumentId::try_from(id).ok()) + .collect(), + }) + } + "create_plan" | "create_documents" | "create_notebook" => { + AIAgentActionType::CreateDocuments(CreateDocumentsRequest { + documents: input + .get("documents") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|document| { + Some(DocumentToCreate { + title: document.get("title")?.as_str()?.to_string(), + content: document.get("content")?.as_str()?.to_string(), + }) + }) + .collect(), + }) + } + "edit_plan" | "edit_documents" | "edit_notebook" => { + AIAgentActionType::EditDocuments(EditDocumentsRequest { + diffs: input + .get("diffs") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|diff| { + Some(DocumentDiff { + document_id: AIDocumentId::try_from(diff.get("document_id")?.as_str()?) + .ok()?, + search: string(diff, "search"), + replace: string(diff, "replace"), + }) + }) + .collect(), + }) + } + "start_agent" => AIAgentActionType::StartAgent { + version: StartAgentVersion::V1, + name: string(input, "name"), + prompt: string(input, "prompt"), + execution_mode: StartAgentExecutionMode::local_with_defaults(), + lifecycle_subscription: None, + }, + "send_message_to_agent" => AIAgentActionType::SendMessageToAgent { + addresses: vec![string(input, "agent_id")], + subject: String::new(), + message: string(input, "message"), + }, + "ask_user_question" => AIAgentActionType::AskUserQuestion { + questions: vec![AskUserQuestionItem { + question_id: Uuid::new_v4().to_string(), + question: string(input, "question"), + question_type: AskUserQuestionType::MultipleChoice { + is_multiselect: false, + options: strings(input, "options") + .into_iter() + .enumerate() + .map(|(index, label)| AskUserQuestionOption { + label, + recommended: index == 0, + }) + .collect(), + supports_other: true, + }, + }], + }, + "read_skill" => { + let skill = string(input, "skill"); + let skill = match input + .get("reference_type") + .and_then(serde_json::Value::as_str) + { + Some("bundled") => SkillReference::BundledSkillId(skill), + Some("path") | Some(_) | None => SkillReference::Path( + skill_path_origin + .location_for_path(skill) + .map_err(|error| error.to_string())?, + ), + }; + AIAgentActionType::ReadSkill(ReadSkillRequest { skill }) + } + "fetch_conversation" => AIAgentActionType::FetchConversation { + conversation_id: string(input, "conversation_id"), + }, + name if name.starts_with("mcp__") => { + let mut parts = name.splitn(3, "__"); + let _prefix = parts.next(); + let server_id = parts.next().and_then(|value| Uuid::parse_str(value).ok()); + let name = parts + .next() + .unwrap_or_else(|| name.strip_prefix("mcp__").unwrap_or(name)) + .to_string(); + AIAgentActionType::CallMCPTool { + server_id, + name, + input: input.clone(), + } + } + name => return Err(format!("unsupported Rig tool proposal: {name}")), + }; + + let tool_name = matches!( + call.name.as_str(), + "read_notebook" | "create_notebook" | "edit_notebook" + ) + .then(|| "notebook".to_string()); + Ok(AIAgentAction { + id: call.id.clone().into(), + task_id: TaskId::new(task_id.to_string()), + action, + requires_result: true, + tool_name, + }) +} + +fn string(input: &serde_json::Value, key: &str) -> String { + input + .get(key) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string() +} + +fn nonempty_string(input: &serde_json::Value, key: &str) -> Option { + let value = string(input, key); + (!value.is_empty()).then_some(value) +} + +fn boolean(input: &serde_json::Value, key: &str) -> bool { + input + .get(key) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) +} + +fn strings(input: &serde_json::Value, key: &str) -> Vec { + input + .get(key) + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(serde_json::Value::as_str) + .map(ToOwned::to_owned) + .collect() +} + +fn nonempty_strings(input: &serde_json::Value, key: &str) -> Option> { + let values = strings(input, key); + (!values.is_empty()).then_some(values) +} + +fn uuid(input: &serde_json::Value, key: &str) -> Option { + input + .get(key) + .and_then(serde_json::Value::as_str) + .and_then(|value| Uuid::parse_str(value).ok()) +} + +fn file_location(file: &serde_json::Value) -> Option { + if let Some(name) = file.as_str() { + return Some(FileLocations { + name: name.to_string(), + lines: Vec::new(), + }); + } + let name = file + .get("path") + .or_else(|| file.get("name"))? + .as_str()? + .to_string(); + let lines = file + .get("line_ranges") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|range| { + let start = usize::try_from(range.get("start")?.as_u64()?).ok()?; + let end = usize::try_from(range.get("end")?.as_u64()?).ok()?; + (start > 0 && end >= start).then_some(start..end) + }) + .collect(); + Some(FileLocations { name, lines }) +} + +fn file_edits(input: &serde_json::Value) -> Vec { + let diffs = input + .get("diffs") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .map(|diff| { + FileEdit::Edit(ParsedDiff::StrReplaceEdit { + file: nonempty_string(diff, "file_path"), + search: nonempty_string(diff, "search"), + replace: nonempty_string(diff, "replace"), + }) + }); + let creates = input + .get("new_files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .map(|file| FileEdit::Create { + file: nonempty_string(file, "file_path"), + content: nonempty_string(file, "content"), + }); + let deletes = input + .get("deleted_files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .map(|file| FileEdit::Delete { + file: file + .as_str() + .map(ToOwned::to_owned) + .or_else(|| nonempty_string(file, "file_path")), + }); + diffs.chain(creates).chain(deletes).collect() +} + +#[cfg(test)] +#[path = "rig_tool_tests.rs"] +mod tests; diff --git a/app/src/ai/runtime/rig_tool_tests.rs b/app/src/ai/runtime/rig_tool_tests.rs new file mode 100644 index 00000000..5921a8be --- /dev/null +++ b/app/src/ai/runtime/rig_tool_tests.rs @@ -0,0 +1,140 @@ +use std::path::PathBuf; + +use ai::diff_validation::ParsedDiff; +use ai::skills::{SkillPathOrigin, SkillReference}; +use galaxy_agent_core::ToolCall; + +use super::action_from_tool_call; +use crate::ai::agent::{AIAgentActionType, FileEdit}; + +fn call(name: &str, arguments: serde_json::Value) -> ToolCall { + ToolCall { + id: "call-1".to_string(), + name: name.to_string(), + arguments, + } +} + +#[test] +fn shell_calls_become_domain_actions_without_a_proto_round_trip() { + let action = action_from_tool_call( + "task-1", + &call( + "run_shell_command", + serde_json::json!({ + "command": "cargo test", + "is_read_only": true, + "is_risky": false + }), + ), + &SkillPathOrigin::Local, + ) + .unwrap(); + + assert_eq!(action.id.to_string(), "call-1"); + assert_eq!(action.task_id.to_string(), "task-1"); + assert!(matches!( + action.action, + AIAgentActionType::RequestCommandOutput { + command, + is_read_only: Some(true), + is_risky: Some(false), + .. + } if command == "cargo test" + )); +} + +#[test] +fn edit_calls_preserve_file_edits_in_the_domain_model() { + let action = action_from_tool_call( + "task-1", + &call( + "apply_file_diffs", + serde_json::json!({ + "summary": "Update greeting", + "diffs": [{ + "file_path": "/tmp/greeting.txt", + "search": "hello", + "replace": "hello galaxy" + }] + }), + ), + &SkillPathOrigin::Local, + ) + .unwrap(); + + let AIAgentActionType::RequestFileEdits { file_edits, title } = action.action else { + panic!("expected file-edit action"); + }; + assert_eq!(title.as_deref(), Some("Update greeting")); + assert!(matches!( + &file_edits[0], + FileEdit::Edit(ParsedDiff::StrReplaceEdit { + file: Some(file), + search: Some(search), + replace: Some(replace), + }) if file == "/tmp/greeting.txt" && search == "hello" && replace == "hello galaxy" + )); +} + +#[test] +fn grouped_mcp_calls_keep_the_installation_uuid_and_json_input() { + let action = action_from_tool_call( + "task-1", + &call( + "mcp__11111111-1111-4111-8111-111111111111__echo", + serde_json::json!({"message": "hello"}), + ), + &SkillPathOrigin::Local, + ) + .unwrap(); + + assert!(matches!( + action.action, + AIAgentActionType::CallMCPTool { + server_id: Some(server_id), + name, + input, + } if server_id.to_string() == "11111111-1111-4111-8111-111111111111" + && name == "echo" + && input == serde_json::json!({"message": "hello"}) + )); +} + +#[test] +fn local_skill_paths_preserve_the_session_origin() { + let action = action_from_tool_call( + "task-1", + &call( + "read_skill", + serde_json::json!({ + "skill": "/tmp/example/SKILL.md", + "reference_type": "path" + }), + ), + &SkillPathOrigin::Local, + ) + .unwrap(); + + assert!(matches!( + action.action, + AIAgentActionType::ReadSkill(request) + if request.skill == SkillReference::Path( + galaxy_util::local_or_remote_path::LocalOrRemotePath::Local(PathBuf::from( + "/tmp/example/SKILL.md" + )) + ) + )); +} + +#[test] +fn unknown_tools_are_rejected_before_the_permission_boundary() { + let error = action_from_tool_call( + "task-1", + &call("invented_tool", serde_json::json!({})), + &SkillPathOrigin::Local, + ) + .unwrap_err(); + + assert!(error.contains("unsupported Rig tool proposal")); +} diff --git a/app/src/integration_testing/agent_mode/assertions.rs b/app/src/integration_testing/agent_mode/assertions.rs index 38966b38..4cb5ee60 100644 --- a/app/src/integration_testing/agent_mode/assertions.rs +++ b/app/src/integration_testing/agent_mode/assertions.rs @@ -612,6 +612,30 @@ pub fn assert_task_is_blocked(conversation_target: ConversationTarget) -> Assert }) } +/// Asserts that the conversation was explicitly cancelled by the user. +pub fn assert_task_is_cancelled(conversation_target: ConversationTarget) -> AssertionCallback { + Box::new(move |app, window_id| { + let terminal_view = terminal_view(app, window_id, 0, 0); + BlocklistAIHistoryModel::handle(app).update(app, |history_model, _| { + let conversation = + match get_conversation(conversation_target, terminal_view.id(), history_model) { + Ok(conversation) => conversation, + Err(assertion) => return assertion, + }; + + match conversation.status() { + ConversationStatus::Cancelled => AssertionOutcome::Success, + status if status.is_in_progress() => { + AssertionOutcome::failure("Task is still in progress".to_owned()) + } + status => AssertionOutcome::immediate_failure(format!( + "Expected task to be cancelled, but status is {status:?}" + )), + } + }) + }) +} + /// Check if a conversation has ended with an API error in its latest exchange. fn check_for_api_error_in_latest_exchange( conversation: &AIConversation, diff --git a/app/src/integration_testing/agent_mode/step.rs b/app/src/integration_testing/agent_mode/step.rs index ed47593c..e7738a1d 100644 --- a/app/src/integration_testing/agent_mode/step.rs +++ b/app/src/integration_testing/agent_mode/step.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::fs::read; use std::io::Cursor; use std::path::Path; @@ -10,6 +11,10 @@ use prost::Message; use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; use crate::ai::execution_profiles::ActionPermission; use crate::ai::llms::{LLMId, LLMPreferences}; +use crate::ai::mcp::{ + JsonTemplate, TemplatableMCPServer, TemplatableMCPServerInstallation, + TemplatableMCPServerManager, +}; use crate::integration_testing::agent_mode::{ assert_latest_task_succeeds_or_blocked, assert_task_is_blocked, ConversationTarget, }; @@ -260,3 +265,119 @@ pub fn set_execution_profile_no_auto_execute() -> TestStep { }, ) } + +/// Sets the execution profile to auto-execute commands. +pub fn set_execution_profile_auto_execute() -> TestStep { + TestStep::new("Set execution profile to auto-execute commands").add_named_assertion( + "Update execution profile", + |app, _window_id| { + AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| { + let default_profile_id = *profiles.default_profile(ctx).id(); + profiles.set_execute_commands( + default_profile_id, + &ActionPermission::AlwaysAllow, + ctx, + ); + }); + async_assert!(true, "Successfully updated execution profile") + }, + ) +} + +/// Sets the execution profile to auto-apply code diffs. +pub fn set_execution_profile_auto_apply_code_diffs() -> TestStep { + TestStep::new("Set execution profile to auto-apply code diffs").add_named_assertion( + "Update execution profile", + |app, _window_id| { + AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| { + let default_profile_id = *profiles.default_profile(ctx).id(); + profiles.set_apply_code_diffs( + default_profile_id, + &ActionPermission::AlwaysAllow, + ctx, + ); + }); + async_assert!(true, "Successfully updated execution profile") + }, + ) +} + +/// Sets the execution profile to auto-execute MCP tools. +pub fn set_execution_profile_auto_execute_mcp_tools() -> TestStep { + TestStep::new("Set execution profile to auto-execute MCP tools").add_named_assertion( + "Update execution profile", + |app, _window_id| { + AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| { + let default_profile_id = *profiles.default_profile(ctx).id(); + profiles.set_mcp_permissions( + default_profile_id, + &ActionPermission::AlwaysAllow, + ctx, + ); + }); + async_assert!(true, "Successfully updated execution profile") + }, + ) +} + +/// Starts an ephemeral MCP stdio server for an integration test. +pub fn start_ephemeral_mcp_server_for_testing( + command: String, + argument: String, + installation_id: &str, + template_id: &str, + server_name: &str, +) -> TestStep { + let installation_id = + uuid::Uuid::parse_str(installation_id).expect("valid MCP installation UUID"); + let template_id = uuid::Uuid::parse_str(template_id).expect("valid MCP template UUID"); + let server_name = server_name.to_string(); + new_step_with_default_assertions("Start ephemeral integration MCP server").add_named_assertion( + "Spawn MCP server", + move |app, _window_id| { + let template_json = serde_json::json!({ + (&server_name): { + "command": command, + "args": [argument], + }, + }) + .to_string(); + let server = TemplatableMCPServer { + uuid: template_id, + name: server_name.clone(), + description: Some("Hermetic MCP server for integration testing".to_string()), + template: JsonTemplate { + json: template_json, + variables: Vec::new(), + }, + version: 1, + gallery_data: None, + }; + let installation = + TemplatableMCPServerInstallation::new(installation_id, server, HashMap::new()); + TemplatableMCPServerManager::handle(app).update(app, |manager, ctx| { + manager.spawn_cli_ephemeral_server(installation, ctx); + }); + async_assert!(true, "Ephemeral MCP server spawn requested") + }, + ) +} + +/// Waits for an ephemeral MCP server to complete its protocol handshake. +pub fn wait_until_mcp_server_is_active_for_testing( + installation_id: &str, + timeout: Duration, +) -> TestStep { + let installation_id = + uuid::Uuid::parse_str(installation_id).expect("valid MCP installation UUID"); + new_step_with_default_assertions("Wait for ephemeral integration MCP server") + .set_timeout(timeout) + .add_named_assertion("MCP server is active", move |app, _window_id| { + let is_active = TemplatableMCPServerManager::handle(app) + .read(app, |manager, _| manager.is_server_active(installation_id)); + async_assert!( + is_active, + "Waiting for ephemeral MCP server to become active" + ) + }) +} diff --git a/crates/integration/src/bin/integration.rs b/crates/integration/src/bin/integration.rs index c448cfe3..00637d25 100644 --- a/crates/integration/src/bin/integration.rs +++ b/crates/integration/src/bin/integration.rs @@ -24,6 +24,10 @@ pub struct Args { } pub fn main() -> Result<()> { + if env::args().nth(1).as_deref() == Some(integration::rig_test_mcp_server::SERVER_ARGUMENT) { + return integration::rig_test_mcp_server::run(); + } + ChannelState::set(ChannelState::new( Channel::Integration, ChannelConfig { @@ -429,6 +433,12 @@ fn register_tests() -> HashMap<&'static str, BoxedBuilderFn> { register_test!(test_agent_mode_pane_minimum_size); register_test!(test_rig_read_tool_round_trip); + register_test!(test_rig_shell_tool_success_round_trip); + register_test!(test_rig_shell_tool_failure_round_trip); + register_test!(test_rig_shell_tool_permission_denial); + register_test!(test_rig_edit_tool_round_trip); + register_test!(test_rig_in_flight_cancellation); + register_test!(test_rig_mcp_tool_round_trip); register_test!(test_git_prompt_chips); // These tests are only invoked manually, and not included in the diff --git a/crates/integration/src/lib.rs b/crates/integration/src/lib.rs index 6b2b4c4b..43eb4d3a 100644 --- a/crates/integration/src/lib.rs +++ b/crates/integration/src/lib.rs @@ -1,4 +1,5 @@ mod builder; +pub mod rig_test_mcp_server; mod step; pub mod test; diff --git a/crates/integration/src/rig_test_mcp_server.rs b/crates/integration/src/rig_test_mcp_server.rs new file mode 100644 index 00000000..5986e995 --- /dev/null +++ b/crates/integration/src/rig_test_mcp_server.rs @@ -0,0 +1,89 @@ +use std::io::{self, BufRead as _, Write as _}; + +use anyhow::Context as _; + +pub const SERVER_ARGUMENT: &str = "--rig-test-mcp-server"; +pub const TOOL_RESULT_PREFIX: &str = "rig-mcp-echo: "; + +pub fn run() -> anyhow::Result<()> { + let stdin = io::stdin(); + let mut stdout = io::stdout().lock(); + + for line in stdin.lock().lines() { + let line = line.context("failed to read MCP request")?; + let request: serde_json::Value = + serde_json::from_str(&line).context("failed to parse MCP request")?; + let Some(id) = request.get("id").cloned() else { + continue; + }; + let method = request + .get("method") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + let result = match method { + "initialize" => { + let protocol_version = request + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| serde_json::json!("2024-11-05")); + serde_json::json!({ + "protocolVersion": protocol_version, + "capabilities": {"tools": {}}, + "serverInfo": { + "name": "rig-integration-mcp", + "version": "1.0.0", + }, + }) + } + "tools/list" => serde_json::json!({ + "tools": [{ + "name": "echo", + "description": "Echo text through the hermetic Rig integration MCP server", + "inputSchema": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + }, + }], + }), + "tools/call" => { + let text = request + .pointer("/params/arguments/text") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + serde_json::json!({ + "content": [{ + "type": "text", + "text": format!("{TOOL_RESULT_PREFIX}{text}"), + }], + "isError": false, + }) + } + "ping" => serde_json::json!({}), + _ => { + let response = serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32601, + "message": format!("Unsupported test MCP method: {method}"), + }, + }); + writeln!(stdout, "{response}").context("failed to write MCP error response")?; + stdout + .flush() + .context("failed to flush MCP error response")?; + continue; + } + }; + let response = serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "result": result, + }); + writeln!(stdout, "{response}").context("failed to write MCP response")?; + stdout.flush().context("failed to flush MCP response")?; + } + + Ok(()) +} diff --git a/crates/integration/src/test/rig_runtime.rs b/crates/integration/src/test/rig_runtime.rs index edc4c392..775053f0 100644 --- a/crates/integration/src/test/rig_runtime.rs +++ b/crates/integration/src/test/rig_runtime.rs @@ -1,48 +1,316 @@ use std::io::{ErrorKind, Read, Write}; use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::path::Path; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::thread::{self, JoinHandle}; use std::time::Duration; +use galaxyui_core::async_assert; use warp::features::FeatureFlag; use warp::integration_testing::agent_mode::{ - assert_latest_exchange_text, enter_agent_view, set_preferred_agent_mode_llm, - submit_ai_query_and_wait_until_done, + assert_latest_exchange_text, assert_task_is_cancelled, enter_agent_view, + set_execution_profile_auto_apply_code_diffs, set_execution_profile_auto_execute, + set_execution_profile_auto_execute_mcp_tools, set_execution_profile_no_auto_execute, + set_preferred_agent_mode_llm, start_ephemeral_mcp_server_for_testing, submit_ai_query, + submit_ai_query_and_wait_until_blocked, submit_ai_query_and_wait_until_done, + wait_until_mcp_server_is_active_for_testing, ConversationTarget, }; use warp::integration_testing::step::new_step_with_default_assertions; -use warp::integration_testing::terminal::wait_until_bootstrapped_single_pane_for_tab; +use warp::integration_testing::terminal::{ + assert_input_is_focused, wait_until_bootstrapped_single_pane_for_tab, +}; use super::new_builder; +use crate::rig_test_mcp_server::{SERVER_ARGUMENT as MCP_SERVER_ARGUMENT, TOOL_RESULT_PREFIX}; use crate::Builder; const MODEL_ID: &str = "integration-rig-model"; -const FINAL_TEXT: &str = "Rig read round trip completed."; +const READ_CALL_ID: &str = "rig-read-call"; +const READ_FINAL_TEXT: &str = "Rig read round trip completed."; const FIXTURE_CONTENT: &str = "content returned through the Galaxy read executor"; +const SHELL_CALL_ID: &str = "rig-shell-call"; +const SHELL_SUCCESS_OUTPUT: &str = "rig-shell-success-output"; +const SHELL_SUCCESS_FINAL_TEXT: &str = "Rig shell success round trip completed."; +const SHELL_FAILURE_OUTPUT: &str = "rig-shell-failure-output"; +const SHELL_FAILURE_FINAL_TEXT: &str = "Rig shell failure round trip completed."; +const SHELL_DENIED_FINAL_TEXT: &str = "Rig shell denial was preserved."; +const EDIT_CALL_ID: &str = "rig-edit-call"; +const EDIT_INITIAL_CONTENT: &str = "before Rig edit\n"; +const EDIT_UPDATED_CONTENT: &str = "after Rig edit\n"; +const EDIT_FINAL_TEXT: &str = "Rig edit round trip completed."; +const MCP_CALL_ID: &str = "rig-mcp-call"; +const MCP_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; +const MCP_TEMPLATE_ID: &str = "22222222-2222-4222-8222-222222222222"; +const MCP_SERVER_NAME: &str = "rig-integration"; +const MCP_TOOL_NAME: &str = "mcp__11111111-1111-4111-8111-111111111111__echo"; +const MCP_INPUT: &str = "hello from Rig"; +const MCP_FINAL_TEXT: &str = "Rig MCP round trip completed."; + +#[derive(Clone)] +enum MockScenario { + Read { + fixture_path: Arc>, + }, + ShellSuccess, + ShellFailure, + ShellDenied { + marker_path: Arc>, + }, + Edit { + fixture_path: Arc>, + }, + Cancellation { + stream_started: Arc, + stream_cancelled: Arc, + }, + Mcp, +} pub fn test_rig_read_tool_round_trip() -> Builder { - FeatureFlag::AgentView.set_enabled(true); let fixture_path = Arc::new(Mutex::new(String::new())); - let stop = Arc::new(AtomicBool::new(false)); - let (address, server_thread) = start_mock_provider(fixture_path.clone(), stop.clone()); - let server_thread = Arc::new(Mutex::new(Some(server_thread))); + rig_builder(MockScenario::Read { fixture_path }) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(enter_agent_view()) + .with_step(submit_ai_query_and_wait_until_done( + "Read the integration fixture and report when the read is complete.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Assert Rig read result reached Agent Mode") + .add_named_assertion( + "Final response follows the real read tool result", + assert_latest_exchange_text(|text| text.contains(READ_FINAL_TEXT)), + ), + ) +} - let setup_fixture_path = fixture_path.clone(); +pub fn test_rig_shell_tool_success_round_trip() -> Builder { + rig_builder(MockScenario::ShellSuccess) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(set_execution_profile_auto_execute()) + .with_step(enter_agent_view()) + .with_step(submit_ai_query_and_wait_until_done( + "Run the requested shell success check.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Assert Rig shell success reached Agent Mode") + .add_named_assertion( + "Final response follows the successful shell result", + assert_latest_exchange_text(|text| text.contains(SHELL_SUCCESS_FINAL_TEXT)), + ), + ) +} + +pub fn test_rig_shell_tool_failure_round_trip() -> Builder { + rig_builder(MockScenario::ShellFailure) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(set_execution_profile_auto_execute()) + .with_step(enter_agent_view()) + .with_step(submit_ai_query_and_wait_until_done( + "Run the requested failing shell check.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Assert Rig shell failure reached Agent Mode") + .add_named_assertion( + "Final response follows the failed shell result", + assert_latest_exchange_text(|text| text.contains(SHELL_FAILURE_FINAL_TEXT)), + ), + ) +} + +pub fn test_rig_shell_tool_permission_denial() -> Builder { + let marker_path = Arc::new(Mutex::new(String::new())); + rig_builder(MockScenario::ShellDenied { marker_path }) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(set_execution_profile_no_auto_execute()) + .with_step(enter_agent_view()) + .with_step(submit_ai_query_and_wait_until_blocked( + "Propose the requested shell command and wait for my decision.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Deny the blocked Rig shell command") + .with_keystrokes(&["ctrl-c"]) + .add_named_assertion( + "Input returns after denying the command", + assert_input_is_focused(), + ), + ) + .with_step(submit_ai_query_and_wait_until_done( + "Continue without running the denied command.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Assert Rig denial reached Agent Mode") + .add_named_assertion( + "Final response follows the explicit denied result", + assert_latest_exchange_text(|text| text.contains(SHELL_DENIED_FINAL_TEXT)), + ), + ) +} + +pub fn test_rig_edit_tool_round_trip() -> Builder { + let fixture_path = Arc::new(Mutex::new(String::new())); + rig_builder(MockScenario::Edit { fixture_path }) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(set_execution_profile_auto_apply_code_diffs()) + .with_step(enter_agent_view()) + .with_step(submit_ai_query_and_wait_until_done( + "Apply the requested edit to the integration fixture.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Assert Rig edit result reached Agent Mode") + .add_named_assertion( + "Final response follows the real edit result", + assert_latest_exchange_text(|text| text.contains(EDIT_FINAL_TEXT)), + ), + ) +} + +pub fn test_rig_in_flight_cancellation() -> Builder { + let stream_started = Arc::new(AtomicBool::new(false)); + let stream_cancelled = Arc::new(AtomicBool::new(false)); + rig_builder(MockScenario::Cancellation { + stream_started: stream_started.clone(), + stream_cancelled: stream_cancelled.clone(), + }) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(enter_agent_view()) + .with_step(submit_ai_query( + "Start the cancellable Rig response.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Wait for the Rig provider stream to start") + .set_timeout(Duration::from_secs(30)) + .add_named_assertion("Provider stream started", move |_app, _window_id| { + async_assert!( + stream_started.load(Ordering::SeqCst), + "Waiting for the provider stream to start" + ) + }), + ) + .with_step( + new_step_with_default_assertions("Cancel the in-flight Rig provider stream") + .set_timeout(Duration::from_secs(30)) + .with_keystrokes(&["ctrl-c"]) + .add_named_assertion( + "Conversation is marked cancelled", + assert_task_is_cancelled(ConversationTarget::Active), + ) + .add_named_assertion( + "Provider connection was dropped", + move |_app, _window_id| { + async_assert!( + stream_cancelled.load(Ordering::SeqCst), + "Waiting for cancellation to close the provider connection" + ) + }, + ) + .add_named_assertion( + "Input returns after cancelling the stream", + assert_input_is_focused(), + ), + ) +} + +pub fn test_rig_mcp_tool_round_trip() -> Builder { + let command = std::env::current_exe() + .expect("integration executable path") + .to_string_lossy() + .into_owned(); + rig_builder(MockScenario::Mcp) + .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) + .with_step(start_ephemeral_mcp_server_for_testing( + command, + MCP_SERVER_ARGUMENT.to_string(), + MCP_INSTALLATION_ID, + MCP_TEMPLATE_ID, + MCP_SERVER_NAME, + )) + .with_step(wait_until_mcp_server_is_active_for_testing( + MCP_INSTALLATION_ID, + Duration::from_secs(30), + )) + .with_step(set_preferred_agent_mode_llm(MODEL_ID)) + .with_step(set_execution_profile_auto_execute_mcp_tools()) + .with_step(enter_agent_view()) + .with_step(submit_ai_query_and_wait_until_done( + "Call the hermetic MCP echo tool.", + Duration::from_secs(60), + )) + .with_step( + new_step_with_default_assertions("Assert Rig MCP result reached Agent Mode") + .add_named_assertion( + "Final response follows the real MCP result", + assert_latest_exchange_text(|text| text.contains(MCP_FINAL_TEXT)), + ), + ) +} + +fn rig_builder(scenario: MockScenario) -> Builder { + FeatureFlag::AgentView.set_enabled(true); + FeatureFlag::MCPGroupedServerContext.set_enabled(true); + let stop = Arc::new(AtomicBool::new(false)); + let setup_scenario = scenario.clone(); + let (address, server_thread) = start_mock_provider(scenario, stop.clone()); + let server_thread = Arc::new(Mutex::new(Some(server_thread))); let cleanup_stop = stop.clone(); let cleanup_thread = server_thread.clone(); + new_builder() .with_setup(move |utils| { - let fixture = utils.test_dir().join("rig-read-fixture.txt"); - std::fs::write(&fixture, FIXTURE_CONTENT) - .expect("should write Rig integration fixture"); - *setup_fixture_path.lock().expect("fixture path lock") = - fixture.to_string_lossy().into_owned(); + write_provider_settings(address); + match &setup_scenario { + MockScenario::Read { fixture_path } => { + let fixture = utils.test_dir().join("rig-read-fixture.txt"); + std::fs::write(&fixture, FIXTURE_CONTENT) + .expect("should write Rig integration fixture"); + *fixture_path.lock().expect("fixture path lock") = + fixture.to_string_lossy().into_owned(); + } + MockScenario::ShellDenied { marker_path } => { + let marker = utils.test_dir().join("rig-denied-command-marker.txt"); + *marker_path.lock().expect("marker path lock") = + marker.to_string_lossy().into_owned(); + } + MockScenario::Edit { fixture_path } => { + let fixture = utils.test_dir().join("rig-edit-fixture.txt"); + std::fs::write(&fixture, EDIT_INITIAL_CONTENT) + .expect("should write Rig edit integration fixture"); + *fixture_path.lock().expect("fixture path lock") = + fixture.to_string_lossy().into_owned(); + } + MockScenario::ShellSuccess + | MockScenario::ShellFailure + | MockScenario::Cancellation { .. } + | MockScenario::Mcp => {} + } + }) + .with_cleanup(move |_utils| { + cleanup_stop.store(true, Ordering::SeqCst); + if let Some(handle) = cleanup_thread.lock().expect("server thread lock").take() { + handle.join().expect("mock provider should stop cleanly"); + } + }) +} - let settings_path = warp::settings::user_preferences_toml_file_path(); - std::fs::create_dir_all(settings_path.parent().expect("settings parent")) - .expect("should create settings directory"); - let settings = format!( - r#"[ai.openai] +fn write_provider_settings(address: SocketAddr) { + let settings_path = warp::settings::user_preferences_toml_file_path(); + std::fs::create_dir_all(settings_path.parent().expect("settings parent")) + .expect("should create settings directory"); + let settings = format!( + r#"[ai.openai] enabled = true [[ai.providers]] @@ -56,33 +324,12 @@ context_size = 128000 use_rig = true supports_system_messages = false "# - ); - std::fs::write(settings_path, settings).expect("should write provider settings"); - }) - .with_cleanup(move |_utils| { - cleanup_stop.store(true, Ordering::SeqCst); - if let Some(handle) = cleanup_thread.lock().expect("server thread lock").take() { - handle.join().expect("mock provider should stop cleanly"); - } - }) - .with_step(wait_until_bootstrapped_single_pane_for_tab(0)) - .with_step(set_preferred_agent_mode_llm(MODEL_ID)) - .with_step(enter_agent_view()) - .with_step(submit_ai_query_and_wait_until_done( - "Read the integration fixture and report when the read is complete.", - Duration::from_secs(60), - )) - .with_step( - new_step_with_default_assertions("Assert Rig read result reached Agent Mode") - .add_named_assertion( - "Final response follows the real read tool result", - assert_latest_exchange_text(|text| text.contains(FINAL_TEXT)), - ), - ) + ); + std::fs::write(settings_path, settings).expect("should write provider settings"); } fn start_mock_provider( - fixture_path: Arc>, + scenario: MockScenario, stop: Arc, ) -> (SocketAddr, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").expect("should bind mock Rig provider"); @@ -95,7 +342,7 @@ fn start_mock_provider( while !stop.load(Ordering::SeqCst) { match listener.accept() { Ok((mut stream, _)) => { - serve_request(&mut stream, &fixture_path, &request_count); + serve_request(&mut stream, &scenario, &request_count, &stop); } Err(error) if error.kind() == ErrorKind::WouldBlock => { thread::sleep(Duration::from_millis(10)); @@ -109,8 +356,9 @@ fn start_mock_provider( fn serve_request( stream: &mut TcpStream, - fixture_path: &Mutex, + scenario: &MockScenario, request_count: &AtomicUsize, + stop: &AtomicBool, ) { stream .set_read_timeout(Some(Duration::from_secs(5))) @@ -130,21 +378,20 @@ fn serve_request( "unexpected mock provider request: {request_line}" ); let turn = request_count.fetch_add(1, Ordering::SeqCst); + if let MockScenario::Cancellation { + stream_started, + stream_cancelled, + } = scenario + { + assert_eq!(turn, 0, "unexpected extra cancellation chat request"); + write_cancellable_response(stream, stream_started, stream_cancelled, stop); + return; + } let body = match turn { - 0 => { - let fixture = fixture_path.lock().expect("fixture path lock").clone(); - tool_call_sse(&fixture) - } + 0 => tool_call_sse(scenario), 1 => { - assert!( - request.contains("rig-read-call"), - "follow-up request should preserve the tool call ID" - ); - assert!( - request.contains(FIXTURE_CONTENT), - "follow-up request should contain the real file contents returned by Galaxy" - ); - final_text_sse() + assert_follow_up_request(scenario, &request); + final_text_sse(final_text(scenario)) } _ => panic!("unexpected extra chat completion request"), }; @@ -186,8 +433,60 @@ fn read_request(stream: &mut TcpStream) -> String { String::from_utf8(request).expect("provider request should be valid UTF-8") } -fn tool_call_sse(fixture_path: &str) -> String { - let arguments = serde_json::json!({"files": [fixture_path]}).to_string(); +fn tool_call_sse(scenario: &MockScenario) -> String { + let (call_id, tool_name, arguments) = match scenario { + MockScenario::Read { fixture_path } => { + let fixture_path = fixture_path.lock().expect("fixture path lock").clone(); + ( + READ_CALL_ID, + "read_files", + serde_json::json!({"files": [fixture_path]}), + ) + } + MockScenario::ShellSuccess => ( + SHELL_CALL_ID, + "run_shell_command", + shell_arguments("printf '%s\\n' 'rig-shell-success-output'"), + ), + MockScenario::ShellFailure => ( + SHELL_CALL_ID, + "run_shell_command", + shell_arguments("(printf '%s\\n' 'rig-shell-failure-output' >&2; exit 7)"), + ), + MockScenario::ShellDenied { marker_path } => { + let marker_path = marker_path.lock().expect("marker path lock").clone(); + ( + SHELL_CALL_ID, + "run_shell_command", + shell_arguments(&format!( + "printf '%s\\n' 'command-must-not-run' > '{marker_path}'" + )), + ) + } + MockScenario::Edit { fixture_path } => { + let fixture_path = fixture_path.lock().expect("fixture path lock").clone(); + ( + EDIT_CALL_ID, + "apply_file_diffs", + serde_json::json!({ + "summary": "Update the Rig edit integration fixture", + "diffs": [{ + "file_path": fixture_path, + "search": EDIT_INITIAL_CONTENT, + "replace": EDIT_UPDATED_CONTENT, + }], + }), + ) + } + MockScenario::Cancellation { .. } => { + unreachable!("cancellation streams do not issue tool calls") + } + MockScenario::Mcp => ( + MCP_CALL_ID, + MCP_TOOL_NAME, + serde_json::json!({"text": MCP_INPUT}), + ), + }; let tool_delta = serde_json::json!({ "id": "rig-integration-1", "model": MODEL_ID, @@ -195,11 +494,11 @@ fn tool_call_sse(fixture_path: &str) -> String { "delta": { "tool_calls": [{ "index": 0, - "id": "rig-read-call", + "id": call_id, "type": "function", "function": { - "name": "read_files", - "arguments": arguments, + "name": tool_name, + "arguments": arguments.to_string(), }, }], }, @@ -220,12 +519,125 @@ fn tool_call_sse(fixture_path: &str) -> String { format!("data: {tool_delta}\n\ndata: {tool_stop}\n\ndata: {usage}\n\ndata: [DONE]\n\n") } -fn final_text_sse() -> String { +fn shell_arguments(command: &str) -> serde_json::Value { + serde_json::json!({ + "command": command, + "is_read_only": false, + "is_risky": false, + "uses_pager": false, + "wait_until_complete": true, + }) +} + +fn assert_follow_up_request(scenario: &MockScenario, request: &str) { + match scenario { + MockScenario::Read { .. } => { + assert!( + request.contains(READ_CALL_ID), + "follow-up request should preserve the read tool call ID" + ); + assert!( + request.contains(FIXTURE_CONTENT), + "follow-up request should contain the real file contents returned by Galaxy" + ); + } + MockScenario::ShellSuccess => { + assert!( + request.contains(SHELL_CALL_ID), + "follow-up request should preserve the shell tool call ID" + ); + assert!( + request.contains(SHELL_SUCCESS_OUTPUT), + "follow-up request should contain successful shell output" + ); + assert!( + request.contains("exit code 0"), + "follow-up request should contain the successful exit code" + ); + } + MockScenario::ShellFailure => { + assert!( + request.contains(SHELL_CALL_ID), + "follow-up request should preserve the failed shell tool call ID" + ); + assert!( + request.contains(SHELL_FAILURE_OUTPUT), + "follow-up request should contain failed shell output" + ); + assert!( + request.contains("exit code 7"), + "follow-up request should contain the failed exit code" + ); + assert!( + request.contains("[ERROR]"), + "failed shell result should remain an explicit model error" + ); + } + MockScenario::ShellDenied { marker_path } => { + assert!( + request.contains(SHELL_CALL_ID), + "follow-up request should preserve the denied shell tool call ID" + ); + assert!( + request.contains("[ERROR] Permission denied by the user"), + "denied shell result should remain an explicit model error" + ); + let marker_path = marker_path.lock().expect("marker path lock"); + assert!( + !Path::new(marker_path.as_str()).exists(), + "denied shell command must not create its marker file" + ); + } + MockScenario::Edit { fixture_path } => { + assert!( + request.contains(EDIT_CALL_ID), + "follow-up request should preserve the edit tool call ID" + ); + assert!( + request.contains(EDIT_UPDATED_CONTENT.trim()), + "follow-up request should contain the updated file content" + ); + let fixture_path = fixture_path.lock().expect("fixture path lock"); + let actual = std::fs::read_to_string(fixture_path.as_str()) + .expect("edited fixture should remain readable"); + assert_eq!(actual, EDIT_UPDATED_CONTENT); + } + MockScenario::Cancellation { .. } => { + unreachable!("cancellation streams do not issue follow-up requests") + } + MockScenario::Mcp => { + assert!( + request.contains(MCP_CALL_ID), + "follow-up request should preserve the MCP tool call ID" + ); + assert!( + request.contains(&format!("{TOOL_RESULT_PREFIX}{MCP_INPUT}")), + "follow-up request should contain the real MCP tool result" + ); + } + } +} + +fn final_text(scenario: &MockScenario) -> &'static str { + match scenario { + MockScenario::Read { .. } => READ_FINAL_TEXT, + MockScenario::ShellSuccess => SHELL_SUCCESS_FINAL_TEXT, + MockScenario::ShellFailure => SHELL_FAILURE_FINAL_TEXT, + MockScenario::ShellDenied { .. } => SHELL_DENIED_FINAL_TEXT, + MockScenario::Edit { .. } => EDIT_FINAL_TEXT, + MockScenario::Cancellation { .. } => { + unreachable!("cancellation streams do not produce final text") + } + MockScenario::Mcp => MCP_FINAL_TEXT, + } +} + +fn final_text_sse(final_text: &str) -> String { let text = serde_json::json!({ "id": "rig-integration-2", "model": MODEL_ID, "choices": [{ - "delta": {"content": FINAL_TEXT, "tool_calls": []}, + "delta": {"content": final_text, "tool_calls": []}, "finish_reason": "stop", }], "usage": null, @@ -246,3 +658,51 @@ fn write_response(stream: &mut TcpStream, content_type: &str, body: &str) { .expect("should write mock provider response"); stream.flush().expect("should flush mock provider response"); } + +fn write_cancellable_response( + stream: &mut TcpStream, + stream_started: &AtomicBool, + stream_cancelled: &AtomicBool, + stop: &AtomicBool, +) { + let text = serde_json::json!({ + "id": "rig-cancellation-1", + "model": MODEL_ID, + "choices": [{ + "delta": {"content": "Rig cancellation stream is active.", "tool_calls": []}, + "finish_reason": null, + }], + "usage": null, + }); + let initial_body = format!("data: {text}\n\n"); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\n{initial_body}" + ) + .expect("should start cancellable provider response"); + stream + .flush() + .expect("should flush cancellable provider response"); + stream_started.store(true, Ordering::SeqCst); + + while !stop.load(Ordering::SeqCst) { + thread::sleep(Duration::from_millis(20)); + if let Err(error) = stream + .write_all(b": keepalive\n\n") + .and_then(|()| stream.flush()) + { + assert!( + matches!( + error.kind(), + ErrorKind::BrokenPipe + | ErrorKind::ConnectionReset + | ErrorKind::ConnectionAborted + | ErrorKind::NotConnected + ), + "unexpected cancellable stream error: {error}" + ); + stream_cancelled.store(true, Ordering::SeqCst); + return; + } + } +} diff --git a/crates/integration/tests/integration/ui_tests.rs b/crates/integration/tests/integration/ui_tests.rs index fcfa1169..58f8b25e 100644 --- a/crates/integration/tests/integration/ui_tests.rs +++ b/crates/integration/tests/integration/ui_tests.rs @@ -312,6 +312,12 @@ integration_tests! { test_middle_click_paste, test_agent_mode_pane_minimum_size, test_rig_read_tool_round_trip, + test_rig_shell_tool_success_round_trip, + test_rig_shell_tool_failure_round_trip, + test_rig_shell_tool_permission_denial, + test_rig_edit_tool_round_trip, + test_rig_in_flight_cancellation, + test_rig_mcp_tool_round_trip, test_rule_creation, test_rule_update, diff --git a/plans/galaxy-local-first-rig.md b/plans/galaxy-local-first-rig.md index 5208ffe9..b66ec4d0 100644 --- a/plans/galaxy-local-first-rig.md +++ b/plans/galaxy-local-first-rig.md @@ -315,18 +315,25 @@ Exit condition: a LiteLLM or local OpenAI-compatible conversation streams throug protobuf round trip. - [x] Separate concise UI result summaries from authoritative model-facing result content so file, code-search, document, skill, and shell results retain their payload without protobuf conversion. -- [ ] Move permission decisions and tool start/result events fully onto the Galaxy domain contract, +- [x] Move permission decisions and tool start/result events fully onto the Galaxy domain contract, removing the temporary Warp protobuf adapter. -- [ ] Add end-to-end integration coverage for representative read, edit, shell, MCP, denial, +- [x] Add end-to-end integration coverage for representative read, edit, shell, MCP, denial, cancellation, and execution-failure flows. - [x] Add a hermetic real-app Rig read-tool round trip covering isolated provider configuration, streamed tool proposal, Galaxy-owned execution, normalized tool result, and model follow-up. +- [x] Add hermetic real-app shell coverage for successful execution, nonzero exit with preserved + stderr/error status, and an `AlwaysAsk` user denial that proves the command never executes. - [x] Port loop prevention, inline `recall_tool_history`, and unknown-tool handling to domain-level policies. Exit condition: representative read, edit, shell, MCP, denial, and failure flows pass integration tests without provider-specific UI code. +Phase 3 is complete. Rig tool proposals now enter the controller as typed `AIAgentAction` values; +Galaxy's action model owns permission and execution lifecycle events; normalized results return to +Rig directly. The legacy Warp response envelope remains only around transcript/init/finished UI +rendering and non-Rig compatibility runtimes, not in Rig's executable tool path. + ### Phase 4 — Bedrock through Rig - Implement Bedrock client construction and model resolution through `rig-bedrock`. @@ -411,8 +418,7 @@ contract is what the UI and persistence observe. ## Immediate next vertical slice -Expand the hermetic Rig integration harness from its passing read-tool round trip to -edit/shell/MCP/denial/cancellation/execution-failure cases. Then move permission decisions and tool -start/result events fully onto the Galaxy domain contract, keeping the outgoing UI response adapter -only until those flows prove that Galaxy-owned events can replace it without changing the -permission UI. +Begin Phase 4 with a focused `rig-bedrock` compatibility spike. Establish client construction and +model/inference-profile resolution first, then add semantic parity fixtures for system prompts, +images, tool schemas, cache controls, reasoning, usage, and context limits before routing any +configured Bedrock model away from the existing compatibility implementation.