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

Complete Rig tool lifecycle migration

Browse Source
This commit is contained in:
rkw6086
2026-08-04 16:00:20 -05:00
parent 91d8bd0381
commit a3c68e9c30
30 changed files with 1494 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/acp/transport.rs
+10 -6
View File
@@ -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)]
app/src/ai/agent/api.rs
+20 -2
View File
@@ -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<warp_multi_agent_api::ResponseEvent, Arc<AIApiError>>;
/// 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<StreamEvent, Arc<AIApiError>>;
pub type LegacyEvent = Result<warp_multi_agent_api::ResponseEvent, Arc<AIApiError>>;
#[cfg(not(target_family = "wasm"))]
pub type ResponseStream = Pin<Box<dyn Stream<Item = Event> + Send + 'static>>;
#[cfg(not(target_family = "wasm"))]
pub type LegacyResponseStream = Pin<Box<dyn Stream<Item = LegacyEvent> + Send + 'static>>;
// The WASM version of this type has no bound on `Send`, which is an unnecessary bound when
// targeting wasm because the browser is single-threaded (and we don't leverage WebWorkers for async
// execution in WoW).
#[cfg(target_family = "wasm")]
pub type ResponseStream = Pin<Box<dyn Stream<Item = Event>>>;
#[cfg(target_family = "wasm")]
pub type LegacyResponseStream = Pin<Box<dyn Stream<Item = LegacyEvent>>>;
#[derive(Debug, Clone)]
pub struct ConversationData {
pub id: AIConversationId,
app/src/ai/agent/api/impl.rs
+7 -3
View File
@@ -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) => {
app/src/ai/agent/conversation.rs
+41
View File
@@ -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<BlocklistAIHistoryModel>,
) -> 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<RequestCost>,
app/src/ai/bedrock/client.rs
+2 -2
View File
@@ -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<Arc<BedrockDiagnosticLogger>>,
messages_sent: Arc<Mutex<Vec<ConversationMessage>>>,
tool_result_archive: Vec<ConversationMessage>,
) -> Result<ResponseStream, BedrockError> {
) -> Result<LegacyResponseStream, BedrockError> {
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)
app/src/ai/bedrock/response_translator.rs
+2 -2
View File
@@ -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<Mutex<Vec<ConversationMessage>>>,
model_id: String,
tool_result_archive: Vec<ConversationMessage>,
) -> BoxStream<'static, Event> {
) -> BoxStream<'static, LegacyEvent> {
let request_id = Uuid::new_v4().to_string();
let conversation_id = Uuid::new_v4().to_string();
app/src/ai/bedrock/translator.rs
+2 -2
View File
@@ -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<ResponseStream, BedrockError> {
) -> Result<LegacyResponseStream, BedrockError> {
let config = params.config.with_external_fallbacks();
let cross_region_inference = config.cross_region_inference;
let bedrock = BedrockClient::from_config(config).await?;
app/src/ai/blocklist/action_model.rs
+7 -1
View File
@@ -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,
}
}
app/src/ai/blocklist/action_model_tests.rs
+2
View File
@@ -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]
app/src/ai/blocklist/controller.rs
+17 -1
View File
@@ -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()
app/src/ai/blocklist/controller/response_stream.rs
+11 -2
View File
@@ -186,7 +186,7 @@ impl ResponseStream {
ctx: &mut ModelContext<Self>,
) {
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)) => {
app/src/ai/blocklist/history_model.rs
+18 -3
View File
@@ -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<Self>,
) -> 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,
app/src/ai/blocklist/passive_suggestions/maa.rs
+1 -1
View File
@@ -486,7 +486,7 @@ async fn extract_suggestion_from_stream(
let mut client_actions: Vec<api::ClientAction> = Vec::new();
let mut server_request_token: Option<String> = 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 {
app/src/ai/openai/response_translator.rs
+2 -2
View File
@@ -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<Item = Result<Bytes, reqwest::Error>> + Send + 'static,
context: OpenAIStreamContext,
) -> BoxStream<'static, Event> {
) -> BoxStream<'static, LegacyEvent> {
use futures::StreamExt;
let OpenAIStreamContext {
app/src/ai/openai/translator.rs
+2 -2
View File
@@ -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<ResponseStream, OpenAIError> {
) -> Result<LegacyResponseStream, OpenAIError> {
let client = OpenAIClient::from_config(params.config.clone());
let PreparedTurn {
task_id,
app/src/ai/runtime/mod.rs
+1
View File
@@ -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;
app/src/ai/runtime/rig.rs
+32 -19
View File
@@ -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<ToolType>,
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) = &current_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) = &current_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<AIAgentAction, String> {
action_from_tool_call(task_id, call, skill_path_origin)
}
fn build_add_reasoning(task_id: &str, message_id: &str, text: &str) -> ResponseEvent {
app/src/ai/runtime/rig_request.rs
+1 -1
View File
@@ -395,7 +395,7 @@ fn tool_definitions(
.collect::<HashSet<_>>();
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,
app/src/ai/runtime/rig_request_tests.rs
+38 -2
View File
@@ -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 = [
app/src/ai/runtime/rig_tests.rs
+30 -53
View File
@@ -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]
app/src/ai/runtime/rig_tool.rs
+319
View File
@@ -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<AIAgentAction, String> {
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<String> {
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<String> {
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<Vec<String>> {
let values = strings(input, key);
(!values.is_empty()).then_some(values)
}
fn uuid(input: &serde_json::Value, key: &str) -> Option<Uuid> {
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<FileLocations> {
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<FileEdit> {
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;
app/src/ai/runtime/rig_tool_tests.rs
+140
View File
@@ -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"));
}
app/src/integration_testing/agent_mode/assertions.rs
+24
View File
@@ -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,
app/src/integration_testing/agent_mode/step.rs
+121
View File
@@ -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"
)
})
}