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

Migrate Rig tool flow to domain runtime

Browse Source
This commit is contained in:
rkw6086
2026-08-04 14:14:51 -05:00
parent 4c7270db8d
commit 91d8bd0381
34 changed files with 2728 additions and 374 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/action_model.rs
+174 -3
View File
@@ -34,6 +34,9 @@ pub use execute::{
StartAgentExecutorEvent, StartAgentRequest, StartAgentRequestId,
};
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 itertools::Itertools;
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 {
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.
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.
///
/// We maintain this so that even though we might process actions in parallel,
/// we can still order the results consistently.
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_action_results: HashMap<AIAgentActionId, Arc<AIAgentActionResult>>,
@@ -266,6 +337,12 @@ impl BlocklistAIActionModel {
ctx.subscribe_to_model(&executor, move |me, _, event, ctx| match event {
BlocklistAIActionExecutorEvent::ExecutingAction { action_id } => {
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 {
result,
@@ -298,10 +375,12 @@ impl BlocklistAIActionModel {
Self {
pending_actions: Default::default(),
finished_action_results: Default::default(),
finished_tool_results: Default::default(),
executor,
past_action_results: HashMap::new(),
running_actions: Default::default(),
action_order: Default::default(),
denied_permissions: Default::default(),
terminal_view_id,
pending_preprocessed_actions: Default::default(),
is_view_only: false,
@@ -533,6 +612,18 @@ impl BlocklistAIActionModel {
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(
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| {
let blocked_action_user_friendly_str = action.action.user_friendly_name();
history_model.update_conversation_status(
@@ -897,6 +999,16 @@ impl BlocklistAIActionModel {
let action_id = action.id.clone();
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
// in-progress update.
let is_wait_for_events = matches!(action.action, AIAgentActionType::WaitForEvents { .. });
@@ -1073,6 +1185,8 @@ impl BlocklistAIActionModel {
reason: CancellationReason,
ctx: &mut ModelContext<Self>,
) {
let status = self.get_action_status(action_id);
let permission_denied = is_permission_denial(reason, status.as_ref());
if self
.running_actions
.get(&conversation_id)
@@ -1092,7 +1206,13 @@ impl BlocklistAIActionModel {
.find_position(|action| action.id == *action_id)
{
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,
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,
pending_action: AIAgentAction,
reason: Option<CancellationReason>,
permission_denied: bool,
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!(
pending_action.action,
AIAgentActionType::RequestComputerUse(_)
@@ -1227,10 +1361,20 @@ impl BlocklistAIActionModel {
.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.
pub(super) fn clear_finished_action_results(&mut self, conversation_id: AIConversationId) {
self.action_order.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,
@@ -1308,10 +1452,31 @@ impl BlocklistAIActionModel {
)
) {
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
.entry(conversation_id)
.or_default()
@@ -1474,6 +1639,11 @@ pub enum BlocklistAIActionEvent {
conversation_id: AIConversationId,
cancellation_reason: Option<CancellationReason>,
},
/// Provider-neutral permission and execution lifecycle event for runtime consumers.
ToolLifecycle {
action_id: AIAgentActionId,
event: ToolEvent,
},
InitProject(AIAgentActionId),
ToggleCodeReview(AIAgentActionId),
InsertCodeReviewComments {
@@ -1491,6 +1661,7 @@ impl BlocklistAIActionEvent {
BlocklistAIActionEvent::ActionBlockedOnUserConfirmation(action_id) => action_id,
BlocklistAIActionEvent::ExecutingAction(action_id) => action_id,
BlocklistAIActionEvent::FinishedAction { action_id, .. } => action_id,
BlocklistAIActionEvent::ToolLifecycle { action_id, .. } => action_id,
BlocklistAIActionEvent::InitProject(action_id) => action_id,
BlocklistAIActionEvent::ToggleCodeReview(action_id) => action_id,
BlocklistAIActionEvent::InsertCodeReviewComments { action_id, .. } => action_id,
app/src/ai/blocklist/action_model_tests.rs
+116 -1
View File
@@ -3,7 +3,9 @@ use std::sync::Arc;
use super::*;
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> {
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 {
let mut current_phase = None;
let mut count = 0;
@@ -100,3 +110,108 @@ fn finished_results_stay_in_original_action_order() {
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),
));
}
app/src/ai/blocklist/block.rs
+2 -1
View File
@@ -4929,7 +4929,8 @@ impl AIBlock {
}
}
BlocklistAIActionEvent::InitProject(_)
BlocklistAIActionEvent::ToolLifecycle { .. }
| BlocklistAIActionEvent::InitProject(_)
| BlocklistAIActionEvent::ToggleCodeReview(_) => {}
}
});
app/src/ai/blocklist/controller.rs
+23 -72
View File
@@ -10,7 +10,7 @@ mod pending_response_streams;
pub mod response_stream;
pub(super) mod shared_session;
mod slash_command;
use std::collections::{HashMap, HashSet, VecDeque};
use std::collections::{HashMap, HashSet};
#[cfg(not(target_family = "wasm"))]
use std::path::PathBuf;
use std::sync::Arc;
@@ -19,6 +19,7 @@ use std::time::Duration;
use ai::skills::SkillPathOrigin;
use anyhow::anyhow;
use chrono::{DateTime, Local};
use galaxy_agent_core::ToolLoopGuard;
use galaxy_core::assertions::safe_assert;
use input_context::{input_context_for_request, parse_context_attachments};
use itertools::Itertools;
@@ -196,60 +197,6 @@ pub enum BlocklistAIControllerEvent {
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)]
pub struct RequestInput {
pub conversation_id: AIConversationId,
@@ -419,7 +366,7 @@ pub struct BlocklistAIController {
pending_passive_follow_ups: HashSet<AIConversationId>,
/// 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.
error_retry_counts: HashMap<AIConversationId, usize>,
/// Passive suggestion results that should be included with the next request
@@ -1992,11 +1939,7 @@ impl BlocklistAIController {
description.hash(&mut hasher);
let input_hash = hasher.finish();
state.record_failure(LoopDetectionEntry {
tool_discriminant: discriminant,
input_hash,
description: description.clone(),
});
state.record_failure(input_hash, description);
} else if result.result.is_successful() {
has_success = true;
}
@@ -2005,23 +1948,21 @@ impl BlocklistAIController {
// If we had at least one success in this batch, clear loop state —
// the agent is making progress.
if has_success {
state.clear();
state.record_success();
return None;
}
// Check for loops
if let Some(looping_entry) = state.detect_loop() {
if let Some(looping_entry) = state.detect_and_reset() {
let warning = format!(
"[SYSTEM] Loop detected: the same action has failed {} or more times consecutively. \
Do NOT repeat this action or any similar approach.\n\n\
Failing action: {}\n\n\
Take a completely different approach to accomplish the goal. \
If you cannot find an alternative, explain to the user what is failing and why.",
LOOP_DETECTION_THRESHOLD,
looping_entry.threshold,
looping_entry.description
);
// Clear the state so we don't keep injecting on every subsequent turn
state.clear();
Some(warning)
} else {
None
@@ -3061,11 +3002,23 @@ impl BlocklistAIController {
query_metadata,
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.agent_name = agent_name;
request_params.bedrock_message_history = bedrock_history;
request_params.bedrock_tool_result_archive = bedrock_tool_result_archive;
request_params.bedrock_progressive_summary = bedrock_progressive_summary;
request_params.message_history = bedrock_history;
request_params.tool_result_archive = bedrock_tool_result_archive;
request_params.progressive_summary = bedrock_progressive_summary;
// 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
@@ -3555,9 +3508,7 @@ impl BlocklistAIController {
// history (input + assistant response) from the Arc back
// into the conversation for the next request cycle.
let new_history = (!response_stream.as_ref(ctx).is_acp())
.then(|| {
response_stream.as_ref(ctx).bedrock_messages_sent().clone()
})
.then(|| response_stream.as_ref(ctx).messages_sent().clone())
.and_then(|messages_sent| {
messages_sent.lock().ok().and_then(|sent| {
if sent.is_empty() {
app/src/ai/blocklist/controller/response_stream.rs
+3 -3
View File
@@ -552,11 +552,11 @@ impl ResponseStream {
}
}
pub fn bedrock_messages_sent(
pub fn messages_sent(
&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.