From f6c6434e5e86de058dd685c743e679b63615caef Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Thu, 22 Oct 2026 14:25:21 +0000 Subject: [PATCH] Fix LiteLLM context budgets and direct computer use Keep efficient context windows separate from expanded maxima throughout discovery, display and provider budgets. Expose correlated computer-use requests and actions with existing permissions, preserve screenshots as images, and retain them across checkpoints, compaction and explicit history recall. Validation: cargo build on Linux passed. Added focused regression coverage; test execution and native GUI verification remain outside this CLI build-only workflow. --- AGENTS.md | 14 +- app/src/ai/agent/api/impl.rs | 8 + app/src/ai/agent/api/impl_tests.rs | 31 ++++ app/src/ai/agent/provider_history.rs | 75 +++++++++ app/src/ai/agent/provider_history_tests.rs | 23 +++ app/src/ai/blocklist/action_model.rs | 55 ++++++- .../action_model/execute/use_computer.rs | 6 +- app/src/ai/blocklist/action_model_tests.rs | 20 +++ app/src/ai/blocklist/controller.rs | 6 +- app/src/ai/llms.rs | 115 ++++++++++---- app/src/ai/llms_tests.rs | 150 ++++++++++++++++++ app/src/ai/openai/response_translator.rs | 2 + app/src/ai/provider/request_translator.rs | 71 +++++++++ app/src/ai/provider/response_translator.rs | 10 ++ app/src/ai/runtime/event_translator_tests.rs | 1 + app/src/ai/runtime/mod.rs | 1 + .../runtime/provider_run_coordinator_tests.rs | 2 + app/src/ai/runtime/rig.rs | 27 ++-- app/src/ai/runtime/rig_request.rs | 18 +-- app/src/ai/runtime/rig_request_tests.rs | 2 + app/src/ai/runtime/rig_tests.rs | 21 +++ app/src/ai/runtime/rig_tool.rs | 124 +++++++++++++++ app/src/ai/runtime/rig_tool_tests.rs | 49 ++++++ app/src/ai/runtime/task_progress_tests.rs | 1 + crates/galaxy_agent_core/src/provider_run.rs | 63 +++++--- .../src/provider_run_tests.rs | 125 ++++++++++++++- crates/galaxy_agent_core/src/tool_policy.rs | 96 +++++++++-- .../src/tool_policy_tests.rs | 114 +++++++++++++ crates/galaxy_agent_core/src/types.rs | 44 +++++ crates/galaxy_agent_core/src/types_tests.rs | 3 + 30 files changed, 1176 insertions(+), 101 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e043e826..0ed85621 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -174,9 +174,10 @@ base_url = "http://localhost:4000/v1" api_key = "sk-..." [[ai.providers.models]] -model_id = "claude-sonnet-4-20250514[1m]" -display_name = "Claude Sonnet 4 (1M)" +model_id = "claude-sonnet-4-20250514" +display_name = "Claude Sonnet 4" context_size = 1000000 +default_context_size = 200000 [[ai.providers]] name = "Ollama (Local)" @@ -189,10 +190,9 @@ context_size = 128000 ``` **OpenAI/LiteLLM model discovery**: -- Models can be auto-fetched from the `/models` endpoint via the Settings > OpenAI / LiteLLM page -- For each model, the system probes `{model_id}[1m]` with a minimal chat completion request -- If the `[1m]` variant is accepted (HTTP 200 or 429), it's used with 1M context window -- Otherwise, the base model ID is used with its reported context size +- Models can be auto-fetched from `/models` and enriched with LiteLLM `/model/info` via the Settings > OpenAI / LiteLLM page. +- Preserve the efficient/default context window separately from the provider's maximum. Ordinary requests, compaction, and context displays use the lower positive default/input/maximum budget, including any explicit profile cap; availability of a larger maximum must never expand that budget automatically. +- Retain lower budgets when merging metadata from both endpoints, apply `effective_context_window_percent` once, and preserve the maximum as metadata. Discovery must not probe or select an extended `[1m]` variant automatically. - Models injected via `ai.providers[]` are routed to their specific endpoint (per-model routing map) - Provider name shown as the description label in the model picker; icon shows OpenAI logo for all OpenAI-compatible providers @@ -211,6 +211,8 @@ Key invariants: - A restored `AwaitingModel` checkpoint has an uncertain remote outcome and must terminate as an explicit restore failure rather than replaying the call; known recoverable failures observed in-process retain the bounded model-retry lifecycle - Same-conversation direct-provider follow-ups queue behind the cancelling generation; the old run keeps the active slot until terminal projection and cleanup, queued intent is checkpointed without credentials for restart recovery, and queued-only restore validates provider ownership and terminalizes the abandoned unprepared exchange from its persisted projection/stream identity before starting the successor; the next generation rebuilds provider history after cleanup so it includes the old generation's final committed output, and stale callbacks are ignored by stream identity - Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs` +- Direct-provider computer use starts with `request_computer_use`; the computer-use profile exposes `use_computer`, including when its model falls back to the main model. Tool availability and every action remain subject to session capabilities and the existing profile permission checks. +- Computer-use screenshots must reach providers as image parts, survive checkpoints and tool-result archival, and remain correlated with their tool call on explicit history recall. Emit every result in a tool batch before its screenshot image groups, and estimate image tokens without counting encoded image bytes as text. - Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose - Direct-provider normal and plan turns expose `update_task_list` for explicit progress updates. Keep the validated current list in durable `ProviderRun` state, inject it into each request, preserve it across compaction and restore, and reconcile restored UI updates idempotently. Finishing a response must never mark unfinished work complete. CLI monitors and completed-command assessments retain their focused tool sets. - Completed local plans are archived intact into `Plans/Archive` after seven days only when their conversation has explicitly completed all checklist items and has no active provider run or child. Draft, unfinished, or reopened plans stay visible; resumed work restores an automatically archived plan, and manual recovery is respected. Archival never deletes plan contents. diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index 00b0d802..df4c1cb2 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -28,12 +28,17 @@ pub(crate) fn prepare_direct_provider_params( .clone() .unwrap_or_else(|| get_supported_tools(params)); remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled); + // Computer use controls this machine, even when the terminal is connected elsewhere. + // Only expose the initial request here; the dedicated profile owns subsequent actions. + supported_tools.retain(|tool| *tool != api::ToolType::UseComputer + && (*tool != api::ToolType::RequestComputerUse || params.computer_use_enabled)); let mut supported_cli_agent_tools = supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(params)); remove_orchestration_tools_if_disabled( &mut supported_cli_agent_tools, params.orchestration_enabled, ); + supported_cli_agent_tools.retain(|tool| !matches!(tool, api::ToolType::UseComputer | api::ToolType::RequestComputerUse)); if params.permission_mode == crate::ai::execution_profiles::PermissionMode::ReadOnly { let permitted = |tool: &api::ToolType| is_read_only_tool(tool) || (!params.command_allowlist.is_empty() && matches!(tool, api::ToolType::RunShellCommand | api::ToolType::ReadShellCommandOutput | api::ToolType::WriteToLongRunningShellCommand)); @@ -73,6 +78,9 @@ fn get_supported_tools(params: &RequestParams) -> Vec { api::ToolType::CreateDocuments, api::ToolType::EditDocuments, ]; + if params.computer_use_enabled { + supported_tools.push(api::ToolType::RequestComputerUse); + } if params.orchestration_enabled { supported_tools.push(api::ToolType::Subagent); diff --git a/app/src/ai/agent/api/impl_tests.rs b/app/src/ai/agent/api/impl_tests.rs index 00b13989..2bfef573 100644 --- a/app/src/ai/agent/api/impl_tests.rs +++ b/app/src/ai/agent/api/impl_tests.rs @@ -91,6 +91,37 @@ fn supported_tools_include_plan_document_capabilities() { assert!(supported_tools.contains(&api::ToolType::EditDocuments)); } +#[test] +fn computer_use_is_explicitly_enabled_and_kept_out_of_command_monitor_tools() { + let mut params = request_params_with_ask_user_question_enabled(false); + params.computer_use_enabled = true; + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + assert!(normal.contains(&api::ToolType::RequestComputerUse)); + assert!(!normal.contains(&api::ToolType::UseComputer)); + assert!(!cli.contains(&api::ToolType::RequestComputerUse)); + assert!(!cli.contains(&api::ToolType::UseComputer)); + + params.computer_use_enabled = false; + let (normal, _) = super::prepare_direct_provider_params(&mut params); + assert!(!normal.contains(&api::ToolType::RequestComputerUse)); +} + +#[test] +fn computer_use_never_expands_explicit_tool_overrides_or_read_only_permissions() { + let mut params = request_params_with_ask_user_question_enabled(false); + params.computer_use_enabled = true; + params.supported_tools_override = Some(Vec::new()); + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + assert!(normal.is_empty()); + assert!(cli.is_empty()); + + params.permission_mode = crate::ai::execution_profiles::PermissionMode::ReadOnly; + params.supported_tools_override = Some(vec![api::ToolType::RequestComputerUse, api::ToolType::UseComputer]); + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + assert!(normal.is_empty()); + assert!(cli.is_empty()); +} + #[test] fn supported_tools_omit_subagents_when_orchestration_is_disabled() { let params = request_params_with_ask_user_question_enabled(false); diff --git a/app/src/ai/agent/provider_history.rs b/app/src/ai/agent/provider_history.rs index 4371903e..3b91c41d 100644 --- a/app/src/ai/agent/provider_history.rs +++ b/app/src/ai/agent/provider_history.rs @@ -8,6 +8,81 @@ use crate::ai::runtime::ProviderActionContext; const TOOL_PREFIX: &str = "galaxy:provider-tool:v1:"; +pub(crate) fn computer_use_tool(action: &super::AIAgentActionType) -> Option { + use api::message::tool_call::{self as tool, use_computer::action as action_api}; + use computer_use::{Action, Key, MouseButton, ScrollDirection, ScrollDistance}; + let coordinates = |point: computer_use::Vector2I| Some(api::Coordinates { x: point.x(), y: point.y() }); + let button = |button: &MouseButton| match button { + MouseButton::Left => action_api::MouseButton::Left as i32, + MouseButton::Right => action_api::MouseButton::Right as i32, + MouseButton::Middle => action_api::MouseButton::Middle as i32, + MouseButton::Back => action_api::MouseButton::Back as i32, + MouseButton::Forward => action_api::MouseButton::Forward as i32, + }; + let key = |key: &Key| Some(action_api::Key { data: Some(match key { + Key::Char(ch) => action_api::key::Data::Char(ch.to_string()), + Key::Keycode(code) => action_api::key::Data::Keycode(*code), + }) }); + match action { + super::AIAgentActionType::RequestComputerUse(request) => Some(tool::Tool::RequestComputerUse(tool::RequestComputerUse { + task_summary: request.task_summary.clone(), + screenshot_params: request.screenshot_params.map(computer_screenshot), + })), + super::AIAgentActionType::UseComputer(request) => Some(tool::Tool::UseComputer(tool::UseComputer { + action_summary: request.action_summary.clone(), + post_actions_screenshot_params: request.screenshot_params.map(computer_screenshot), + actions: request.actions.iter().map(|targeted| tool::use_computer::Action { + target: computer_target(targeted.target), + r#type: Some(match &targeted.action { + Action::MouseMove { to } => action_api::Type::MouseMove(action_api::MouseMove { to: coordinates(*to) }), + Action::MouseDown { button: b, at } => action_api::Type::MouseDown(action_api::MouseDown { button: button(b), at: coordinates(*at) }), + Action::MouseUp { button: b } => action_api::Type::MouseUp(action_api::MouseUp { button: button(b) }), + Action::MouseWheel { at, direction, distance } => action_api::Type::MouseWheel(action_api::MouseWheel { + at: coordinates(*at), + direction: match direction { + ScrollDirection::Up => action_api::mouse_wheel::Direction::Up as i32, + ScrollDirection::Down => action_api::mouse_wheel::Direction::Down as i32, + ScrollDirection::Left => action_api::mouse_wheel::Direction::Left as i32, + ScrollDirection::Right => action_api::mouse_wheel::Direction::Right as i32, + }, + distance: Some(match distance { + ScrollDistance::Pixels(value) => action_api::mouse_wheel::Distance::Pixels(*value), + ScrollDistance::Clicks(value) => action_api::mouse_wheel::Distance::Clicks(*value), + }), + }), + Action::TypeText { text } => action_api::Type::TypeText(action_api::TypeText { text: text.clone() }), + Action::KeyDown { key: k } => action_api::Type::KeyDown(action_api::KeyDown { key: key(k) }), + Action::KeyUp { key: k } => action_api::Type::KeyUp(action_api::KeyUp { key: key(k) }), + Action::Wait(duration) => action_api::Type::Wait(action_api::Wait { duration: Some(prost_types::Duration { + seconds: duration.as_secs() as i64, nanos: duration.subsec_nanos() as i32, + }) }), + }), + }).collect(), + })), + _ => None, + } +} + +fn computer_target(target: computer_use::Target) -> Option { + use api::message::tool_call::{computer_use_target, ComputerUseTarget}; + match target { + computer_use::Target::Screen => None, + computer_use::Target::Window { window_id, pid } => Some(ComputerUseTarget { target: Some(computer_use_target::Target::Window(computer_use_target::Window { window_id: window_id.to_string(), pid })) }), + } +} + +fn computer_screenshot(params: computer_use::ScreenshotParams) -> api::message::tool_call::ScreenshotParams { + api::message::tool_call::ScreenshotParams { + max_long_edge_px: params.max_long_edge_px.unwrap_or_default() as i32, + max_total_px: params.max_total_px.unwrap_or_default() as i32, + target: computer_target(params.target), + region: params.region.map(|region| api::message::tool_call::screenshot_params::Region { + top_left: Some(api::Coordinates { x: region.top_left.x(), y: region.top_left.y() }), + bottom_right: Some(api::Coordinates { x: region.bottom_right.x(), y: region.bottom_right.y() }), + }), + } +} + #[derive(Serialize, Deserialize)] struct SavedTool { context: ProviderActionContext, diff --git a/app/src/ai/agent/provider_history_tests.rs b/app/src/ai/agent/provider_history_tests.rs index 0963bd74..bb146de2 100644 --- a/app/src/ai/agent/provider_history_tests.rs +++ b/app/src/ai/agent/provider_history_tests.rs @@ -24,6 +24,29 @@ fn query(id: &str) -> api::Message { } } +#[test] +fn computer_tool_display_and_native_restoration_preserve_the_same_actions() { + let context = ProviderActionContext::new_for_test("task"); + let call = ToolCall { + id: "desktop-call".into(), name: "use_computer".into(), arguments: json!({ + "action_summary": "Type and press Return", "actions": [ + { "type": "type_text", "text": "hello" }, + { "type": "key_down", "keycode": 36 }, + { "type": "key_up", "keycode": 36 }, + ], + }), + }; + let action = context.action_from_tool_call(&call).unwrap(); + let message = tool_message(&context, &call, &action).unwrap(); + let Some(api::message::Message::ToolCall(display)) = message.message else { panic!("expected tool display"); }; + let Some(api::message::tool_call::Tool::UseComputer(display)) = display.tool else { panic!("expected computer display"); }; + let display_action = crate::ai::agent::AIAgentActionType::try_from(display).unwrap(); + assert_eq!(display_action, action.action); + let restored = super::restore_tool(&message.server_message_data).unwrap().unwrap(); + assert_eq!(restored.action, action.action); + assert_eq!(restored.id, action.id); +} + #[test] fn saved_tools_and_results_restore_as_individual_panes_across_turns() { let context = ProviderActionContext::new_for_test("task"); diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 55b0bf08..e79c9b5a 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -418,19 +418,72 @@ fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bo } else { ToolResultStatus::Success }; - let content = action_result.result.model_content(); + let (content, images) = computer_use_model_content(&action_result.result) + .unwrap_or_else(|| (action_result.result.model_content(), Vec::new())); let content = if permission_denied { format!("Permission denied by the user. {content}") } else { content }; ToolResult { + images: if permission_denied { Vec::new() } else { images }, call_id: action_result.id.to_string(), content, status, } } +fn computer_use_model_content( + result: &AIAgentActionResultType, +) -> Option<(String, Vec)> { + let (screenshot, mut metadata) = match result { + AIAgentActionResultType::RequestComputerUse(RequestComputerUseResult::Approved { + screenshot, + platform, + windows, + }) => ( + Some(screenshot), + serde_json::json!({ + "status": "approved", + "platform": format!("{platform:?}"), + "windows": windows, + }), + ), + AIAgentActionResultType::UseComputer(UseComputerResult::Success(result)) => ( + result.screenshot.as_ref(), + serde_json::json!({ + "status": "success", + "cursor_position": result.cursor_position.map(|point| serde_json::json!({ + "x": point.x(), "y": point.y(), + })), + "windows": result.windows, + "captured_window": result.captured_window.map(|window| serde_json::json!({ + "window_id": window.window_id, + "width_px": window.width_px, + "height_px": window.height_px, + })), + }), + ), + _ => return None, + }; + let images = screenshot + .map(|screenshot| { + metadata["screenshot"] = serde_json::json!({ + "width": screenshot.width, + "height": screenshot.height, + "original_width": screenshot.original_width, + "original_height": screenshot.original_height, + "coordinate_space": "Use original physical pixels; scale image coordinates by original_width/width and original_height/height.", + }); + vec![galaxy_agent_core::ToolResultImage { + data: screenshot.data.clone(), + mime_type: screenshot.mime_type.to_string(), + }] + }) + .unwrap_or_default(); + Some((metadata.to_string(), images)) +} + #[cfg(not(target_family = "wasm"))] fn action_tool_name(action: &AIAgentAction) -> String { action diff --git a/app/src/ai/blocklist/action_model/execute/use_computer.rs b/app/src/ai/blocklist/action_model/execute/use_computer.rs index 66ed3223..ba5d085c 100644 --- a/app/src/ai/blocklist/action_model/execute/use_computer.rs +++ b/app/src/ai/blocklist/action_model/execute/use_computer.rs @@ -24,10 +24,8 @@ impl UseComputerExecutor { return false; }; - // We unconditionally return true here because this action is only executed by - // the computer use subagent, which cannot begin without the user approving it via - // a `RequestComputerUse` action, and the approval extends to all computer use - // actions within that computer use subagent. + // The shared action boundary applies the active permission mode to every action. + // Once that check passes, the executor has no additional approval requirement. true } diff --git a/app/src/ai/blocklist/action_model_tests.rs b/app/src/ai/blocklist/action_model_tests.rs index d8ce5253..38f29821 100644 --- a/app/src/ai/blocklist/action_model_tests.rs +++ b/app/src/ai/blocklist/action_model_tests.rs @@ -366,6 +366,26 @@ fn domain_read_result_contains_the_file_contents_for_the_next_model_turn() { assert!(result.content.contains("pub fn answer() -> u8 { 42 }")); } +#[test] +fn computer_results_keep_real_image_bytes_and_coordinate_metadata() { + let screenshot = computer_use::Screenshot { + width: 500, height: 300, original_width: 1000, original_height: 600, + data: vec![1, 2, 3, 4], mime_type: "image/png".into(), + }; + let action = action_result("desktop", AIAgentActionResultType::RequestComputerUse(RequestComputerUseResult::Approved { + screenshot, platform: computer_use::Platform::Mac, windows: Vec::new(), + })); + let result = domain_tool_result(&action, false); + assert_eq!(result.images.len(), 1); + assert_eq!(result.images[0].data, vec![1, 2, 3, 4]); + assert_eq!(result.images[0].mime_type, "image/png"); + let metadata: serde_json::Value = serde_json::from_str(&result.content).unwrap(); + assert_eq!(metadata["platform"], "Mac"); + assert_eq!(metadata["screenshot"]["width"], 500); + assert_eq!(metadata["screenshot"]["original_width"], 1000); + assert!(domain_tool_result(&action, true).images.is_empty()); +} + #[test] fn action_permission_kinds_match_the_safety_boundary() { assert_eq!( diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 985e0fe3..46d948b9 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -1039,6 +1039,7 @@ fn convert_provider_tool_batch( matches!(pending.state, PendingToolCallState::RecoveryPending), )), Err(message) => invalid_results.push(ToolResult { + images: Vec::new(), call_id: pending.call.id.clone(), content: format!("Invalid {} tool input: {message}", pending.call.name), status: ToolResultStatus::Error, @@ -8460,7 +8461,10 @@ impl BlocklistAIController { matches!( pending.call.name.as_str(), "request_computer_use" | "use_computer" - ) + ) && pending.state.result().or_else(|| match event { + galaxy_agent_core::ToolEvent::Completed { result } if result.call_id == pending.call.id => Some(result), + _ => None, + }).is_some_and(|result| result.status == galaxy_agent_core::ToolResultStatus::Success) }); used_computer.then_some(COMPUTER_USE_PROVIDER_PROFILE) } diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 5fb6c1a6..45ecb63a 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -2819,9 +2819,12 @@ pub(crate) fn openai_model_context_size( .max_input_tokens .unwrap_or(model.context_size) .min(model.context_size); - if provider_kind == OpenAIProviderKind::ChatGPTSubscription { - // The catalog's default window is the supported session budget. The - // advertised maximum must not silently opt subscriptions into long context. + if matches!( + provider_kind, + OpenAIProviderKind::ChatGPTSubscription | OpenAIProviderKind::LiteLLM + ) { + // Use the efficient session window for both subscriptions and proxies. + // An advertised maximum does not opt a connection into long context. // Older saved catalogs kept the default in max_input_tokens instead. return model .default_context_size @@ -2959,7 +2962,7 @@ fn openai_model_context_window( is_configurable: max_context_size > default_context_size, min: default_context_size, max: max_context_size, - // Subscription limits are already capped to the catalog's default above. + // Subscription and LiteLLM limits are capped to the catalog default above. default_max: max_context_size, } } @@ -2971,6 +2974,53 @@ fn u32_from_json_any(value: &serde_json::Value, keys: &[&str]) -> Option { .and_then(|value| u32::try_from(value).ok()) } +/// Keep an endpoint's recommended window separate from its optional expanded +/// maximum. A bare model ID carries no token metadata, so it must not replace +/// richer limits when the two catalog endpoints are combined. +#[cfg(not(target_family = "wasm"))] +fn openai_context_limits(metadata: &serde_json::Value) -> (u32, Option, Option) { + let positive_limit = |keys: &[&str]| { + keys.iter().find_map(|key| { + metadata[*key] + .as_u64() + .and_then(|value| u32::try_from(value).ok()) + .filter(|value| *value > 0) + }) + }; + let max_input_tokens = + positive_limit(&["max_input_tokens", "input_token_limit", "max_prompt_tokens"]); + let default_context_size = positive_limit(&[ + "default_context_size", + "default_context_window", + "context_window", + ]); + let maximum_context_size = positive_limit(&[ + "max_context_window", + "max_model_len", + "context_size", + "token_size", + ]); + let effective_context_percent = metadata["effective_context_window_percent"] + .as_u64() + .filter(|value| (1..=100).contains(value)) + .unwrap_or(100); + // A proxy may already reduce max_input_tokens by this percentage. Apply it + // only to the original session window, never to that input limit again. + let default_context_size = default_context_size + .map(|value| { + u32::try_from(u64::from(value) * effective_context_percent / 100) + .unwrap_or(value) + .max(1) + }) + .or(max_input_tokens) + .or(maximum_context_size); + let context_size = maximum_context_size + .or(default_context_size) + .or(max_input_tokens) + .unwrap_or(DEFAULT_DISCOVERED_MODEL_CONTEXT_SIZE); + (context_size, max_input_tokens, default_context_size) +} + #[cfg(not(target_family = "wasm"))] fn normalize_codex_release_version(version: &str) -> Option { let version = version.trim(); @@ -3209,19 +3259,11 @@ async fn fetch_from_litellm_model_info( let model_name = entry["model_name"].as_str()?; let model_info = &entry["model_info"]; - let max_input_tokens = model_info["max_input_tokens"] - .as_u64() - .and_then(|v| u32::try_from(v).ok()); + let (context_size, max_input_tokens, default_context_size) = + openai_context_limits(model_info); let max_output_tokens = model_info["max_output_tokens"] .as_u64() .and_then(|v| u32::try_from(v).ok()); - let context_size = u32_from_json_any( - model_info, - &["max_context_window", "context_window", "max_model_len"], - ) - .or(max_input_tokens) - .unwrap_or(200_000); - let vision_supported = model_info["supports_vision"].as_bool().unwrap_or(true); let display_name = model_name.replace(['-', '_'], " "); @@ -3271,7 +3313,7 @@ async fn fetch_from_litellm_model_info( vision_supported, context_size, max_input_tokens, - default_context_size: None, + default_context_size, max_output_tokens, provider, supports_system_messages: if model_name.starts_with("codex-gpt-") { @@ -3318,7 +3360,30 @@ async fn fetch_litellm_catalog( metadata .iter() .find(|info| info.model_id == model.model_id) - .cloned() + .map(|info| { + let mut enriched = info.clone(); + // /model/info supplies capabilities, but must not erase a + // smaller session budget supplied by /models (or vice versa). + enriched.default_context_size = model + .default_context_size + .into_iter() + .chain(info.default_context_size) + .min(); + enriched.max_input_tokens = model + .max_input_tokens + .into_iter() + .chain(info.max_input_tokens) + .min(); + enriched.context_size = match ( + model.default_context_size, + info.default_context_size, + ) { + (Some(_), Some(_)) => model.context_size.max(info.context_size), + (Some(_), None) => model.context_size, + _ => info.context_size, + }; + enriched + }) .unwrap_or(model) }) .collect()) @@ -3369,21 +3434,7 @@ async fn fetch_from_openai_models( .iter() .filter_map(|m| { let id = m["id"].as_str()?; - let max_input_tokens = u32_from_any( - m, - &["max_input_tokens", "input_token_limit", "max_prompt_tokens"], - ); - let context_size = u32_from_any( - m, - &[ - "max_context_window", - "max_model_len", - "context_window", - "token_size", - ], - ) - .or(max_input_tokens) - .unwrap_or(200_000); + let (context_size, max_input_tokens, default_context_size) = openai_context_limits(m); let max_output_tokens = u32_from_any( m, &[ @@ -3430,7 +3481,7 @@ async fn fetch_from_openai_models( .unwrap_or(true), context_size, max_input_tokens, - default_context_size: None, + default_context_size, max_output_tokens, provider, supports_system_messages: if id.starts_with("codex-gpt-") { diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index c0a83ce8..e8488faf 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -1,3 +1,4 @@ +use serde_json::json; use warpui::App; use super::*; @@ -1189,6 +1190,155 @@ fn chatgpt_default_window_limits_routing_without_a_separate_input_limit() { ); } +#[test] +fn litellm_efficient_windows_limit_model_picker_and_provider_routing() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + let models = [("terminal-model", 272_000), ("main-model", 500_000)] + .into_iter() + .map(|(id, efficient_window)| { + let mut model = openai_model(id); + model.context_size = 872_000; + model.max_input_tokens = Some(872_000); + model.default_context_size = Some(efficient_window); + model + }) + .collect(); + AISettings::handle(&app).update(&mut app, |settings, ctx| { + settings.openai_enabled.set_value(true, ctx).unwrap(); + settings + .openai_providers + .set_value( + vec![OpenAIProviderConfig { + kind: OpenAIProviderKind::LiteLLM, + enabled: true, + name: "LiteLLM".to_owned(), + base_url: "https://proxy.test/v1".to_owned(), + api_key: None, + project_id: None, + location: None, + models, + }], + ctx, + ) + .unwrap(); + }); + let mut preferences = empty_preferences(); + app.read(|ctx| preferences.inject_openai_models(ctx)); + for (id, efficient_window) in [("terminal-model", 272_000), ("main-model", 500_000)] { + let route = preferences.openai_client_config_for_model(id).unwrap(); + assert_eq!(route.max_input_tokens, Some(efficient_window)); + let model = preferences + .models_by_feature + .agent_mode + .choices + .iter() + .find(|model| model.id.as_str() == id) + .unwrap(); + assert_eq!(model.context_window.max, efficient_window); + assert_eq!(model.context_window.default_max, efficient_window); + assert!(!model.context_window.is_configurable); + } + }); +} + +#[tokio::test] +async fn litellm_discovery_preserves_efficient_windows_across_catalog_endpoints() { + let mut server = mockito::Server::new_async().await; + for (model_metadata, rich_metadata, maximum, efficient_window, request_budget) in [ + ( + json!({"id": "model"}), + Some(json!({ + "context_window": 500000, + "max_context_window": 872000, + "max_input_tokens": 872000 + })), + 872_000, + 500_000, + 500_000, + ), + ( + json!({"id": "model", "context_window": 272000, "max_context_window": 872000}), + Some(json!({"max_context_window": 872000})), + 872_000, + 272_000, + 272_000, + ), + ( + json!({"id": "model", "context_window": 272000, "max_context_window": 872000}), + None, + 872_000, + 272_000, + 272_000, + ), + ( + json!({"id": "model", "context_window": 500000, "max_context_window": 872000}), + Some(json!({"supports_vision": true})), + 872_000, + 500_000, + 500_000, + ), + ( + json!({"id": "model"}), + Some(json!({ + "context_window": 272000, + "max_context_window": 872000, + "max_input_tokens": 258400, + "effective_context_window_percent": 95 + })), + 872_000, + 258_400, + 258_400, + ), + ( + json!({"id": "model"}), + Some(json!({ + "context_window": 272000, + "max_context_window": 872000, + "max_input_tokens": 128000 + })), + 872_000, + 272_000, + 128_000, + ), + ] { + let catalog = server + .mock("GET", "/models") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(json!({"data": [model_metadata]}).to_string()) + .create_async() + .await; + let metadata = server + .mock("GET", "/model/info") + .with_status(if rich_metadata.is_some() { 200 } else { 404 }) + .with_header("content-type", "application/json") + .with_body( + json!({"data": [{"model_name": "model", "model_info": rich_metadata}]}).to_string(), + ) + .create_async() + .await; + let models = fetch_litellm_catalog(&server.url(), None, &reqwest::Client::new()) + .await + .unwrap(); + assert_eq!(models.len(), 1); + assert_eq!(models[0].context_size, maximum); + assert_eq!(models[0].default_context_size, Some(efficient_window)); + assert_eq!( + openai_model_context_size(&models[0], OpenAIProviderKind::LiteLLM), + request_budget + ); + let window = openai_model_context_window(&models[0], OpenAIProviderKind::LiteLLM); + assert_eq!(window.default_max, request_budget); + assert_eq!(window.max, request_budget); + assert!(!window.is_configurable); + catalog.assert_async().await; + metadata.assert_async().await; + catalog.remove_async().await; + metadata.remove_async().await; + } +} + #[test] fn litellm_refresh_replaces_membership_and_preserves_preferences() { let mut retained = openai_model("retained"); diff --git a/app/src/ai/openai/response_translator.rs b/app/src/ai/openai/response_translator.rs index 18f491d4..ab6b50ef 100644 --- a/app/src/ai/openai/response_translator.rs +++ b/app/src/ai/openai/response_translator.rs @@ -663,6 +663,8 @@ fn estimate_cost_cents( } const KNOWN_TOOLS: &[&str] = &[ + "request_computer_use", + "use_computer", "run_shell_command", "read_files", "apply_file_diffs", diff --git a/app/src/ai/provider/request_translator.rs b/app/src/ai/provider/request_translator.rs index cded1d37..6f830adc 100644 --- a/app/src/ai/provider/request_translator.rs +++ b/app/src/ai/provider/request_translator.rs @@ -1714,6 +1714,8 @@ pub(crate) fn tool_name_is_supported(name: &str, supported: &HashSet has(ToolType::AskUserQuestion), "read_skill" => has(ToolType::ReadSkill), "fetch_conversation" => has(ToolType::FetchConversation), + "request_computer_use" => has(ToolType::RequestComputerUse), + "use_computer" => has(ToolType::UseComputer), // This tool is implemented entirely inside the direct-provider response // translator, so it does not need a client ToolType capability bit. "recall_tool_history" => true, @@ -1730,8 +1732,77 @@ pub(crate) const IMPLEMENTATION_PROGRESS_GUIDANCE: &str = "## Incremental Implem pub(crate) const SHELL_EXECUTION_GUIDANCE: &str = "## Shell Execution\nCommands use a persistent terminal. Never exit, log out of, replace with exec, or kill that shell or Galaxy. Do not append `exit $status` to a command to report its result; let the command's exit status be returned naturally. If a script needs exit, run it in an explicit child shell such as `sh -c '...'`. Stop a running agent command with interrupt_shell_command, not broad kill/pkill commands.\nRun builds, tests, installs, and workspace checks as separate shell tool calls so each has its own output and exit status. Keep stdout and stderr live: do not hide them in subprocess capture_output/PIPE, redirect everything to a file, or buffer a whole validation suite and print only its final tail. If a log file is needed, stream through tee while preserving the original command's exit status. Use non-interactive, non-watch modes and disable pagers. Give intentionally quiet commands a descriptive tool title so the user knows what is running; silence alone is not evidence of a hang. If repeated snapshots show no progress, investigate whether the command is waiting for input, watching, or buffering; do not blindly poll forever or start duplicate commands. Interrupt and correct an unintended watcher or blocked command before retrying.\n\n"; +fn computer_target_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "type": { "type": "string", "enum": ["screen", "window"], "default": "screen" }, + "window_id": { "type": "integer", "minimum": 1 }, + "pid": { "type": "integer", "minimum": 1 }, + }, + "required": ["type"], + "description": "Omit to target the main screen. Use window only with a window_id and pid from the latest returned windows list; requires background computer use support.", + }) +} + +fn computer_screenshot_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "max_long_edge_px": { "type": "integer", "minimum": 1, "maximum": 4096, "default": 1568 }, + "max_total_px": { "type": "integer", "minimum": 1, "maximum": 4194304, "default": 1150000 }, + "target": computer_target_schema(), + }, + "description": "Screenshot size limits and target. Omit for a main-screen screenshot with bounded image size.", + }) +} + pub fn default_tool_definitions() -> Vec { vec![ + ToolDefinition { + name: "request_computer_use".to_string(), + description: "Request desktop computer use on the machine running Galaxy and capture its initial screenshot. This uses the user's local desktop even when the terminal is connected by SSH. Call before interacting with the desktop; subsequent turns use the configured computer-use model and use_computer tool. Existing profile permissions apply. Screen Recording permission is needed to capture the screen; Accessibility permission is needed to control input on macOS.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "task_summary": { "type": "string", "description": "Specific desktop task and application to operate." }, + "screenshot_params": computer_screenshot_schema(), + }, + "required": ["task_summary"], + }), + }, + ToolDefinition { + name: "use_computer".to_string(), + description: "Perform a bounded batch of desktop actions after request_computer_use succeeds, then receive a screenshot. Use an empty actions array to inspect the desktop. Coordinates use original physical screenshot pixels, so scale coordinates from the displayed image using the returned original_width/width and original_height/height. Screen targets use main-display coordinates; window targets use coordinates relative to the exact enumerated window. Prefer small action batches and inspect the result. Named keys use OS keycodes; key accepts a single character. Button/key down and up are separate actions.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "action_summary": { "type": "string" }, + "actions": { + "type": "array", "maxItems": 20, + "items": { + "type": "object", + "properties": { + "type": { "type": "string", "enum": ["wait", "mouse_move", "mouse_down", "mouse_up", "mouse_wheel", "type_text", "key_down", "key_up"] }, + "x": { "type": "integer" }, "y": { "type": "integer" }, + "button": { "type": "string", "enum": ["left", "right", "middle", "back", "forward"] }, + "text": { "type": "string", "maxLength": 100000 }, + "key": { "type": "string", "minLength": 1, "maxLength": 1 }, + "keycode": { "type": "integer", "description": "Platform keycode for modifier/navigation keys, such as macOS Return 36, Escape 53, Command 55, Shift 56, Option 58, Control 59." }, + "seconds": { "type": "number", "minimum": 0, "maximum": 10 }, + "direction": { "type": "string", "enum": ["up", "down", "left", "right"] }, + "pixels": { "type": "integer", "minimum": 0 }, + "target": computer_target_schema(), + }, + "required": ["type"], + }, + }, + "take_screenshot": { "type": "boolean", "default": true }, + "screenshot_params": computer_screenshot_schema(), + }, + "required": ["action_summary", "actions"], + }), + }, ToolDefinition { name: "run_shell_command".to_string(), description: "Execute a shell command in the user's terminal and return its output. Use for builds, tests, git, package managers, and other shell work. Commands run in this conversation's persistent shell and environment, starting in the supplied working directory. Successful cd commands affect subsequent calls in this shell; subshell directory changes do not. Avoid prefixing commands with cd to the directory you are already in, and use pwd only when the current location is uncertain. Never exit, log out of, replace, or kill this persistent shell. Run each build/test/install separately and keep its output live; avoid buffered subprocess wrappers and redirecting all output to files. Set wait_until_complete=false for commands that might run longer than a few seconds; Galaxy will return a command_id that can be monitored with read_shell_command_output.".to_string(), diff --git a/app/src/ai/provider/response_translator.rs b/app/src/ai/provider/response_translator.rs index fd470747..df6b11cd 100644 --- a/app/src/ai/provider/response_translator.rs +++ b/app/src/ai/provider/response_translator.rs @@ -839,6 +839,14 @@ pub fn build_tool_call_message( serde_json::from_str(tool_input_json).unwrap_or(serde_json::json!({})); let tool = match tool_name { + "request_computer_use" | "use_computer" => { + crate::ai::runtime::action_from_tool_call( + task_id, + &galaxy_agent_core::ToolCall { id: tool_use_id.to_string(), name: tool_name.to_string(), arguments: input.clone() }, + &ai::skills::SkillPathOrigin::Local, + &std::collections::HashMap::new(), + ).ok().and_then(|action| crate::ai::agent::provider_history::computer_use_tool(&action.action)) + } "run_shell_command" => { let command = input .get("command") @@ -1596,6 +1604,8 @@ fn run_agents_harness_from_str(harness_type: &str) -> Option { /// Built-in tools that Galaxy knows how to execute directly. const KNOWN_TOOLS: &[&str] = &[ + "request_computer_use", + "use_computer", "run_shell_command", "read_files", "apply_file_diffs", diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index d36af971..bba27bfd 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -104,6 +104,7 @@ fn restored_task_progress_replay_does_not_duplicate_the_persisted_checklist() { arguments: serde_json::to_value(&snapshot).unwrap(), }, state: galaxy_agent_core::PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { + images: Vec::new(), call_id: "progress".to_owned(), status: galaxy_agent_core::ToolResultStatus::Success, content: snapshot.to_tool_result_content(), }}, diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index 39c64c12..7abefc26 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -4,6 +4,7 @@ mod provider_run_coordinator; mod rig; mod rig_request; mod rig_tool; +pub(crate) use rig_tool::action_from_tool_call; mod task_progress; pub(crate) use child_context::add_parent_context; diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index a4059281..b535d26e 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -461,6 +461,7 @@ async fn coordinator_awaiting_tools( fn tool_result(call_id: &str, content: &str, status: ToolResultStatus) -> ToolEvent { ToolEvent::Completed { result: ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content: content.to_string(), status, @@ -514,6 +515,7 @@ async fn one_run_drives_model_tool_and_followup_turns_with_atomic_history() { .complete_tool( &batch.work_id, ToolResult { + images: Vec::new(), call_id: "read-1".to_string(), content: "manifest contents".to_string(), status: ToolResultStatus::Success, diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index c4593a09..94e8876f 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -133,20 +133,25 @@ pub(crate) async fn prepare_provider_run( let skill_path_origin = params.session_context.skill_path_origin(); let mut cli_params = params.clone(); let computer_use_model = params.computer_use_model.as_str(); - let computer_use_profile = if params.computer_use_enabled - && !computer_use_model.is_empty() - && computer_use_model != "computer-use-agent-auto" - && computer_use_model != params.model.as_str() - && !matches!( - &computer_use_provider_config, - crate::ai::provider::ProviderConfig::None - ) { + let computer_use_profile = if supported_tools.contains(&ToolType::RequestComputerUse) { let mut computer_params = params.clone(); - computer_params.model = params.computer_use_model.clone(); + let computer_provider = if computer_use_model.is_empty() + || computer_use_model == "computer-use-agent-auto" + || computer_use_model == params.model.as_str() + || matches!(&computer_use_provider_config, crate::ai::provider::ProviderConfig::None) + { + base_provider_config.clone() + } else { + computer_params.model = params.computer_use_model.clone(); + computer_use_provider_config + }; + let mut computer_tools = supported_tools.clone(); + computer_tools.retain(|tool| *tool != ToolType::RequestComputerUse); + computer_tools.push(ToolType::UseComputer); let (runtime, prepared) = prepare_provider_profile( - computer_use_provider_config, + computer_provider, computer_params, - supported_tools.clone(), + computer_tools, supported_cli_agent_tools.clone(), None, ) diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index 766e0e07..8650961a 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -389,17 +389,7 @@ fn input_messages( ) -> Vec { let mut messages = Vec::new(); if !tool_results.is_empty() { - let mut parts = tool_results - .into_iter() - .map(|result| { - let is_error = result.is_error(); - ContentPart::ToolResult { - tool_use_id: result.call_id, - content: result.content, - is_error, - } - }) - .collect::>(); + let mut parts = ToolResult::into_batch_content_parts(tool_results); let content = if parts.len() == 1 { let ContentPart::ToolResult { tool_use_id, @@ -1043,6 +1033,12 @@ fn build_system_prompt( { prompt.push_str(CODING_DELEGATION_GUIDANCE); } + if tools.iter().any(|tool| tool.name == "request_computer_use") { + prompt.push_str("## Desktop Computer Use\nUse request_computer_use for tasks that require seeing or interacting with the user's desktop. This controls the machine running Galaxy, including when the terminal is connected to another host. Request the specific authorized task and wait for the screenshot. A successful request switches to the configured computer-use model; rejected or failed requests do not grant control.\n\n"); + } + if tools.iter().any(|tool| tool.name == "use_computer") { + prompt.push_str("## Desktop Computer Use\nInspect the latest returned screenshot before acting. Use use_computer with small action batches, then inspect the new screenshot. Use original physical pixel coordinates: multiply displayed image x by original_width/width and y by original_height/height. Every button/key press requires both down and up actions. For window targeting, select an exact window_id/pid from the latest windows list; never invent a window. Stop on errors or denied permission and report the actual cause. Galaxy's Full Disk Access is separate from Screen Recording and Accessibility permissions. When desktop work is complete and terminal/file work remains, call the appropriate normal tool and Galaxy resumes the base model. Only finish the response when the whole user's task is complete.\n\n"); + } prompt.push_str("## Available Tools\n"); if tools.is_empty() { prompt.push_str("No tools are available. Do not invent tool calls.\n"); diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index cce253c0..f14f8dad 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -533,6 +533,7 @@ fn lrc_snapshot_follow_up_uses_the_cli_monitor_prompt_and_tools() { ), }; let snapshot_tool_result = ToolResult { + images: Vec::new(), call_id: "run-call".to_string(), content: result.result.model_content(), status: ToolResultStatus::Success, @@ -842,6 +843,7 @@ fn normalized_tool_outcomes_are_the_only_action_results_sent_to_rig() { let tool_results = statuses .iter() .map(|(call_id, status, _)| ToolResult { + images: Vec::new(), call_id: (*call_id).to_string(), content: format!("normalized-{call_id}"), status: *status, diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index e685ff1c..b157252a 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -87,6 +87,27 @@ async fn selected_computer_use_model_gets_its_own_profile() { ); } +#[tokio::test] +async fn automatic_and_same_model_computer_use_keep_a_dedicated_profile() { + for selected in ["computer-use-agent-auto", "base-selection", "unroutable-model"] { + let mut params = RequestParams::new_for_test(); + params.model = LLMId::from("base-selection"); + params.computer_use_model = LLMId::from(selected); + params.computer_use_enabled = true; + let prepared = prepare_provider_run( + ProviderConfig::OpenAI(openai_config("base-provider-model")), + ProviderConfig::None, ProviderConfig::None, params, Vec::new(), + ).await.unwrap(); + assert!(prepared.base_profile.request.tools.iter().any(|tool| tool.name == "request_computer_use")); + assert!(!prepared.base_profile.request.tools.iter().any(|tool| tool.name == "use_computer")); + let computer = prepared.computer_use_profile.unwrap(); + assert_eq!(computer.request.model.as_str(), "base-provider-model"); + assert!(computer.request.tools.iter().any(|tool| tool.name == "use_computer")); + assert!(!computer.request.tools.iter().any(|tool| tool.name == "request_computer_use")); + assert!(computer.request.system_prompt.unwrap().contains("Inspect the latest returned screenshot before acting")); + } +} + #[test] fn subscription_context_budget_uses_catalog_limit_and_clamps_stale_overrides() { let mut config = openai_config("subscription-model"); diff --git a/app/src/ai/runtime/rig_tool.rs b/app/src/ai/runtime/rig_tool.rs index 50d50be4..f87ba6dc 100644 --- a/app/src/ai/runtime/rig_tool.rs +++ b/app/src/ai/runtime/rig_tool.rs @@ -55,6 +55,23 @@ pub(crate) fn action_from_tool_call( rationale: optional_string(input, "rationale")?, citations: Vec::new(), }, + "request_computer_use" => AIAgentActionType::RequestComputerUse( + crate::ai::agent::RequestComputerUseRequest { + task_summary: required_nonempty_string(input, "task_summary")?, + screenshot_params: Some(computer_screenshot_params(input)?), + }, + ), + "use_computer" => AIAgentActionType::UseComputer( + crate::ai::agent::UseComputerRequest { + action_summary: required_nonempty_string(input, "action_summary")?, + actions: computer_actions(input)?, + screenshot_params: if optional_boolean(input, "take_screenshot")?.unwrap_or(true) { + Some(computer_screenshot_params(input)?) + } else { + None + }, + }, + ), "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest { locations: limited_required_array(input, "files", MAX_BATCH_READ_ITEMS)? .iter() @@ -472,6 +489,113 @@ fn optional_nonnegative_i32(input: &serde_json::Value, key: &str) -> Result Result { + input.get(key).and_then(serde_json::Value::as_i64) + .and_then(|value| i32::try_from(value).ok()) + .ok_or_else(|| format!("invalid field {key:?}: expected a 32-bit integer")) +} + +fn computer_target(input: &serde_json::Value) -> Result { + let Some(target) = input.get("target") else { return Ok(computer_use::Target::Screen); }; + require_object(target, "target")?; + match required_string(target, "type")?.as_str() { + "screen" => Ok(computer_use::Target::Screen), + "window" => { + let window_id = optional_bounded_u64(target, "window_id", u64::from(u32::MAX))? + .filter(|id| *id > 0).ok_or("window target requires a non-zero window_id")? as u32; + let pid = computer_i32(target, "pid")?; + if pid <= 0 { return Err("window target requires a positive pid".to_string()); } + if !crate::features::FeatureFlag::BackgroundComputerUse.is_enabled() + || !computer_use::background_supported() + { + return Err("Background window targeting is unavailable; use the screen target.".to_string()); + } + Ok(computer_use::Target::Window { window_id, pid }) + } + _ => Err("target.type must be screen or window".to_string()), + } +} + +fn computer_screenshot_params(input: &serde_json::Value) -> Result { + let empty = serde_json::json!({}); + let params = input.get("screenshot_params").unwrap_or(&empty); + require_object(params, "screenshot_params")?; + let positive = |key, maximum, default| -> Result { + let value = optional_bounded_u64(params, key, maximum)?.unwrap_or(default); + if value == 0 { return Err(format!("{key} must be positive")); } + Ok(value as usize) + }; + Ok(computer_use::ScreenshotParams { + max_long_edge_px: Some(positive("max_long_edge_px", 4096, 1568)?), + max_total_px: Some(positive("max_total_px", 4_194_304, 1_150_000)?), + region: None, + target: computer_target(params)?, + }) +} + +fn computer_actions(input: &serde_json::Value) -> Result, String> { + use computer_use::{Action, Key, MouseButton, ScrollDirection, ScrollDistance, TargetedAction, Vector2I}; + let mut wait_seconds = 0.0; + limited_required_array(input, "actions", 20)?.iter().map(|input| { + require_object(input, "action")?; + let position = || -> Result { + Ok(Vector2I::new(computer_i32(input, "x")?, computer_i32(input, "y")?)) + }; + let button = || -> Result { + match required_string(input, "button")?.as_str() { + "left" => Ok(MouseButton::Left), "right" => Ok(MouseButton::Right), + "middle" => Ok(MouseButton::Middle), "back" => Ok(MouseButton::Back), + "forward" => Ok(MouseButton::Forward), + _ => Err("button must be left, right, middle, back, or forward".to_string()), + } + }; + let key = || -> Result { + match (input.get("key"), input.get("keycode")) { + (Some(_), None) => { + let text = required_string(input, "key")?; + let mut chars = text.chars(); + let key = chars.next().ok_or("key must contain exactly one character")?; + if chars.next().is_some() { return Err("key must contain exactly one character".to_string()); } + Ok(Key::Char(key)) + } + (None, Some(_)) => Ok(Key::Keycode(computer_i32(input, "keycode")?)), + _ => Err("provide exactly one of key or keycode".to_string()), + } + }; + let action = match required_string(input, "type")?.as_str() { + "mouse_move" => Action::MouseMove { to: position()? }, + "mouse_down" => Action::MouseDown { button: button()?, at: position()? }, + "mouse_up" => Action::MouseUp { button: button()? }, + "mouse_wheel" => Action::MouseWheel { + at: position()?, + direction: match required_string(input, "direction")?.as_str() { + "up" => ScrollDirection::Up, "down" => ScrollDirection::Down, + "left" => ScrollDirection::Left, "right" => ScrollDirection::Right, + _ => return Err("direction must be up, down, left, or right".to_string()), + }, + distance: ScrollDistance::Pixels(optional_nonnegative_i32(input, "pixels")?.ok_or("mouse_wheel requires pixels")?), + }, + "type_text" => { + let text = required_string(input, "text")?; + if text.len() > 100_000 { return Err("text must be no longer than 100000 bytes".to_string()); } + Action::TypeText { text } + } + "key_down" => Action::KeyDown { key: key()? }, + "key_up" => Action::KeyUp { key: key()? }, + "wait" => { + let seconds = input.get("seconds").and_then(serde_json::Value::as_f64) + .filter(|seconds| seconds.is_finite() && *seconds >= 0.0 && *seconds <= 10.0) + .ok_or("seconds must be between 0 and 10")?; + wait_seconds += seconds; + if wait_seconds > 10.0 { return Err("a computer-use batch may wait at most 10 seconds".to_string()); } + Action::Wait(Duration::from_secs_f64(seconds)) + } + _ => return Err("unknown computer action type".to_string()), + }; + Ok(TargetedAction { action, target: computer_target(input)? }) + }).collect() +} + fn skill_references( input: &serde_json::Value, skill_path_origin: &SkillPathOrigin, diff --git a/app/src/ai/runtime/rig_tool_tests.rs b/app/src/ai/runtime/rig_tool_tests.rs index 71a7d1fb..7b66189a 100644 --- a/app/src/ai/runtime/rig_tool_tests.rs +++ b/app/src/ai/runtime/rig_tool_tests.rs @@ -16,6 +16,55 @@ fn call(name: &str, arguments: serde_json::Value) -> ToolCall { } } +#[test] +fn computer_use_request_always_captures_a_bounded_initial_screenshot() { + let action = action_from_tool_call("task", &call("request_computer_use", serde_json::json!({ + "task_summary": "Inspect the browser", + })), &SkillPathOrigin::Local, &HashMap::new()).unwrap(); + let AIAgentActionType::RequestComputerUse(request) = action.action else { panic!("expected computer-use request"); }; + let screenshot = request.screenshot_params.unwrap(); + assert_eq!(screenshot.max_long_edge_px, Some(1568)); + assert_eq!(screenshot.max_total_px, Some(1_150_000)); + assert_eq!(screenshot.target, computer_use::Target::Screen); +} + +#[test] +fn computer_actions_preserve_coordinates_keys_and_screenshot_choices() { + let action = action_from_tool_call("task", &call("use_computer", serde_json::json!({ + "action_summary": "Select the field and type", + "actions": [ + { "type": "mouse_down", "button": "left", "x": 420, "y": 310 }, + { "type": "mouse_up", "button": "left" }, + { "type": "type_text", "text": "hello" }, + { "type": "key_down", "keycode": 36 }, + { "type": "key_up", "keycode": 36 }, + ], + "take_screenshot": false, + })), &SkillPathOrigin::Local, &HashMap::new()).unwrap(); + let AIAgentActionType::UseComputer(request) = action.action else { panic!("expected computer actions"); }; + assert_eq!(request.actions.len(), 5); + assert_eq!(request.actions[0].action, computer_use::Action::MouseDown { + button: computer_use::MouseButton::Left, at: computer_use::Vector2I::new(420, 310), + }); + assert_eq!(request.actions[3].action, computer_use::Action::KeyDown { key: computer_use::Key::Keycode(36) }); + assert_eq!(request.screenshot_params, None); +} + +#[test] +fn computer_actions_reject_invalid_targets_keys_and_unbounded_waiting() { + for actions in [ + serde_json::json!([{ "type": "key_down", "key": "Return" }]), + serde_json::json!([{ "type": "key_up", "key": "a", "keycode": 0 }]), + serde_json::json!([{ "type": "mouse_move", "x": 1, "y": 2, "target": { "type": "window", "window_id": 0, "pid": 12 } }]), + serde_json::json!([{ "type": "wait", "seconds": 6 }, { "type": "wait", "seconds": 6 }]), + serde_json::json!(vec![serde_json::json!({ "type": "wait", "seconds": 0 }); 21]), + ] { + assert!(action_from_tool_call("task", &call("use_computer", serde_json::json!({ + "action_summary": "Invalid request", "actions": actions, + })), &SkillPathOrigin::Local, &HashMap::new()).is_err()); + } +} + #[test] fn shell_calls_become_domain_actions_without_a_proto_round_trip() { let action = action_from_tool_call( diff --git a/app/src/ai/runtime/task_progress_tests.rs b/app/src/ai/runtime/task_progress_tests.rs index 0c794d6b..9b69c172 100644 --- a/app/src/ai/runtime/task_progress_tests.rs +++ b/app/src/ai/runtime/task_progress_tests.rs @@ -106,6 +106,7 @@ fn rejected_plan_creation_cannot_seed_visible_or_durable_tasks() { arguments: serde_json::json!({"documents": [{"content": "## Tasks\n- [ ] Inspect\n- [ ] Verify"}]}), }, state: PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { + images: Vec::new(), call_id: "plan".to_owned(), content: "Not advertised".to_owned(), status: galaxy_agent_core::ToolResultStatus::Error, diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 16a8a7dc..9c99f7db 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -4,6 +4,7 @@ use std::fmt; use serde::{Deserialize, Serialize}; +use crate::types::tool_image_group_call_id; use crate::{ AgentError, ContentPart, ConversationMessage, MessageContent, MessageRole, PermissionDecision, PermissionRequest, StopReason, TaskListSnapshot, ToolCall, ToolCallDecision, ToolPolicy, @@ -860,6 +861,7 @@ impl ProviderRun { .push(pending.call.id.clone()); pending.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: pending.call.id.clone(), content: "Tool execution was interrupted by application restart and was not replayed." .to_string(), @@ -1161,6 +1163,7 @@ impl ProviderRun { let content = reason.unwrap_or_else(|| "Tool permission was denied.".to_string()); call.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content, status: ToolResultStatus::Denied, @@ -1244,6 +1247,7 @@ impl ProviderRun { | PendingToolCallState::RecoveryPending => { call.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content: reason.into(), status: ToolResultStatus::Cancelled, @@ -1468,6 +1472,7 @@ impl ProviderRun { let state = if has_launch && !is_local_agent_launch(&call) { PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: "Not executed: launching a child group yields this turn. Wait for the group's completion before scheduling dependent work.".to_string(), status: ToolResultStatus::Error, @@ -1476,6 +1481,7 @@ impl ProviderRun { } else if !call.arguments.is_object() { PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: format!( "Error: '{}' received malformed arguments; expected a JSON object.", @@ -1587,6 +1593,7 @@ impl ProviderRun { if call.state.result().is_none() { call.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call.call.id.clone(), content: result_content.clone(), status, @@ -1782,16 +1789,7 @@ fn assistant_message(turn: &CompletedModelTurn) -> ConversationMessage { fn tool_result_message(results: Vec) -> ConversationMessage { ConversationMessage { role: MessageRole::User, - content: MessageContent::MultiPart( - results - .into_iter() - .map(|result| ContentPart::ToolResult { - tool_use_id: result.call_id, - content: result.content, - is_error: !matches!(result.status, ToolResultStatus::Success), - }) - .collect(), - ), + content: MessageContent::MultiPart(ToolResult::into_batch_content_parts(results)), } } @@ -1862,17 +1860,42 @@ pub fn archive_tool_results( seen_results.insert(tool_use_id.clone()) } MessageContent::MultiPart(parts) => { - parts.retain(|part| match part { - ContentPart::ToolUse { tool_use_id, .. } => { - seen_uses.insert(tool_use_id.clone()) + let mut retained_results = HashSet::new(); + let mut keep_parts: Vec<_> = parts + .iter() + .map(|part| match part { + ContentPart::ToolUse { tool_use_id, .. } => { + seen_uses.insert(tool_use_id.clone()) + } + ContentPart::ToolResult { tool_use_id, .. } => { + let keep = seen_results.insert(tool_use_id.clone()); + if keep { + retained_results.insert(tool_use_id.clone()); + } + keep + } + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } => false, + }) + .collect(); + let mut keep_image_group = false; + for (index, part) in parts.iter().enumerate() { + match part { + ContentPart::Text(_) => { + keep_image_group = tool_image_group_call_id(part) + .is_some_and(|call_id| retained_results.contains(call_id)) + && matches!(parts.get(index + 1), Some(ContentPart::Image { .. })); + keep_parts[index] = keep_image_group; + } + ContentPart::Image { .. } => keep_parts[index] = keep_image_group, + ContentPart::Reasoning { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => keep_image_group = false, } - ContentPart::ToolResult { tool_use_id, .. } => { - seen_results.insert(tool_use_id.clone()) - } - ContentPart::Text(_) - | ContentPart::Reasoning { .. } - | ContentPart::Image { .. } => false, - }); + } + let mut retained = keep_parts.into_iter(); + parts.retain(|_| retained.next().unwrap_or(false)); !parts.is_empty() } MessageContent::Text(_) => false, diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index da299550..b7b391ce 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -1,7 +1,7 @@ use serde_json::json; use super::*; -use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus}; +use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus, ToolResultImage}; fn initial_messages() -> Vec { vec![ConversationMessage { @@ -88,6 +88,7 @@ fn accept_tool_turn(run: &mut ProviderRun, turn: CompletedModelTurn) -> PendingT fn successful_result(call_id: &str, content: &str) -> ToolResult { ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content: content.to_string(), status: ToolResultStatus::Success, @@ -1673,6 +1674,127 @@ fn multipart_archive_is_not_multiplied_and_legacy_duplicates_are_repaired() { assert_eq!(archive, expected); } +#[test] +fn screenshot_results_survive_checkpoint_compaction_and_archive_deduplication() { + let mut run = ProviderRun::new( + "screenshots", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(vec![ContentPart::Image { + data: vec![99], + mime_type: "image/png".to_string(), + }]), + }], + "base", + ProviderRunLimits::default(), + ); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("capture", "use_computer"), + tool_call("capture-more", "use_computer"), + ], + &["use_computer"], + ), + ); + for (call_id, bytes) in [("capture", vec![1, 2]), ("capture-more", vec![3, 4])] { + let mut result = successful_result(call_id, "screenshot captured"); + result.images.push(ToolResultImage { + data: bytes, + mime_type: "image/png".to_string(), + }); + run.complete_tool(&batch.work_id, result).unwrap(); + } + // Pending typed results and their image payloads must survive a checkpoint. + assert_serialization_round_trip(&run); + run = serde_json::from_str(&serde_json::to_string(&run).unwrap()).unwrap(); + run.commit_tool_batch(&batch.work_id).unwrap(); + let transcript = run.transcript().to_vec(); + let summary = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Screenshots inspected.".to_string()), + }]; + run.compact_transcript_at_model_boundary(transcript.len(), summary.clone()) + .unwrap(); + assert_eq!(run.transcript(), summary); + assert_eq!(run.tool_result_archive().len(), 2); + let MessageContent::MultiPart(parts) = &run.tool_result_archive()[1].content else { + panic!("expected archived screenshot result batch"); + }; + assert_eq!(parts.len(), 6); + assert!(matches!( + &parts[2], ContentPart::Text(text) if text == "Images returned by tool call capture:" + )); + assert!(matches!( + &parts[3], ContentPart::Image { data, .. } if data == &[1, 2] + )); + assert!(matches!( + &parts[4], ContentPart::Text(text) if text == "Images returned by tool call capture-more:" + )); + assert!(matches!( + &parts[5], ContentPart::Image { data, .. } if data == &[3, 4] + )); + let expected = run.tool_result_archive().to_vec(); + let duplicated = expected.iter().cloned().chain(transcript).collect(); + run.replace_tool_result_archive(duplicated); + assert_eq!(run.tool_result_archive(), expected); + assert_serialization_round_trip(&run); +} + +#[test] +fn archives_only_images_labelled_for_new_results_in_the_same_message() { + let result = successful_result("capture", "captured"); + let mut parts = ToolResult::into_batch_content_parts(vec![result]); + parts.extend([ + ContentPart::Image { + data: vec![99], + mime_type: "image/png".to_string(), + }, + ContentPart::Text("Images returned by tool call missing:".to_string()), + ContentPart::Image { + data: vec![88], + mime_type: "image/png".to_string(), + }, + ContentPart::Text("Images returned by tool call capture:".to_string()), + ContentPart::Image { + data: vec![1], + mime_type: "image/png".to_string(), + }, + ContentPart::Text("Unrelated user upload".to_string()), + ContentPart::Image { + data: vec![77], + mime_type: "image/png".to_string(), + }, + ]); + let mut archive = Vec::new(); + let message = ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(parts), + }; + archive_tool_results(&mut archive, vec![message.clone()]); + let MessageContent::MultiPart(parts) = &archive[0].content else { + panic!("expected correlated tool result with one image"); + }; + assert_eq!(parts.len(), 3); + assert!(matches!(&parts[2], ContentPart::Image { data, .. } if data == &[1])); + let expected = archive.clone(); + archive_tool_results(&mut archive, vec![message]); + assert_eq!(archive, expected); +} + +#[test] +fn legacy_tool_results_without_images_still_restore() { + let result: ToolResult = serde_json::from_value(json!({ + "call_id": "legacy", + "content": "old text result", + "status": "Success" + })) + .unwrap(); + assert!(result.images.is_empty()); + assert_eq!(result.content, "old text result"); +} + #[test] fn local_child_launch_yields_and_does_not_execute_batched_follow_up_tools() { let mut run = run(); @@ -1718,6 +1840,7 @@ fn failed_child_launch_allows_a_corrective_model_turn() { run.complete_tool( &batch.work_id, ToolResult { + images: Vec::new(), call_id: "children".to_string(), content: "Launch failed".to_string(), status: ToolResultStatus::Error, diff --git a/crates/galaxy_agent_core/src/tool_policy.rs b/crates/galaxy_agent_core/src/tool_policy.rs index acf9eec9..b97c9fc7 100644 --- a/crates/galaxy_agent_core/src/tool_policy.rs +++ b/crates/galaxy_agent_core/src/tool_policy.rs @@ -2,9 +2,10 @@ use std::collections::{BTreeSet, HashMap, VecDeque}; use serde_json::Value as JsonValue; +use crate::types::tool_image_group_call_id; use crate::{ ContentPart, ConversationMessage, MessageContent, TaskListSnapshot, ToolCall, ToolDefinition, - ToolResult, ToolResultStatus, UPDATE_TASK_LIST_NAME, + ToolResult, ToolResultImage, ToolResultStatus, UPDATE_TASK_LIST_NAME, }; pub const RECALL_TOOL_HISTORY_NAME: &str = "recall_tool_history"; @@ -124,6 +125,7 @@ impl ToolPolicy { ) }; return ToolCallDecision::Reject(ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: format!( "Error: '{}' is not a valid tool for this request; {available}. Do not invent tool names.", @@ -134,13 +136,15 @@ impl ToolPolicy { } if call.name == RECALL_TOOL_HISTORY_NAME { + let (content, images) = recall_tool_history_with_images( + messages, + archive, + ToolHistoryQuery::from_arguments(&call.arguments), + ); return ToolCallDecision::Inline(ToolResult { + images, call_id: call.id.clone(), - content: recall_tool_history( - messages, - archive, - ToolHistoryQuery::from_arguments(&call.arguments), - ), + content, status: ToolResultStatus::Success, }); } @@ -148,11 +152,13 @@ impl ToolPolicy { if call.name == UPDATE_TASK_LIST_NAME { return match TaskListSnapshot::from_arguments(&call.arguments) { Ok(snapshot) => ToolCallDecision::Inline(ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: snapshot.to_tool_result_content(), status: ToolResultStatus::Success, }), Err(error) => ToolCallDecision::Reject(ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: format!("Error: {error}"), status: ToolResultStatus::Error, @@ -201,6 +207,14 @@ pub fn recall_tool_history( archive: &[ConversationMessage], query: ToolHistoryQuery<'_>, ) -> String { + recall_tool_history_with_images(messages, archive, query).0 +} + +fn recall_tool_history_with_images( + messages: &[ConversationMessage], + archive: &[ConversationMessage], + query: ToolHistoryQuery<'_>, +) -> (String, Vec) { let mut entries = Vec::new(); collect_tool_entries(archive, &mut entries); collect_tool_entries(messages, &mut entries); @@ -222,26 +236,43 @@ pub fn recall_tool_history( .len() .checked_sub(1usize.saturating_add(query.offset_from_end)) else { - return "No matching tool calls found in conversation history.".to_string(); + return ( + "No matching tool calls found in conversation history.".to_string(), + Vec::new(), + ); }; 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 + ( + format!( + "Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{result}", + entry.name, entry.tool_use_id, entry.input + ), + entry + .images + .iter() + .map(|(data, mime_type)| ToolResultImage { + data: data.to_vec(), + mime_type: (*mime_type).to_string(), + }) + .collect(), ) } -#[derive(Clone, Debug, PartialEq, Eq)] -struct ToolHistoryEntry { +#[derive(Clone, Debug, PartialEq)] +struct ToolHistoryEntry<'a> { tool_use_id: String, name: String, input: String, result: String, + images: Vec<(&'a [u8], &'a str)>, } -fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec) { +fn collect_tool_entries<'a>( + messages: &'a [ConversationMessage], + entries: &mut Vec>, +) { let mut pending = Vec::new(); for message in messages { match &message.content { @@ -256,6 +287,7 @@ fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec pair_result(tool_use_id, content, &mut pending, entries), MessageContent::MultiPart(parts) => { + let entries_start = entries.len(); for part in parts { match part { ContentPart::ToolUse { @@ -275,6 +307,41 @@ fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec {} } } + // ToolResult::into_batch_content_parts places labelled image groups + // after every correlated result. Only attach images to results from + // this same message, never a nearby user upload or another call. + let result_ids: BTreeSet<_> = parts + .iter() + .filter_map(|part| match part { + ContentPart::ToolResult { tool_use_id, .. } => Some(tool_use_id.as_str()), + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } + | ContentPart::ToolUse { .. } => None, + }) + .collect(); + let mut image_call_id = None; + for part in parts { + match part { + ContentPart::Text(_) => { + image_call_id = tool_image_group_call_id(part) + .filter(|call_id| result_ids.contains(call_id)); + } + ContentPart::Image { data, mime_type } => { + if let Some(entry) = image_call_id.and_then(|call_id| { + entries[entries_start..] + .iter_mut() + .rev() + .find(|entry| entry.tool_use_id == call_id) + }) { + entry.images.push((data, mime_type)); + } + } + ContentPart::Reasoning { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => image_call_id = None, + } + } } MessageContent::Text(_) => {} } @@ -285,7 +352,7 @@ fn pair_result( tool_use_id: &str, content: &str, pending: &mut Vec<(String, String, String)>, - entries: &mut Vec, + entries: &mut Vec>, ) { let Some(index) = pending.iter().position(|(id, _, _)| id == tool_use_id) else { return; @@ -296,6 +363,7 @@ fn pair_result( name, input, result: content.to_string(), + images: Vec::new(), }); } diff --git a/crates/galaxy_agent_core/src/tool_policy_tests.rs b/crates/galaxy_agent_core/src/tool_policy_tests.rs index 5e876110..1758b90b 100644 --- a/crates/galaxy_agent_core/src/tool_policy_tests.rs +++ b/crates/galaxy_agent_core/src/tool_policy_tests.rs @@ -212,6 +212,120 @@ fn recall_never_returns_a_prior_recall_call() { ); } +fn screenshot_exchange(id: &str, bytes: Vec) -> Vec { + let mut messages = tool_exchange(id, "use_computer", serde_json::json!({}), "captured"); + messages[1].content = MessageContent::MultiPart(ToolResult::into_batch_content_parts(vec![ + ToolResult { + call_id: id.to_string(), + content: "captured".to_string(), + status: ToolResultStatus::Success, + images: vec![ToolResultImage { + data: bytes, + mime_type: "image/png".to_string(), + }], + }, + ])); + messages +} + +#[test] +fn recall_attaches_only_the_exact_selected_screenshots_from_archived_or_live_results() { + let archive = screenshot_exchange("capture", vec![1, 2]); + let mut messages = screenshot_exchange("capture-more", vec![3, 4]); + messages.push(ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(vec![ + ContentPart::Text("Images returned by tool call capture:".to_string()), + ContentPart::Image { + data: vec![99], + mime_type: "image/png".to_string(), + }, + ]), + }); + let policy = ToolPolicy::new(&[definition(RECALL_TOOL_HISTORY_NAME)]); + for (arguments, original_id, bytes) in [ + ( + serde_json::json!({"tool_use_id": "capture"}), + "capture", + vec![1, 2], + ), + ( + serde_json::json!({"tool_name": "use_computer"}), + "capture-more", + vec![3, 4], + ), + ( + serde_json::json!({"tool_name": "use_computer", "offset_from_end": 1}), + "capture", + vec![1, 2], + ), + ] { + let call = ToolCall { + id: "recall-image".to_string(), + name: RECALL_TOOL_HISTORY_NAME.to_string(), + arguments, + }; + let ToolCallDecision::Inline(result) = policy.decide(&call, &messages, &archive) else { + panic!("recall should execute inline"); + }; + assert_eq!(result.call_id, call.id); + assert!(result + .content + .contains(&format!("Tool Use ID: {original_id}\n"))); + assert_eq!( + result.images, + vec![ToolResultImage { + data: bytes, + mime_type: "image/png".to_string(), + }] + ); + let parts = ToolResult::into_batch_content_parts(vec![result]); + assert!(matches!( + &parts[0], ContentPart::ToolResult { tool_use_id, .. } if tool_use_id == "recall-image" + )); + assert!(matches!( + &parts[1], ContentPart::Text(text) if text == "Images returned by tool call recall-image:" + )); + assert!(matches!(&parts[2], ContentPart::Image { .. })); + } + let call = ToolCall { + id: "missing".to_string(), + name: RECALL_TOOL_HISTORY_NAME.to_string(), + arguments: serde_json::json!({"tool_use_id": "capture-missing"}), + }; + let ToolCallDecision::Inline(result) = policy.decide(&call, &messages, &archive) else { + panic!("missing recall should execute inline"); + }; + assert!(result.images.is_empty()); + assert_eq!( + result.content, + "No matching tool calls found in conversation history." + ); +} + +#[test] +fn recall_does_not_reattach_images_from_an_earlier_recall_result() { + let mut messages = screenshot_exchange("recall-previous", vec![99]); + let MessageContent::ToolUse { name, .. } = &mut messages[0].content else { + panic!("expected tool call"); + }; + *name = RECALL_TOOL_HISTORY_NAME.to_string(); + let policy = ToolPolicy::new(&[definition(RECALL_TOOL_HISTORY_NAME)]); + let call = ToolCall { + id: "recall-now".to_string(), + name: RECALL_TOOL_HISTORY_NAME.to_string(), + arguments: serde_json::json!({}), + }; + let ToolCallDecision::Inline(result) = policy.decide(&call, &messages, &[]) else { + panic!("recall should execute inline"); + }; + assert!(result.images.is_empty()); + assert_eq!( + result.content, + "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); diff --git a/crates/galaxy_agent_core/src/types.rs b/crates/galaxy_agent_core/src/types.rs index 8f5c720d..e17f88dc 100644 --- a/crates/galaxy_agent_core/src/types.rs +++ b/crates/galaxy_agent_core/src/types.rs @@ -190,9 +190,53 @@ pub struct ToolResult { pub call_id: String, pub content: String, pub status: ToolResultStatus, + /// Images returned by a tool, kept outside the text so provider adapters receive real images. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub images: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ToolResultImage { + pub data: Vec, + pub mime_type: String, +} + +pub(crate) fn tool_image_group_call_id(part: &ContentPart) -> Option<&str> { + let ContentPart::Text(text) = part else { + return None; + }; + text.strip_prefix("Images returned by tool call ")? + .strip_suffix(':') + .filter(|call_id| !call_id.is_empty()) } impl ToolResult { + pub fn into_batch_content_parts(results: Vec) -> Vec { + let mut parts = Vec::new(); + let mut image_parts = Vec::new(); + for result in results { + let is_error = !matches!(result.status, ToolResultStatus::Success); + if !result.images.is_empty() { + image_parts.push(ContentPart::Text(format!( + "Images returned by tool call {}:", result.call_id + ))); + image_parts.extend(result.images.into_iter().map(|image| ContentPart::Image { + data: image.data, + mime_type: image.mime_type, + })); + } + parts.push(ContentPart::ToolResult { + tool_use_id: result.call_id, + content: result.content, + is_error, + }); + } + // OpenAI requires every tool result before the next user message. Its adapter splits + // adjacent images into a user message, so never interleave them with a result batch. + parts.extend(image_parts); + parts + } + pub fn is_error(&self) -> bool { matches!( self.status, diff --git a/crates/galaxy_agent_core/src/types_tests.rs b/crates/galaxy_agent_core/src/types_tests.rs index 54e88198..62a65df0 100644 --- a/crates/galaxy_agent_core/src/types_tests.rs +++ b/crates/galaxy_agent_core/src/types_tests.rs @@ -169,6 +169,7 @@ fn tool_events_keep_one_call_id_across_permission_and_execution() { }, ToolEvent::Completed { result: ToolResult { + images: Vec::new(), call_id: "call-1".to_string(), content: "ok".to_string(), status: ToolResultStatus::Success, @@ -182,11 +183,13 @@ fn tool_events_keep_one_call_id_across_permission_and_execution() { #[test] fn denied_results_are_errors_but_cancelled_results_are_distinct() { let denied = ToolResult { + images: Vec::new(), call_id: "denied".to_string(), content: "permission denied".to_string(), status: ToolResultStatus::Denied, }; let cancelled = ToolResult { + images: Vec::new(), call_id: "cancelled".to_string(), content: "cancelled".to_string(), status: ToolResultStatus::Cancelled,