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

Complete Rig tool lifecycle migration

Browse Source
This commit is contained in:
rkw6086
2026-08-04 16:00:20 -05:00
parent 91d8bd0381
commit a3c68e9c30
30 changed files with 1494 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/integration/src/bin/integration.rs
+10
View File
@@ -24,6 +24,10 @@ pub struct Args {
}
pub fn main() -> Result<()> {
if env::args().nth(1).as_deref() == Some(integration::rig_test_mcp_server::SERVER_ARGUMENT) {
return integration::rig_test_mcp_server::run();
}
ChannelState::set(ChannelState::new(
Channel::Integration,
ChannelConfig {
@@ -429,6 +433,12 @@ fn register_tests() -> HashMap<&'static str, BoxedBuilderFn> {
register_test!(test_agent_mode_pane_minimum_size);
register_test!(test_rig_read_tool_round_trip);
register_test!(test_rig_shell_tool_success_round_trip);
register_test!(test_rig_shell_tool_failure_round_trip);
register_test!(test_rig_shell_tool_permission_denial);
register_test!(test_rig_edit_tool_round_trip);
register_test!(test_rig_in_flight_cancellation);
register_test!(test_rig_mcp_tool_round_trip);
register_test!(test_git_prompt_chips);
// These tests are only invoked manually, and not included in the
crates/integration/src/lib.rs
+1
View File
@@ -1,4 +1,5 @@
mod builder;
pub mod rig_test_mcp_server;
mod step;
pub mod test;
crates/integration/src/rig_test_mcp_server.rs
+89
View File
@@ -0,0 +1,89 @@
use std::io::{self, BufRead as _, Write as _};
use anyhow::Context as _;
pub const SERVER_ARGUMENT: &str = "--rig-test-mcp-server";
pub const TOOL_RESULT_PREFIX: &str = "rig-mcp-echo: ";
pub fn run() -> anyhow::Result<()> {
let stdin = io::stdin();
let mut stdout = io::stdout().lock();
for line in stdin.lock().lines() {
let line = line.context("failed to read MCP request")?;
let request: serde_json::Value =
serde_json::from_str(&line).context("failed to parse MCP request")?;
let Some(id) = request.get("id").cloned() else {
continue;
};
let method = request
.get("method")
.and_then(serde_json::Value::as_str)
.unwrap_or_default();
let result = match method {
"initialize" => {
let protocol_version = request
.pointer("/params/protocolVersion")
.cloned()
.unwrap_or_else(|| serde_json::json!("2024-11-05"));
serde_json::json!({
"protocolVersion": protocol_version,
"capabilities": {"tools": {}},
"serverInfo": {
"name": "rig-integration-mcp",
"version": "1.0.0",
},
})
}
"tools/list" => serde_json::json!({
"tools": [{
"name": "echo",
"description": "Echo text through the hermetic Rig integration MCP server",
"inputSchema": {
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"],
},
}],
}),
"tools/call" => {
let text = request
.pointer("/params/arguments/text")
.and_then(serde_json::Value::as_str)
.unwrap_or_default();
serde_json::json!({
"content": [{
"type": "text",
"text": format!("{TOOL_RESULT_PREFIX}{text}"),
}],
"isError": false,
})
}
"ping" => serde_json::json!({}),
_ => {
let response = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"error": {
"code": -32601,
"message": format!("Unsupported test MCP method: {method}"),
},
});
writeln!(stdout, "{response}").context("failed to write MCP error response")?;
stdout
.flush()
.context("failed to flush MCP error response")?;
continue;
}
};
let response = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"result": result,
});
writeln!(stdout, "{response}").context("failed to write MCP response")?;
stdout.flush().context("failed to flush MCP response")?;
}
Ok(())
}
crates/integration/src/test/rig_runtime.rs
+525 -65
View File
@@ -1,48 +1,316 @@
use std::io::{ErrorKind, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;
use galaxyui_core::async_assert;
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,
assert_latest_exchange_text, assert_task_is_cancelled, enter_agent_view,
set_execution_profile_auto_apply_code_diffs, set_execution_profile_auto_execute,
set_execution_profile_auto_execute_mcp_tools, set_execution_profile_no_auto_execute,
set_preferred_agent_mode_llm, start_ephemeral_mcp_server_for_testing, submit_ai_query,
submit_ai_query_and_wait_until_blocked, submit_ai_query_and_wait_until_done,
wait_until_mcp_server_is_active_for_testing, ConversationTarget,
};
use warp::integration_testing::step::new_step_with_default_assertions;
use warp::integration_testing::terminal::wait_until_bootstrapped_single_pane_for_tab;
use warp::integration_testing::terminal::{
assert_input_is_focused, wait_until_bootstrapped_single_pane_for_tab,
};
use super::new_builder;
use crate::rig_test_mcp_server::{SERVER_ARGUMENT as MCP_SERVER_ARGUMENT, TOOL_RESULT_PREFIX};
use crate::Builder;
const MODEL_ID: &str = "integration-rig-model";
const FINAL_TEXT: &str = "Rig read round trip completed.";
const READ_CALL_ID: &str = "rig-read-call";
const READ_FINAL_TEXT: &str = "Rig read round trip completed.";
const FIXTURE_CONTENT: &str = "content returned through the Galaxy read executor";
const SHELL_CALL_ID: &str = "rig-shell-call";
const SHELL_SUCCESS_OUTPUT: &str = "rig-shell-success-output";
const SHELL_SUCCESS_FINAL_TEXT: &str = "Rig shell success round trip completed.";
const SHELL_FAILURE_OUTPUT: &str = "rig-shell-failure-output";
const SHELL_FAILURE_FINAL_TEXT: &str = "Rig shell failure round trip completed.";
const SHELL_DENIED_FINAL_TEXT: &str = "Rig shell denial was preserved.";
const EDIT_CALL_ID: &str = "rig-edit-call";
const EDIT_INITIAL_CONTENT: &str = "before Rig edit\n";
const EDIT_UPDATED_CONTENT: &str = "after Rig edit\n";
const EDIT_FINAL_TEXT: &str = "Rig edit round trip completed.";
const MCP_CALL_ID: &str = "rig-mcp-call";
const MCP_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
const MCP_TEMPLATE_ID: &str = "22222222-2222-4222-8222-222222222222";
const MCP_SERVER_NAME: &str = "rig-integration";
const MCP_TOOL_NAME: &str = "mcp__11111111-1111-4111-8111-111111111111__echo";
const MCP_INPUT: &str = "hello from Rig";
const MCP_FINAL_TEXT: &str = "Rig MCP round trip completed.";
#[derive(Clone)]
enum MockScenario {
Read {
fixture_path: Arc<Mutex<String>>,
},
ShellSuccess,
ShellFailure,
ShellDenied {
marker_path: Arc<Mutex<String>>,
},
Edit {
fixture_path: Arc<Mutex<String>>,
},
Cancellation {
stream_started: Arc<AtomicBool>,
stream_cancelled: Arc<AtomicBool>,
},
Mcp,
}
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)));
rig_builder(MockScenario::Read { fixture_path })
.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(READ_FINAL_TEXT)),
),
)
}
let setup_fixture_path = fixture_path.clone();
pub fn test_rig_shell_tool_success_round_trip() -> Builder {
rig_builder(MockScenario::ShellSuccess)
.with_step(wait_until_bootstrapped_single_pane_for_tab(0))
.with_step(set_preferred_agent_mode_llm(MODEL_ID))
.with_step(set_execution_profile_auto_execute())
.with_step(enter_agent_view())
.with_step(submit_ai_query_and_wait_until_done(
"Run the requested shell success check.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Assert Rig shell success reached Agent Mode")
.add_named_assertion(
"Final response follows the successful shell result",
assert_latest_exchange_text(|text| text.contains(SHELL_SUCCESS_FINAL_TEXT)),
),
)
}
pub fn test_rig_shell_tool_failure_round_trip() -> Builder {
rig_builder(MockScenario::ShellFailure)
.with_step(wait_until_bootstrapped_single_pane_for_tab(0))
.with_step(set_preferred_agent_mode_llm(MODEL_ID))
.with_step(set_execution_profile_auto_execute())
.with_step(enter_agent_view())
.with_step(submit_ai_query_and_wait_until_done(
"Run the requested failing shell check.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Assert Rig shell failure reached Agent Mode")
.add_named_assertion(
"Final response follows the failed shell result",
assert_latest_exchange_text(|text| text.contains(SHELL_FAILURE_FINAL_TEXT)),
),
)
}
pub fn test_rig_shell_tool_permission_denial() -> Builder {
let marker_path = Arc::new(Mutex::new(String::new()));
rig_builder(MockScenario::ShellDenied { marker_path })
.with_step(wait_until_bootstrapped_single_pane_for_tab(0))
.with_step(set_preferred_agent_mode_llm(MODEL_ID))
.with_step(set_execution_profile_no_auto_execute())
.with_step(enter_agent_view())
.with_step(submit_ai_query_and_wait_until_blocked(
"Propose the requested shell command and wait for my decision.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Deny the blocked Rig shell command")
.with_keystrokes(&["ctrl-c"])
.add_named_assertion(
"Input returns after denying the command",
assert_input_is_focused(),
),
)
.with_step(submit_ai_query_and_wait_until_done(
"Continue without running the denied command.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Assert Rig denial reached Agent Mode")
.add_named_assertion(
"Final response follows the explicit denied result",
assert_latest_exchange_text(|text| text.contains(SHELL_DENIED_FINAL_TEXT)),
),
)
}
pub fn test_rig_edit_tool_round_trip() -> Builder {
let fixture_path = Arc::new(Mutex::new(String::new()));
rig_builder(MockScenario::Edit { fixture_path })
.with_step(wait_until_bootstrapped_single_pane_for_tab(0))
.with_step(set_preferred_agent_mode_llm(MODEL_ID))
.with_step(set_execution_profile_auto_apply_code_diffs())
.with_step(enter_agent_view())
.with_step(submit_ai_query_and_wait_until_done(
"Apply the requested edit to the integration fixture.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Assert Rig edit result reached Agent Mode")
.add_named_assertion(
"Final response follows the real edit result",
assert_latest_exchange_text(|text| text.contains(EDIT_FINAL_TEXT)),
),
)
}
pub fn test_rig_in_flight_cancellation() -> Builder {
let stream_started = Arc::new(AtomicBool::new(false));
let stream_cancelled = Arc::new(AtomicBool::new(false));
rig_builder(MockScenario::Cancellation {
stream_started: stream_started.clone(),
stream_cancelled: stream_cancelled.clone(),
})
.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(
"Start the cancellable Rig response.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Wait for the Rig provider stream to start")
.set_timeout(Duration::from_secs(30))
.add_named_assertion("Provider stream started", move |_app, _window_id| {
async_assert!(
stream_started.load(Ordering::SeqCst),
"Waiting for the provider stream to start"
)
}),
)
.with_step(
new_step_with_default_assertions("Cancel the in-flight Rig provider stream")
.set_timeout(Duration::from_secs(30))
.with_keystrokes(&["ctrl-c"])
.add_named_assertion(
"Conversation is marked cancelled",
assert_task_is_cancelled(ConversationTarget::Active),
)
.add_named_assertion(
"Provider connection was dropped",
move |_app, _window_id| {
async_assert!(
stream_cancelled.load(Ordering::SeqCst),
"Waiting for cancellation to close the provider connection"
)
},
)
.add_named_assertion(
"Input returns after cancelling the stream",
assert_input_is_focused(),
),
)
}
pub fn test_rig_mcp_tool_round_trip() -> Builder {
let command = std::env::current_exe()
.expect("integration executable path")
.to_string_lossy()
.into_owned();
rig_builder(MockScenario::Mcp)
.with_step(wait_until_bootstrapped_single_pane_for_tab(0))
.with_step(start_ephemeral_mcp_server_for_testing(
command,
MCP_SERVER_ARGUMENT.to_string(),
MCP_INSTALLATION_ID,
MCP_TEMPLATE_ID,
MCP_SERVER_NAME,
))
.with_step(wait_until_mcp_server_is_active_for_testing(
MCP_INSTALLATION_ID,
Duration::from_secs(30),
))
.with_step(set_preferred_agent_mode_llm(MODEL_ID))
.with_step(set_execution_profile_auto_execute_mcp_tools())
.with_step(enter_agent_view())
.with_step(submit_ai_query_and_wait_until_done(
"Call the hermetic MCP echo tool.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Assert Rig MCP result reached Agent Mode")
.add_named_assertion(
"Final response follows the real MCP result",
assert_latest_exchange_text(|text| text.contains(MCP_FINAL_TEXT)),
),
)
}
fn rig_builder(scenario: MockScenario) -> Builder {
FeatureFlag::AgentView.set_enabled(true);
FeatureFlag::MCPGroupedServerContext.set_enabled(true);
let stop = Arc::new(AtomicBool::new(false));
let setup_scenario = scenario.clone();
let (address, server_thread) = start_mock_provider(scenario, stop.clone());
let server_thread = Arc::new(Mutex::new(Some(server_thread)));
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();
write_provider_settings(address);
match &setup_scenario {
MockScenario::Read { fixture_path } => {
let fixture = utils.test_dir().join("rig-read-fixture.txt");
std::fs::write(&fixture, FIXTURE_CONTENT)
.expect("should write Rig integration fixture");
*fixture_path.lock().expect("fixture path lock") =
fixture.to_string_lossy().into_owned();
}
MockScenario::ShellDenied { marker_path } => {
let marker = utils.test_dir().join("rig-denied-command-marker.txt");
*marker_path.lock().expect("marker path lock") =
marker.to_string_lossy().into_owned();
}
MockScenario::Edit { fixture_path } => {
let fixture = utils.test_dir().join("rig-edit-fixture.txt");
std::fs::write(&fixture, EDIT_INITIAL_CONTENT)
.expect("should write Rig edit integration fixture");
*fixture_path.lock().expect("fixture path lock") =
fixture.to_string_lossy().into_owned();
}
MockScenario::ShellSuccess
| MockScenario::ShellFailure
| MockScenario::Cancellation { .. }
| MockScenario::Mcp => {}
}
})
.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");
}
})
}
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]
fn write_provider_settings(address: SocketAddr) {
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]]
@@ -56,33 +324,12 @@ 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)),
),
)
);
std::fs::write(settings_path, settings).expect("should write provider settings");
}
fn start_mock_provider(
fixture_path: Arc<Mutex<String>>,
scenario: MockScenario,
stop: Arc<AtomicBool>,
) -> (SocketAddr, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").expect("should bind mock Rig provider");
@@ -95,7 +342,7 @@ fn start_mock_provider(
while !stop.load(Ordering::SeqCst) {
match listener.accept() {
Ok((mut stream, _)) => {
serve_request(&mut stream, &fixture_path, &request_count);
serve_request(&mut stream, &scenario, &request_count, &stop);
}
Err(error) if error.kind() == ErrorKind::WouldBlock => {
thread::sleep(Duration::from_millis(10));
@@ -109,8 +356,9 @@ fn start_mock_provider(
fn serve_request(
stream: &mut TcpStream,
fixture_path: &Mutex<String>,
scenario: &MockScenario,
request_count: &AtomicUsize,
stop: &AtomicBool,
) {
stream
.set_read_timeout(Some(Duration::from_secs(5)))
@@ -130,21 +378,20 @@ fn serve_request(
"unexpected mock provider request: {request_line}"
);
let turn = request_count.fetch_add(1, Ordering::SeqCst);
if let MockScenario::Cancellation {
stream_started,
stream_cancelled,
} = scenario
{
assert_eq!(turn, 0, "unexpected extra cancellation chat request");
write_cancellable_response(stream, stream_started, stream_cancelled, stop);
return;
}
let body = match turn {
0 => {
let fixture = fixture_path.lock().expect("fixture path lock").clone();
tool_call_sse(&fixture)
}
0 => tool_call_sse(scenario),
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()
assert_follow_up_request(scenario, &request);
final_text_sse(final_text(scenario))
}
_ => panic!("unexpected extra chat completion request"),
};
@@ -186,8 +433,60 @@ fn read_request(stream: &mut TcpStream) -> String {
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();
fn tool_call_sse(scenario: &MockScenario) -> String {
let (call_id, tool_name, arguments) = match scenario {
MockScenario::Read { fixture_path } => {
let fixture_path = fixture_path.lock().expect("fixture path lock").clone();
(
READ_CALL_ID,
"read_files",
serde_json::json!({"files": [fixture_path]}),
)
}
MockScenario::ShellSuccess => (
SHELL_CALL_ID,
"run_shell_command",
shell_arguments("printf '%s\\n' 'rig-shell-success-output'"),
),
MockScenario::ShellFailure => (
SHELL_CALL_ID,
"run_shell_command",
shell_arguments("(printf '%s\\n' 'rig-shell-failure-output' >&2; exit 7)"),
),
MockScenario::ShellDenied { marker_path } => {
let marker_path = marker_path.lock().expect("marker path lock").clone();
(
SHELL_CALL_ID,
"run_shell_command",
shell_arguments(&format!(
"printf '%s\\n' 'command-must-not-run' > '{marker_path}'"
)),
)
}
MockScenario::Edit { fixture_path } => {
let fixture_path = fixture_path.lock().expect("fixture path lock").clone();
(
EDIT_CALL_ID,
"apply_file_diffs",
serde_json::json!({
"summary": "Update the Rig edit integration fixture",
"diffs": [{
"file_path": fixture_path,
"search": EDIT_INITIAL_CONTENT,
"replace": EDIT_UPDATED_CONTENT,
}],
}),
)
}
MockScenario::Cancellation { .. } => {
unreachable!("cancellation streams do not issue tool calls")
}
MockScenario::Mcp => (
MCP_CALL_ID,
MCP_TOOL_NAME,
serde_json::json!({"text": MCP_INPUT}),
),
};
let tool_delta = serde_json::json!({
"id": "rig-integration-1",
"model": MODEL_ID,
@@ -195,11 +494,11 @@ fn tool_call_sse(fixture_path: &str) -> String {
"delta": {
"tool_calls": [{
"index": 0,
"id": "rig-read-call",
"id": call_id,
"type": "function",
"function": {
"name": "read_files",
"arguments": arguments,
"name": tool_name,
"arguments": arguments.to_string(),
},
}],
},
@@ -220,12 +519,125 @@ fn tool_call_sse(fixture_path: &str) -> String {
format!("data: {tool_delta}\n\ndata: {tool_stop}\n\ndata: {usage}\n\ndata: [DONE]\n\n")
}
fn final_text_sse() -> String {
fn shell_arguments(command: &str) -> serde_json::Value {
serde_json::json!({
"command": command,
"is_read_only": false,
"is_risky": false,
"uses_pager": false,
"wait_until_complete": true,
})
}
fn assert_follow_up_request(scenario: &MockScenario, request: &str) {
match scenario {
MockScenario::Read { .. } => {
assert!(
request.contains(READ_CALL_ID),
"follow-up request should preserve the read tool call ID"
);
assert!(
request.contains(FIXTURE_CONTENT),
"follow-up request should contain the real file contents returned by Galaxy"
);
}
MockScenario::ShellSuccess => {
assert!(
request.contains(SHELL_CALL_ID),
"follow-up request should preserve the shell tool call ID"
);
assert!(
request.contains(SHELL_SUCCESS_OUTPUT),
"follow-up request should contain successful shell output"
);
assert!(
request.contains("exit code 0"),
"follow-up request should contain the successful exit code"
);
}
MockScenario::ShellFailure => {
assert!(
request.contains(SHELL_CALL_ID),
"follow-up request should preserve the failed shell tool call ID"
);
assert!(
request.contains(SHELL_FAILURE_OUTPUT),
"follow-up request should contain failed shell output"
);
assert!(
request.contains("exit code 7"),
"follow-up request should contain the failed exit code"
);
assert!(
request.contains("[ERROR]"),
"failed shell result should remain an explicit model error"
);
}
MockScenario::ShellDenied { marker_path } => {
assert!(
request.contains(SHELL_CALL_ID),
"follow-up request should preserve the denied shell tool call ID"
);
assert!(
request.contains("[ERROR] Permission denied by the user"),
"denied shell result should remain an explicit model error"
);
let marker_path = marker_path.lock().expect("marker path lock");
assert!(
!Path::new(marker_path.as_str()).exists(),
"denied shell command must not create its marker file"
);
}
MockScenario::Edit { fixture_path } => {
assert!(
request.contains(EDIT_CALL_ID),
"follow-up request should preserve the edit tool call ID"
);
assert!(
request.contains(EDIT_UPDATED_CONTENT.trim()),
"follow-up request should contain the updated file content"
);
let fixture_path = fixture_path.lock().expect("fixture path lock");
let actual = std::fs::read_to_string(fixture_path.as_str())
.expect("edited fixture should remain readable");
assert_eq!(actual, EDIT_UPDATED_CONTENT);
}
MockScenario::Cancellation { .. } => {
unreachable!("cancellation streams do not issue follow-up requests")
}
MockScenario::Mcp => {
assert!(
request.contains(MCP_CALL_ID),
"follow-up request should preserve the MCP tool call ID"
);
assert!(
request.contains(&format!("{TOOL_RESULT_PREFIX}{MCP_INPUT}")),
"follow-up request should contain the real MCP tool result"
);
}
}
}
fn final_text(scenario: &MockScenario) -> &'static str {
match scenario {
MockScenario::Read { .. } => READ_FINAL_TEXT,
MockScenario::ShellSuccess => SHELL_SUCCESS_FINAL_TEXT,
MockScenario::ShellFailure => SHELL_FAILURE_FINAL_TEXT,
MockScenario::ShellDenied { .. } => SHELL_DENIED_FINAL_TEXT,
MockScenario::Edit { .. } => EDIT_FINAL_TEXT,
MockScenario::Cancellation { .. } => {
unreachable!("cancellation streams do not produce final text")
}
MockScenario::Mcp => MCP_FINAL_TEXT,
}
}
fn final_text_sse(final_text: &str) -> String {
let text = serde_json::json!({
"id": "rig-integration-2",
"model": MODEL_ID,
"choices": [{
"delta": {"content": FINAL_TEXT, "tool_calls": []},
"delta": {"content": final_text, "tool_calls": []},
"finish_reason": "stop",
}],
"usage": null,
@@ -246,3 +658,51 @@ fn write_response(stream: &mut TcpStream, content_type: &str, body: &str) {
.expect("should write mock provider response");
stream.flush().expect("should flush mock provider response");
}
fn write_cancellable_response(
stream: &mut TcpStream,
stream_started: &AtomicBool,
stream_cancelled: &AtomicBool,
stop: &AtomicBool,
) {
let text = serde_json::json!({
"id": "rig-cancellation-1",
"model": MODEL_ID,
"choices": [{
"delta": {"content": "Rig cancellation stream is active.", "tool_calls": []},
"finish_reason": null,
}],
"usage": null,
});
let initial_body = format!("data: {text}\n\n");
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\n{initial_body}"
)
.expect("should start cancellable provider response");
stream
.flush()
.expect("should flush cancellable provider response");
stream_started.store(true, Ordering::SeqCst);
while !stop.load(Ordering::SeqCst) {
thread::sleep(Duration::from_millis(20));
if let Err(error) = stream
.write_all(b": keepalive\n\n")
.and_then(|()| stream.flush())
{
assert!(
matches!(
error.kind(),
ErrorKind::BrokenPipe
| ErrorKind::ConnectionReset
| ErrorKind::ConnectionAborted
| ErrorKind::NotConnected
),
"unexpected cancellable stream error: {error}"
);
stream_cancelled.store(true, Ordering::SeqCst);
return;
}
}
}