Migrate Rig tool flow to domain runtime

This commit is contained in:
2026-08-04 14:14:51 -05:00
parent 4c7270db8d
commit 91d8bd0381
34 changed files with 2728 additions and 374 deletions
+1 -1
View File
@@ -53,7 +53,7 @@ fn canonical_json_value(value: &serde_json::Value) -> String {
), ),
serde_json::Value::Object(values) => { serde_json::Value::Object(values) => {
let mut entries = values.iter().collect::<Vec<_>>(); let mut entries = values.iter().collect::<Vec<_>>();
entries.sort_by(|(left, _), (right, _)| left.cmp(right)); entries.sort_by_key(|(key, _)| *key);
format!( format!(
"{{{}}}", "{{{}}}",
entries entries
+2 -4
View File
@@ -164,10 +164,8 @@ impl AcpRuntimeModel {
) -> BTreeMap<String, serde_json::Value> { ) -> BTreeMap<String, serde_json::Value> {
options options
.iter() .iter()
.filter_map(|option| { .filter(|option| !option.current_value.is_null())
(!option.current_value.is_null()) .map(|option| (option.id.clone(), option.current_value.clone()))
.then(|| (option.id.clone(), option.current_value.clone()))
})
.collect() .collect()
} }
+23 -19
View File
@@ -14,6 +14,7 @@ pub use convert_from::{
MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError, MaybeAIAgentOutputMessage, MessageToAIAgentOutputMessageError,
}; };
use futures_lite::Stream; use futures_lite::Stream;
use galaxy_agent_core::ToolResult;
use galaxy_core::channel::ChannelState; use galaxy_core::channel::ChannelState;
use galaxy_core::execution_mode::AppExecutionMode; use galaxy_core::execution_mode::AppExecutionMode;
use galaxy_core::features::FeatureFlag; use galaxy_core::features::FeatureFlag;
@@ -96,6 +97,8 @@ pub struct RequestParams {
/// locally so ACP-provided Galaxy tools can be pinned to the exact pane. /// locally so ACP-provided Galaxy tools can be pinned to the exact pane.
pub terminal_view_id: Option<EntityId>, pub terminal_view_id: Option<EntityId>,
pub input: Vec<AIAgentInput>, pub input: Vec<AIAgentInput>,
/// Normalized results consumed directly by Rig-selected models.
pub tool_results: Vec<ToolResult>,
pub conversation_token: Option<ServerConversationToken>, pub conversation_token: Option<ServerConversationToken>,
pub forked_from_conversation_token: Option<ServerConversationToken>, pub forked_from_conversation_token: Option<ServerConversationToken>,
pub ambient_agent_task_id: Option<AmbientAgentTaskId>, pub ambient_agent_task_id: Option<AmbientAgentTaskId>,
@@ -140,21 +143,20 @@ pub struct RequestParams {
pub parent_agent_id: Option<String>, pub parent_agent_id: Option<String>,
/// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator. /// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator.
pub agent_name: Option<String>, pub agent_name: Option<String>,
/// Full Bedrock conversation history for direct Bedrock calls. /// Provider-neutral conversation history for direct model calls.
/// When present, the Bedrock path uses this instead of extracting from task_context. pub message_history: Vec<crate::ai::provider::types::ConversationMessage>,
pub bedrock_message_history: Vec<crate::ai::bedrock::convert::ConversationMessage>,
/// Progressive summary of older conversation history. Prepended as the first /// Progressive summary of older conversation history. Prepended as the first
/// message pair in the messages array sent to Bedrock. /// message pair in the messages array sent to the model.
pub bedrock_progressive_summary: Option<String>, pub progressive_summary: Option<String>,
/// Archived tool_use/tool_result pairs from previous summarization drains. /// Archived tool_use/tool_result pairs from previous summarization drains.
/// Passed to the Bedrock translator so `recall_tool_history` can search archived /// Kept separately so `recall_tool_history` can search archived results even after
/// results even after they've been summarized away from live history. /// they've been summarized away from live history.
pub bedrock_tool_result_archive: Vec<crate::ai::bedrock::convert::ConversationMessage>, pub tool_result_archive: Vec<crate::ai::provider::types::ConversationMessage>,
/// Populated by the Bedrock path after building the message list. /// Populated by direct-provider paths after building the message list.
/// Contains the full messages sent (old history + new input) so the controller /// Contains the full messages sent (old history + new input) so the controller
/// can store them back into the conversation for the next request cycle. /// can store them back into the conversation for the next request cycle.
pub bedrock_messages_sent: pub messages_sent:
std::sync::Arc<std::sync::Mutex<Vec<crate::ai::bedrock::convert::ConversationMessage>>>, std::sync::Arc<std::sync::Mutex<Vec<crate::ai::provider::types::ConversationMessage>>>,
/// Global rules (name, content) from the local CloudModel (AIFact/AIMemory). /// Global rules (name, content) from the local CloudModel (AIFact/AIMemory).
/// Injected into the system prompt when `is_memory_enabled` is true. /// Injected into the system prompt when `is_memory_enabled` is true.
pub global_rules: Vec<(String, String)>, pub global_rules: Vec<(String, String)>,
@@ -187,6 +189,7 @@ impl RequestParams {
Self { Self {
terminal_view_id: None, terminal_view_id: None,
input: vec![], input: vec![],
tool_results: vec![],
conversation_token: None, conversation_token: None,
forked_from_conversation_token: None, forked_from_conversation_token: None,
ambient_agent_task_id: None, ambient_agent_task_id: None,
@@ -218,10 +221,10 @@ impl RequestParams {
parent_agent_id: None, parent_agent_id: None,
agent_name: None, agent_name: None,
root_task_id: None, root_task_id: None,
bedrock_message_history: vec![], message_history: vec![],
bedrock_progressive_summary: None, progressive_summary: None,
bedrock_tool_result_archive: vec![], tool_result_archive: vec![],
bedrock_messages_sent: std::sync::Arc::new(std::sync::Mutex::new(vec![])), messages_sent: std::sync::Arc::new(std::sync::Mutex::new(vec![])),
global_rules: vec![], global_rules: vec![],
} }
} }
@@ -391,6 +394,7 @@ impl RequestParams {
Self { Self {
terminal_view_id, terminal_view_id,
input: request_input.all_inputs().cloned().collect(), input: request_input.all_inputs().cloned().collect(),
tool_results: Vec::new(),
conversation_token: conversation.server_conversation_token, conversation_token: conversation.server_conversation_token,
forked_from_conversation_token: conversation.forked_from_conversation_token, forked_from_conversation_token: conversation.forked_from_conversation_token,
ambient_agent_task_id: conversation.ambient_agent_task_id, ambient_agent_task_id: conversation.ambient_agent_task_id,
@@ -426,10 +430,10 @@ impl RequestParams {
.map(|id| id.to_string()), .map(|id| id.to_string()),
parent_agent_id: None, parent_agent_id: None,
agent_name: None, agent_name: None,
bedrock_message_history: Vec::new(), message_history: Vec::new(),
bedrock_progressive_summary: None, progressive_summary: None,
bedrock_tool_result_archive: Vec::new(), tool_result_archive: Vec::new(),
bedrock_messages_sent: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), messages_sent: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
global_rules: if is_memory_enabled { global_rules: if is_memory_enabled {
Self::load_global_rules(app) Self::load_global_rules(app)
} else { } else {
+24 -26
View File
@@ -24,6 +24,22 @@ pub async fn generate_multi_agent_output(
.unwrap_or_else(|| get_supported_tools(&params)); .unwrap_or_else(|| get_supported_tools(&params));
let supported_cli_agent_tools = let supported_cli_agent_tools =
supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(&params)); supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(&params));
if params.should_redact_secrets {
redaction::redact_inputs(&mut params.input);
}
if let ProviderConfig::OpenAI(config) = &provider_config {
if config.use_rig {
return Ok(crate::ai::runtime::rig_openai_response_stream(
config.clone(),
params,
supported_tools,
supported_cli_agent_tools,
cancellation_rx,
));
}
}
let mut logging_metadata = HashMap::new(); let mut logging_metadata = HashMap::new();
if let Some(ref metadata) = params.metadata { if let Some(ref metadata) = params.metadata {
logging_metadata.insert( logging_metadata.insert(
@@ -52,16 +68,6 @@ pub async fn generate_multi_agent_output(
); );
} }
if params.should_redact_secrets {
redaction::redact_inputs(&mut params.input);
}
let rig_params = matches!(
&provider_config,
ProviderConfig::OpenAI(config) if config.use_rig
)
.then(|| params.clone());
let mut request = api::Request { let mut request = api::Request {
task_context: Some(api::request::TaskContext { task_context: Some(api::request::TaskContext {
tasks: params.tasks, tasks: params.tasks,
@@ -144,23 +150,15 @@ pub async fn generate_multi_agent_output(
}; };
match provider_config { match provider_config {
ProviderConfig::OpenAI(config) if config.use_rig => {
Ok(crate::ai::runtime::rig_openai_response_stream(
config,
rig_params.expect("Rig request parameters should be retained for a Rig model"),
&mut request,
cancellation_rx,
))
}
ProviderConfig::OpenAI(config) => { ProviderConfig::OpenAI(config) => {
let translator_request = openai_translator::TranslatorRequest { let translator_request = openai_translator::TranslatorRequest {
config, config,
model_id: params.model.as_str().to_string(), model_id: params.model.as_str().to_string(),
root_task_id: params.root_task_id.clone(), root_task_id: params.root_task_id.clone(),
message_history: params.bedrock_message_history.clone(), message_history: params.message_history.clone(),
tool_result_archive: params.bedrock_tool_result_archive.clone(), tool_result_archive: params.tool_result_archive.clone(),
progressive_summary: params.bedrock_progressive_summary.clone(), progressive_summary: params.progressive_summary.clone(),
messages_sent: params.bedrock_messages_sent.clone(), messages_sent: params.messages_sent.clone(),
global_rules: params.global_rules.clone(), global_rules: params.global_rules.clone(),
}; };
@@ -189,10 +187,10 @@ pub async fn generate_multi_agent_output(
config, config,
model_id: params.model.as_str().to_string(), model_id: params.model.as_str().to_string(),
root_task_id: params.root_task_id.clone(), root_task_id: params.root_task_id.clone(),
bedrock_message_history: params.bedrock_message_history.clone(), bedrock_message_history: params.message_history.clone(),
bedrock_tool_result_archive: params.bedrock_tool_result_archive.clone(), bedrock_tool_result_archive: params.tool_result_archive.clone(),
bedrock_progressive_summary: params.bedrock_progressive_summary.clone(), bedrock_progressive_summary: params.progressive_summary.clone(),
bedrock_messages_sent: params.bedrock_messages_sent.clone(), bedrock_messages_sent: params.messages_sent.clone(),
global_rules: params.global_rules.clone(), global_rules: params.global_rules.clone(),
}; };
+5 -4
View File
@@ -14,6 +14,7 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool
RequestParams { RequestParams {
terminal_view_id: None, terminal_view_id: None,
input: vec![], input: vec![],
tool_results: vec![],
conversation_token: None, conversation_token: None,
forked_from_conversation_token: None, forked_from_conversation_token: None,
ambient_agent_task_id: None, ambient_agent_task_id: None,
@@ -45,10 +46,10 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool
root_task_id: None, root_task_id: None,
parent_agent_id: None, parent_agent_id: None,
agent_name: None, agent_name: None,
bedrock_message_history: Vec::new(), message_history: Vec::new(),
bedrock_progressive_summary: None, progressive_summary: None,
bedrock_tool_result_archive: Vec::new(), tool_result_archive: Vec::new(),
bedrock_messages_sent: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), messages_sent: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
global_rules: Vec::new(), global_rules: Vec::new(),
} }
} }
+1 -1
View File
@@ -1617,7 +1617,7 @@ fn supported_tool_types(request: &api::Request) -> Option<HashSet<api::ToolType>
) )
} }
fn tool_name_is_supported(name: &str, supported: &HashSet<api::ToolType>) -> bool { pub(crate) fn tool_name_is_supported(name: &str, supported: &HashSet<api::ToolType>) -> bool {
use api::ToolType; use api::ToolType;
let has = |tool| supported.contains(&tool); let has = |tool| supported.contains(&tool);
+8 -139
View File
@@ -8,6 +8,7 @@ use aws_sdk_bedrockruntime::types::{
ReasoningContentBlockDelta, StopReason, ReasoningContentBlockDelta, StopReason,
}; };
use futures::stream::BoxStream; use futures::stream::BoxStream;
use galaxy_agent_core::{recall_tool_history, ToolHistoryQuery};
use uuid::Uuid; use uuid::Uuid;
use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
@@ -231,13 +232,15 @@ pub fn bedrock_stream_to_response_events(
.unwrap_or(0) as usize; .unwrap_or(0) as usize;
let recall_result = match messages_sent.lock() { let recall_result = match messages_sent.lock() {
Ok(sent) => recall_from_history( Ok(sent) => recall_tool_history(
&sent, &sent,
&tool_result_archive, &tool_result_archive,
search_query, ToolHistoryQuery {
tool_name_filter, search_query,
tool_use_id, tool_name: tool_name_filter,
offset, tool_use_id,
offset_from_end: offset,
},
), ),
Err(_) => "Error: could not access conversation history.".to_string(), Err(_) => "Error: could not access conversation history.".to_string(),
}; };
@@ -1504,137 +1507,3 @@ pub(super) fn is_known_tool(name: &str) -> bool {
fn is_notebook_tool(name: &str) -> bool { fn is_notebook_tool(name: &str) -> bool {
matches!(name, "create_notebook" | "read_notebook" | "edit_notebook") matches!(name, "create_notebook" | "read_notebook" | "edit_notebook")
} }
/// Searches conversation message history for tool call results matching the given criteria.
pub(crate) fn recall_from_history(
messages: &[ConversationMessage],
archive: &[ConversationMessage],
search_query: &str,
tool_name_filter: &str,
tool_use_id: &str,
offset_from_end: usize,
) -> String {
use super::convert::{ContentPart, MessageContent};
struct ToolEntry {
tool_use_id: String,
name: String,
input: String,
result: String,
}
let mut tool_entries: Vec<ToolEntry> = Vec::new();
let mut pending_tool_uses: Vec<(String, String, String)> = Vec::new(); // (id, name, input)
for msg in messages.iter().chain(archive.iter()) {
match &msg.content {
MessageContent::ToolUse {
tool_use_id,
name,
input,
} => {
pending_tool_uses.push((tool_use_id.clone(), name.clone(), input.to_string()));
}
MessageContent::ToolResult {
tool_use_id,
content,
..
} => {
if let Some(pos) = pending_tool_uses
.iter()
.position(|(id, _, _)| id == tool_use_id)
{
let (tuid, name, input) = pending_tool_uses.remove(pos);
tool_entries.push(ToolEntry {
tool_use_id: tuid,
name,
input,
result: content.clone(),
});
}
}
MessageContent::MultiPart(parts) => {
for part in parts {
match part {
ContentPart::ToolUse {
tool_use_id,
name,
input,
} => {
pending_tool_uses.push((
tool_use_id.clone(),
name.clone(),
input.to_string(),
));
}
ContentPart::ToolResult {
tool_use_id,
content,
..
} => {
if let Some(pos) = pending_tool_uses
.iter()
.position(|(id, _, _)| id == tool_use_id)
{
let (tuid, name, input) = pending_tool_uses.remove(pos);
tool_entries.push(ToolEntry {
tool_use_id: tuid,
name,
input,
result: content.clone(),
});
}
}
_ => {}
}
}
}
_ => {}
}
}
let filtered: Vec<&ToolEntry> = tool_entries
.iter()
.filter(|entry| {
if !tool_use_id.is_empty() && entry.tool_use_id != tool_use_id {
return false;
}
if !tool_name_filter.is_empty() && entry.name != tool_name_filter {
return false;
}
if !search_query.is_empty() {
let haystack = format!("{} {} {}", entry.name, entry.input, entry.result);
let query_lower = search_query.to_lowercase();
if !haystack.to_lowercase().contains(&query_lower) {
return false;
}
}
true
})
.collect();
if filtered.is_empty() {
return "No matching tool calls found in conversation history.".to_string();
}
// Get the entry at offset_from_end (0 = most recent)
let idx = if offset_from_end >= filtered.len() {
0
} else {
filtered.len() - 1 - offset_from_end
};
let entry = &filtered[idx];
let result_display = if entry.result.len() > 50000 {
let trunc = entry.result.chars().take(50000).collect::<String>();
format!("{trunc}... [truncated, {} total chars]", entry.result.len())
} else {
entry.result.clone()
};
format!(
"Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{}",
entry.name, entry.tool_use_id, entry.input, result_display
)
}
+174 -3
View File
@@ -34,6 +34,9 @@ pub use execute::{
StartAgentExecutorEvent, StartAgentRequest, StartAgentRequestId, StartAgentExecutorEvent, StartAgentRequest, StartAgentRequestId,
}; };
use futures::future::{join_all, BoxFuture}; use futures::future::{join_all, BoxFuture};
use galaxy_agent_core::{
PermissionDecision, PermissionKind, PermissionRequest, ToolEvent, ToolResult, ToolResultStatus,
};
use galaxyui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity}; use galaxyui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonEntity};
use itertools::Itertools; use itertools::Itertools;
use parking_lot::FairMutex; use parking_lot::FairMutex;
@@ -210,6 +213,67 @@ fn can_start_action_with_current_phase(
} }
} }
fn permission_request_id(action_id: &AIAgentActionId) -> String {
format!("permission:{action_id}")
}
fn is_permission_denial(reason: CancellationReason, status: Option<&AIActionStatus>) -> bool {
matches!(reason, CancellationReason::ManuallyCancelled)
&& matches!(status, Some(AIActionStatus::Blocked))
}
fn permission_kind_for_action(action: &AIAgentActionType) -> PermissionKind {
match action {
AIAgentActionType::ReadFiles(_)
| AIAgentActionType::SearchCodebase(_)
| AIAgentActionType::Grep { .. }
| AIAgentActionType::FileGlob { .. }
| AIAgentActionType::FileGlobV2 { .. }
| AIAgentActionType::ReadMCPResource { .. }
| AIAgentActionType::ReadDocuments(_)
| AIAgentActionType::ReadSkill(_)
| AIAgentActionType::FetchConversation { .. }
| AIAgentActionType::WaitForEvents { .. } => PermissionKind::Read,
AIAgentActionType::RequestFileEdits { .. }
| AIAgentActionType::EditDocuments(_)
| AIAgentActionType::CreateDocuments(_)
| AIAgentActionType::InitProject
| AIAgentActionType::InsertCodeReviewComments { .. } => PermissionKind::Write,
AIAgentActionType::RequestCommandOutput { .. }
| AIAgentActionType::WriteToLongRunningShellCommand { .. }
| AIAgentActionType::ReadShellCommandOutput { .. }
| AIAgentActionType::UseComputer(_)
| AIAgentActionType::RequestComputerUse(_)
| AIAgentActionType::TransferShellCommandControlToUser { .. }
| AIAgentActionType::OpenCodeReview => PermissionKind::Execute,
AIAgentActionType::UploadArtifact(_) => PermissionKind::Network,
AIAgentActionType::CallMCPTool { .. }
| AIAgentActionType::SuggestNewConversation { .. }
| AIAgentActionType::SuggestPrompt(_)
| AIAgentActionType::StartAgent { .. }
| AIAgentActionType::SendMessageToAgent { .. }
| AIAgentActionType::AskUserQuestion { .. }
| AIAgentActionType::RunAgents(_) => PermissionKind::ExternalTool,
}
}
fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bool) -> ToolResult {
let status = if permission_denied {
ToolResultStatus::Denied
} else if action_result.result.is_cancelled() {
ToolResultStatus::Cancelled
} else if action_result.result.is_failed() {
ToolResultStatus::Error
} else {
ToolResultStatus::Success
};
ToolResult {
call_id: action_result.id.to_string(),
content: action_result.result.model_content(),
status,
}
}
pub struct BlocklistAIActionModel { pub struct BlocklistAIActionModel {
executor: ModelHandle<BlocklistAIActionExecutor>, executor: ModelHandle<BlocklistAIActionExecutor>,
@@ -224,12 +288,19 @@ pub struct BlocklistAIActionModel {
/// Map from conversation ID to actions received in the most recent AI output that are finished. /// Map from conversation ID to actions received in the most recent AI output that are finished.
finished_action_results: HashMap<AIConversationId, Vec<Arc<AIAgentActionResult>>>, finished_action_results: HashMap<AIConversationId, Vec<Arc<AIAgentActionResult>>>,
/// Provider-neutral results for the same finished actions. Rig consumes these directly rather
/// than reconstructing them from the legacy request protobuf.
finished_tool_results: HashMap<AIConversationId, Vec<ToolResult>>,
/// Original order for the current batch of actions. /// Original order for the current batch of actions.
/// ///
/// We maintain this so that even though we might process actions in parallel, /// We maintain this so that even though we might process actions in parallel,
/// we can still order the results consistently. /// we can still order the results consistently.
action_order: HashMap<AIConversationId, HashMap<AIAgentActionId, usize>>, action_order: HashMap<AIConversationId, HashMap<AIAgentActionId, usize>>,
/// Permission-card rejections that still need a correlated completion event.
denied_permissions: HashSet<(AIConversationId, AIAgentActionId)>,
/// Past actions and their corresponding statuses from previous AI exchanges. /// Past actions and their corresponding statuses from previous AI exchanges.
past_action_results: HashMap<AIAgentActionId, Arc<AIAgentActionResult>>, past_action_results: HashMap<AIAgentActionId, Arc<AIAgentActionResult>>,
@@ -266,6 +337,12 @@ impl BlocklistAIActionModel {
ctx.subscribe_to_model(&executor, move |me, _, event, ctx| match event { ctx.subscribe_to_model(&executor, move |me, _, event, ctx| match event {
BlocklistAIActionExecutorEvent::ExecutingAction { action_id } => { BlocklistAIActionExecutorEvent::ExecutingAction { action_id } => {
ctx.emit(BlocklistAIActionEvent::ExecutingAction(action_id.clone())); ctx.emit(BlocklistAIActionEvent::ExecutingAction(action_id.clone()));
ctx.emit(BlocklistAIActionEvent::ToolLifecycle {
action_id: action_id.clone(),
event: ToolEvent::Started {
call_id: action_id.to_string(),
},
});
} }
BlocklistAIActionExecutorEvent::FinishedAction { BlocklistAIActionExecutorEvent::FinishedAction {
result, result,
@@ -298,10 +375,12 @@ impl BlocklistAIActionModel {
Self { Self {
pending_actions: Default::default(), pending_actions: Default::default(),
finished_action_results: Default::default(), finished_action_results: Default::default(),
finished_tool_results: Default::default(),
executor, executor,
past_action_results: HashMap::new(), past_action_results: HashMap::new(),
running_actions: Default::default(), running_actions: Default::default(),
action_order: Default::default(), action_order: Default::default(),
denied_permissions: Default::default(),
terminal_view_id, terminal_view_id,
pending_preprocessed_actions: Default::default(), pending_preprocessed_actions: Default::default(),
is_view_only: false, is_view_only: false,
@@ -533,6 +612,18 @@ impl BlocklistAIActionModel {
action_order.get(&result.id).copied().unwrap_or(usize::MAX) action_order.get(&result.id).copied().unwrap_or(usize::MAX)
}); });
} }
if let Some(tool_results) = self.finished_tool_results.get_mut(&conversation_id) {
let tool_order = action_order
.iter()
.map(|(id, index)| (id.to_string(), *index))
.collect::<HashMap<_, _>>();
tool_results.sort_by_key(|result| {
tool_order
.get(&result.call_id)
.copied()
.unwrap_or(usize::MAX)
});
}
} }
} }
@@ -833,6 +924,17 @@ impl BlocklistAIActionModel {
ctx.emit(BlocklistAIActionEvent::ActionBlockedOnUserConfirmation( ctx.emit(BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(
action.id.clone(), action.id.clone(),
)); ));
ctx.emit(BlocklistAIActionEvent::ToolLifecycle {
action_id: action.id.clone(),
event: ToolEvent::PermissionRequested {
request: PermissionRequest {
id: permission_request_id(&action.id),
call_id: action.id.to_string(),
kind: permission_kind_for_action(&action.action),
reason: Some(action.action.user_friendly_name()),
},
},
});
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
let blocked_action_user_friendly_str = action.action.user_friendly_name(); let blocked_action_user_friendly_str = action.action.user_friendly_name();
history_model.update_conversation_status( history_model.update_conversation_status(
@@ -897,6 +999,16 @@ impl BlocklistAIActionModel {
let action_id = action.id.clone(); let action_id = action.id.clone();
let phase = self.action_phase_for_action(&action, ctx); let phase = self.action_phase_for_action(&action, ctx);
if is_user_initiated {
ctx.emit(BlocklistAIActionEvent::ToolLifecycle {
action_id: action_id.clone(),
event: ToolEvent::PermissionResolved {
request_id: permission_request_id(&action_id),
call_id: action_id.to_string(),
decision: PermissionDecision::AllowOnce,
},
});
}
// WaitForEvents owns its own status transition; skip the default // WaitForEvents owns its own status transition; skip the default
// in-progress update. // in-progress update.
let is_wait_for_events = matches!(action.action, AIAgentActionType::WaitForEvents { .. }); let is_wait_for_events = matches!(action.action, AIAgentActionType::WaitForEvents { .. });
@@ -1073,6 +1185,8 @@ impl BlocklistAIActionModel {
reason: CancellationReason, reason: CancellationReason,
ctx: &mut ModelContext<Self>, ctx: &mut ModelContext<Self>,
) { ) {
let status = self.get_action_status(action_id);
let permission_denied = is_permission_denial(reason, status.as_ref());
if self if self
.running_actions .running_actions
.get(&conversation_id) .get(&conversation_id)
@@ -1092,7 +1206,13 @@ impl BlocklistAIActionModel {
.find_position(|action| action.id == *action_id) .find_position(|action| action.id == *action_id)
{ {
if let Some(action) = pending_actions_for_conversation.remove(idx) { if let Some(action) = pending_actions_for_conversation.remove(idx) {
self.cancel_pending_action(conversation_id, action, Some(reason), ctx); self.cancel_pending_action(
conversation_id,
action,
Some(reason),
permission_denied,
ctx,
);
} }
} }
} }
@@ -1140,7 +1260,7 @@ impl BlocklistAIActionModel {
reason, reason,
std::backtrace::Backtrace::force_capture() std::backtrace::Backtrace::force_capture()
); );
self.cancel_pending_action(conversation_id, action, reason, ctx); self.cancel_pending_action(conversation_id, action, reason, false, ctx);
} }
} }
@@ -1177,8 +1297,22 @@ impl BlocklistAIActionModel {
conversation_id: AIConversationId, conversation_id: AIConversationId,
pending_action: AIAgentAction, pending_action: AIAgentAction,
reason: Option<CancellationReason>, reason: Option<CancellationReason>,
permission_denied: bool,
ctx: &mut ModelContext<Self>, ctx: &mut ModelContext<Self>,
) { ) {
if permission_denied {
self.denied_permissions
.insert((conversation_id, pending_action.id.clone()));
ctx.emit(BlocklistAIActionEvent::ToolLifecycle {
action_id: pending_action.id.clone(),
event: ToolEvent::PermissionResolved {
request_id: permission_request_id(&pending_action.id),
call_id: pending_action.id.to_string(),
decision: PermissionDecision::Denied { reason: None },
},
});
}
if matches!( if matches!(
pending_action.action, pending_action.action,
AIAgentActionType::RequestComputerUse(_) AIAgentActionType::RequestComputerUse(_)
@@ -1227,10 +1361,20 @@ impl BlocklistAIActionModel {
.collect_vec() .collect_vec()
} }
pub(super) fn drain_finished_tool_results(
&mut self,
conversation_id: AIConversationId,
) -> Vec<ToolResult> {
self.finished_tool_results
.remove(&conversation_id)
.unwrap_or_default()
}
/// Clears finished action results for a conversation. Used when reverting. /// Clears finished action results for a conversation. Used when reverting.
pub(super) fn clear_finished_action_results(&mut self, conversation_id: AIConversationId) { pub(super) fn clear_finished_action_results(&mut self, conversation_id: AIConversationId) {
self.action_order.remove(&conversation_id); self.action_order.remove(&conversation_id);
self.finished_action_results.remove(&conversation_id); self.finished_action_results.remove(&conversation_id);
self.finished_tool_results.remove(&conversation_id);
} }
/// The control flow for initiating cancellations across suggested plans, requested commands, /// The control flow for initiating cancellations across suggested plans, requested commands,
@@ -1308,10 +1452,31 @@ impl BlocklistAIActionModel {
) )
) { ) {
for action in self.drain_pending_request_command_actions(conversation_id) { for action in self.drain_pending_request_command_actions(conversation_id) {
self.cancel_pending_action(conversation_id, action, cancellation_reason, ctx); self.cancel_pending_action(
conversation_id,
action,
cancellation_reason,
false,
ctx,
);
} }
} }
let permission_denied = self
.denied_permissions
.remove(&(conversation_id, action_result.id.clone()));
let tool_result = domain_tool_result(&action_result, permission_denied);
self.finished_tool_results
.entry(conversation_id)
.or_default()
.push(tool_result.clone());
ctx.emit(BlocklistAIActionEvent::ToolLifecycle {
action_id: action_result.id.clone(),
event: ToolEvent::Completed {
result: tool_result,
},
});
self.finished_action_results self.finished_action_results
.entry(conversation_id) .entry(conversation_id)
.or_default() .or_default()
@@ -1474,6 +1639,11 @@ pub enum BlocklistAIActionEvent {
conversation_id: AIConversationId, conversation_id: AIConversationId,
cancellation_reason: Option<CancellationReason>, cancellation_reason: Option<CancellationReason>,
}, },
/// Provider-neutral permission and execution lifecycle event for runtime consumers.
ToolLifecycle {
action_id: AIAgentActionId,
event: ToolEvent,
},
InitProject(AIAgentActionId), InitProject(AIAgentActionId),
ToggleCodeReview(AIAgentActionId), ToggleCodeReview(AIAgentActionId),
InsertCodeReviewComments { InsertCodeReviewComments {
@@ -1491,6 +1661,7 @@ impl BlocklistAIActionEvent {
BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(action_id) => action_id, BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(action_id) => action_id,
BlocklistAIActionEvent::ExecutingAction(action_id) => action_id, BlocklistAIActionEvent::ExecutingAction(action_id) => action_id,
BlocklistAIActionEvent::FinishedAction { action_id, .. } => action_id, BlocklistAIActionEvent::FinishedAction { action_id, .. } => action_id,
BlocklistAIActionEvent::ToolLifecycle { action_id, .. } => action_id,
BlocklistAIActionEvent::InitProject(action_id) => action_id, BlocklistAIActionEvent::InitProject(action_id) => action_id,
BlocklistAIActionEvent::ToggleCodeReview(action_id) => action_id, BlocklistAIActionEvent::ToggleCodeReview(action_id) => action_id,
BlocklistAIActionEvent::InsertCodeReviewComments { action_id, .. } => action_id, BlocklistAIActionEvent::InsertCodeReviewComments { action_id, .. } => action_id,
+116 -1
View File
@@ -3,7 +3,9 @@ use std::sync::Arc;
use super::*; use super::*;
use crate::ai::agent::task::TaskId; use crate::ai::agent::task::TaskId;
use crate::ai::agent::AIAgentActionResultType; use crate::ai::agent::{
AIAgentActionResultType, AnyFileContent, FileContext, GrepResult, ReadFilesResult,
};
fn make_action_result(id: &str) -> Arc<AIAgentActionResult> { fn make_action_result(id: &str) -> Arc<AIAgentActionResult> {
Arc::new(AIAgentActionResult { Arc::new(AIAgentActionResult {
@@ -13,6 +15,14 @@ fn make_action_result(id: &str) -> Arc<AIAgentActionResult> {
}) })
} }
fn action_result(id: &str, result: AIAgentActionResultType) -> AIAgentActionResult {
AIAgentActionResult {
id: AIAgentActionId::from(id.to_owned()),
task_id: TaskId::new("task".to_owned()),
result,
}
}
fn count_startable_actions_for_pass(phases: &[(RunningActionPhase, bool)]) -> usize { fn count_startable_actions_for_pass(phases: &[(RunningActionPhase, bool)]) -> usize {
let mut current_phase = None; let mut current_phase = None;
let mut count = 0; let mut count = 0;
@@ -100,3 +110,108 @@ fn finished_results_stay_in_original_action_order() {
AIAgentActionId::from("third".to_owned()) AIAgentActionId::from("third".to_owned())
); );
} }
#[test]
fn domain_tool_results_preserve_success_failure_cancellation_and_denial() {
let success = domain_tool_result(
&action_result("success", AIAgentActionResultType::InitProject),
false,
);
let failure = domain_tool_result(
&action_result(
"failure",
AIAgentActionResultType::Grep(GrepResult::Error("boom".to_string())),
),
false,
);
let cancelled_result = action_result(
"cancelled",
AIAgentActionResultType::Grep(GrepResult::Cancelled),
);
let cancelled = domain_tool_result(&cancelled_result, false);
let denied = domain_tool_result(&cancelled_result, true);
assert_eq!(success.status, ToolResultStatus::Success);
assert_eq!(failure.status, ToolResultStatus::Error);
assert_eq!(cancelled.status, ToolResultStatus::Cancelled);
assert_eq!(denied.status, ToolResultStatus::Denied);
assert_eq!(success.call_id, "success");
assert_eq!(failure.call_id, "failure");
assert_eq!(cancelled.call_id, "cancelled");
assert_eq!(denied.call_id, "cancelled");
}
#[test]
fn domain_read_result_contains_the_file_contents_for_the_next_model_turn() {
let result = action_result(
"read-call",
AIAgentActionResultType::ReadFiles(ReadFilesResult::Success {
files: vec![FileContext::new(
"/workspace/src/lib.rs".to_string(),
AnyFileContent::StringContent("pub fn answer() -> u8 { 42 }".to_string()),
None,
None,
)],
}),
);
let result = domain_tool_result(&result, false);
assert_eq!(result.status, ToolResultStatus::Success);
assert_eq!(result.call_id, "read-call");
assert!(result.content.contains("/workspace/src/lib.rs"));
assert!(result.content.contains("pub fn answer() -> u8 { 42 }"));
}
#[test]
fn action_permission_kinds_match_the_safety_boundary() {
assert_eq!(
permission_kind_for_action(&AIAgentActionType::Grep {
queries: vec!["needle".to_string()],
path: ".".to_string(),
}),
PermissionKind::Read
);
assert_eq!(
permission_kind_for_action(&AIAgentActionType::InitProject),
PermissionKind::Write
);
assert_eq!(
permission_kind_for_action(&AIAgentActionType::RequestCommandOutput {
command: "cargo test".to_string(),
is_read_only: Some(true),
is_risky: Some(false),
wait_until_completion: true,
uses_pager: Some(false),
rationale: None,
citations: Vec::new(),
}),
PermissionKind::Execute
);
assert_eq!(
permission_kind_for_action(&AIAgentActionType::CallMCPTool {
server_id: None,
name: "tool".to_string(),
input: serde_json::json!({}),
}),
PermissionKind::ExternalTool
);
}
#[test]
fn only_rejecting_a_blocked_action_is_a_permission_denial() {
assert!(is_permission_denial(
CancellationReason::ManuallyCancelled,
Some(&AIActionStatus::Blocked),
));
assert!(!is_permission_denial(
CancellationReason::ManuallyCancelled,
Some(&AIActionStatus::Queued),
));
assert!(!is_permission_denial(
CancellationReason::FollowUpSubmitted {
is_for_same_conversation: true,
},
Some(&AIActionStatus::Blocked),
));
}
+2 -1
View File
@@ -4929,7 +4929,8 @@ impl AIBlock {
} }
} }
BlocklistAIActionEvent::InitProject(_) BlocklistAIActionEvent::ToolLifecycle { .. }
| BlocklistAIActionEvent::InitProject(_)
| BlocklistAIActionEvent::ToggleCodeReview(_) => {} | BlocklistAIActionEvent::ToggleCodeReview(_) => {}
} }
}); });
+23 -72
View File
@@ -10,7 +10,7 @@ mod pending_response_streams;
pub mod response_stream; pub mod response_stream;
pub(super) mod shared_session; pub(super) mod shared_session;
mod slash_command; mod slash_command;
use std::collections::{HashMap, HashSet, VecDeque}; use std::collections::{HashMap, HashSet};
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::Arc; use std::sync::Arc;
@@ -19,6 +19,7 @@ use std::time::Duration;
use ai::skills::SkillPathOrigin; use ai::skills::SkillPathOrigin;
use anyhow::anyhow; use anyhow::anyhow;
use chrono::{DateTime, Local}; use chrono::{DateTime, Local};
use galaxy_agent_core::ToolLoopGuard;
use galaxy_core::assertions::safe_assert; use galaxy_core::assertions::safe_assert;
use input_context::{input_context_for_request, parse_context_attachments}; use input_context::{input_context_for_request, parse_context_attachments};
use itertools::Itertools; use itertools::Itertools;
@@ -196,60 +197,6 @@ pub enum BlocklistAIControllerEvent {
FreeTierLimitCheckTriggered, FreeTierLimitCheckTriggered,
} }
/// Tracks recent failed action signatures for loop detection.
/// When the same tool+input pattern fails repeatedly, we inject
/// corrective instructions to break the cycle.
#[derive(Debug, Clone)]
struct LoopDetectionEntry {
/// Discriminant of the action result type (e.g. RequestCommandOutput, ApplyFileDiffs)
tool_discriminant: std::mem::Discriminant<AIAgentActionResultType>,
/// Hash of the action's identifying input (command string, file paths, etc.)
input_hash: u64,
/// Human-readable description of what failed
description: String,
}
#[derive(Debug, Default, Clone)]
struct LoopDetectionState {
recent_failures: VecDeque<LoopDetectionEntry>,
}
const LOOP_DETECTION_WINDOW: usize = 10;
const LOOP_DETECTION_THRESHOLD: usize = 3;
impl LoopDetectionState {
fn record_failure(&mut self, entry: LoopDetectionEntry) {
self.recent_failures.push_back(entry);
if self.recent_failures.len() > LOOP_DETECTION_WINDOW {
self.recent_failures.pop_front();
}
}
fn detect_loop(&self) -> Option<&LoopDetectionEntry> {
use std::collections::HashMap as CountMap;
let mut counts: CountMap<
(std::mem::Discriminant<AIAgentActionResultType>, u64),
(usize, usize),
> = CountMap::new();
for (idx, entry) in self.recent_failures.iter().enumerate() {
let key = (entry.tool_discriminant, entry.input_hash);
let counter = counts.entry(key).or_insert((0, 0));
counter.0 += 1;
counter.1 = idx; // Track most recent occurrence
}
for ((_disc, _hash), (count, latest_idx)) in &counts {
if *count >= LOOP_DETECTION_THRESHOLD {
return self.recent_failures.get(*latest_idx);
}
}
None
}
fn clear(&mut self) {
self.recent_failures.clear();
}
}
#[derive(Debug)] #[derive(Debug)]
pub struct RequestInput { pub struct RequestInput {
pub conversation_id: AIConversationId, pub conversation_id: AIConversationId,
@@ -419,7 +366,7 @@ pub struct BlocklistAIController {
pending_passive_follow_ups: HashSet<AIConversationId>, pending_passive_follow_ups: HashSet<AIConversationId>,
/// Per-conversation loop detection state for preventing recursive tool failures. /// Per-conversation loop detection state for preventing recursive tool failures.
loop_detection: HashMap<AIConversationId, LoopDetectionState>, loop_detection: HashMap<AIConversationId, ToolLoopGuard>,
/// Per-conversation error retry count for injecting corrective messages on failure. /// Per-conversation error retry count for injecting corrective messages on failure.
error_retry_counts: HashMap<AIConversationId, usize>, error_retry_counts: HashMap<AIConversationId, usize>,
/// Passive suggestion results that should be included with the next request /// Passive suggestion results that should be included with the next request
@@ -1992,11 +1939,7 @@ impl BlocklistAIController {
description.hash(&mut hasher); description.hash(&mut hasher);
let input_hash = hasher.finish(); let input_hash = hasher.finish();
state.record_failure(LoopDetectionEntry { state.record_failure(input_hash, description);
tool_discriminant: discriminant,
input_hash,
description: description.clone(),
});
} else if result.result.is_successful() { } else if result.result.is_successful() {
has_success = true; has_success = true;
} }
@@ -2005,23 +1948,21 @@ impl BlocklistAIController {
// If we had at least one success in this batch, clear loop state — // If we had at least one success in this batch, clear loop state —
// the agent is making progress. // the agent is making progress.
if has_success { if has_success {
state.clear(); state.record_success();
return None; return None;
} }
// Check for loops // Check for loops
if let Some(looping_entry) = state.detect_loop() { if let Some(looping_entry) = state.detect_and_reset() {
let warning = format!( let warning = format!(
"[SYSTEM] Loop detected: the same action has failed {} or more times consecutively. \ "[SYSTEM] Loop detected: the same action has failed {} or more times consecutively. \
Do NOT repeat this action or any similar approach.\n\n\ Do NOT repeat this action or any similar approach.\n\n\
Failing action: {}\n\n\ Failing action: {}\n\n\
Take a completely different approach to accomplish the goal. \ Take a completely different approach to accomplish the goal. \
If you cannot find an alternative, explain to the user what is failing and why.", If you cannot find an alternative, explain to the user what is failing and why.",
LOOP_DETECTION_THRESHOLD, looping_entry.threshold,
looping_entry.description looping_entry.description
); );
// Clear the state so we don't keep injecting on every subsequent turn
state.clear();
Some(warning) Some(warning)
} else { } else {
None None
@@ -3061,11 +3002,23 @@ impl BlocklistAIController {
query_metadata, query_metadata,
ctx, ctx,
); );
let action_result_ids = request_input
.all_inputs()
.filter_map(AIAgentInput::action_result)
.map(|result| result.id.to_string())
.collect::<HashSet<_>>();
request_params.tool_results = self.action_model.update(ctx, |action_model, _| {
action_model
.drain_finished_tool_results(conversation_id)
.into_iter()
.filter(|result| action_result_ids.contains(&result.call_id))
.collect()
});
request_params.parent_agent_id = parent_agent_id; request_params.parent_agent_id = parent_agent_id;
request_params.agent_name = agent_name; request_params.agent_name = agent_name;
request_params.bedrock_message_history = bedrock_history; request_params.message_history = bedrock_history;
request_params.bedrock_tool_result_archive = bedrock_tool_result_archive; request_params.tool_result_archive = bedrock_tool_result_archive;
request_params.bedrock_progressive_summary = bedrock_progressive_summary; request_params.progressive_summary = bedrock_progressive_summary;
// For the Bedrock path, when this is the first request in a new conversation // For the Bedrock path, when this is the first request in a new conversation
// (no tasks established yet), use the conversation's root task ID so the // (no tasks established yet), use the conversation's root task ID so the
@@ -3555,9 +3508,7 @@ impl BlocklistAIController {
// history (input + assistant response) from the Arc back // history (input + assistant response) from the Arc back
// into the conversation for the next request cycle. // into the conversation for the next request cycle.
let new_history = (!response_stream.as_ref(ctx).is_acp()) let new_history = (!response_stream.as_ref(ctx).is_acp())
.then(|| { .then(|| response_stream.as_ref(ctx).messages_sent().clone())
response_stream.as_ref(ctx).bedrock_messages_sent().clone()
})
.and_then(|messages_sent| { .and_then(|messages_sent| {
messages_sent.lock().ok().and_then(|sent| { messages_sent.lock().ok().and_then(|sent| {
if sent.is_empty() { if sent.is_empty() {
@@ -552,11 +552,11 @@ impl ResponseStream {
} }
} }
pub fn bedrock_messages_sent( pub fn messages_sent(
&self, &self,
) -> &std::sync::Arc<std::sync::Mutex<Vec<crate::ai::bedrock::convert::ConversationMessage>>> ) -> &std::sync::Arc<std::sync::Mutex<Vec<crate::ai::provider::types::ConversationMessage>>>
{ {
&self.params.bedrock_messages_sent &self.params.messages_sent
} }
/// Returns the model ID associated with this response stream's request. /// Returns the model ID associated with this response stream's request.
+9 -6
View File
@@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex};
use bytes::Bytes; use bytes::Bytes;
use futures::stream::BoxStream; use futures::stream::BoxStream;
use futures::Stream; use futures::Stream;
use galaxy_agent_core::{recall_tool_history, ToolHistoryQuery};
use serde_json::Value as JsonValue; use serde_json::Value as JsonValue;
use uuid::Uuid; use uuid::Uuid;
use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::response_event::stream_finished;
@@ -10,7 +11,7 @@ use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
use crate::ai::agent::api::Event; use crate::ai::agent::api::Event;
use crate::ai::bedrock::response_translator::{ use crate::ai::bedrock::response_translator::{
build_create_task, build_stream_init, context_window_for_model, recall_from_history, build_create_task, build_stream_init, context_window_for_model,
}; };
use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent, MessageRole}; use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent, MessageRole};
use crate::server::server_api::AIApiError; use crate::server::server_api::AIApiError;
@@ -276,13 +277,15 @@ pub fn openai_stream_to_response_events(
.and_then(|value| value.as_u64()) .and_then(|value| value.as_u64())
.unwrap_or(0) as usize; .unwrap_or(0) as usize;
let recall_result = match messages_sent.lock() { let recall_result = match messages_sent.lock() {
Ok(sent) => recall_from_history( Ok(sent) => recall_tool_history(
&sent, &sent,
&tool_result_archive, &tool_result_archive,
search_query, ToolHistoryQuery {
tool_name_filter, search_query,
tool_use_id, tool_name: tool_name_filter,
offset, tool_use_id,
offset_from_end: offset,
},
), ),
Err(_) => "Error: could not access conversation history.".to_string(), Err(_) => "Error: could not access conversation history.".to_string(),
}; };
+1
View File
@@ -1,5 +1,6 @@
mod provider; mod provider;
mod rig; mod rig;
mod rig_request;
pub(crate) use provider::ProviderRuntime; pub(crate) use provider::ProviderRuntime;
pub(crate) use rig::rig_openai_response_stream; pub(crate) use rig::rig_openai_response_stream;
+139 -61
View File
@@ -4,13 +4,14 @@ use futures::channel::oneshot;
use futures::{FutureExt, StreamExt}; use futures::{FutureExt, StreamExt};
use galaxy_agent_core::{ use galaxy_agent_core::{
turn_control, AgentError, AgentEvent, AgentRuntime, MessageContent, MessageRole, StopReason, turn_control, AgentError, AgentEvent, AgentRuntime, MessageContent, MessageRole, StopReason,
TurnCommand, TurnRequest, Usage, ToolCall, ToolCallDecision, ToolEvent, ToolPolicy, ToolResult, TurnCommand, Usage,
}; };
use galaxy_agent_rig::{OpenAICompatibleRuntime, OpenAICompatibleRuntimeConfig}; use galaxy_agent_rig::{OpenAICompatibleRuntime, OpenAICompatibleRuntimeConfig};
use uuid::Uuid; use uuid::Uuid;
use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; 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 crate::ai::agent::api::{Event, RequestParams, ResponseStream};
use crate::ai::bedrock::response_translator::{ use crate::ai::bedrock::response_translator::{
build_add_agent_output_message, build_append_text, build_create_task, build_stream_init, build_add_agent_output_message, build_append_text, build_create_task, build_stream_init,
@@ -18,56 +19,30 @@ use crate::ai::bedrock::response_translator::{
}; };
use crate::ai::openai::client::OpenAIClientConfig; use crate::ai::openai::client::OpenAIClientConfig;
use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage}; use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage};
use crate::ai::openai::translator::{prepare_turn, PreparedTurn, TranslatorRequest}; use crate::ai::provider::types::{ContentPart, ConversationMessage};
use crate::ai::provider::types::ConversationMessage;
use crate::server::server_api::AIApiError; use crate::server::server_api::AIApiError;
pub(crate) fn rig_openai_response_stream( pub(crate) fn rig_openai_response_stream(
config: OpenAIClientConfig, config: OpenAIClientConfig,
params: RequestParams, params: RequestParams,
request: &mut api::Request, supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
cancellation_rx: oneshot::Receiver<()>, cancellation_rx: oneshot::Receiver<()>,
) -> ResponseStream { ) -> ResponseStream {
let translator_request = TranslatorRequest { let PreparedRigTurn {
config: config.clone(),
model_id: params.model.as_str().to_string(),
root_task_id: params.root_task_id,
message_history: params.bedrock_message_history,
tool_result_archive: params.bedrock_tool_result_archive,
progressive_summary: params.bedrock_progressive_summary,
messages_sent: params.bedrock_messages_sent,
global_rules: params.global_rules,
};
let PreparedTurn {
task_id, task_id,
needs_create_task, needs_create_task,
user_query, user_query,
messages, request: turn_request,
system_prompt, persistent_messages,
tools: _, tool_result_archive,
model_id, messages_sent,
persistent_message_count, } = prepare_rig_turn(&config, params, supported_tools, supported_cli_agent_tools);
} = prepare_turn(&translator_request, request); store_messages_sent(&messages_sent, &persistent_messages);
store_messages_sent( let conversation_id = turn_request.conversation_id.clone();
&translator_request.messages_sent, let model_id = turn_request.model.as_str().to_string();
&messages, let tool_policy = ToolPolicy::new(&turn_request.tools);
persistent_message_count,
);
let conversation_id = request
.metadata
.as_ref()
.map(|metadata| metadata.conversation_id.clone())
.filter(|id| !id.is_empty());
let mut turn_request = TurnRequest::new(model_id.clone(), messages);
turn_request.conversation_id = conversation_id.clone();
turn_request.system_prompt = system_prompt;
// Phase 2 deliberately validates the model streaming seam. Galaxy tool
// execution moves behind AgentRuntime in Phase 3; exposing the legacy tool
// list here would split ownership across both systems.
turn_request.tools = Vec::new();
turn_request.max_output_tokens = config.max_output_tokens.map(u64::from);
let runtime = OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig { let runtime = OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig {
base_url: config.base_url, base_url: config.base_url,
@@ -76,7 +51,6 @@ pub(crate) fn rig_openai_response_stream(
max_output_tokens: config.max_output_tokens.map(u64::from), max_output_tokens: config.max_output_tokens.map(u64::from),
supports_system_messages: config.supports_system_messages, supports_system_messages: config.supports_system_messages,
}); });
let messages_sent = translator_request.messages_sent;
let max_context_tokens = config.max_input_tokens; let max_context_tokens = config.max_input_tokens;
let stream = async_stream::stream! { let stream = async_stream::stream! {
let (control_sender, control) = turn_control(); let (control_sender, control) = turn_control();
@@ -110,6 +84,8 @@ pub(crate) fn rig_openai_response_stream(
let mut current_text_message_id: Option<String> = None; let mut current_text_message_id: Option<String> = None;
let mut current_reasoning_message_id: Option<String> = None; let mut current_reasoning_message_id: Option<String> = None;
let mut full_text = String::new(); let mut full_text = String::new();
let mut proposed_tools = Vec::new();
let mut assistant_history_index = None;
let mut usage = Usage::default(); let mut usage = Usage::default();
loop { loop {
@@ -163,11 +139,57 @@ pub(crate) fn rig_openai_response_stream(
} }
} }
AgentEvent::UsageUpdated { usage: updated } => usage = updated, AgentEvent::UsageUpdated { usage: updated } => usage = updated,
AgentEvent::Tool {
event: ToolEvent::Proposed { call },
} => {
proposed_tools.push(call.clone());
sync_assistant_turn(
&messages_sent,
&full_text,
&proposed_tools,
&mut assistant_history_index,
);
let history = messages_sent
.lock()
.map(|sent| sent.clone())
.unwrap_or_default();
match tool_policy.decide(&call, &history, &tool_result_archive) {
ToolCallDecision::Execute => {
yield Ok(build_tool_proposed(&task_id, &call));
}
ToolCallDecision::Inline(result) => {
append_tool_result(&messages_sent, result);
}
ToolCallDecision::Reject(result) => {
log::warn!(
"Rig model called unavailable tool '{}' (id={})",
call.name,
call.id
);
let error_display = format!(
"Failed tool call: `{}`\n\n{}",
call.name, result.content
);
append_tool_result(&messages_sent, result);
let message_id = Uuid::new_v4().to_string();
yield Ok(build_add_agent_output_message(
&task_id,
&message_id,
&error_display,
));
}
}
}
AgentEvent::TurnStopped { reason } => { AgentEvent::TurnStopped { reason } => {
if !initialized { if !initialized {
yield Ok(build_stream_init(&request_id, &conversation_id)); yield Ok(build_stream_init(&request_id, &conversation_id));
} }
store_assistant_text(&messages_sent, full_text); sync_assistant_turn(
&messages_sent,
&full_text,
&proposed_tools,
&mut assistant_history_index,
);
yield Ok(build_stream_finished( yield Ok(build_stream_finished(
map_stop_reason(reason), map_stop_reason(reason),
StreamUsage { StreamUsage {
@@ -184,13 +206,10 @@ pub(crate) fn rig_openai_response_stream(
)); ));
return; return;
} }
AgentEvent::ToolProposed { .. } AgentEvent::Tool { .. } => {
| AgentEvent::PermissionRequested { .. }
| AgentEvent::ToolStarted { .. }
| AgentEvent::ToolCompleted { .. } => {
yield Err(agent_error(AgentError::new( yield Err(agent_error(AgentError::new(
galaxy_agent_core::AgentErrorKind::Protocol, galaxy_agent_core::AgentErrorKind::Protocol,
"the Phase 2 Rig runtime emitted a tool event while tools are disabled", "the provider runtime attempted to execute a tool outside Galaxy's permission boundary",
))); )));
return; return;
} }
@@ -206,31 +225,90 @@ pub(crate) fn rig_openai_response_stream(
fn store_messages_sent( fn store_messages_sent(
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>, messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
messages: &[ConversationMessage], messages: &[ConversationMessage],
persistent_message_count: usize,
) { ) {
let Ok(mut sent) = messages_sent.lock() else { let Ok(mut sent) = messages_sent.lock() else {
return; return;
}; };
if persistent_message_count > 0 && messages.len() >= persistent_message_count { *sent = messages.to_vec();
*sent = messages[messages.len() - persistent_message_count..].to_vec(); }
} else {
*sent = messages.to_vec(); fn append_tool_result(
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
result: ToolResult,
) {
let is_error = result.is_error();
let message = ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: result.call_id,
content: result.content,
is_error,
},
};
if let Ok(mut sent) = messages_sent.lock() {
sent.push(message);
} }
} }
fn store_assistant_text( fn sync_assistant_turn(
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>, messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
text: String, text: &str,
tool_calls: &[ToolCall],
history_index: &mut Option<usize>,
) { ) {
if text.is_empty() { let mut parts = Vec::with_capacity(usize::from(!text.is_empty()) + tool_calls.len());
if !text.is_empty() {
parts.push(ContentPart::Text(text.to_string()));
}
parts.extend(tool_calls.iter().map(|call| ContentPart::ToolUse {
tool_use_id: call.id.clone(),
name: call.name.clone(),
input: call.arguments.clone(),
}));
if parts.is_empty() {
return; return;
} }
if let Ok(mut sent) = messages_sent.lock() {
sent.push(ConversationMessage { let content = if parts.len() == 1 {
role: MessageRole::Assistant, match parts.pop().unwrap() {
content: MessageContent::Text(text), ContentPart::Text(text) => MessageContent::Text(text),
}); ContentPart::ToolUse {
tool_use_id,
name,
input,
} => MessageContent::ToolUse {
tool_use_id,
name,
input,
},
ContentPart::Image { .. } | ContentPart::ToolResult { .. } => unreachable!(),
}
} else {
MessageContent::MultiPart(parts)
};
let message = ConversationMessage {
role: MessageRole::Assistant,
content,
};
let Ok(mut sent) = messages_sent.lock() else {
return;
};
if let Some(index) = *history_index {
if index < sent.len() {
sent[index] = message;
return;
}
} }
*history_index = Some(sent.len());
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_add_reasoning(task_id: &str, message_id: &str, text: &str) -> ResponseEvent { fn build_add_reasoning(task_id: &str, message_id: &str, text: &str) -> ResponseEvent {
+625
View File
@@ -0,0 +1,625 @@
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use ai::agent::action_result::AnyFileContent;
use ai::skills::SkillReference;
use base64::engine::general_purpose;
use base64::Engine as _;
use galaxy_agent_core::{
ContentPart, ConversationMessage, MessageContent, MessageRole, ToolDefinition, ToolResult,
TurnRequest,
};
use warp_multi_agent_api::ToolType;
use crate::ai::agent::api::RequestParams;
use crate::ai::agent::{AIAgentContext, AIAgentInput, MCPContext, UserQueryMode};
use crate::ai::bedrock::request_translator::{default_tool_definitions, tool_name_is_supported};
use crate::ai::openai::client::OpenAIClientConfig;
use crate::ai::openai::request_translator::sanitize_messages_for_openai;
pub(crate) struct PreparedRigTurn {
pub task_id: String,
pub needs_create_task: bool,
pub user_query: Option<String>,
pub request: TurnRequest,
pub persistent_messages: Vec<ConversationMessage>,
pub tool_result_archive: Vec<ConversationMessage>,
pub messages_sent: Arc<Mutex<Vec<ConversationMessage>>>,
}
pub(crate) fn prepare_rig_turn(
config: &OpenAIClientConfig,
params: RequestParams,
supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
) -> PreparedRigTurn {
let RequestParams {
input,
tool_results,
conversation_token,
tasks,
model,
root_task_id,
message_history,
progressive_summary,
tool_result_archive,
messages_sent,
global_rules,
mcp_context,
..
} = params;
let task_id = root_task_id
.or_else(|| tasks.first().map(|task| task.id.clone()))
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let needs_create_task = tasks.is_empty();
let user_query = input.iter().find_map(input_user_query);
let mode = request_mode(&input);
let available_tools = match mode {
RigRequestMode::Cli => supported_cli_agent_tools,
RigRequestMode::Normal | RigRequestMode::Plan | RigRequestMode::Orchestrate => {
supported_tools
}
};
let tools = tool_definitions(&available_tools, mcp_context.as_ref());
let system_prompt = build_system_prompt(&input, &tools, &global_rules, mode);
let mut new_messages = input_messages(input, tool_results);
let mut persistent_messages = message_history;
persistent_messages.append(&mut new_messages);
for message in &mut persistent_messages {
message.truncate_tool_results_for_provider_request();
}
sanitize_messages_for_openai(&mut persistent_messages);
let mut turn_messages = Vec::new();
if let Some(summary) = progressive_summary {
turn_messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(format!(
"<conversation-history-summary>\n{summary}\n</conversation-history-summary>\n\n\
The above summarizes earlier conversation history. The detailed messages below are the most recent exchanges."
)),
});
turn_messages.push(ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::Text(
"Understood, I have the prior context. Continuing with the recent conversation."
.to_string(),
),
});
}
turn_messages.extend(persistent_messages.clone());
let model_id = config
.model
.clone()
.filter(|model| !model.is_empty() && model != "auto")
.unwrap_or_else(|| model.as_str().to_string());
let mut request = TurnRequest::new(model_id, turn_messages);
request.conversation_id = conversation_token.map(|token| token.as_str().to_string());
request.system_prompt = Some(system_prompt);
request.tools = tools;
request.max_output_tokens = config.max_output_tokens.map(u64::from);
PreparedRigTurn {
task_id,
needs_create_task,
user_query,
request,
persistent_messages,
tool_result_archive,
messages_sent,
}
}
fn input_messages(
inputs: Vec<AIAgentInput>,
tool_results: Vec<ToolResult>,
) -> Vec<ConversationMessage> {
let mut messages = Vec::new();
if !tool_results.is_empty() {
let mut parts = tool_results
.into_iter()
.map(|result| {
let is_error = result.is_error();
ContentPart::ToolResult {
tool_use_id: result.call_id,
content: result.content,
is_error,
}
})
.collect::<Vec<_>>();
let content = if parts.len() == 1 {
let ContentPart::ToolResult {
tool_use_id,
content,
is_error,
} = parts.pop().expect("one tool result exists")
else {
unreachable!()
};
MessageContent::ToolResult {
tool_use_id,
content,
is_error,
}
} else {
MessageContent::MultiPart(parts)
};
messages.push(ConversationMessage {
role: MessageRole::User,
content,
});
}
messages.extend(inputs.into_iter().filter_map(input_message));
messages
}
fn input_message(input: AIAgentInput) -> Option<ConversationMessage> {
let (text, images) = match input {
AIAgentInput::UserQuery {
query,
context,
running_command,
..
} => {
let text = if let Some(command) = running_command {
format!(
"[Running command: {}]\n[Command ID: {}]\n[Terminal output:\n{}\n]\n{}",
command.command, command.block_id, command.grid_contents, query
)
} else {
query
};
(text, image_parts(&context))
}
AIAgentInput::ActionResult { .. } => return None,
AIAgentInput::AutoCodeDiffQuery { query, .. } => (query, Vec::new()),
AIAgentInput::ResumeConversation { .. } => (
"Continue where we left off. Review the conversation history and proceed with the next steps."
.to_string(),
Vec::new(),
),
AIAgentInput::InitProjectRules { .. } => (
"Initialize this project. Analyze the codebase structure and files, generate an AGENTS.md file documenting project conventions and setup instructions, and offer to create a development environment configuration. Use the available tools to inspect the project before responding."
.to_string(),
Vec::new(),
),
AIAgentInput::CreateEnvironment { repo_paths, .. } => (
format!(
"Create a development environment for this project. Set up necessary dependencies, configuration files, and tooling. Repositories: {}",
repo_paths.join(", ")
),
Vec::new(),
),
AIAgentInput::TriggerPassiveSuggestion { .. } => (
"Suggest a useful next action based on the current project context.".to_string(),
Vec::new(),
),
AIAgentInput::CreateNewProject { query, .. } => {
(format!("Create a new project: {query}"), Vec::new())
}
AIAgentInput::CloneRepository {
clone_repo_url, ..
} => (
format!(
"Clone the repository at {} and set it up for development.",
clone_repo_url.into_url()
),
Vec::new(),
),
AIAgentInput::CodeReview { .. } => (
"Review the provided code changes and address the review comments.".to_string(),
Vec::new(),
),
AIAgentInput::FetchReviewComments { repo_path, .. } => (
format!("Fetch and review the pull-request comments for {repo_path}."),
Vec::new(),
),
AIAgentInput::SummarizeConversation { prompt, .. } => (
prompt.unwrap_or_else(|| {
"Summarize this conversation, preserving decisions, changes, and context needed to continue."
.to_string()
}),
Vec::new(),
),
AIAgentInput::InvokeSkill {
skill, user_query, ..
} => {
let suffix = user_query
.map(|query| query.query)
.filter(|query| !query.is_empty())
.map(|query| format!("\n\nAdditional context from user: {query}"))
.unwrap_or_default();
(
format!(
"Execute the following skill: {}\n\n<skill-instructions>\n{}\n</skill-instructions>{suffix}",
skill.name, skill.content
),
Vec::new(),
)
}
AIAgentInput::StartFromAmbientRunPrompt { ambient_run_id, .. } => (
format!("Continue the configured ambient-agent run {ambient_run_id}."),
Vec::new(),
),
AIAgentInput::MessagesReceivedFromAgents { messages } => (
messages
.into_iter()
.map(|message| {
format!(
"Message from {} ({})\nSubject: {}\n{}",
message.sender_agent_id,
message.addresses.join(", "),
message.subject,
message.message_body
)
})
.collect::<Vec<_>>()
.join("\n\n"),
Vec::new(),
),
AIAgentInput::EventsFromAgents { events } => (
format!("Agent lifecycle events:\n{events:#?}"),
Vec::new(),
),
AIAgentInput::PassiveSuggestionResult { suggestion, .. } => (
format!("The user responded to a passive suggestion: {suggestion:?}"),
Vec::new(),
),
AIAgentInput::OrchestrationConfigUpdate {
plan_id,
config,
status,
} => (
format!(
"Orchestration configuration updated for plan {plan_id}: status={status:?}, config={config:?}"
),
Vec::new(),
),
};
let content = if images.is_empty() {
MessageContent::Text(text)
} else {
let mut parts = Vec::with_capacity(images.len() + 1);
parts.push(ContentPart::Text(text));
parts.extend(images);
MessageContent::MultiPart(parts)
};
Some(ConversationMessage {
role: MessageRole::User,
content,
})
}
fn image_parts(context: &[AIAgentContext]) -> Vec<ContentPart> {
context
.iter()
.filter_map(|context| {
let AIAgentContext::Image(image) = context else {
return None;
};
let data = match general_purpose::STANDARD.decode(&image.data) {
Ok(data) => data,
Err(error) => {
log::warn!("Skipping invalid base64 image supplied to Rig: {error}");
return None;
}
};
Some(ContentPart::Image {
data,
mime_type: image.mime_type.clone(),
})
})
.collect()
}
fn input_user_query(input: &AIAgentInput) -> Option<String> {
match input {
AIAgentInput::UserQuery { query, .. } => Some(query.clone()),
AIAgentInput::InvokeSkill { skill, .. } => Some(format!("/{}", skill.name)),
AIAgentInput::AutoCodeDiffQuery { .. }
| AIAgentInput::ResumeConversation { .. }
| AIAgentInput::InitProjectRules { .. }
| AIAgentInput::CreateEnvironment { .. }
| AIAgentInput::TriggerPassiveSuggestion { .. }
| AIAgentInput::CreateNewProject { .. }
| AIAgentInput::CloneRepository { .. }
| AIAgentInput::CodeReview { .. }
| AIAgentInput::FetchReviewComments { .. }
| AIAgentInput::SummarizeConversation { .. }
| AIAgentInput::StartFromAmbientRunPrompt { .. }
| AIAgentInput::ActionResult { .. }
| AIAgentInput::MessagesReceivedFromAgents { .. }
| AIAgentInput::EventsFromAgents { .. }
| AIAgentInput::PassiveSuggestionResult { .. }
| AIAgentInput::OrchestrationConfigUpdate { .. } => None,
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum RigRequestMode {
Normal,
Plan,
Orchestrate,
Cli,
}
fn request_mode(inputs: &[AIAgentInput]) -> RigRequestMode {
for input in inputs {
if matches!(
input,
AIAgentInput::UserQuery {
running_command: Some(_),
..
}
) {
return RigRequestMode::Cli;
}
if let AIAgentInput::UserQuery {
user_query_mode, ..
} = input
{
match user_query_mode {
UserQueryMode::Normal => {}
UserQueryMode::Plan => return RigRequestMode::Plan,
UserQueryMode::Orchestrate => return RigRequestMode::Orchestrate,
}
}
}
RigRequestMode::Normal
}
fn tool_definitions(
supported_tools: &[ToolType],
mcp_context: Option<&MCPContext>,
) -> Vec<ToolDefinition> {
let supported = supported_tools.iter().copied().collect::<HashSet<_>>();
let mut tools = default_tool_definitions()
.into_iter()
.filter(|tool| tool_name_is_supported(&tool.name, &supported))
.collect::<Vec<_>>();
if !supported.contains(&ToolType::CallMcpTool) {
return tools;
}
let Some(mcp_context) = mcp_context else {
return tools;
};
let mut seen = tools
.iter()
.map(|tool| tool.name.clone())
.collect::<HashSet<_>>();
for server in &mcp_context.servers {
for tool in &server.tools {
let name = format!("mcp__{}__{}", server.name, tool.name);
if seen.insert(name.clone()) {
tools.push(ToolDefinition {
name,
description: tool
.description
.as_deref()
.map(str::to_string)
.unwrap_or_else(|| format!("MCP tool from {} server", server.name)),
input_schema: serde_json::Value::Object(tool.input_schema.as_ref().clone()),
});
}
}
}
#[allow(deprecated)]
for tool in &mcp_context.tools {
let name = format!("mcp__{}", tool.name);
if seen.insert(name.clone()) {
tools.push(ToolDefinition {
name,
description: tool
.description
.as_deref()
.map(str::to_string)
.unwrap_or_else(|| "MCP tool".to_string()),
input_schema: serde_json::Value::Object(tool.input_schema.as_ref().clone()),
});
}
}
tools
}
fn build_system_prompt(
inputs: &[AIAgentInput],
tools: &[ToolDefinition],
global_rules: &[(String, String)],
mode: RigRequestMode,
) -> String {
let mut prompt = String::from(
"You are Galaxy, a local-first software-engineering and terminal agent. Complete the user's task through inspection, implementation, and proportionate validation. Galaxy owns tool permissions and execution; use only the tools advertised in this request and treat every result as authoritative evidence.\n\n",
);
let contexts = inputs.iter().filter_map(AIAgentInput::context).flatten();
let mut environment = Vec::new();
let mut project_rules = Vec::new();
let mut available_skills = Vec::new();
let mut attached_context = Vec::new();
for context in contexts {
match context {
AIAgentContext::Directory {
pwd,
home_dir,
are_file_symbols_indexed,
} => {
if let Some(pwd) = pwd {
environment.push(format!("Working directory: {pwd}"));
}
if let Some(home_dir) = home_dir {
environment.push(format!("Home directory: {home_dir}"));
}
environment.push(format!(
"Working-directory file symbols indexed: {are_file_symbols_indexed}"
));
}
AIAgentContext::ExecutionEnvironment(execution) => {
let shell_version = execution
.shell_version
.as_deref()
.map(|version| format!(" {version}"))
.unwrap_or_default();
environment.push(format!("Shell: {}{shell_version}", execution.shell_name));
if let Some(os) = &execution.os.category {
environment.push(format!("OS: {os}"));
}
if let Some(distribution) = &execution.os.distribution {
environment.push(format!("OS distribution: {distribution}"));
}
}
AIAgentContext::ProjectRules {
root_path,
active_rules,
additional_rule_paths,
} => {
for rule in active_rules {
if let AnyFileContent::StringContent(content) = &rule.content {
project_rules.push((root_path.clone(), content.clone()));
}
}
if !additional_rule_paths.is_empty() {
environment.push(format!(
"Additional project rule paths: {}",
additional_rule_paths.join(", ")
));
}
}
AIAgentContext::Git { head, branch } => {
environment.push(format!("Git HEAD: {head}"));
if let Some(branch) = branch {
environment.push(format!("Git branch: {branch}"));
}
}
AIAgentContext::Skills { skills } => {
for skill in skills {
let (reference_type, reference) = match &skill.reference {
SkillReference::Path(path) => ("path", path.display_path()),
SkillReference::BundledSkillId(id) => ("bundled", id.clone()),
};
available_skills.push(format!(
"- name={:?}; reference_type={reference_type:?}; skill={reference:?}; description={:?}",
skill.name, skill.description
));
}
}
AIAgentContext::SelectedText(text) => {
attached_context.push(("Selected text".to_string(), text.clone()));
}
AIAgentContext::CurrentTime { current_time } => {
environment.push(format!("Current time: {current_time}"));
}
AIAgentContext::Codebase { path, name } => {
environment.push(format!("Indexed codebase: {name} ({path})"));
}
AIAgentContext::File(file) => match &file.content {
AnyFileContent::StringContent(content) => {
attached_context.push((format!("Attached file: {file}"), content.clone()));
}
AnyFileContent::BinaryContent(_) => {
environment.push(format!("Attached binary file (content omitted): {file}"));
}
},
AIAgentContext::Repository { name, owner } => {
let owner = owner
.as_deref()
.map(|owner| format!("{owner}/"))
.unwrap_or_default();
environment.push(format!("Repository: {owner}{name}"));
}
AIAgentContext::PullRequest {
number,
state,
draft,
base_branch,
} => {
environment.push(format!(
"Pull request: #{number}; state={state}; draft={draft}; base={base_branch}"
));
}
AIAgentContext::Block(block) => {
let details = format!(
"Command: {}\nExit code: {}\nOutput:\n{}",
block.command, block.exit_code, block.output
);
attached_context.push((format!("Terminal block {}", block.id), details));
}
AIAgentContext::Image(_) => {}
}
}
if !environment.is_empty() {
prompt.push_str("## Environment\n");
for item in environment {
prompt.push_str("- ");
prompt.push_str(&item);
prompt.push('\n');
}
prompt.push('\n');
}
if !project_rules.is_empty() {
prompt.push_str("## Project Rules\n");
for (root, content) in project_rules {
prompt.push_str(&format!("### Rules from {root}\n{content}\n"));
}
prompt.push('\n');
}
if !attached_context.is_empty() {
prompt.push_str("## Attached Context\n");
for (label, content) in attached_context {
prompt.push_str(&format!(
"<context label={label:?}>\n{content}\n</context>\n"
));
}
prompt.push('\n');
}
if !available_skills.is_empty() && tools.iter().any(|tool| tool.name == "read_skill") {
prompt.push_str("## Available Skills\n");
prompt.push_str(&available_skills.join("\n"));
prompt.push_str("\n\n");
}
if !global_rules.is_empty() {
prompt.push_str("## Global Rules\n");
for (name, content) in global_rules {
if !name.is_empty() {
prompt.push_str(&format!("### {name}\n"));
}
prompt.push_str(content);
prompt.push_str("\n\n");
}
}
match mode {
RigRequestMode::Normal => {}
RigRequestMode::Plan => prompt.push_str(
"## Plan Mode\nInspect and produce an implementation-ready plan. Do not edit files or perform state-changing actions.\n\n",
),
RigRequestMode::Orchestrate => prompt.push_str(
"## Orchestration Mode\nDelegate only independent, bounded work where parallelism materially helps, then synthesize the results.\n\n",
),
RigRequestMode::Cli => prompt.push_str(
"## Running Command Monitor\nMonitor the existing command by its command ID. Never start a duplicate command. Poll briefly, respect stop conditions, and report only verified outcomes.\n\n",
),
}
prompt.push_str("## Available Tools\n");
if tools.is_empty() {
prompt.push_str("No tools are available. Do not invent tool calls.\n");
} else {
prompt.push_str("Use only these tools: ");
prompt.push_str(
&tools
.iter()
.map(|tool| tool.name.as_str())
.collect::<Vec<_>>()
.join(", "),
);
prompt.push_str(".\nNever invent tool names or parameters. Check command exit codes and tool error results before claiming success.\n");
}
prompt
}
#[cfg(test)]
#[path = "rig_request_tests.rs"]
mod tests;
+170
View File
@@ -0,0 +1,170 @@
use std::collections::HashMap;
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 crate::ai::agent::api::RequestParams;
use crate::ai::agent::{AIAgentContext, AIAgentInput, AnyFileContent, FileContext, UserQueryMode};
use crate::ai::llms::LLMId;
use crate::ai::openai::client::OpenAIClientConfig;
fn config() -> OpenAIClientConfig {
OpenAIClientConfig {
base_url: "http://localhost:4000/v1".to_string(),
api_key: None,
model: Some("provider-model".to_string()),
max_input_tokens: Some(128_000),
max_output_tokens: Some(8_192),
use_rig: true,
supports_system_messages: true,
}
}
fn user_query(query: &str) -> AIAgentInput {
user_query_with_context(query, Vec::new())
}
fn user_query_with_context(query: &str, context: Vec<AIAgentContext>) -> AIAgentInput {
AIAgentInput::UserQuery {
query: query.to_string(),
context: Arc::from(context),
static_query_type: None,
referenced_attachments: HashMap::new(),
user_query_mode: UserQueryMode::Normal,
running_command: None,
intended_agent: None,
}
}
#[test]
fn native_context_reaches_rig_without_a_proto_context_conversion() {
let mut params = RequestParams::new_for_test();
params.input = vec![user_query_with_context(
"Explain the selected implementation",
vec![
AIAgentContext::SelectedText("prepare_rig_turn(params)".to_string()),
AIAgentContext::File(FileContext::new(
"/repo/src/runtime.rs".to_string(),
AnyFileContent::StringContent("fn prepare_rig_turn() {}".to_string()),
None,
None,
)),
AIAgentContext::Codebase {
path: "/repo".to_string(),
name: "galaxy".to_string(),
},
],
)];
let prepared = prepare_rig_turn(&config(), params, Vec::new(), Vec::new());
let prompt = prepared.request.system_prompt.expect("system prompt");
assert!(prompt.contains("prepare_rig_turn(params)"));
assert!(prompt.contains("fn prepare_rig_turn() {}"));
assert!(prompt.contains("Indexed codebase: galaxy (/repo)"));
}
#[test]
fn builds_a_rig_turn_directly_from_galaxy_request_state() {
let mut params = RequestParams::new_for_test();
params.model = LLMId::from("selected-model");
params.root_task_id = Some("task-1".to_string());
params.input = vec![user_query("Inspect this repository")];
let prepared = prepare_rig_turn(
&config(),
params,
vec![ToolType::ReadFiles, ToolType::RunShellCommand],
Vec::new(),
);
assert_eq!(prepared.task_id, "task-1");
assert_eq!(
prepared.user_query.as_deref(),
Some("Inspect this repository")
);
assert_eq!(prepared.request.model.as_str(), "provider-model");
assert_eq!(prepared.request.max_output_tokens, Some(8_192));
assert_eq!(prepared.request.messages, prepared.persistent_messages);
assert!(prepared
.request
.tools
.iter()
.any(|tool| tool.name == "read_files"));
assert!(prepared
.request
.tools
.iter()
.any(|tool| tool.name == "run_shell_command"));
assert!(prepared
.request
.system_prompt
.as_deref()
.is_some_and(|prompt| prompt.contains("Galaxy owns tool permissions and execution")));
assert!(matches!(
&prepared.request.messages[0],
galaxy_agent_core::ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(text),
} if text == "Inspect this repository"
));
}
#[test]
fn normalized_tool_outcomes_are_the_only_action_results_sent_to_rig() {
let statuses = [
("read", ToolResultStatus::Success, false),
("shell", ToolResultStatus::Error, true),
("denied", ToolResultStatus::Denied, true),
("cancelled", ToolResultStatus::Cancelled, false),
];
let tool_results = statuses
.iter()
.map(|(call_id, status, _)| ToolResult {
call_id: (*call_id).to_string(),
content: format!("normalized-{call_id}"),
status: *status,
})
.collect();
let messages = input_messages(Vec::new(), tool_results);
assert_eq!(messages.len(), 1);
let MessageContent::MultiPart(parts) = &messages[0].content else {
panic!("expected normalized tool results to remain in one user turn");
};
for ((call_id, _, expected_error), part) in statuses.iter().zip(parts) {
assert!(matches!(
part,
ContentPart::ToolResult {
tool_use_id,
content,
is_error,
} if tool_use_id == call_id
&& content == &format!("normalized-{call_id}")
&& is_error == expected_error
));
}
}
#[test]
fn progressive_summary_is_provider_context_not_persistent_history() {
let mut params = RequestParams::new_for_test();
params.input = vec![user_query("Continue")];
params.progressive_summary = Some("Earlier work was validated.".to_string());
let prepared = prepare_rig_turn(&config(), params, Vec::new(), Vec::new());
assert_eq!(prepared.persistent_messages.len(), 1);
assert_eq!(prepared.request.messages.len(), 3);
assert!(matches!(
&prepared.request.messages[0].content,
MessageContent::Text(text) if text.contains("Earlier work was validated.")
));
assert!(matches!(
&prepared.request.messages[2].content,
MessageContent::Text(text) if text == "Continue"
));
}
+175 -2
View File
@@ -1,7 +1,14 @@
use galaxy_agent_core::StopReason; use std::sync::{Arc, Mutex};
use galaxy_agent_core::{
MessageContent, MessageRole, StopReason, ToolCall, ToolResult, ToolResultStatus,
};
use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::response_event::stream_finished;
use super::{build_add_reasoning, build_append_reasoning, map_stop_reason, saturating_i32}; use super::{
append_tool_result, build_add_reasoning, build_append_reasoning, build_tool_proposed,
map_stop_reason, saturating_i32, sync_assistant_turn,
};
#[test] #[test]
fn stop_reasons_map_to_the_existing_ui_contract() { fn stop_reasons_map_to_the_existing_ui_contract() {
@@ -57,3 +64,169 @@ fn reasoning_events_match_the_existing_ui_message_contract() {
["agent_reasoning.reasoning"] ["agent_reasoning.reasoning"]
); );
} }
#[test]
fn tool_proposal_matches_the_existing_permission_ui_contract() {
let event = build_tool_proposed(
"task",
&ToolCall {
id: "call-1".to_string(),
name: "run_shell_command".to_string(),
arguments: serde_json::json!({
"command": "cargo test",
"is_read_only": true
}),
},
);
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);
}
#[test]
fn mcp_tool_proposal_routes_through_the_existing_mcp_executor_contract() {
let event = build_tool_proposed(
"task",
&ToolCall {
id: "call-mcp".to_string(),
name: "mcp__filesystem__read_file".to_string(),
arguments: serde_json::json!({"path": "Cargo.toml"}),
},
);
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());
}
#[test]
fn assistant_history_is_updated_before_fast_tool_execution_can_continue() {
let messages = Arc::new(Mutex::new(Vec::new()));
let mut history_index = None;
let first_call = ToolCall {
id: "call-1".to_string(),
name: "read_files".to_string(),
arguments: serde_json::json!({"files": ["Cargo.toml"]}),
};
let second_call = ToolCall {
id: "call-2".to_string(),
name: "grep".to_string(),
arguments: serde_json::json!({"queries": ["rig"]}),
};
sync_assistant_turn(
&messages,
"I'll inspect both.",
std::slice::from_ref(&first_call),
&mut history_index,
);
sync_assistant_turn(
&messages,
"I'll inspect both.",
&[first_call, second_call],
&mut history_index,
);
let messages = messages.lock().unwrap();
assert_eq!(messages.len(), 1);
let MessageContent::MultiPart(parts) = &messages[0].content else {
panic!("expected combined assistant content");
};
assert_eq!(parts.len(), 3);
assert!(
matches!(&parts[0], galaxy_agent_core::ContentPart::Text(text) if text == "I'll inspect both.")
);
assert!(
matches!(&parts[1], galaxy_agent_core::ContentPart::ToolUse { tool_use_id, .. } if tool_use_id == "call-1")
);
assert!(
matches!(&parts[2], galaxy_agent_core::ContentPart::ToolUse { tool_use_id, .. } if tool_use_id == "call-2")
);
}
#[test]
fn rejected_tool_result_is_paired_with_the_assistant_call_in_history() {
let messages = Arc::new(Mutex::new(Vec::new()));
let mut history_index = None;
let call = ToolCall {
id: "call-unknown".to_string(),
name: "invented_tool".to_string(),
arguments: serde_json::json!({}),
};
sync_assistant_turn(
&messages,
"",
std::slice::from_ref(&call),
&mut history_index,
);
append_tool_result(
&messages,
ToolResult {
call_id: call.id.clone(),
content: "tool is unavailable".to_string(),
status: ToolResultStatus::Error,
},
);
let messages = messages.lock().unwrap();
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].role, MessageRole::Assistant);
assert!(matches!(
&messages[0].content,
MessageContent::ToolUse {
tool_use_id,
name,
..
} if tool_use_id == "call-unknown" && name == "invented_tool"
));
assert_eq!(messages[1].role, MessageRole::User);
assert!(matches!(
&messages[1].content,
MessageContent::ToolResult {
tool_use_id,
content,
is_error: true,
} if tool_use_id == "call-unknown" && content == "tool is unavailable"
));
}
@@ -25,8 +25,13 @@ use super::hydrate_ai_conversation_assertion;
/// Assumes that the terminal input is currently not in AI input mode. /// Assumes that the terminal input is currently not in AI input mode.
pub fn enter_agent_view() -> TestStep { pub fn enter_agent_view() -> TestStep {
let keystroke = if cfg!(target_os = "macos") {
"cmd-enter"
} else {
"ctrl-shift-enter"
};
new_step_with_default_assertions("Enter Agent View") new_step_with_default_assertions("Enter Agent View")
.with_keystrokes(&["ctrl-shift-enter"]) .with_keystrokes(&[keystroke])
.add_named_assertion( .add_named_assertion(
"Assert that we are in Agent View and AI input mode", "Assert that we are in Agent View and AI input mode",
move |app, window_id| { move |app, window_id| {
+5 -5
View File
@@ -356,9 +356,9 @@ fn initial_litellm_provider_maps_codex_model_to_rig_without_a_committed_key() {
assert_eq!(provider.models.len(), 1); assert_eq!(provider.models.len(), 1);
let model = &provider.models[0]; let model = &provider.models[0];
assert_eq!(model.model_id, INITIAL_RIG_MODEL_ID); assert_eq!(model.model_id, INITIAL_RIG_MODEL_ID);
assert_eq!(model.use_rig, true); assert!(model.use_rig);
assert_eq!(model.supports_system_messages, Some(false)); assert_eq!(model.supports_system_messages, Some(false));
assert_eq!(model.supports_system_messages(), false); assert!(!model.supports_system_messages());
} }
#[test] #[test]
@@ -366,14 +366,14 @@ fn codex_litellm_model_infers_missing_system_message_capability() {
let mut model = default_openai_providers().remove(0).models.remove(0); let mut model = default_openai_providers().remove(0).models.remove(0);
model.supports_system_messages = None; model.supports_system_messages = None;
assert_eq!(model.supports_system_messages(), false); assert!(!model.supports_system_messages());
model.model_id = "gpt-4o".to_string(); model.model_id = "gpt-4o".to_string();
assert_eq!(model.supports_system_messages(), true); assert!(model.supports_system_messages());
model.model_id = INITIAL_RIG_MODEL_ID.to_string(); model.model_id = INITIAL_RIG_MODEL_ID.to_string();
model.supports_system_messages = Some(true); model.supports_system_messages = Some(true);
assert_eq!(model.supports_system_messages(), true); assert!(model.supports_system_messages());
} }
#[test] #[test]
+2 -1
View File
@@ -7324,7 +7324,8 @@ impl TerminalView {
); );
} }
} }
BlocklistAIActionEvent::QueuedAction(_) => {} BlocklistAIActionEvent::QueuedAction(_)
| BlocklistAIActionEvent::ToolLifecycle { .. } => {}
} }
} }
+210
View File
@@ -135,6 +135,216 @@ impl AIAgentActionResultType {
_ => None, _ => None,
} }
} }
/// Returns the authoritative result content to send back to a model.
///
/// `Display` is intentionally concise for UI summaries, so content-bearing
/// results must not use it directly when constructing the next model turn.
pub fn model_content(&self) -> String {
match self {
Self::RequestCommandOutput(result) => match result {
RequestCommandOutputResult::Completed {
command,
output,
exit_code,
..
} => command_result_content(Some(command), output, exit_code.value()),
RequestCommandOutputResult::LongRunningCommandSnapshot {
command,
grid_contents,
cursor,
is_alt_screen_active,
..
} => shell_snapshot_content(
Some(command),
grid_contents,
cursor,
*is_alt_screen_active,
None,
),
RequestCommandOutputResult::CancelledBeforeExecution
| RequestCommandOutputResult::Denylisted { .. } => result.to_string(),
},
Self::WriteToLongRunningShellCommand(result) => match result {
WriteToLongRunningShellCommandResult::Snapshot {
grid_contents,
cursor,
is_alt_screen_active,
is_preempted,
..
} => shell_snapshot_content(
None,
grid_contents,
cursor,
*is_alt_screen_active,
Some(*is_preempted),
),
WriteToLongRunningShellCommandResult::CommandFinished {
output, exit_code, ..
} => command_result_content(None, output, exit_code.value()),
WriteToLongRunningShellCommandResult::Cancelled
| WriteToLongRunningShellCommandResult::Error(_) => result.to_string(),
},
Self::ReadFiles(result) => match result {
ReadFilesResult::Success { files } => file_contexts_content(files),
ReadFilesResult::Error(_) | ReadFilesResult::Cancelled => result.to_string(),
},
Self::SearchCodebase(result) => match result {
SearchCodebaseResult::Success { files } => file_contexts_content(files),
SearchCodebaseResult::Failed { .. } | SearchCodebaseResult::Cancelled => {
result.to_string()
}
},
Self::ReadSkill(result) => match result {
ReadSkillResult::Success { content } => file_context_content(content),
ReadSkillResult::Error(_) | ReadSkillResult::Cancelled => result.to_string(),
},
Self::ReadDocuments(result) => match result {
ReadDocumentsResult::Success { documents } => document_contexts_content(documents),
ReadDocumentsResult::Error(_) | ReadDocumentsResult::Cancelled => {
result.to_string()
}
},
Self::EditDocuments(result) => match result {
EditDocumentsResult::Success { updated_documents } => {
document_contexts_content(updated_documents)
}
EditDocumentsResult::Error(_) | EditDocumentsResult::Cancelled => {
result.to_string()
}
},
Self::CreateDocuments(result) => match result {
CreateDocumentsResult::Success { created_documents } => {
document_contexts_content(created_documents)
}
CreateDocumentsResult::Error(_) | CreateDocumentsResult::Cancelled => {
result.to_string()
}
},
Self::ReadShellCommandOutput(result) => match result {
ReadShellCommandOutputResult::CommandFinished {
command,
output,
exit_code,
..
} => command_result_content(Some(command), output, exit_code.value()),
ReadShellCommandOutputResult::LongRunningCommandSnapshot {
command,
grid_contents,
cursor,
is_alt_screen_active,
is_preempted,
..
} => shell_snapshot_content(
Some(command),
grid_contents,
cursor,
*is_alt_screen_active,
Some(*is_preempted),
),
ReadShellCommandOutputResult::Cancelled
| ReadShellCommandOutputResult::Error(_) => result.to_string(),
},
Self::TransferShellCommandControlToUser(result) => match result {
TransferShellCommandControlToUserResult::Snapshot {
grid_contents,
cursor,
is_alt_screen_active,
is_preempted,
..
} => format!(
"{}\nControl has been transferred to the user. Do not write to the command until control is returned.",
shell_snapshot_content(
None,
grid_contents,
cursor,
*is_alt_screen_active,
Some(*is_preempted),
)
),
TransferShellCommandControlToUserResult::CommandFinished {
output, exit_code, ..
} => command_result_content(None, output, exit_code.value()),
TransferShellCommandControlToUserResult::Cancelled
| TransferShellCommandControlToUserResult::Error(_) => result.to_string(),
},
Self::RequestFileEdits(_)
| Self::UploadArtifact(_)
| Self::Grep(_)
| Self::FileGlob(_)
| Self::FileGlobV2(_)
| Self::ReadMCPResource(_)
| Self::CallMCPTool(_)
| Self::SuggestNewConversation(_)
| Self::SuggestPrompt(_)
| Self::OpenCodeReview
| Self::InitProject
| Self::UseComputer(_)
| Self::InsertReviewComments(_)
| Self::RequestComputerUse(_)
| Self::FetchConversation(_)
| Self::StartAgent(_)
| Self::SendMessageToAgent(_)
| Self::AskUserQuestion(_)
| Self::RunAgents(_)
| Self::WaitForEvents(_) => self.to_string(),
}
}
}
fn command_result_content(command: Option<&str>, output: &str, exit_code: i32) -> String {
let command = command
.map(|command| format!("Command: {command}\n"))
.unwrap_or_default();
let output = if output.is_empty() {
"(no output)"
} else {
output
};
format!("{command}Command finished with exit code {exit_code}.\nOutput:\n{output}")
}
fn shell_snapshot_content(
command: Option<&str>,
grid_contents: &str,
cursor: &str,
is_alt_screen_active: bool,
is_preempted: Option<bool>,
) -> String {
let command = command
.map(|command| format!("Command: {command}\n"))
.unwrap_or_default();
let preempted = is_preempted
.map(|is_preempted| format!("\nPreempted: {is_preempted}"))
.unwrap_or_default();
format!(
"{command}Command is still running.\nCurrent output:\n{grid_contents}\nCursor: {cursor}\nAlt screen active: {is_alt_screen_active}{preempted}"
)
}
fn file_contexts_content(files: &[FileContext]) -> String {
files
.iter()
.map(file_context_content)
.collect::<Vec<_>>()
.join("\n\n")
}
fn file_context_content(file: &FileContext) -> String {
match &file.content {
AnyFileContent::StringContent(content) => format!("{file}:\n{content}"),
AnyFileContent::BinaryContent(content) => {
format!("{file}:\n[binary file, {} bytes]", content.len())
}
}
}
fn document_contexts_content(documents: &[DocumentContext]) -> String {
documents
.iter()
.map(|document| format!("{document}:\n{}", document.content))
.collect::<Vec<_>>()
.join("\n\n")
} }
#[cfg(test)] #[cfg(test)]
+2
View File
@@ -5,7 +5,9 @@
//! depend on GalaxyUI, provider SDKs, persistence, or Warp wire protocols. //! depend on GalaxyUI, provider SDKs, persistence, or Warp wire protocols.
mod runtime; mod runtime;
mod tool_policy;
mod types; mod types;
pub use runtime::*; pub use runtime::*;
pub use tool_policy::*;
pub use types::*; pub use types::*;
+294
View File
@@ -0,0 +1,294 @@
use std::collections::{BTreeSet, HashMap, VecDeque};
use serde_json::Value as JsonValue;
use crate::{
ContentPart, ConversationMessage, MessageContent, ToolCall, ToolDefinition, ToolResult,
ToolResultStatus,
};
pub const RECALL_TOOL_HISTORY_NAME: &str = "recall_tool_history";
const MAX_RECALLED_RESULT_CHARS: usize = 50_000;
const DEFAULT_LOOP_WINDOW: usize = 10;
const DEFAULT_LOOP_THRESHOLD: usize = 3;
#[derive(Clone, Debug, PartialEq)]
pub enum ToolCallDecision {
Execute,
Inline(ToolResult),
Reject(ToolResult),
}
#[derive(Clone, Debug, Default)]
pub struct ToolPolicy {
advertised_tools: BTreeSet<String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct ToolFailureRecord {
signature: u64,
description: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ToolLoopDetected {
pub description: String,
pub threshold: usize,
}
#[derive(Clone, Debug)]
pub struct ToolLoopGuard {
recent_failures: VecDeque<ToolFailureRecord>,
window: usize,
threshold: usize,
}
impl Default for ToolLoopGuard {
fn default() -> Self {
Self::new(DEFAULT_LOOP_WINDOW, DEFAULT_LOOP_THRESHOLD)
}
}
impl ToolLoopGuard {
pub fn new(window: usize, threshold: usize) -> Self {
Self {
recent_failures: VecDeque::new(),
window: window.max(1),
threshold: threshold.max(1),
}
}
pub fn record_failure(&mut self, signature: u64, description: impl Into<String>) {
self.recent_failures.push_back(ToolFailureRecord {
signature,
description: description.into(),
});
if self.recent_failures.len() > self.window {
self.recent_failures.pop_front();
}
}
pub fn record_success(&mut self) {
self.recent_failures.clear();
}
pub fn detect_and_reset(&mut self) -> Option<ToolLoopDetected> {
let mut counts = HashMap::new();
for (index, failure) in self.recent_failures.iter().enumerate() {
let count = counts.entry(failure.signature).or_insert((0usize, 0usize));
count.0 += 1;
count.1 = index;
}
let latest_index = counts
.values()
.filter(|(count, _)| *count >= self.threshold)
.map(|(_, index)| *index)
.max()?;
let detected = ToolLoopDetected {
description: self.recent_failures[latest_index].description.clone(),
threshold: self.threshold,
};
self.recent_failures.clear();
Some(detected)
}
}
impl ToolPolicy {
pub fn new(tools: &[ToolDefinition]) -> Self {
Self {
advertised_tools: tools.iter().map(|tool| tool.name.clone()).collect(),
}
}
pub fn decide(
&self,
call: &ToolCall,
messages: &[ConversationMessage],
archive: &[ConversationMessage],
) -> ToolCallDecision {
if !self.advertised_tools.contains(&call.name) {
let available = if self.advertised_tools.is_empty() {
"no tools are available".to_string()
} else {
format!(
"available tools are: {}",
self.advertised_tools
.iter()
.cloned()
.collect::<Vec<_>>()
.join(", ")
)
};
return ToolCallDecision::Reject(ToolResult {
call_id: call.id.clone(),
content: format!(
"Error: '{}' is not a valid tool for this request; {available}. Do not invent tool names.",
call.name
),
status: ToolResultStatus::Error,
});
}
if call.name == RECALL_TOOL_HISTORY_NAME {
return ToolCallDecision::Inline(ToolResult {
call_id: call.id.clone(),
content: recall_tool_history(
messages,
archive,
ToolHistoryQuery::from_arguments(&call.arguments),
),
status: ToolResultStatus::Success,
});
}
ToolCallDecision::Execute
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ToolHistoryQuery<'a> {
pub search_query: &'a str,
pub tool_name: &'a str,
pub tool_use_id: &'a str,
pub offset_from_end: usize,
}
impl<'a> ToolHistoryQuery<'a> {
fn from_arguments(arguments: &'a JsonValue) -> Self {
Self {
search_query: arguments
.get("search_query")
.and_then(JsonValue::as_str)
.unwrap_or_default(),
tool_name: arguments
.get("tool_name")
.and_then(JsonValue::as_str)
.unwrap_or_default(),
tool_use_id: arguments
.get("tool_use_id")
.and_then(JsonValue::as_str)
.unwrap_or_default(),
offset_from_end: arguments
.get("offset_from_end")
.and_then(JsonValue::as_u64)
.and_then(|offset| usize::try_from(offset).ok())
.unwrap_or_default(),
}
}
}
pub fn recall_tool_history(
messages: &[ConversationMessage],
archive: &[ConversationMessage],
query: ToolHistoryQuery<'_>,
) -> String {
let mut entries = Vec::new();
collect_tool_entries(archive, &mut entries);
collect_tool_entries(messages, &mut entries);
let filtered = entries
.iter()
.filter(|entry| {
(query.tool_use_id.is_empty() || entry.tool_use_id == query.tool_use_id)
&& (query.tool_name.is_empty() || entry.name == query.tool_name)
&& (query.search_query.is_empty()
|| format!("{} {} {}", entry.name, entry.input, entry.result)
.to_lowercase()
.contains(&query.search_query.to_lowercase()))
})
.collect::<Vec<_>>();
let Some(index) = filtered
.len()
.checked_sub(1usize.saturating_add(query.offset_from_end))
else {
return "No matching tool calls found in conversation history.".to_string();
};
let entry = filtered[index];
let result = truncate_recalled_result(&entry.result);
format!(
"Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{result}",
entry.name, entry.tool_use_id, entry.input
)
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct ToolHistoryEntry {
tool_use_id: String,
name: String,
input: String,
result: String,
}
fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec<ToolHistoryEntry>) {
let mut pending = Vec::new();
for message in messages {
match &message.content {
MessageContent::ToolUse {
tool_use_id,
name,
input,
} => pending.push((tool_use_id.clone(), name.clone(), input.to_string())),
MessageContent::ToolResult {
tool_use_id,
content,
..
} => pair_result(tool_use_id, content, &mut pending, entries),
MessageContent::MultiPart(parts) => {
for part in parts {
match part {
ContentPart::ToolUse {
tool_use_id,
name,
input,
} => {
pending.push((tool_use_id.clone(), name.clone(), input.to_string()));
}
ContentPart::ToolResult {
tool_use_id,
content,
..
} => pair_result(tool_use_id, content, &mut pending, entries),
ContentPart::Text(_) | ContentPart::Image { .. } => {}
}
}
}
MessageContent::Text(_) => {}
}
}
}
fn pair_result(
tool_use_id: &str,
content: &str,
pending: &mut Vec<(String, String, String)>,
entries: &mut Vec<ToolHistoryEntry>,
) {
let Some(index) = pending.iter().position(|(id, _, _)| id == tool_use_id) else {
return;
};
let (tool_use_id, name, input) = pending.remove(index);
entries.push(ToolHistoryEntry {
tool_use_id,
name,
input,
result: content.to_string(),
});
}
fn truncate_recalled_result(result: &str) -> String {
let char_count = result.chars().count();
if char_count <= MAX_RECALLED_RESULT_CHARS {
return result.to_string();
}
let truncated = result
.chars()
.take(MAX_RECALLED_RESULT_CHARS)
.collect::<String>();
format!("{truncated}... [truncated, {char_count} total chars]")
}
#[cfg(test)]
#[path = "tool_policy_tests.rs"]
mod tests;
@@ -0,0 +1,171 @@
use super::*;
use crate::MessageRole;
fn definition(name: &str) -> ToolDefinition {
ToolDefinition {
name: name.to_string(),
description: String::new(),
input_schema: serde_json::json!({"type": "object"}),
}
}
fn tool_exchange(id: &str, name: &str, input: JsonValue, result: &str) -> Vec<ConversationMessage> {
vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: id.to_string(),
name: name.to_string(),
input,
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: id.to_string(),
content: result.to_string(),
is_error: false,
},
},
]
}
#[test]
fn executes_only_tools_advertised_for_this_turn() {
let policy = ToolPolicy::new(&[definition("read_files")]);
let read = ToolCall {
id: "read-1".to_string(),
name: "read_files".to_string(),
arguments: serde_json::json!({"files": ["Cargo.toml"]}),
};
let shell = ToolCall {
id: "shell-1".to_string(),
name: "run_shell_command".to_string(),
arguments: serde_json::json!({"command": "pwd"}),
};
assert_eq!(policy.decide(&read, &[], &[]), ToolCallDecision::Execute);
let ToolCallDecision::Reject(result) = policy.decide(&shell, &[], &[]) else {
panic!("unadvertised tool should be rejected");
};
assert_eq!(result.call_id, "shell-1");
assert_eq!(result.status, ToolResultStatus::Error);
assert!(result.content.contains("read_files"));
assert!(
!result
.content
.contains("available tools are: run_shell_command")
);
}
#[test]
fn recall_searches_archived_and_live_results_with_live_results_most_recent() {
let policy = ToolPolicy::new(&[definition(RECALL_TOOL_HISTORY_NAME)]);
let archive = tool_exchange(
"archived-read",
"read_files",
serde_json::json!({"files": ["old.txt"]}),
"old contents",
);
let messages = tool_exchange(
"live-read",
"read_files",
serde_json::json!({"files": ["new.txt"]}),
"new contents",
);
let latest = ToolCall {
id: "recall-latest".to_string(),
name: RECALL_TOOL_HISTORY_NAME.to_string(),
arguments: serde_json::json!({"tool_name": "read_files"}),
};
let previous = ToolCall {
id: "recall-previous".to_string(),
name: RECALL_TOOL_HISTORY_NAME.to_string(),
arguments: serde_json::json!({
"tool_name": "read_files",
"offset_from_end": 1,
}),
};
let ToolCallDecision::Inline(latest_result) = policy.decide(&latest, &messages, &archive)
else {
panic!("recall should execute inline");
};
assert!(latest_result.content.contains("Tool Use ID: live-read"));
assert!(latest_result.content.contains("new contents"));
let ToolCallDecision::Inline(previous_result) = policy.decide(&previous, &messages, &archive)
else {
panic!("recall should execute inline");
};
assert!(
previous_result
.content
.contains("Tool Use ID: archived-read")
);
assert!(previous_result.content.contains("old contents"));
}
#[test]
fn recall_supports_exact_call_id_and_case_insensitive_text_search() {
let messages = tool_exchange(
"shell-7",
"run_shell_command",
serde_json::json!({"command": "cargo test"}),
"ALL TESTS PASSED",
);
let exact = recall_tool_history(
&messages,
&[],
ToolHistoryQuery {
tool_use_id: "shell-7",
search_query: "all tests",
..Default::default()
},
);
let missing = recall_tool_history(
&messages,
&[],
ToolHistoryQuery {
tool_use_id: "missing",
..Default::default()
},
);
assert!(exact.contains("cargo test"));
assert!(exact.contains("ALL TESTS PASSED"));
assert_eq!(
missing,
"No matching tool calls found in conversation history."
);
}
#[test]
fn loop_guard_detects_repeated_failures_and_resets_after_detection() {
let mut guard = ToolLoopGuard::new(5, 3);
guard.record_failure(7, "cargo test failed");
guard.record_failure(11, "another command failed");
guard.record_failure(7, "cargo test failed again");
assert_eq!(guard.detect_and_reset(), None);
guard.record_failure(7, "cargo test failed a third time");
assert_eq!(
guard.detect_and_reset(),
Some(ToolLoopDetected {
description: "cargo test failed a third time".to_string(),
threshold: 3,
})
);
assert_eq!(guard.detect_and_reset(), None);
}
#[test]
fn loop_guard_clears_failures_when_a_tool_makes_progress() {
let mut guard = ToolLoopGuard::new(5, 2);
guard.record_failure(7, "first failure");
guard.record_success();
guard.record_failure(7, "failure after success");
assert_eq!(guard.detect_and_reset(), None);
}
+61 -6
View File
@@ -126,7 +126,24 @@ pub struct ToolCall {
pub struct ToolResult { pub struct ToolResult {
pub call_id: String, pub call_id: String,
pub content: String, pub content: String,
pub is_error: bool, pub status: ToolResultStatus,
}
impl ToolResult {
pub fn is_error(&self) -> bool {
matches!(
self.status,
ToolResultStatus::Error | ToolResultStatus::Denied
)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ToolResultStatus {
Success,
Error,
Denied,
Cancelled,
} }
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -141,11 +158,52 @@ pub enum PermissionKind {
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PermissionRequest { pub struct PermissionRequest {
pub id: String, pub id: String,
pub tool_call: ToolCall, pub call_id: String,
pub kind: PermissionKind, pub kind: PermissionKind,
pub reason: Option<String>, pub reason: Option<String>,
} }
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum PermissionDecision {
AllowOnce,
AlwaysAllow,
Denied { reason: Option<String> },
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum ToolEvent {
Proposed {
call: ToolCall,
},
PermissionRequested {
request: PermissionRequest,
},
PermissionResolved {
request_id: String,
call_id: String,
decision: PermissionDecision,
},
Started {
call_id: String,
},
Completed {
result: ToolResult,
},
}
impl ToolEvent {
pub fn call_id(&self) -> &str {
match self {
ToolEvent::Proposed { call } => &call.id,
ToolEvent::PermissionRequested { request } => &request.call_id,
ToolEvent::PermissionResolved { call_id, .. } | ToolEvent::Started { call_id } => {
call_id
}
ToolEvent::Completed { result } => &result.call_id,
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage { pub struct Usage {
pub input_tokens: u64, pub input_tokens: u64,
@@ -176,10 +234,7 @@ pub enum AgentEvent {
TurnStarted { runtime_request_id: String }, TurnStarted { runtime_request_id: String },
TextDelta { text: String }, TextDelta { text: String },
ReasoningDelta { text: String }, ReasoningDelta { text: String },
ToolProposed { call: ToolCall }, Tool { event: ToolEvent },
PermissionRequested { request: PermissionRequest },
ToolStarted { call: ToolCall },
ToolCompleted { result: ToolResult },
UsageUpdated { usage: Usage }, UsageUpdated { usage: Usage },
TurnStopped { reason: StopReason }, TurnStopped { reason: StopReason },
} }
@@ -36,3 +36,58 @@ fn usage_total_excludes_cached_breakdown_to_avoid_double_counting() {
assert_eq!(usage.total_tokens(), 125); assert_eq!(usage.total_tokens(), 125);
} }
#[test]
fn tool_events_keep_one_call_id_across_permission_and_execution() {
let events = [
ToolEvent::Proposed {
call: ToolCall {
id: "call-1".to_string(),
name: "run_shell_command".to_string(),
arguments: serde_json::json!({"command": "cargo test"}),
},
},
ToolEvent::PermissionRequested {
request: PermissionRequest {
id: "permission:call-1".to_string(),
call_id: "call-1".to_string(),
kind: PermissionKind::Execute,
reason: Some("Run a command".to_string()),
},
},
ToolEvent::PermissionResolved {
request_id: "permission:call-1".to_string(),
call_id: "call-1".to_string(),
decision: PermissionDecision::AllowOnce,
},
ToolEvent::Started {
call_id: "call-1".to_string(),
},
ToolEvent::Completed {
result: ToolResult {
call_id: "call-1".to_string(),
content: "ok".to_string(),
status: ToolResultStatus::Success,
},
},
];
assert!(events.iter().all(|event| event.call_id() == "call-1"));
}
#[test]
fn denied_results_are_errors_but_cancelled_results_are_distinct() {
let denied = ToolResult {
call_id: "denied".to_string(),
content: "permission denied".to_string(),
status: ToolResultStatus::Denied,
};
let cancelled = ToolResult {
call_id: "cancelled".to_string(),
content: "cancelled".to_string(),
status: ToolResultStatus::Cancelled,
};
assert!(denied.is_error());
assert!(!cancelled.is_error());
}
@@ -180,12 +180,14 @@ where
} }
} }
Ok(StreamedAssistantContent::ToolCall { tool_call, .. }) => { Ok(StreamedAssistantContent::ToolCall { tool_call, .. }) => {
yield Ok(AgentEvent::ToolProposed { yield Ok(AgentEvent::Tool {
call: ToolCall { event: galaxy_agent_core::ToolEvent::Proposed {
call: ToolCall {
id: tool_call.id, id: tool_call.id,
name: tool_call.function.name, name: tool_call.function.name,
arguments: tool_call.function.arguments, arguments: tool_call.function.arguments,
}, },
},
}); });
} }
Ok(StreamedAssistantContent::ToolCallDelta { .. }) => { Ok(StreamedAssistantContent::ToolCallDelta { .. }) => {
@@ -298,10 +300,10 @@ fn user_content(content: MessageContent) -> Result<OneOrMany<UserContent>, Agent
MessageContent::ToolResult { MessageContent::ToolResult {
tool_use_id, tool_use_id,
content, content,
.. is_error,
} => vec![UserContent::tool_result( } => vec![UserContent::tool_result(
tool_use_id, tool_use_id,
OneOrMany::one(ToolResultContent::text(content)), OneOrMany::one(ToolResultContent::text(tool_result_text(content, is_error))),
)], )],
MessageContent::MultiPart(parts) => parts MessageContent::MultiPart(parts) => parts
.into_iter() .into_iter()
@@ -344,10 +346,10 @@ fn convert_user_part(part: ContentPart) -> Result<UserContent, AgentError> {
ContentPart::ToolResult { ContentPart::ToolResult {
tool_use_id, tool_use_id,
content, content,
.. is_error,
} => Ok(UserContent::tool_result( } => Ok(UserContent::tool_result(
tool_use_id, tool_use_id,
OneOrMany::one(ToolResultContent::text(content)), OneOrMany::one(ToolResultContent::text(tool_result_text(content, is_error))),
)), )),
ContentPart::ToolUse { .. } => Err(invalid_role("tool use", "user")), ContentPart::ToolUse { .. } => Err(invalid_role("tool use", "user")),
} }
@@ -371,6 +373,14 @@ fn convert_assistant_part(part: ContentPart) -> Result<AssistantContent, AgentEr
} }
} }
fn tool_result_text(content: String, is_error: bool) -> String {
if is_error {
format!("[ERROR] {content}")
} else {
content
}
}
fn one_or_many<T: Clone>(parts: Vec<T>, role: &str) -> Result<OneOrMany<T>, AgentError> { fn one_or_many<T: Clone>(parts: Vec<T>, role: &str) -> Result<OneOrMany<T>, AgentError> {
OneOrMany::many(parts).map_err(|_| { OneOrMany::many(parts).map_err(|_| {
AgentError::new( AgentError::new(
@@ -1,6 +1,6 @@
use futures::StreamExt; use futures::StreamExt;
use galaxy_agent_core::{ use galaxy_agent_core::{
AgentEvent, AgentRuntime, ConversationMessage, MessageContent, MessageRole, AgentEvent, AgentRuntime, ConversationMessage, MessageContent, MessageRole, ToolEvent,
}; };
use rig_core::client::CompletionClient; use rig_core::client::CompletionClient;
use rig_core::providers::openai; use rig_core::providers::openai;
@@ -152,6 +152,71 @@ async fn usage_at_the_requested_limit_maps_to_max_tokens() {
); );
} }
#[tokio::test]
async fn rig_stream_maps_complete_tool_call_without_executing_it() {
let http_client = MockStreamingClient {
sse_bytes: sse(&[
r#"{"id":"cmpl-1","model":"test-model","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"read_files","arguments":"{\"files\":[\"Cargo.toml\"]}"}}]},"finish_reason":null}],"usage":null}"#,
r#"{"id":"cmpl-1","model":"test-model","choices":[{"delta":{"tool_calls":[]},"finish_reason":"tool_calls"}],"usage":null}"#,
r#"{"choices":[],"usage":{"prompt_tokens":8,"completion_tokens":4,"total_tokens":12}}"#,
"[DONE]",
]),
};
let client = openai::CompletionsClient::builder()
.api_key("test-key")
.base_url("http://localhost/v1")
.http_client(http_client)
.build()
.unwrap();
let model = client.completion_model("test-model");
let (_, control) = galaxy_agent_core::turn_control();
let mut request = text_request();
request.tools.push(galaxy_agent_core::ToolDefinition {
name: "read_files".to_string(),
description: "Read files".to_string(),
input_schema: serde_json::json!({"type": "object"}),
});
let events = start_model_turn(model, request, control, None, true)
.await
.unwrap()
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert!(events.iter().any(|event| {
matches!(
event,
AgentEvent::Tool {
event: ToolEvent::Proposed { call },
}
if call.id == "call-1"
&& call.name == "read_files"
&& call.arguments == serde_json::json!({"files": ["Cargo.toml"]})
)
}));
assert_eq!(
events.last(),
Some(&AgentEvent::TurnStopped {
reason: StopReason::Completed,
})
);
assert_eq!(
events
.iter()
.filter(|event| matches!(
event,
AgentEvent::Tool {
event: ToolEvent::Started { .. } | ToolEvent::Completed { .. },
}
))
.count(),
0
);
}
#[test] #[test]
fn request_conversion_preserves_history_tools_and_limits() { fn request_conversion_preserves_history_tools_and_limits() {
let mut request = text_request(); let mut request = text_request();
@@ -175,6 +240,57 @@ fn request_conversion_preserves_history_tools_and_limits() {
)); ));
} }
#[test]
fn request_conversion_preserves_tool_call_and_denied_result_for_the_next_turn() {
let request = TurnRequest::new(
"test-model",
vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "call-1".to_string(),
name: "run_shell_command".to_string(),
input: serde_json::json!({"command": "cargo test"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "call-1".to_string(),
content: "Command not executed — permission denied.".to_string(),
is_error: true,
},
},
],
);
let converted = build_completion_request(request, None, true).unwrap();
let messages = converted.chat_history.iter().collect::<Vec<_>>();
let Message::Assistant { content, .. } = messages[0] else {
panic!("expected assistant tool call");
};
let Some(AssistantContent::ToolCall(call)) = content.iter().next() else {
panic!("expected assistant tool call content");
};
assert_eq!(call.id, "call-1");
assert_eq!(call.function.name, "run_shell_command");
let Message::User { content } = messages[1] else {
panic!("expected user tool result");
};
let Some(UserContent::ToolResult(result)) = content.iter().next() else {
panic!("expected user tool result content");
};
assert_eq!(result.id, "call-1");
let Some(ToolResultContent::Text(text)) = result.content.iter().next() else {
panic!("expected text tool result");
};
assert_eq!(
text.text,
"[ERROR] Command not executed — permission denied."
);
}
#[test] #[test]
fn request_conversion_places_system_prompt_in_user_message_when_system_role_is_unsupported() { fn request_conversion_places_system_prompt_in_user_message_when_system_role_is_unsupported() {
let mut request = text_request(); let mut request = text_request();
@@ -428,6 +428,7 @@ fn register_tests() -> HashMap<&'static str, BoxedBuilderFn> {
register_test!(test_restored_ai_block_renders_mermaid_and_local_images); register_test!(test_restored_ai_block_renders_mermaid_and_local_images);
register_test!(test_agent_mode_pane_minimum_size); register_test!(test_agent_mode_pane_minimum_size);
register_test!(test_rig_read_tool_round_trip);
register_test!(test_git_prompt_chips); register_test!(test_git_prompt_chips);
// These tests are only invoked manually, and not included in the // These tests are only invoked manually, and not included in the
+2
View File
@@ -21,6 +21,7 @@ mod pane_restoration;
mod preview_config_migration; mod preview_config_migration;
mod remote_server; mod remote_server;
mod rich_input_ctrl_enter; mod rich_input_ctrl_enter;
mod rig_runtime;
mod rules; mod rules;
mod secrets; mod secrets;
mod session_restoration; mod session_restoration;
@@ -77,6 +78,7 @@ use pathfinder_geometry::vector::Vector2F;
pub use preview_config_migration::*; pub use preview_config_migration::*;
pub use remote_server::*; pub use remote_server::*;
pub use rich_input_ctrl_enter::*; pub use rich_input_ctrl_enter::*;
pub use rig_runtime::*;
pub use rules::*; pub use rules::*;
use rust_embed::RustEmbed; use rust_embed::RustEmbed;
pub use secrets::*; pub use secrets::*;
+248
View File
@@ -0,0 +1,248 @@
use std::io::{ErrorKind, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;
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,
};
use warp::integration_testing::step::new_step_with_default_assertions;
use warp::integration_testing::terminal::wait_until_bootstrapped_single_pane_for_tab;
use super::new_builder;
use crate::Builder;
const MODEL_ID: &str = "integration-rig-model";
const FINAL_TEXT: &str = "Rig read round trip completed.";
const FIXTURE_CONTENT: &str = "content returned through the Galaxy read executor";
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)));
let setup_fixture_path = fixture_path.clone();
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();
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]]
name = "Rig Integration"
base_url = "http://{address}/v1"
[[ai.providers.models]]
model_id = "{MODEL_ID}"
display_name = "Rig Integration Model"
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)),
),
)
}
fn start_mock_provider(
fixture_path: Arc<Mutex<String>>,
stop: Arc<AtomicBool>,
) -> (SocketAddr, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").expect("should bind mock Rig provider");
let address = listener.local_addr().expect("mock provider address");
listener
.set_nonblocking(true)
.expect("should make mock provider nonblocking");
let request_count = AtomicUsize::new(0);
let thread = thread::spawn(move || {
while !stop.load(Ordering::SeqCst) {
match listener.accept() {
Ok((mut stream, _)) => {
serve_request(&mut stream, &fixture_path, &request_count);
}
Err(error) if error.kind() == ErrorKind::WouldBlock => {
thread::sleep(Duration::from_millis(10));
}
Err(error) => panic!("mock Rig provider accept failed: {error}"),
}
}
});
(address, thread)
}
fn serve_request(
stream: &mut TcpStream,
fixture_path: &Mutex<String>,
request_count: &AtomicUsize,
) {
stream
.set_read_timeout(Some(Duration::from_secs(5)))
.expect("should set request timeout");
let request = read_request(stream);
let request_line = request.lines().next().unwrap_or_default();
if request_line.contains("/models") {
let body =
format!(r#"{{"object":"list","data":[{{"id":"{MODEL_ID}","object":"model"}}]}}"#);
write_response(stream, "application/json", &body);
return;
}
assert!(
request_line.contains("/chat/completions"),
"unexpected mock provider request: {request_line}"
);
let turn = request_count.fetch_add(1, Ordering::SeqCst);
let body = match turn {
0 => {
let fixture = fixture_path.lock().expect("fixture path lock").clone();
tool_call_sse(&fixture)
}
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()
}
_ => panic!("unexpected extra chat completion request"),
};
write_response(stream, "text/event-stream", &body);
}
fn read_request(stream: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut chunk = [0; 8 * 1024];
loop {
let bytes_read = stream
.read(&mut chunk)
.expect("should read provider request");
if bytes_read == 0 {
break;
}
request.extend_from_slice(&chunk[..bytes_read]);
assert!(
request.len() <= 1024 * 1024,
"mock provider request exceeded 1 MiB"
);
let Some(headers_end) = request.windows(4).position(|window| window == b"\r\n\r\n") else {
continue;
};
let body_start = headers_end + 4;
let headers = String::from_utf8_lossy(&request[..headers_end]);
let content_length = headers.lines().find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
});
match content_length {
Some(content_length) if request.len() < body_start + content_length => continue,
Some(_) | None => break,
}
}
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();
let tool_delta = serde_json::json!({
"id": "rig-integration-1",
"model": MODEL_ID,
"choices": [{
"delta": {
"tool_calls": [{
"index": 0,
"id": "rig-read-call",
"type": "function",
"function": {
"name": "read_files",
"arguments": arguments,
},
}],
},
"finish_reason": null,
}],
"usage": null,
});
let tool_stop = serde_json::json!({
"id": "rig-integration-1",
"model": MODEL_ID,
"choices": [{"delta": {"tool_calls": []}, "finish_reason": "tool_calls"}],
"usage": null,
});
let usage = serde_json::json!({
"choices": [],
"usage": {"prompt_tokens": 20, "completion_tokens": 8, "total_tokens": 28},
});
format!("data: {tool_delta}\n\ndata: {tool_stop}\n\ndata: {usage}\n\ndata: [DONE]\n\n")
}
fn final_text_sse() -> String {
let text = serde_json::json!({
"id": "rig-integration-2",
"model": MODEL_ID,
"choices": [{
"delta": {"content": FINAL_TEXT, "tool_calls": []},
"finish_reason": "stop",
}],
"usage": null,
});
let usage = serde_json::json!({
"choices": [],
"usage": {"prompt_tokens": 30, "completion_tokens": 6, "total_tokens": 36},
});
format!("data: {text}\n\ndata: {usage}\n\ndata: [DONE]\n\n")
}
fn write_response(stream: &mut TcpStream, content_type: &str, body: &str) {
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
)
.expect("should write mock provider response");
stream.flush().expect("should flush mock provider response");
}
@@ -311,6 +311,7 @@ integration_tests! {
#[cfg(any(target_os = "linux", target_os = "freebsd"))] #[cfg(any(target_os = "linux", target_os = "freebsd"))]
test_middle_click_paste, test_middle_click_paste,
test_agent_mode_pane_minimum_size, test_agent_mode_pane_minimum_size,
test_rig_read_tool_round_trip,
test_rule_creation, test_rule_creation,
test_rule_update, test_rule_update,
+36 -11
View File
@@ -286,19 +286,43 @@ use_rig = true
supports_system_messages = false supports_system_messages = false
``` ```
Phase 2 intentionally does not expose Galaxy's legacy tool list to Rig. That ownership moves as a At Phase 2 completion, Galaxy intentionally did not expose its legacy tool list to Rig. Phase 3
unit in Phase 3; until then, the opt-in slice validates text conversation streaming without two then moves that ownership behind the Galaxy safety boundary without introducing a second tool
competing tool executors. executor.
Exit condition: a LiteLLM or local OpenAI-compatible conversation streams through Rig without Exit condition: a LiteLLM or local OpenAI-compatible conversation streams through Rig without
`warp_multi_agent_api::Request` on the provider side. `warp_multi_agent_api::Request` on the provider side.
### Phase 3 — Tools, permissions, MCP, and multi-turn behavior ### Phase 3 — Tools, permissions, MCP, and multi-turn behavior
- Bridge the core Galaxy tools into Rig. - [x] Advertise Galaxy's current core tool definitions to Rig and translate streamed
- Preserve permission cards, denial, cancellation, parallel-call ordering, and error visibility. `ToolProposed` events into the existing permission/action UI contract.
- Bridge current MCP tools through Rig's `rmcp` support or a single Galaxy tool-server adapter. - [x] Keep Galaxy's action model as the sole execution authority; the provider runtime cannot
- Port loop prevention and unknown-tool handling to domain-level policies. execute shell, file, or MCP tools itself.
- [x] Persist assistant tool calls before exposing them to the executor, preserving parallel-call
order and preventing fast results from outrunning conversation history.
- [x] Preserve denied and failed tool results as explicit errors when building the next Rig turn.
- [x] Route current MCP tool proposals through Galaxy's existing MCP executor adapter.
- [x] Define one provider-neutral `ToolEvent` lifecycle with stable call IDs, permission request and
resolution, execution start, and success/error/denied/cancelled completion states.
- [x] Emit the normalized permission and execution lifecycle from Galaxy's existing action model
while keeping legacy UI events as a temporary compatibility layer.
- [x] Build Rig `TurnRequest`s directly from Galaxy request state before the legacy
`warp_multi_agent_api::Request` boundary; Bedrock and legacy OpenAI alone retain that request
adapter.
- [x] Feed normalized tool results directly into the next Rig turn, preserving success, failure,
denial, cancellation, call IDs, ordering, and persistent assistant tool-call history without a
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,
removing the temporary Warp protobuf adapter.
- [ ] 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] 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 Exit condition: representative read, edit, shell, MCP, denial, and failure flows pass integration
tests without provider-specific UI code. tests without provider-specific UI code.
@@ -387,7 +411,8 @@ contract is what the UI and persistence observe.
## Immediate next vertical slice ## Immediate next vertical slice
After the Phase 0 egress guard and UI ledger are verified, the next implementation change is a small Expand the hermetic Rig integration harness from its passing read-tool round trip to
`galaxy_agent_core` crate plus a legacy adapter. It should move only provider-neutral message/event edit/shell/MCP/denial/cancellation/execution-failure cases. Then move permission decisions and tool
types and runtime selection. Adding Rig before this seam would couple the UI to a new framework and start/result events fully onto the Galaxy domain contract, keeping the outgoing UI response adapter
repeat the current mistake with a different name. only until those flows prove that Galaxy-owned events can replace it without changing the
permission UI.