e2e testing

This commit is contained in:
2026-08-10 07:22:56 -05:00
parent fa43f723a5
commit 88ad290c7e
29 changed files with 1695 additions and 99 deletions
+64 -1
View File
@@ -294,8 +294,8 @@ impl AIAgentActionResultType {
| Self::StartAgent(_)
| Self::SendMessageToAgent(_)
| Self::AskUserQuestion(_)
| Self::RunAgents(_)
| Self::WaitForEvents(_) => self.to_string(),
Self::RunAgents(result) => result.model_content(),
}
}
}
@@ -1629,6 +1629,69 @@ pub enum RunAgentsAgentOutcomeKind {
Failed { error: String },
}
impl RunAgentsResult {
fn model_content(&self) -> String {
let value = match self {
Self::Launched {
model_id,
harness_type,
execution_mode,
agents,
} => {
let execution_mode = match execution_mode {
RunAgentsLaunchedExecutionMode::Local => serde_json::json!({
"type": "local",
}),
RunAgentsLaunchedExecutionMode::Remote {
environment_id,
worker_host,
computer_use_enabled,
} => serde_json::json!({
"type": "remote",
"environment_id": environment_id,
"worker_host": worker_host,
"computer_use_enabled": computer_use_enabled,
}),
};
let agents = agents
.iter()
.map(|agent| match &agent.kind {
RunAgentsAgentOutcomeKind::Launched { agent_id } => serde_json::json!({
"name": agent.name,
"status": "launched",
"agent_id": agent_id,
}),
RunAgentsAgentOutcomeKind::Failed { error } => serde_json::json!({
"name": agent.name,
"status": "failed",
"error": error,
}),
})
.collect::<Vec<_>>();
serde_json::json!({
"status": "launched",
"model_id": model_id,
"harness_type": harness_type,
"execution_mode": execution_mode,
"agents": agents,
})
}
Self::Denied { reason } => serde_json::json!({
"status": "denied",
"reason": reason,
}),
Self::Failure { error } => serde_json::json!({
"status": "failure",
"error": error,
}),
Self::Cancelled => serde_json::json!({
"status": "cancelled",
}),
};
value.to_string()
}
}
impl Display for RunAgentsResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
+91 -1
View File
@@ -1,4 +1,7 @@
use super::{StartAgentResult, StartAgentVersion};
use super::{
AIAgentActionResultType, RunAgentsAgentOutcome, RunAgentsAgentOutcomeKind,
RunAgentsLaunchedExecutionMode, RunAgentsResult, StartAgentResult, StartAgentVersion,
};
#[test]
fn deserializes_legacy_start_agent_success_without_version_as_v1() {
@@ -42,3 +45,90 @@ fn deserializes_legacy_start_agent_cancelled_without_version_as_v1() {
}
);
}
#[test]
fn run_agents_model_content_contains_resolved_config_and_agent_outcomes() {
let result = AIAgentActionResultType::RunAgents(RunAgentsResult::Launched {
model_id: "resolved-model".to_string(),
harness_type: "oz".to_string(),
execution_mode: RunAgentsLaunchedExecutionMode::Remote {
environment_id: "env-1".to_string(),
worker_host: "worker.example".to_string(),
computer_use_enabled: true,
},
agents: vec![
RunAgentsAgentOutcome {
name: "research".to_string(),
kind: RunAgentsAgentOutcomeKind::Launched {
agent_id: "agent-1".to_string(),
},
},
RunAgentsAgentOutcome {
name: "tests".to_string(),
kind: RunAgentsAgentOutcomeKind::Failed {
error: "capacity exhausted".to_string(),
},
},
],
});
let content: serde_json::Value = serde_json::from_str(&result.model_content())
.expect("run-agents model content should be valid JSON");
assert_eq!(
content,
serde_json::json!({
"status": "launched",
"model_id": "resolved-model",
"harness_type": "oz",
"execution_mode": {
"type": "remote",
"environment_id": "env-1",
"worker_host": "worker.example",
"computer_use_enabled": true,
},
"agents": [
{
"name": "research",
"status": "launched",
"agent_id": "agent-1",
},
{
"name": "tests",
"status": "failed",
"error": "capacity exhausted",
},
],
})
);
assert_eq!(
result.to_string(),
"Orchestrate launched (1/2 agents started)"
);
}
#[test]
fn run_agents_model_content_serializes_terminal_non_launch_outcomes() {
for (result, expected) in [
(
RunAgentsResult::Denied {
reason: "not approved".to_string(),
},
serde_json::json!({ "status": "denied", "reason": "not approved" }),
),
(
RunAgentsResult::Failure {
error: "invalid request".to_string(),
},
serde_json::json!({ "status": "failure", "error": "invalid request" }),
),
(
RunAgentsResult::Cancelled,
serde_json::json!({ "status": "cancelled" }),
),
] {
let result = AIAgentActionResultType::RunAgents(result);
let content: serde_json::Value = serde_json::from_str(&result.model_content())
.expect("run-agents model content should be valid JSON");
assert_eq!(content, expected);
}
}
@@ -0,0 +1,69 @@
use std::fs::File;
use std::io::BufReader;
use std::path::{Path, PathBuf};
use rig_core::client::CompletionClient;
use rig_core::completion::CompletionModel;
use rig_core::providers::chatgpt::{self, ChatGPTAuth};
fn codex_auth_path() -> PathBuf {
if let Some(codex_home) = std::env::var_os("CODEX_HOME") {
return PathBuf::from(codex_home).join("auth.json");
}
let home = std::env::var_os("HOME").expect("HOME must be set to locate ~/.codex/auth.json");
PathBuf::from(home).join(".codex").join("auth.json")
}
fn load_codex_auth(path: &Path) -> ChatGPTAuth {
let file = File::open(path)
.unwrap_or_else(|error| panic!("failed to open {}: {error}", path.display()));
let document: serde_json::Value = serde_json::from_reader(BufReader::new(file))
.unwrap_or_else(|error| panic!("failed to parse {}: {error}", path.display()));
let tokens = document
.get("tokens")
.unwrap_or_else(|| panic!("{} does not contain a tokens object", path.display()));
let access_token = tokens
.get("access_token")
.and_then(serde_json::Value::as_str)
.filter(|token| !token.is_empty())
.unwrap_or_else(|| panic!("{} does not contain an access token", path.display()));
let account_id = tokens
.get("account_id")
.and_then(serde_json::Value::as_str)
.filter(|account_id| !account_id.is_empty())
.map(str::to_string);
ChatGPTAuth::AccessToken {
access_token: access_token.to_string(),
account_id,
}
}
#[tokio::test(flavor = "current_thread")]
#[ignore = "makes a live ChatGPT request using local Codex credentials"]
async fn live_chatgpt_backend_via_rig_records_full_response() {
let auth_path = codex_auth_path();
let client = chatgpt::Client::builder()
.api_key(load_codex_auth(&auth_path))
.allow_device_flow(false)
.build()
.expect("Rig ChatGPT client should build");
let model_id =
std::env::var("GALAXY_CHATGPT_LIVE_MODEL").unwrap_or_else(|_| chatgpt::GPT_5_4.to_string());
let prompt = std::env::var("GALAXY_CHATGPT_LIVE_PROMPT").unwrap_or_else(|_| {
"Reply with exactly two short sentences explaining what a live backend smoke test verifies."
.to_string()
});
let model = client.completion_model(&model_id);
let request = model.completion_request(prompt).build();
let response = model
.completion(request)
.await
.expect("live ChatGPT completion should succeed");
let recorded = serde_json::to_string_pretty(&response)
.expect("the normalized Rig response should serialize");
println!("CHATGPT_LIVE_RESPONSE_BEGIN\n{recorded}\nCHATGPT_LIVE_RESPONSE_END");
}
@@ -435,10 +435,12 @@ fn register_tests() -> HashMap<&'static str, BoxedBuilderFn> {
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_long_running_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_rig_local_run_agents_round_trip);
register_test!(test_git_prompt_chips);
// These tests are only invoked manually, and not included in the
+407 -25
View File
@@ -1,7 +1,7 @@
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::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;
@@ -9,10 +9,12 @@ use std::time::Duration;
use galaxyui_core::async_assert;
use warp::features::FeatureFlag;
use warp::integration_testing::agent_mode::{
assert_latest_exchange_text, assert_task_is_cancelled, enter_agent_view,
assert_any_exchange_text, assert_latest_exchange_text,
assert_single_hidden_child_agent_succeeds, 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,
set_execution_profile_auto_execute_mcp_tools, set_execution_profile_auto_run_agents,
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,
};
@@ -35,6 +37,11 @@ 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 LONG_RUNNING_CALL_ID: &str = "rig-long-running-call";
const LONG_RUNNING_POLL_CALL_ID: &str = "rig-long-running-poll-call";
const LONG_RUNNING_START_OUTPUT: &str = "rig-long-running-start";
const LONG_RUNNING_COMPLETE_OUTPUT: &str = "rig-long-running-complete";
const LONG_RUNNING_FINAL_TEXT: &str = "Rig long-running shell round trip completed.";
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";
@@ -46,6 +53,12 @@ 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.";
const RUN_AGENTS_CALL_ID: &str = "rig-run-agents-call";
const RUN_AGENTS_CHILD_NAME: &str = "rig-child";
const RUN_AGENTS_CHILD_PROMPT: &str =
"Return the deterministic Rig child completion marker without calling tools.";
const RUN_AGENTS_CHILD_OUTPUT: &str = "Rig child agent completed.";
const RUN_AGENTS_FINAL_TEXT: &str = "Rig local orchestration round trip completed.";
#[derive(Clone)]
enum MockScenario {
@@ -54,6 +67,7 @@ enum MockScenario {
},
ShellSuccess,
ShellFailure,
ShellLongRunning,
ShellDenied {
marker_path: Arc<Mutex<String>>,
},
@@ -65,6 +79,7 @@ enum MockScenario {
stream_cancelled: Arc<AtomicBool>,
},
Mcp,
LocalRunAgents,
}
pub fn test_rig_read_tool_round_trip() -> Builder {
@@ -124,6 +139,25 @@ pub fn test_rig_shell_tool_failure_round_trip() -> Builder {
)
}
pub fn test_rig_shell_long_running_round_trip() -> Builder {
rig_builder(MockScenario::ShellLongRunning)
.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 long-running shell monitor check.",
Duration::from_secs(60),
))
.with_step(
new_step_with_default_assertions("Assert Rig long-running shell reached Agent Mode")
.add_named_assertion(
"Final response follows the completed shell poll",
assert_any_exchange_text(|text| text.contains(LONG_RUNNING_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 })
@@ -258,6 +292,32 @@ pub fn test_rig_mcp_tool_round_trip() -> Builder {
)
}
pub fn test_rig_local_run_agents_round_trip() -> Builder {
rig_builder(MockScenario::LocalRunAgents)
.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_run_agents())
.with_step(enter_agent_view())
.with_step(submit_ai_query_and_wait_until_done(
"Launch the deterministic local child agent.",
Duration::from_secs(90),
))
.with_step(
new_step_with_default_assertions("Assert Rig local orchestration reached Agent Mode")
.add_named_assertion(
"Final response follows the structured run-agents result",
assert_latest_exchange_text(|text| text.contains(RUN_AGENTS_FINAL_TEXT)),
)
.add_named_assertion(
"Hidden child completed as a leaf worker",
assert_single_hidden_child_agent_succeeds(
RUN_AGENTS_CHILD_NAME,
RUN_AGENTS_CHILD_OUTPUT,
),
),
)
}
fn rig_builder(scenario: MockScenario) -> Builder {
FeatureFlag::AgentView.set_enabled(true);
FeatureFlag::MCPGroupedServerContext.set_enabled(true);
@@ -293,8 +353,10 @@ fn rig_builder(scenario: MockScenario) -> Builder {
}
MockScenario::ShellSuccess
| MockScenario::ShellFailure
| MockScenario::ShellLongRunning
| MockScenario::Cancellation { .. }
| MockScenario::Mcp => {}
| MockScenario::Mcp
| MockScenario::LocalRunAgents => {}
}
})
.with_cleanup(move |_utils| {
@@ -337,12 +399,11 @@ fn start_mock_provider(
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, &scenario, &request_count, &stop);
serve_request(&mut stream, &scenario, &stop);
}
Err(error) if error.kind() == ErrorKind::WouldBlock => {
thread::sleep(Duration::from_millis(10));
@@ -354,12 +415,7 @@ fn start_mock_provider(
(address, thread)
}
fn serve_request(
stream: &mut TcpStream,
scenario: &MockScenario,
request_count: &AtomicUsize,
stop: &AtomicBool,
) {
fn serve_request(stream: &mut TcpStream, scenario: &MockScenario, stop: &AtomicBool) {
stream
.set_read_timeout(Some(Duration::from_secs(5)))
.expect("should set request timeout");
@@ -377,27 +433,137 @@ fn serve_request(
request_line.contains("/chat/completions"),
"unexpected mock provider request: {request_line}"
);
let turn = request_count.fetch_add(1, Ordering::SeqCst);
let request_body = parse_request_body(&request);
if matches!(scenario, MockScenario::ShellLongRunning) {
let body = long_running_shell_sse(&request_body);
write_response(stream, "text/event-stream", &body);
return;
}
if matches!(scenario, MockScenario::LocalRunAgents) {
let body = local_run_agents_sse(&request_body);
write_response(stream, "text/event-stream", &body);
return;
}
if let MockScenario::Cancellation {
stream_started,
stream_cancelled,
} = scenario
{
assert_eq!(turn, 0, "unexpected extra cancellation chat request");
assert!(
!request_has_tool_result(&request_body),
"cancellation scenario should not issue a follow-up tool result"
);
write_cancellable_response(stream, stream_started, stream_cancelled, stop);
return;
}
let body = match turn {
0 => tool_call_sse(scenario),
1 => {
assert_follow_up_request(scenario, &request);
final_text_sse(final_text(scenario))
}
_ => panic!("unexpected extra chat completion request"),
let call_id = scenario_call_id(scenario).expect("non-cancellation scenario should call a tool");
let body = if let Some(content) = tool_result_content(&request_body, call_id) {
assert!(
!content.is_empty(),
"follow-up tool result should contain model-facing content"
);
assert_follow_up_request(scenario, &request);
final_text_sse(final_text(scenario))
} else {
assert!(
!request_has_tool_result(&request_body),
"unexpected tool result in initial scenario request"
);
let expected_tool = scenario_tool_name(scenario)
.expect("non-cancellation scenario should advertise its tool");
assert!(
advertised_tool_names(&request_body).contains(&expected_tool),
"initial request should advertise {expected_tool}"
);
tool_call_sse(scenario)
};
write_response(stream, "text/event-stream", &body);
}
fn parse_request_body(request: &str) -> serde_json::Value {
let (_, body) = request
.split_once("\r\n\r\n")
.expect("provider request should contain an HTTP body");
serde_json::from_str(body).expect("provider request body should be valid JSON")
}
fn request_has_tool_result(request: &serde_json::Value) -> bool {
request
.get("messages")
.and_then(serde_json::Value::as_array)
.is_some_and(|messages| {
messages.iter().any(|message| {
message.get("role").and_then(serde_json::Value::as_str) == Some("tool")
&& message.get("tool_call_id").is_some()
})
})
}
fn tool_result_content(request: &serde_json::Value, call_id: &str) -> Option<String> {
request
.get("messages")?
.as_array()?
.iter()
.find(|message| {
message.get("role").and_then(serde_json::Value::as_str) == Some("tool")
&& message
.get("tool_call_id")
.and_then(serde_json::Value::as_str)
== Some(call_id)
})
.and_then(|message| message.get("content"))
.map(|content| {
content
.as_str()
.map(ToOwned::to_owned)
.unwrap_or_else(|| content.to_string())
})
}
fn scenario_call_id(scenario: &MockScenario) -> Option<&'static str> {
match scenario {
MockScenario::Read { .. } => Some(READ_CALL_ID),
MockScenario::ShellSuccess
| MockScenario::ShellFailure
| MockScenario::ShellDenied { .. } => Some(SHELL_CALL_ID),
MockScenario::ShellLongRunning => Some(LONG_RUNNING_CALL_ID),
MockScenario::Edit { .. } => Some(EDIT_CALL_ID),
MockScenario::Cancellation { .. } => None,
MockScenario::Mcp => Some(MCP_CALL_ID),
MockScenario::LocalRunAgents => Some(RUN_AGENTS_CALL_ID),
}
}
fn scenario_tool_name(scenario: &MockScenario) -> Option<&'static str> {
match scenario {
MockScenario::Read { .. } => Some("read_files"),
MockScenario::ShellSuccess
| MockScenario::ShellFailure
| MockScenario::ShellLongRunning
| MockScenario::ShellDenied { .. } => Some("run_shell_command"),
MockScenario::Edit { .. } => Some("apply_file_diffs"),
MockScenario::Cancellation { .. } => None,
MockScenario::Mcp => Some(MCP_TOOL_NAME),
MockScenario::LocalRunAgents => Some("run_agents"),
}
}
fn advertised_tool_names(request: &serde_json::Value) -> Vec<&str> {
request
.get("tools")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(|tool| {
tool.get("function")
.and_then(|function| function.get("name"))
.or_else(|| tool.get("name"))
.and_then(serde_json::Value::as_str)
})
.collect()
}
fn read_request(stream: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut chunk = [0; 8 * 1024];
@@ -453,6 +619,9 @@ fn tool_call_sse(scenario: &MockScenario) -> String {
"run_shell_command",
shell_arguments("(printf '%s\\n' 'rig-shell-failure-output' >&2; exit 7)"),
),
MockScenario::ShellLongRunning => {
unreachable!("long-running shell requests use request-aware routing")
}
MockScenario::ShellDenied { marker_path } => {
let marker_path = marker_path.lock().expect("marker path lock").clone();
(
@@ -486,7 +655,14 @@ fn tool_call_sse(scenario: &MockScenario) -> String {
MCP_TOOL_NAME,
serde_json::json!({"text": MCP_INPUT}),
),
MockScenario::LocalRunAgents => {
unreachable!("local orchestration requests use request-aware routing")
}
};
tool_call_sse_for(call_id, tool_name, arguments)
}
fn tool_call_sse_for(call_id: &str, tool_name: &str, arguments: serde_json::Value) -> String {
let tool_delta = serde_json::json!({
"id": "rig-integration-1",
"model": MODEL_ID,
@@ -519,13 +695,205 @@ fn tool_call_sse(scenario: &MockScenario) -> String {
format!("data: {tool_delta}\n\ndata: {tool_stop}\n\ndata: {usage}\n\ndata: [DONE]\n\n")
}
fn long_running_shell_sse(request: &serde_json::Value) -> String {
if let Some(poll_result) = tool_result_content(request, LONG_RUNNING_POLL_CALL_ID) {
assert!(
poll_result.contains(LONG_RUNNING_COMPLETE_OUTPUT),
"completed poll should contain the command's final output"
);
assert!(
poll_result.contains("exit code 0"),
"completed poll should contain the successful exit code"
);
let initial_result = tool_result_content(request, LONG_RUNNING_CALL_ID)
.expect("completed poll request should preserve the initial running snapshot");
let command_id = command_id_from_result(&initial_result);
let poll_arguments = tool_call_arguments(request, LONG_RUNNING_POLL_CALL_ID)
.expect("completed poll request should preserve the polling tool call");
assert_eq!(
poll_arguments
.get("command_id")
.and_then(serde_json::Value::as_str),
Some(command_id.as_str()),
"poll must reuse the dynamic command ID returned by Galaxy"
);
return final_text_sse(LONG_RUNNING_FINAL_TEXT);
}
if let Some(initial_result) = tool_result_content(request, LONG_RUNNING_CALL_ID) {
assert!(
initial_result.contains("Command is still running"),
"non-blocking command should first return a running snapshot"
);
assert!(
initial_result.contains(LONG_RUNNING_START_OUTPUT),
"running snapshot should contain real intermediate output"
);
let command_id = command_id_from_result(&initial_result);
assert_ne!(
command_id, LONG_RUNNING_CALL_ID,
"the command ID should be the real terminal block ID, not the tool call ID"
);
assert!(
advertised_tool_names(request).contains(&"read_shell_command_output"),
"running snapshot follow-up should advertise the shell polling tool"
);
return tool_call_sse_for(
LONG_RUNNING_POLL_CALL_ID,
"read_shell_command_output",
serde_json::json!({"command_id": command_id, "wait_seconds": 5}),
);
}
assert!(
!request_has_tool_result(request),
"initial long-running shell request should not contain tool results"
);
assert!(
advertised_tool_names(request).contains(&"run_shell_command"),
"initial long-running request should advertise the shell tool"
);
let command = format!(
"printf '%s\\n' '{LONG_RUNNING_START_OUTPUT}'; sleep 4; printf '%s\\n' '{LONG_RUNNING_COMPLETE_OUTPUT}'"
);
tool_call_sse_for(
LONG_RUNNING_CALL_ID,
"run_shell_command",
shell_arguments_with_wait(&command, false),
)
}
fn local_run_agents_sse(request: &serde_json::Value) -> String {
if let Some(result) = tool_result_content(request, RUN_AGENTS_CALL_ID) {
let result: serde_json::Value =
serde_json::from_str(&result).expect("run-agents result should be structured JSON");
assert_eq!(
result.get("status").and_then(serde_json::Value::as_str),
Some("launched")
);
assert_eq!(
result
.pointer("/execution_mode/type")
.and_then(serde_json::Value::as_str),
Some("local")
);
let agents = result
.get("agents")
.and_then(serde_json::Value::as_array)
.expect("run-agents result should contain agent outcomes");
assert_eq!(agents.len(), 1, "exactly one child should be launched");
assert_eq!(
agents[0].get("name").and_then(serde_json::Value::as_str),
Some(RUN_AGENTS_CHILD_NAME)
);
assert_eq!(
agents[0].get("status").and_then(serde_json::Value::as_str),
Some("launched")
);
assert!(
agents[0]
.get("agent_id")
.and_then(serde_json::Value::as_str)
.is_some_and(|id| !id.is_empty()),
"launched child should return its real conversation ID"
);
return final_text_sse(RUN_AGENTS_FINAL_TEXT);
}
if request_messages_contain(request, RUN_AGENTS_CHILD_PROMPT) {
assert!(
!request_has_tool_result(request),
"child's initial request should not contain tool results"
);
let tools = advertised_tool_names(request);
for delegation_tool in [
"run_agents",
"start_agent",
"send_message_to_agent",
"wait_for_events",
] {
assert!(
!tools.contains(&delegation_tool),
"leaf child must not advertise {delegation_tool}"
);
}
return final_text_sse(RUN_AGENTS_CHILD_OUTPUT);
}
assert!(
!request_has_tool_result(request),
"root's initial orchestration request should not contain tool results"
);
let tools = advertised_tool_names(request);
assert!(
tools.contains(&"run_agents"),
"root request should advertise modern run_agents"
);
assert!(
tools.contains(&"start_agent"),
"root request should retain legacy start_agent compatibility"
);
tool_call_sse_for(
RUN_AGENTS_CALL_ID,
"run_agents",
serde_json::json!({
"summary": "Launch one deterministic local child",
"base_prompt": "Complete the assigned task directly.",
"agent_run_configs": [{
"name": RUN_AGENTS_CHILD_NAME,
"prompt": RUN_AGENTS_CHILD_PROMPT,
"title": "Rig child agent"
}]
}),
)
}
fn request_messages_contain(request: &serde_json::Value, expected: &str) -> bool {
request
.get("messages")
.is_some_and(|messages| messages.to_string().contains(expected))
}
fn command_id_from_result(result: &str) -> String {
result
.lines()
.find_map(|line| line.strip_prefix("Command ID: "))
.map(str::trim)
.filter(|id| !id.is_empty())
.map(ToOwned::to_owned)
.expect("long-running result should include a dynamic command ID")
}
fn tool_call_arguments(request: &serde_json::Value, call_id: &str) -> Option<serde_json::Value> {
request
.get("messages")?
.as_array()?
.iter()
.filter(|message| {
message.get("role").and_then(serde_json::Value::as_str) == Some("assistant")
})
.filter_map(|message| message.get("tool_calls")?.as_array())
.flatten()
.find(|call| call.get("id").and_then(serde_json::Value::as_str) == Some(call_id))
.and_then(|call| call.get("function"))
.and_then(|function| function.get("arguments"))
.and_then(|arguments| match arguments {
serde_json::Value::String(arguments) => serde_json::from_str(arguments).ok(),
arguments => Some(arguments.clone()),
})
}
fn shell_arguments(command: &str) -> serde_json::Value {
shell_arguments_with_wait(command, true)
}
fn shell_arguments_with_wait(command: &str, wait_until_complete: bool) -> serde_json::Value {
serde_json::json!({
"command": command,
"is_read_only": false,
"is_risky": false,
"uses_pager": false,
"wait_until_complete": true,
"wait_until_complete": wait_until_complete,
})
}
@@ -573,6 +941,9 @@ fn assert_follow_up_request(scenario: &MockScenario, request: &str) {
"failed shell result should remain an explicit model error"
);
}
MockScenario::ShellLongRunning => {
unreachable!("long-running shell follow-ups use request-aware routing")
}
MockScenario::ShellDenied { marker_path } => {
assert!(
request.contains(SHELL_CALL_ID),
@@ -615,6 +986,9 @@ fn assert_follow_up_request(scenario: &MockScenario, request: &str) {
"follow-up request should contain the real MCP tool result"
);
}
MockScenario::LocalRunAgents => {
unreachable!("local orchestration follow-ups use request-aware routing")
}
}
}
@@ -623,12 +997,14 @@ fn final_text(scenario: &MockScenario) -> &'static str {
MockScenario::Read { .. } => READ_FINAL_TEXT,
MockScenario::ShellSuccess => SHELL_SUCCESS_FINAL_TEXT,
MockScenario::ShellFailure => SHELL_FAILURE_FINAL_TEXT,
MockScenario::ShellLongRunning => LONG_RUNNING_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,
MockScenario::LocalRunAgents => RUN_AGENTS_FINAL_TEXT,
}
}
@@ -638,15 +1014,21 @@ fn final_text_sse(final_text: &str) -> String {
"model": MODEL_ID,
"choices": [{
"delta": {"content": final_text, "tool_calls": []},
"finish_reason": "stop",
"finish_reason": null,
}],
"usage": null,
});
let stop = serde_json::json!({
"id": "rig-integration-2",
"model": MODEL_ID,
"choices": [{"delta": {"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")
format!("data: {text}\n\ndata: {stop}\n\ndata: {usage}\n\ndata: [DONE]\n\n")
}
fn write_response(stream: &mut TcpStream, content_type: &str, body: &str) {
@@ -314,10 +314,12 @@ integration_tests! {
test_rig_read_tool_round_trip,
test_rig_shell_tool_success_round_trip,
test_rig_shell_tool_failure_round_trip,
test_rig_shell_long_running_round_trip,
test_rig_shell_tool_permission_denial,
test_rig_edit_tool_round_trip,
test_rig_in_flight_cancellation,
test_rig_mcp_tool_round_trip,
test_rig_local_run_agents_round_trip,
test_rule_creation,
test_rule_update,