384 lines
16 KiB
Rust
384 lines
16 KiB
Rust
use std::collections::HashMap;
|
|
use std::fs::read;
|
|
use std::io::Cursor;
|
|
use std::path::Path;
|
|
use std::time::Duration;
|
|
|
|
use galaxyui::integration::TestStep;
|
|
use galaxyui::{async_assert, SingletonEntity};
|
|
use prost::Message;
|
|
|
|
use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel;
|
|
use crate::ai::execution_profiles::ActionPermission;
|
|
use crate::ai::llms::{LLMId, LLMPreferences};
|
|
use crate::ai::mcp::{
|
|
JsonTemplate, TemplatableMCPServer, TemplatableMCPServerInstallation,
|
|
TemplatableMCPServerManager,
|
|
};
|
|
use crate::integration_testing::agent_mode::{
|
|
assert_latest_task_succeeds_or_blocked, assert_task_is_blocked, ConversationTarget,
|
|
};
|
|
use crate::integration_testing::step::{
|
|
new_step_with_default_assertions, new_step_with_default_assertions_for_pane,
|
|
};
|
|
use crate::integration_testing::terminal::assert_input_is_focused;
|
|
use crate::integration_testing::view_getters::terminal_view;
|
|
|
|
pub const AGENT_MODE_RUNNING_STEP_GROUP_NAME: &str = "Agent mode running";
|
|
|
|
use super::hydrate_ai_conversation_assertion;
|
|
|
|
/// Assumes that the terminal input is currently not in AI input mode.
|
|
pub fn enter_agent_view() -> TestStep {
|
|
let keystroke = if cfg!(target_os = "macos") {
|
|
"cmd-enter"
|
|
} else {
|
|
"ctrl-shift-enter"
|
|
};
|
|
new_step_with_default_assertions("Enter Agent View")
|
|
.with_keystrokes(&[keystroke])
|
|
.add_named_assertion(
|
|
"Assert that we are in Agent View and AI input mode",
|
|
move |app, window_id| {
|
|
let terminal_view = terminal_view(app, window_id, 0, 0);
|
|
terminal_view.read(app, |terminal_view, app| {
|
|
let is_ai_input_mode = terminal_view
|
|
.input()
|
|
.read(app, |input, app| input.input_type(app).is_ai());
|
|
let agent_view_state = {
|
|
let model = terminal_view.model.lock();
|
|
model.block_list().agent_view_state().clone()
|
|
};
|
|
async_assert!(
|
|
is_ai_input_mode && agent_view_state.is_active(),
|
|
"Expected fullscreen Agent View + AI input mode, got agent_view_state={agent_view_state:?}, is_ai_input_mode={is_ai_input_mode}"
|
|
)
|
|
})
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Assumes that the terminal input is currently in AI input mode.
|
|
pub fn exit_agent_view() -> TestStep {
|
|
new_step_with_default_assertions("Exit Agent View")
|
|
.with_keystrokes(&["escape"])
|
|
.add_named_assertion(
|
|
"Assert that we exited Agent View and are not in AI input mode",
|
|
move |app, window_id| {
|
|
let terminal_view = terminal_view(app, window_id, 0, 0);
|
|
terminal_view.read(app, |terminal_view, app| {
|
|
let is_ai_input_mode = terminal_view
|
|
.input()
|
|
.read(app, |input, app| input.input_type(app).is_ai());
|
|
let agent_view_state = {
|
|
let model = terminal_view.model.lock();
|
|
model.block_list().agent_view_state().clone()
|
|
};
|
|
async_assert!(
|
|
!is_ai_input_mode && !agent_view_state.is_active(),
|
|
"Expected inactive Agent View + non-AI input mode, got agent_view_state={agent_view_state:?}, is_ai_input_mode={is_ai_input_mode}"
|
|
)
|
|
})
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Hydrates a conversation from a protobuf file.
|
|
/// File should be generated into the `input_data` directory.
|
|
/// See the agent_mode_eval README for more details.
|
|
pub fn hydrate_ai_conversation(file_name: &str) -> TestStep {
|
|
let file_bytes = get_input_data(file_name);
|
|
let Ok(request) = warp_multi_agent_api::Request::decode(file_bytes) else {
|
|
panic!("Failed to decode request from protobuf");
|
|
};
|
|
|
|
let tasks = request
|
|
.task_context
|
|
.map(|ctx| ctx.tasks)
|
|
.unwrap_or_default();
|
|
|
|
new_step_with_default_assertions("Hydrate AI conversation").add_named_assertion(
|
|
"Assert that conversation was hydrated successfully",
|
|
hydrate_ai_conversation_assertion(tasks),
|
|
)
|
|
}
|
|
|
|
/// Attach the latest block in the blocklist (command + output) to the AI query.
|
|
pub fn attach_recent_block_as_context() -> TestStep {
|
|
TestStep::new("Attach last block as context").add_named_assertion(
|
|
"Attach last block as context",
|
|
|app, window_id| {
|
|
let terminal_view = terminal_view(app, window_id, 0, 0);
|
|
terminal_view.update(app, |view, ctx| {
|
|
let last_index = {
|
|
let model = view.model.lock();
|
|
model.block_list().last_non_hidden_block_by_index()
|
|
};
|
|
if let Some(idx) = last_index {
|
|
view.integration_test_change_block_selection_to_single(idx, ctx);
|
|
}
|
|
});
|
|
|
|
terminal_view.read(app, |view, ctx| {
|
|
let count = view
|
|
.ai_context_model()
|
|
.as_ref(ctx)
|
|
.pending_context_block_ids()
|
|
.len();
|
|
async_assert!(
|
|
count == 1,
|
|
"Expected exactly 1 attached context block, got {count}"
|
|
)
|
|
})
|
|
},
|
|
)
|
|
}
|
|
|
|
// This will fail immediately on any error responses.
|
|
pub fn submit_ai_query_and_wait_until_done(query: &str, timeout: Duration) -> TestStep {
|
|
submit_ai_query(query, timeout)
|
|
.add_named_assertion(
|
|
"Assert the agent task is complete",
|
|
assert_latest_task_succeeds_or_blocked(ConversationTarget::Active, None),
|
|
)
|
|
.add_named_assertion(
|
|
"Assert that that input has been returned to the user",
|
|
assert_input_is_focused(),
|
|
)
|
|
}
|
|
|
|
/// Submits an AI query and waits until the task is blocked (waiting for user approval).
|
|
/// This is useful for tests where auto-execution is disabled and you want to verify
|
|
/// the command that would be executed without actually running it.
|
|
pub fn submit_ai_query_and_wait_until_blocked(query: &str, timeout: Duration) -> TestStep {
|
|
submit_ai_query(query, timeout).add_named_assertion(
|
|
"Assert the agent task is blocked",
|
|
assert_task_is_blocked(ConversationTarget::Active),
|
|
)
|
|
}
|
|
|
|
// Runs an AI query without waiting for anything.
|
|
// This is useful if you expect a specific sequence of responses (e.g. expect a certain command to be requested immediately),
|
|
// since it lets you make assertions on responses as they become ready and fail early instead of waiting for the agent to finish all its turns.
|
|
pub fn submit_ai_query(query: &str, timeout: Duration) -> TestStep {
|
|
new_step_with_default_assertions_for_pane(&format!("Enter AI query: {query}"), 0, 0)
|
|
.set_timeout(timeout)
|
|
.set_step_group_name(AGENT_MODE_RUNNING_STEP_GROUP_NAME)
|
|
.with_typed_characters(&[query])
|
|
.with_keystrokes(&["enter"])
|
|
.add_named_assertion(
|
|
"Print conversation ID to stdout",
|
|
print_conversation_id_assertion(),
|
|
)
|
|
}
|
|
|
|
/// Returns an assertion that prints the conversation ID to stdout once available.
|
|
/// This assertion will poll until the conversation token is received from the server.
|
|
fn print_conversation_id_assertion(
|
|
) -> impl FnMut(&mut galaxyui::App, galaxyui::WindowId) -> galaxyui::integration::AssertionOutcome {
|
|
|app, window_id| {
|
|
use galaxyui::integration::AssertionOutcome;
|
|
|
|
use crate::BlocklistAIHistoryModel;
|
|
let terminal_view = terminal_view(app, window_id, 0, 0);
|
|
BlocklistAIHistoryModel::handle(app).read(app, |history_model, _| {
|
|
if let Some(conversation) = history_model.active_conversation(terminal_view.id()) {
|
|
if let Some(token) = conversation.server_conversation_token() {
|
|
// The debug link within the container will be using host.docker.internal, but we're opening
|
|
// from outside the container.
|
|
let debug_link = token
|
|
.debug_link()
|
|
.replace("host.docker.internal", "localhost");
|
|
println!("Conversation ID (debug link): {debug_link}");
|
|
return AssertionOutcome::Success;
|
|
}
|
|
}
|
|
// If we don't have a conversation token yet, keep polling
|
|
AssertionOutcome::failure("Waiting for conversation token to be available".to_string())
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Sets the preferred agent mode LLM. This is the base model for agent and inline AI conversations.
|
|
pub fn set_preferred_agent_mode_llm(llm_id: &str) -> TestStep {
|
|
let llm_id = LLMId::from(llm_id);
|
|
TestStep::new(&format!("Set preferred agent mode LLM to {llm_id}")).add_named_assertion(
|
|
"Update preferred agent mode LLM",
|
|
move |app, window_id| {
|
|
let llm_id = llm_id.clone();
|
|
let terminal_view_id = terminal_view(app, window_id, 0, 0).id();
|
|
LLMPreferences::handle(app).update(app, |llm_preferences, ctx| {
|
|
// Validate that the LLM ID is actually available. We only do this
|
|
// for the base model, since the coding and planning models are
|
|
// currently unused in the product.
|
|
assert!(
|
|
llm_preferences.is_available_agent_mode_llm(&llm_id),
|
|
"LLM ID '{llm_id}' is not a valid agent mode LLM",
|
|
);
|
|
llm_preferences.update_preferred_agent_mode_llm(&llm_id, terminal_view_id, ctx);
|
|
});
|
|
async_assert!(true, "Successfully updated preferred agent mode LLM")
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Sets the preferred coding LLM. Note that the server currently ignores this.
|
|
pub fn set_preferred_coding_llm(llm_id: &str) -> TestStep {
|
|
let llm_id = LLMId::from(llm_id);
|
|
TestStep::new(&format!("Set preferred coding LLM to {llm_id}")).add_named_assertion(
|
|
"Update preferred coding LLM",
|
|
move |app, window_id| {
|
|
let llm_id = llm_id.clone();
|
|
let terminal_view_id = terminal_view(app, window_id, 0, 0).id();
|
|
LLMPreferences::handle(app).update(app, |llm_preferences, ctx| {
|
|
llm_preferences.update_preferred_coding_llm(&llm_id, Some(terminal_view_id), ctx);
|
|
});
|
|
async_assert!(true, "Successfully updated preferred coding LLM")
|
|
},
|
|
)
|
|
}
|
|
|
|
fn get_input_data(file_name: &str) -> Cursor<Vec<u8>> {
|
|
let input_data_dir = std::env::var("INPUT_DATA_DIR").expect(
|
|
"INPUT_DATA_DIR is not set. This is needed to hydrate conversations from eval tests.",
|
|
);
|
|
let path = Path::new(&input_data_dir).join(file_name);
|
|
Cursor::new(read(&path).expect("Failed to read binary input data"))
|
|
}
|
|
|
|
/// Sets the execution profile to not auto-execute commands.
|
|
/// This changes the `execute_commands` permission from `AlwaysAllow` to `AlwaysAsk`,
|
|
/// which means commands will be proposed but not automatically executed.
|
|
pub fn set_execution_profile_no_auto_execute() -> TestStep {
|
|
TestStep::new("Set execution profile to not auto-execute commands").add_named_assertion(
|
|
"Update execution profile",
|
|
|app, _window_id| {
|
|
AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| {
|
|
let default_profile_id = *profiles.default_profile(ctx).id();
|
|
profiles.set_execute_commands(
|
|
default_profile_id,
|
|
&ActionPermission::AlwaysAsk,
|
|
ctx,
|
|
);
|
|
});
|
|
async_assert!(true, "Successfully updated execution profile")
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Sets the execution profile to auto-execute commands.
|
|
pub fn set_execution_profile_auto_execute() -> TestStep {
|
|
TestStep::new("Set execution profile to auto-execute commands").add_named_assertion(
|
|
"Update execution profile",
|
|
|app, _window_id| {
|
|
AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| {
|
|
let default_profile_id = *profiles.default_profile(ctx).id();
|
|
profiles.set_execute_commands(
|
|
default_profile_id,
|
|
&ActionPermission::AlwaysAllow,
|
|
ctx,
|
|
);
|
|
});
|
|
async_assert!(true, "Successfully updated execution profile")
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Sets the execution profile to auto-apply code diffs.
|
|
pub fn set_execution_profile_auto_apply_code_diffs() -> TestStep {
|
|
TestStep::new("Set execution profile to auto-apply code diffs").add_named_assertion(
|
|
"Update execution profile",
|
|
|app, _window_id| {
|
|
AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| {
|
|
let default_profile_id = *profiles.default_profile(ctx).id();
|
|
profiles.set_apply_code_diffs(
|
|
default_profile_id,
|
|
&ActionPermission::AlwaysAllow,
|
|
ctx,
|
|
);
|
|
});
|
|
async_assert!(true, "Successfully updated execution profile")
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Sets the execution profile to auto-execute MCP tools.
|
|
pub fn set_execution_profile_auto_execute_mcp_tools() -> TestStep {
|
|
TestStep::new("Set execution profile to auto-execute MCP tools").add_named_assertion(
|
|
"Update execution profile",
|
|
|app, _window_id| {
|
|
AIExecutionProfilesModel::handle(app).update(app, |profiles, ctx| {
|
|
let default_profile_id = *profiles.default_profile(ctx).id();
|
|
profiles.set_mcp_permissions(
|
|
default_profile_id,
|
|
&ActionPermission::AlwaysAllow,
|
|
ctx,
|
|
);
|
|
});
|
|
async_assert!(true, "Successfully updated execution profile")
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Starts an ephemeral MCP stdio server for an integration test.
|
|
pub fn start_ephemeral_mcp_server_for_testing(
|
|
command: String,
|
|
argument: String,
|
|
installation_id: &str,
|
|
template_id: &str,
|
|
server_name: &str,
|
|
) -> TestStep {
|
|
let installation_id =
|
|
uuid::Uuid::parse_str(installation_id).expect("valid MCP installation UUID");
|
|
let template_id = uuid::Uuid::parse_str(template_id).expect("valid MCP template UUID");
|
|
let server_name = server_name.to_string();
|
|
new_step_with_default_assertions("Start ephemeral integration MCP server").add_named_assertion(
|
|
"Spawn MCP server",
|
|
move |app, _window_id| {
|
|
let template_json = serde_json::json!({
|
|
(&server_name): {
|
|
"command": command,
|
|
"args": [argument],
|
|
},
|
|
})
|
|
.to_string();
|
|
let server = TemplatableMCPServer {
|
|
uuid: template_id,
|
|
name: server_name.clone(),
|
|
description: Some("Hermetic MCP server for integration testing".to_string()),
|
|
template: JsonTemplate {
|
|
json: template_json,
|
|
variables: Vec::new(),
|
|
},
|
|
version: 1,
|
|
gallery_data: None,
|
|
};
|
|
let installation =
|
|
TemplatableMCPServerInstallation::new(installation_id, server, HashMap::new());
|
|
TemplatableMCPServerManager::handle(app).update(app, |manager, ctx| {
|
|
manager.spawn_cli_ephemeral_server(installation, ctx);
|
|
});
|
|
async_assert!(true, "Ephemeral MCP server spawn requested")
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Waits for an ephemeral MCP server to complete its protocol handshake.
|
|
pub fn wait_until_mcp_server_is_active_for_testing(
|
|
installation_id: &str,
|
|
timeout: Duration,
|
|
) -> TestStep {
|
|
let installation_id =
|
|
uuid::Uuid::parse_str(installation_id).expect("valid MCP installation UUID");
|
|
new_step_with_default_assertions("Wait for ephemeral integration MCP server")
|
|
.set_timeout(timeout)
|
|
.add_named_assertion("MCP server is active", move |app, _window_id| {
|
|
let is_active = TemplatableMCPServerManager::handle(app)
|
|
.read(app, |manager, _| manager.is_server_active(installation_id));
|
|
async_assert!(
|
|
is_active,
|
|
"Waiting for ephemeral MCP server to become active"
|
|
)
|
|
})
|
|
}
|