Migrate Rig tool flow to domain runtime
This commit is contained in:
@@ -5,7 +5,9 @@
|
||||
//! depend on GalaxyUI, provider SDKs, persistence, or Warp wire protocols.
|
||||
|
||||
mod runtime;
|
||||
mod tool_policy;
|
||||
mod types;
|
||||
|
||||
pub use runtime::*;
|
||||
pub use tool_policy::*;
|
||||
pub use types::*;
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -126,7 +126,24 @@ pub struct ToolCall {
|
||||
pub struct ToolResult {
|
||||
pub call_id: 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)]
|
||||
@@ -141,11 +158,52 @@ pub enum PermissionKind {
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct PermissionRequest {
|
||||
pub id: String,
|
||||
pub tool_call: ToolCall,
|
||||
pub call_id: String,
|
||||
pub kind: PermissionKind,
|
||||
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)]
|
||||
pub struct Usage {
|
||||
pub input_tokens: u64,
|
||||
@@ -176,10 +234,7 @@ pub enum AgentEvent {
|
||||
TurnStarted { runtime_request_id: String },
|
||||
TextDelta { text: String },
|
||||
ReasoningDelta { text: String },
|
||||
ToolProposed { call: ToolCall },
|
||||
PermissionRequested { request: PermissionRequest },
|
||||
ToolStarted { call: ToolCall },
|
||||
ToolCompleted { result: ToolResult },
|
||||
Tool { event: ToolEvent },
|
||||
UsageUpdated { usage: Usage },
|
||||
TurnStopped { reason: StopReason },
|
||||
}
|
||||
|
||||
@@ -36,3 +36,58 @@ fn usage_total_excludes_cached_breakdown_to_avoid_double_counting() {
|
||||
|
||||
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());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user