Migrate Rig tool flow to domain runtime
This commit is contained in:
@@ -135,6 +135,216 @@ impl AIAgentActionResultType {
|
||||
_ => 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)]
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
@@ -180,12 +180,14 @@ where
|
||||
}
|
||||
}
|
||||
Ok(StreamedAssistantContent::ToolCall { tool_call, .. }) => {
|
||||
yield Ok(AgentEvent::ToolProposed {
|
||||
call: ToolCall {
|
||||
yield Ok(AgentEvent::Tool {
|
||||
event: galaxy_agent_core::ToolEvent::Proposed {
|
||||
call: ToolCall {
|
||||
id: tool_call.id,
|
||||
name: tool_call.function.name,
|
||||
arguments: tool_call.function.arguments,
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
Ok(StreamedAssistantContent::ToolCallDelta { .. }) => {
|
||||
@@ -298,10 +300,10 @@ fn user_content(content: MessageContent) -> Result<OneOrMany<UserContent>, Agent
|
||||
MessageContent::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
..
|
||||
is_error,
|
||||
} => vec![UserContent::tool_result(
|
||||
tool_use_id,
|
||||
OneOrMany::one(ToolResultContent::text(content)),
|
||||
OneOrMany::one(ToolResultContent::text(tool_result_text(content, is_error))),
|
||||
)],
|
||||
MessageContent::MultiPart(parts) => parts
|
||||
.into_iter()
|
||||
@@ -344,10 +346,10 @@ fn convert_user_part(part: ContentPart) -> Result<UserContent, AgentError> {
|
||||
ContentPart::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
..
|
||||
is_error,
|
||||
} => Ok(UserContent::tool_result(
|
||||
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")),
|
||||
}
|
||||
@@ -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> {
|
||||
OneOrMany::many(parts).map_err(|_| {
|
||||
AgentError::new(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use futures::StreamExt;
|
||||
use galaxy_agent_core::{
|
||||
AgentEvent, AgentRuntime, ConversationMessage, MessageContent, MessageRole,
|
||||
AgentEvent, AgentRuntime, ConversationMessage, MessageContent, MessageRole, ToolEvent,
|
||||
};
|
||||
use rig_core::client::CompletionClient;
|
||||
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]
|
||||
fn request_conversion_preserves_history_tools_and_limits() {
|
||||
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]
|
||||
fn request_conversion_places_system_prompt_in_user_message_when_system_role_is_unsupported() {
|
||||
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_agent_mode_pane_minimum_size);
|
||||
register_test!(test_rig_read_tool_round_trip);
|
||||
register_test!(test_git_prompt_chips);
|
||||
|
||||
// These tests are only invoked manually, and not included in the
|
||||
|
||||
@@ -21,6 +21,7 @@ mod pane_restoration;
|
||||
mod preview_config_migration;
|
||||
mod remote_server;
|
||||
mod rich_input_ctrl_enter;
|
||||
mod rig_runtime;
|
||||
mod rules;
|
||||
mod secrets;
|
||||
mod session_restoration;
|
||||
@@ -77,6 +78,7 @@ use pathfinder_geometry::vector::Vector2F;
|
||||
pub use preview_config_migration::*;
|
||||
pub use remote_server::*;
|
||||
pub use rich_input_ctrl_enter::*;
|
||||
pub use rig_runtime::*;
|
||||
pub use rules::*;
|
||||
use rust_embed::RustEmbed;
|
||||
pub use secrets::*;
|
||||
|
||||
@@ -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"))]
|
||||
test_middle_click_paste,
|
||||
test_agent_mode_pane_minimum_size,
|
||||
test_rig_read_tool_round_trip,
|
||||
|
||||
test_rule_creation,
|
||||
test_rule_update,
|
||||
|
||||
Reference in New Issue
Block a user