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.
This commit is contained in:
Ryan Ward
2026-10-22 14:25:21 +00:00
parent d5fa711e5c
commit f6c6434e5e
30 changed files with 1176 additions and 101 deletions
+8 -6
View File
@@ -174,9 +174,10 @@ base_url = "http://localhost:4000/v1"
api_key = "sk-..." api_key = "sk-..."
[[ai.providers.models]] [[ai.providers.models]]
model_id = "claude-sonnet-4-20250514[1m]" model_id = "claude-sonnet-4-20250514"
display_name = "Claude Sonnet 4 (1M)" display_name = "Claude Sonnet 4"
context_size = 1000000 context_size = 1000000
default_context_size = 200000
[[ai.providers]] [[ai.providers]]
name = "Ollama (Local)" name = "Ollama (Local)"
@@ -189,10 +190,9 @@ context_size = 128000
``` ```
**OpenAI/LiteLLM model discovery**: **OpenAI/LiteLLM model discovery**:
- Models can be auto-fetched from the `/models` endpoint via the Settings > OpenAI / LiteLLM page - Models can be auto-fetched from `/models` and enriched with LiteLLM `/model/info` via the Settings > OpenAI / LiteLLM page.
- For each model, the system probes `{model_id}[1m]` with a minimal chat completion request - 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.
- If the `[1m]` variant is accepted (HTTP 200 or 429), it's used with 1M context window - 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.
- Otherwise, the base model ID is used with its reported context size
- Models injected via `ai.providers[]` are routed to their specific endpoint (per-model routing map) - 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 - 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 - 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 - 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` - 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 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. - 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. - 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.
+8
View File
@@ -28,12 +28,17 @@ pub(crate) fn prepare_direct_provider_params(
.clone() .clone()
.unwrap_or_else(|| get_supported_tools(params)); .unwrap_or_else(|| get_supported_tools(params));
remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled); 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 = let mut supported_cli_agent_tools =
supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(params)); supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(params));
remove_orchestration_tools_if_disabled( remove_orchestration_tools_if_disabled(
&mut supported_cli_agent_tools, &mut supported_cli_agent_tools,
params.orchestration_enabled, 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 { if params.permission_mode == crate::ai::execution_profiles::PermissionMode::ReadOnly {
let permitted = |tool: &api::ToolType| is_read_only_tool(tool) 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)); || (!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> {
api::ToolType::CreateDocuments, api::ToolType::CreateDocuments,
api::ToolType::EditDocuments, api::ToolType::EditDocuments,
]; ];
if params.computer_use_enabled {
supported_tools.push(api::ToolType::RequestComputerUse);
}
if params.orchestration_enabled { if params.orchestration_enabled {
supported_tools.push(api::ToolType::Subagent); supported_tools.push(api::ToolType::Subagent);
+31
View File
@@ -91,6 +91,37 @@ fn supported_tools_include_plan_document_capabilities() {
assert!(supported_tools.contains(&api::ToolType::EditDocuments)); 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] #[test]
fn supported_tools_omit_subagents_when_orchestration_is_disabled() { fn supported_tools_omit_subagents_when_orchestration_is_disabled() {
let params = request_params_with_ask_user_question_enabled(false); let params = request_params_with_ask_user_question_enabled(false);
+75
View File
@@ -8,6 +8,81 @@ use crate::ai::runtime::ProviderActionContext;
const TOOL_PREFIX: &str = "galaxy:provider-tool:v1:"; const TOOL_PREFIX: &str = "galaxy:provider-tool:v1:";
pub(crate) fn computer_use_tool(action: &super::AIAgentActionType) -> Option<api::message::tool_call::Tool> {
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<api::message::tool_call::ComputerUseTarget> {
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)] #[derive(Serialize, Deserialize)]
struct SavedTool { struct SavedTool {
context: ProviderActionContext, context: ProviderActionContext,
@@ -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] #[test]
fn saved_tools_and_results_restore_as_individual_panes_across_turns() { fn saved_tools_and_results_restore_as_individual_panes_across_turns() {
let context = ProviderActionContext::new_for_test("task"); let context = ProviderActionContext::new_for_test("task");
+54 -1
View File
@@ -418,19 +418,72 @@ fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bo
} else { } else {
ToolResultStatus::Success 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 { let content = if permission_denied {
format!("Permission denied by the user. {content}") format!("Permission denied by the user. {content}")
} else { } else {
content content
}; };
ToolResult { ToolResult {
images: if permission_denied { Vec::new() } else { images },
call_id: action_result.id.to_string(), call_id: action_result.id.to_string(),
content, content,
status, status,
} }
} }
fn computer_use_model_content(
result: &AIAgentActionResultType,
) -> Option<(String, Vec<galaxy_agent_core::ToolResultImage>)> {
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"))] #[cfg(not(target_family = "wasm"))]
fn action_tool_name(action: &AIAgentAction) -> String { fn action_tool_name(action: &AIAgentAction) -> String {
action action
@@ -24,10 +24,8 @@ impl UseComputerExecutor {
return false; return false;
}; };
// We unconditionally return true here because this action is only executed by // The shared action boundary applies the active permission mode to every action.
// the computer use subagent, which cannot begin without the user approving it via // Once that check passes, the executor has no additional approval requirement.
// a `RequestComputerUse` action, and the approval extends to all computer use
// actions within that computer use subagent.
true true
} }
@@ -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 }")); 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] #[test]
fn action_permission_kinds_match_the_safety_boundary() { fn action_permission_kinds_match_the_safety_boundary() {
assert_eq!( assert_eq!(
+5 -1
View File
@@ -1039,6 +1039,7 @@ fn convert_provider_tool_batch(
matches!(pending.state, PendingToolCallState::RecoveryPending), matches!(pending.state, PendingToolCallState::RecoveryPending),
)), )),
Err(message) => invalid_results.push(ToolResult { Err(message) => invalid_results.push(ToolResult {
images: Vec::new(),
call_id: pending.call.id.clone(), call_id: pending.call.id.clone(),
content: format!("Invalid {} tool input: {message}", pending.call.name), content: format!("Invalid {} tool input: {message}", pending.call.name),
status: ToolResultStatus::Error, status: ToolResultStatus::Error,
@@ -8460,7 +8461,10 @@ impl BlocklistAIController {
matches!( matches!(
pending.call.name.as_str(), pending.call.name.as_str(),
"request_computer_use" | "use_computer" "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) used_computer.then_some(COMPUTER_USE_PROVIDER_PROFILE)
} }
+83 -32
View File
@@ -2819,9 +2819,12 @@ pub(crate) fn openai_model_context_size(
.max_input_tokens .max_input_tokens
.unwrap_or(model.context_size) .unwrap_or(model.context_size)
.min(model.context_size); .min(model.context_size);
if provider_kind == OpenAIProviderKind::ChatGPTSubscription { if matches!(
// The catalog's default window is the supported session budget. The provider_kind,
// advertised maximum must not silently opt subscriptions into long context. 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. // Older saved catalogs kept the default in max_input_tokens instead.
return model return model
.default_context_size .default_context_size
@@ -2959,7 +2962,7 @@ fn openai_model_context_window(
is_configurable: max_context_size > default_context_size, is_configurable: max_context_size > default_context_size,
min: default_context_size, min: default_context_size,
max: max_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, default_max: max_context_size,
} }
} }
@@ -2971,6 +2974,53 @@ fn u32_from_json_any(value: &serde_json::Value, keys: &[&str]) -> Option<u32> {
.and_then(|value| u32::try_from(value).ok()) .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<u32>, Option<u32>) {
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"))] #[cfg(not(target_family = "wasm"))]
fn normalize_codex_release_version(version: &str) -> Option<String> { fn normalize_codex_release_version(version: &str) -> Option<String> {
let version = version.trim(); 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_name = entry["model_name"].as_str()?;
let model_info = &entry["model_info"]; let model_info = &entry["model_info"];
let max_input_tokens = model_info["max_input_tokens"] let (context_size, max_input_tokens, default_context_size) =
.as_u64() openai_context_limits(model_info);
.and_then(|v| u32::try_from(v).ok());
let max_output_tokens = model_info["max_output_tokens"] let max_output_tokens = model_info["max_output_tokens"]
.as_u64() .as_u64()
.and_then(|v| u32::try_from(v).ok()); .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 vision_supported = model_info["supports_vision"].as_bool().unwrap_or(true);
let display_name = model_name.replace(['-', '_'], " "); let display_name = model_name.replace(['-', '_'], " ");
@@ -3271,7 +3313,7 @@ async fn fetch_from_litellm_model_info(
vision_supported, vision_supported,
context_size, context_size,
max_input_tokens, max_input_tokens,
default_context_size: None, default_context_size,
max_output_tokens, max_output_tokens,
provider, provider,
supports_system_messages: if model_name.starts_with("codex-gpt-") { supports_system_messages: if model_name.starts_with("codex-gpt-") {
@@ -3318,7 +3360,30 @@ async fn fetch_litellm_catalog(
metadata metadata
.iter() .iter()
.find(|info| info.model_id == model.model_id) .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) .unwrap_or(model)
}) })
.collect()) .collect())
@@ -3369,21 +3434,7 @@ async fn fetch_from_openai_models(
.iter() .iter()
.filter_map(|m| { .filter_map(|m| {
let id = m["id"].as_str()?; let id = m["id"].as_str()?;
let max_input_tokens = u32_from_any( let (context_size, max_input_tokens, default_context_size) = openai_context_limits(m);
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 max_output_tokens = u32_from_any( let max_output_tokens = u32_from_any(
m, m,
&[ &[
@@ -3430,7 +3481,7 @@ async fn fetch_from_openai_models(
.unwrap_or(true), .unwrap_or(true),
context_size, context_size,
max_input_tokens, max_input_tokens,
default_context_size: None, default_context_size,
max_output_tokens, max_output_tokens,
provider, provider,
supports_system_messages: if id.starts_with("codex-gpt-") { supports_system_messages: if id.starts_with("codex-gpt-") {
+150
View File
@@ -1,3 +1,4 @@
use serde_json::json;
use warpui::App; use warpui::App;
use super::*; 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] #[test]
fn litellm_refresh_replaces_membership_and_preserves_preferences() { fn litellm_refresh_replaces_membership_and_preserves_preferences() {
let mut retained = openai_model("retained"); let mut retained = openai_model("retained");
+2
View File
@@ -663,6 +663,8 @@ fn estimate_cost_cents(
} }
const KNOWN_TOOLS: &[&str] = &[ const KNOWN_TOOLS: &[&str] = &[
"request_computer_use",
"use_computer",
"run_shell_command", "run_shell_command",
"read_files", "read_files",
"apply_file_diffs", "apply_file_diffs",
+71
View File
@@ -1714,6 +1714,8 @@ pub(crate) fn tool_name_is_supported(name: &str, supported: &HashSet<api::ToolTy
"ask_user_question" => has(ToolType::AskUserQuestion), "ask_user_question" => has(ToolType::AskUserQuestion),
"read_skill" => has(ToolType::ReadSkill), "read_skill" => has(ToolType::ReadSkill),
"fetch_conversation" => has(ToolType::FetchConversation), "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 // This tool is implemented entirely inside the direct-provider response
// translator, so it does not need a client ToolType capability bit. // translator, so it does not need a client ToolType capability bit.
"recall_tool_history" => true, "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"; 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<ToolDefinition> { pub fn default_tool_definitions() -> Vec<ToolDefinition> {
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 { ToolDefinition {
name: "run_shell_command".to_string(), 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(), 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(),
@@ -839,6 +839,14 @@ pub fn build_tool_call_message(
serde_json::from_str(tool_input_json).unwrap_or(serde_json::json!({})); serde_json::from_str(tool_input_json).unwrap_or(serde_json::json!({}));
let tool = match tool_name { 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" => { "run_shell_command" => {
let command = input let command = input
.get("command") .get("command")
@@ -1596,6 +1604,8 @@ fn run_agents_harness_from_str(harness_type: &str) -> Option<api::Harness> {
/// Built-in tools that Galaxy knows how to execute directly. /// Built-in tools that Galaxy knows how to execute directly.
const KNOWN_TOOLS: &[&str] = &[ const KNOWN_TOOLS: &[&str] = &[
"request_computer_use",
"use_computer",
"run_shell_command", "run_shell_command",
"read_files", "read_files",
"apply_file_diffs", "apply_file_diffs",
@@ -104,6 +104,7 @@ fn restored_task_progress_replay_does_not_duplicate_the_persisted_checklist() {
arguments: serde_json::to_value(&snapshot).unwrap(), arguments: serde_json::to_value(&snapshot).unwrap(),
}, },
state: galaxy_agent_core::PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { 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, call_id: "progress".to_owned(), status: galaxy_agent_core::ToolResultStatus::Success,
content: snapshot.to_tool_result_content(), content: snapshot.to_tool_result_content(),
}}, }},
+1
View File
@@ -4,6 +4,7 @@ mod provider_run_coordinator;
mod rig; mod rig;
mod rig_request; mod rig_request;
mod rig_tool; mod rig_tool;
pub(crate) use rig_tool::action_from_tool_call;
mod task_progress; mod task_progress;
pub(crate) use child_context::add_parent_context; pub(crate) use child_context::add_parent_context;
@@ -461,6 +461,7 @@ async fn coordinator_awaiting_tools(
fn tool_result(call_id: &str, content: &str, status: ToolResultStatus) -> ToolEvent { fn tool_result(call_id: &str, content: &str, status: ToolResultStatus) -> ToolEvent {
ToolEvent::Completed { ToolEvent::Completed {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: call_id.to_string(), call_id: call_id.to_string(),
content: content.to_string(), content: content.to_string(),
status, status,
@@ -514,6 +515,7 @@ async fn one_run_drives_model_tool_and_followup_turns_with_atomic_history() {
.complete_tool( .complete_tool(
&batch.work_id, &batch.work_id,
ToolResult { ToolResult {
images: Vec::new(),
call_id: "read-1".to_string(), call_id: "read-1".to_string(),
content: "manifest contents".to_string(), content: "manifest contents".to_string(),
status: ToolResultStatus::Success, status: ToolResultStatus::Success,
+15 -10
View File
@@ -133,20 +133,25 @@ pub(crate) async fn prepare_provider_run(
let skill_path_origin = params.session_context.skill_path_origin(); let skill_path_origin = params.session_context.skill_path_origin();
let mut cli_params = params.clone(); let mut cli_params = params.clone();
let computer_use_model = params.computer_use_model.as_str(); let computer_use_model = params.computer_use_model.as_str();
let computer_use_profile = if params.computer_use_enabled let computer_use_profile = if supported_tools.contains(&ToolType::RequestComputerUse) {
&& !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 mut computer_params = params.clone(); let mut computer_params = params.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_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( let (runtime, prepared) = prepare_provider_profile(
computer_use_provider_config, computer_provider,
computer_params, computer_params,
supported_tools.clone(), computer_tools,
supported_cli_agent_tools.clone(), supported_cli_agent_tools.clone(),
None, None,
) )
+7 -11
View File
@@ -389,17 +389,7 @@ fn input_messages(
) -> Vec<ConversationMessage> { ) -> Vec<ConversationMessage> {
let mut messages = Vec::new(); let mut messages = Vec::new();
if !tool_results.is_empty() { if !tool_results.is_empty() {
let mut parts = tool_results let mut parts = ToolResult::into_batch_content_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::<Vec<_>>();
let content = if parts.len() == 1 { let content = if parts.len() == 1 {
let ContentPart::ToolResult { let ContentPart::ToolResult {
tool_use_id, tool_use_id,
@@ -1043,6 +1033,12 @@ fn build_system_prompt(
{ {
prompt.push_str(CODING_DELEGATION_GUIDANCE); 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"); prompt.push_str("## Available Tools\n");
if tools.is_empty() { if tools.is_empty() {
prompt.push_str("No tools are available. Do not invent tool calls.\n"); prompt.push_str("No tools are available. Do not invent tool calls.\n");
+2
View File
@@ -533,6 +533,7 @@ fn lrc_snapshot_follow_up_uses_the_cli_monitor_prompt_and_tools() {
), ),
}; };
let snapshot_tool_result = ToolResult { let snapshot_tool_result = ToolResult {
images: Vec::new(),
call_id: "run-call".to_string(), call_id: "run-call".to_string(),
content: result.result.model_content(), content: result.result.model_content(),
status: ToolResultStatus::Success, status: ToolResultStatus::Success,
@@ -842,6 +843,7 @@ fn normalized_tool_outcomes_are_the_only_action_results_sent_to_rig() {
let tool_results = statuses let tool_results = statuses
.iter() .iter()
.map(|(call_id, status, _)| ToolResult { .map(|(call_id, status, _)| ToolResult {
images: Vec::new(),
call_id: (*call_id).to_string(), call_id: (*call_id).to_string(),
content: format!("normalized-{call_id}"), content: format!("normalized-{call_id}"),
status: *status, status: *status,
+21
View File
@@ -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] #[test]
fn subscription_context_budget_uses_catalog_limit_and_clamps_stale_overrides() { fn subscription_context_budget_uses_catalog_limit_and_clamps_stale_overrides() {
let mut config = openai_config("subscription-model"); let mut config = openai_config("subscription-model");
+124
View File
@@ -55,6 +55,23 @@ pub(crate) fn action_from_tool_call(
rationale: optional_string(input, "rationale")?, rationale: optional_string(input, "rationale")?,
citations: Vec::new(), 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 { "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest {
locations: limited_required_array(input, "files", MAX_BATCH_READ_ITEMS)? locations: limited_required_array(input, "files", MAX_BATCH_READ_ITEMS)?
.iter() .iter()
@@ -472,6 +489,113 @@ fn optional_nonnegative_i32(input: &serde_json::Value, key: &str) -> Result<Opti
.transpose() .transpose()
} }
fn computer_i32(input: &serde_json::Value, key: &str) -> Result<i32, String> {
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<computer_use::Target, String> {
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<computer_use::ScreenshotParams, String> {
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<usize, String> {
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<Vec<computer_use::TargetedAction>, 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<Vector2I, String> {
Ok(Vector2I::new(computer_i32(input, "x")?, computer_i32(input, "y")?))
};
let button = || -> Result<MouseButton, String> {
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<Key, String> {
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( fn skill_references(
input: &serde_json::Value, input: &serde_json::Value,
skill_path_origin: &SkillPathOrigin, skill_path_origin: &SkillPathOrigin,
+49
View File
@@ -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] #[test]
fn shell_calls_become_domain_actions_without_a_proto_round_trip() { fn shell_calls_become_domain_actions_without_a_proto_round_trip() {
let action = action_from_tool_call( let action = action_from_tool_call(
@@ -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"}]}), arguments: serde_json::json!({"documents": [{"content": "## Tasks\n- [ ] Inspect\n- [ ] Verify"}]}),
}, },
state: PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { state: PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult {
images: Vec::new(),
call_id: "plan".to_owned(), call_id: "plan".to_owned(),
content: "Not advertised".to_owned(), content: "Not advertised".to_owned(),
status: galaxy_agent_core::ToolResultStatus::Error, status: galaxy_agent_core::ToolResultStatus::Error,
+36 -13
View File
@@ -4,6 +4,7 @@ use std::fmt;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::types::tool_image_group_call_id;
use crate::{ use crate::{
AgentError, ContentPart, ConversationMessage, MessageContent, MessageRole, PermissionDecision, AgentError, ContentPart, ConversationMessage, MessageContent, MessageRole, PermissionDecision,
PermissionRequest, StopReason, TaskListSnapshot, ToolCall, ToolCallDecision, ToolPolicy, PermissionRequest, StopReason, TaskListSnapshot, ToolCall, ToolCallDecision, ToolPolicy,
@@ -860,6 +861,7 @@ impl ProviderRun {
.push(pending.call.id.clone()); .push(pending.call.id.clone());
pending.state = PendingToolCallState::Resolved { pending.state = PendingToolCallState::Resolved {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: pending.call.id.clone(), call_id: pending.call.id.clone(),
content: "Tool execution was interrupted by application restart and was not replayed." content: "Tool execution was interrupted by application restart and was not replayed."
.to_string(), .to_string(),
@@ -1161,6 +1163,7 @@ impl ProviderRun {
let content = reason.unwrap_or_else(|| "Tool permission was denied.".to_string()); let content = reason.unwrap_or_else(|| "Tool permission was denied.".to_string());
call.state = PendingToolCallState::Resolved { call.state = PendingToolCallState::Resolved {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: call_id.to_string(), call_id: call_id.to_string(),
content, content,
status: ToolResultStatus::Denied, status: ToolResultStatus::Denied,
@@ -1244,6 +1247,7 @@ impl ProviderRun {
| PendingToolCallState::RecoveryPending => { | PendingToolCallState::RecoveryPending => {
call.state = PendingToolCallState::Resolved { call.state = PendingToolCallState::Resolved {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: call_id.to_string(), call_id: call_id.to_string(),
content: reason.into(), content: reason.into(),
status: ToolResultStatus::Cancelled, status: ToolResultStatus::Cancelled,
@@ -1468,6 +1472,7 @@ impl ProviderRun {
let state = if has_launch && !is_local_agent_launch(&call) { let state = if has_launch && !is_local_agent_launch(&call) {
PendingToolCallState::Resolved { PendingToolCallState::Resolved {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: call.id.clone(), 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(), 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, status: ToolResultStatus::Error,
@@ -1476,6 +1481,7 @@ impl ProviderRun {
} else if !call.arguments.is_object() { } else if !call.arguments.is_object() {
PendingToolCallState::Resolved { PendingToolCallState::Resolved {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: call.id.clone(), call_id: call.id.clone(),
content: format!( content: format!(
"Error: '{}' received malformed arguments; expected a JSON object.", "Error: '{}' received malformed arguments; expected a JSON object.",
@@ -1587,6 +1593,7 @@ impl ProviderRun {
if call.state.result().is_none() { if call.state.result().is_none() {
call.state = PendingToolCallState::Resolved { call.state = PendingToolCallState::Resolved {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: call.call.id.clone(), call_id: call.call.id.clone(),
content: result_content.clone(), content: result_content.clone(),
status, status,
@@ -1782,16 +1789,7 @@ fn assistant_message(turn: &CompletedModelTurn) -> ConversationMessage {
fn tool_result_message(results: Vec<ToolResult>) -> ConversationMessage { fn tool_result_message(results: Vec<ToolResult>) -> ConversationMessage {
ConversationMessage { ConversationMessage {
role: MessageRole::User, role: MessageRole::User,
content: MessageContent::MultiPart( content: MessageContent::MultiPart(ToolResult::into_batch_content_parts(results)),
results
.into_iter()
.map(|result| ContentPart::ToolResult {
tool_use_id: result.call_id,
content: result.content,
is_error: !matches!(result.status, ToolResultStatus::Success),
})
.collect(),
),
} }
} }
@@ -1862,17 +1860,42 @@ pub fn archive_tool_results(
seen_results.insert(tool_use_id.clone()) seen_results.insert(tool_use_id.clone())
} }
MessageContent::MultiPart(parts) => { MessageContent::MultiPart(parts) => {
parts.retain(|part| match part { let mut retained_results = HashSet::new();
let mut keep_parts: Vec<_> = parts
.iter()
.map(|part| match part {
ContentPart::ToolUse { tool_use_id, .. } => { ContentPart::ToolUse { tool_use_id, .. } => {
seen_uses.insert(tool_use_id.clone()) seen_uses.insert(tool_use_id.clone())
} }
ContentPart::ToolResult { tool_use_id, .. } => { ContentPart::ToolResult { tool_use_id, .. } => {
seen_results.insert(tool_use_id.clone()) let keep = seen_results.insert(tool_use_id.clone());
if keep {
retained_results.insert(tool_use_id.clone());
}
keep
} }
ContentPart::Text(_) ContentPart::Text(_)
| ContentPart::Reasoning { .. } | ContentPart::Reasoning { .. }
| ContentPart::Image { .. } => false, | 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,
}
}
let mut retained = keep_parts.into_iter();
parts.retain(|_| retained.next().unwrap_or(false));
!parts.is_empty() !parts.is_empty()
} }
MessageContent::Text(_) => false, MessageContent::Text(_) => false,
@@ -1,7 +1,7 @@
use serde_json::json; use serde_json::json;
use super::*; use super::*;
use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus}; use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus, ToolResultImage};
fn initial_messages() -> Vec<ConversationMessage> { fn initial_messages() -> Vec<ConversationMessage> {
vec![ConversationMessage { 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 { fn successful_result(call_id: &str, content: &str) -> ToolResult {
ToolResult { ToolResult {
images: Vec::new(),
call_id: call_id.to_string(), call_id: call_id.to_string(),
content: content.to_string(), content: content.to_string(),
status: ToolResultStatus::Success, status: ToolResultStatus::Success,
@@ -1673,6 +1674,127 @@ fn multipart_archive_is_not_multiplied_and_legacy_duplicates_are_repaired() {
assert_eq!(archive, expected); 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] #[test]
fn local_child_launch_yields_and_does_not_execute_batched_follow_up_tools() { fn local_child_launch_yields_and_does_not_execute_batched_follow_up_tools() {
let mut run = run(); let mut run = run();
@@ -1718,6 +1840,7 @@ fn failed_child_launch_allows_a_corrective_model_turn() {
run.complete_tool( run.complete_tool(
&batch.work_id, &batch.work_id,
ToolResult { ToolResult {
images: Vec::new(),
call_id: "children".to_string(), call_id: "children".to_string(),
content: "Launch failed".to_string(), content: "Launch failed".to_string(),
status: ToolResultStatus::Error, status: ToolResultStatus::Error,
+78 -10
View File
@@ -2,9 +2,10 @@ use std::collections::{BTreeSet, HashMap, VecDeque};
use serde_json::Value as JsonValue; use serde_json::Value as JsonValue;
use crate::types::tool_image_group_call_id;
use crate::{ use crate::{
ContentPart, ConversationMessage, MessageContent, TaskListSnapshot, ToolCall, ToolDefinition, 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"; pub const RECALL_TOOL_HISTORY_NAME: &str = "recall_tool_history";
@@ -124,6 +125,7 @@ impl ToolPolicy {
) )
}; };
return ToolCallDecision::Reject(ToolResult { return ToolCallDecision::Reject(ToolResult {
images: Vec::new(),
call_id: call.id.clone(), call_id: call.id.clone(),
content: format!( content: format!(
"Error: '{}' is not a valid tool for this request; {available}. Do not invent tool names.", "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 { if call.name == RECALL_TOOL_HISTORY_NAME {
return ToolCallDecision::Inline(ToolResult { let (content, images) = recall_tool_history_with_images(
call_id: call.id.clone(),
content: recall_tool_history(
messages, messages,
archive, archive,
ToolHistoryQuery::from_arguments(&call.arguments), ToolHistoryQuery::from_arguments(&call.arguments),
), );
return ToolCallDecision::Inline(ToolResult {
images,
call_id: call.id.clone(),
content,
status: ToolResultStatus::Success, status: ToolResultStatus::Success,
}); });
} }
@@ -148,11 +152,13 @@ impl ToolPolicy {
if call.name == UPDATE_TASK_LIST_NAME { if call.name == UPDATE_TASK_LIST_NAME {
return match TaskListSnapshot::from_arguments(&call.arguments) { return match TaskListSnapshot::from_arguments(&call.arguments) {
Ok(snapshot) => ToolCallDecision::Inline(ToolResult { Ok(snapshot) => ToolCallDecision::Inline(ToolResult {
images: Vec::new(),
call_id: call.id.clone(), call_id: call.id.clone(),
content: snapshot.to_tool_result_content(), content: snapshot.to_tool_result_content(),
status: ToolResultStatus::Success, status: ToolResultStatus::Success,
}), }),
Err(error) => ToolCallDecision::Reject(ToolResult { Err(error) => ToolCallDecision::Reject(ToolResult {
images: Vec::new(),
call_id: call.id.clone(), call_id: call.id.clone(),
content: format!("Error: {error}"), content: format!("Error: {error}"),
status: ToolResultStatus::Error, status: ToolResultStatus::Error,
@@ -201,6 +207,14 @@ pub fn recall_tool_history(
archive: &[ConversationMessage], archive: &[ConversationMessage],
query: ToolHistoryQuery<'_>, query: ToolHistoryQuery<'_>,
) -> String { ) -> String {
recall_tool_history_with_images(messages, archive, query).0
}
fn recall_tool_history_with_images(
messages: &[ConversationMessage],
archive: &[ConversationMessage],
query: ToolHistoryQuery<'_>,
) -> (String, Vec<ToolResultImage>) {
let mut entries = Vec::new(); let mut entries = Vec::new();
collect_tool_entries(archive, &mut entries); collect_tool_entries(archive, &mut entries);
collect_tool_entries(messages, &mut entries); collect_tool_entries(messages, &mut entries);
@@ -222,26 +236,43 @@ pub fn recall_tool_history(
.len() .len()
.checked_sub(1usize.saturating_add(query.offset_from_end)) .checked_sub(1usize.saturating_add(query.offset_from_end))
else { 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 entry = filtered[index];
let result = truncate_recalled_result(&entry.result); let result = truncate_recalled_result(&entry.result);
(
format!( format!(
"Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{result}", "Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{result}",
entry.name, entry.tool_use_id, entry.input 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)] #[derive(Clone, Debug, PartialEq)]
struct ToolHistoryEntry { struct ToolHistoryEntry<'a> {
tool_use_id: String, tool_use_id: String,
name: String, name: String,
input: String, input: String,
result: String, result: String,
images: Vec<(&'a [u8], &'a str)>,
} }
fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec<ToolHistoryEntry>) { fn collect_tool_entries<'a>(
messages: &'a [ConversationMessage],
entries: &mut Vec<ToolHistoryEntry<'a>>,
) {
let mut pending = Vec::new(); let mut pending = Vec::new();
for message in messages { for message in messages {
match &message.content { match &message.content {
@@ -256,6 +287,7 @@ fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec<Tool
.. ..
} => pair_result(tool_use_id, content, &mut pending, entries), } => pair_result(tool_use_id, content, &mut pending, entries),
MessageContent::MultiPart(parts) => { MessageContent::MultiPart(parts) => {
let entries_start = entries.len();
for part in parts { for part in parts {
match part { match part {
ContentPart::ToolUse { ContentPart::ToolUse {
@@ -275,6 +307,41 @@ fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec<Tool
| ContentPart::Image { .. } => {} | ContentPart::Image { .. } => {}
} }
} }
// 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(_) => {} MessageContent::Text(_) => {}
} }
@@ -285,7 +352,7 @@ fn pair_result(
tool_use_id: &str, tool_use_id: &str,
content: &str, content: &str,
pending: &mut Vec<(String, String, String)>, pending: &mut Vec<(String, String, String)>,
entries: &mut Vec<ToolHistoryEntry>, entries: &mut Vec<ToolHistoryEntry<'_>>,
) { ) {
let Some(index) = pending.iter().position(|(id, _, _)| id == tool_use_id) else { let Some(index) = pending.iter().position(|(id, _, _)| id == tool_use_id) else {
return; return;
@@ -296,6 +363,7 @@ fn pair_result(
name, name,
input, input,
result: content.to_string(), result: content.to_string(),
images: Vec::new(),
}); });
} }
@@ -212,6 +212,120 @@ fn recall_never_returns_a_prior_recall_call() {
); );
} }
fn screenshot_exchange(id: &str, bytes: Vec<u8>) -> Vec<ConversationMessage> {
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] #[test]
fn loop_guard_detects_repeated_failures_and_resets_after_detection() { fn loop_guard_detects_repeated_failures_and_resets_after_detection() {
let mut guard = ToolLoopGuard::new(5, 3); let mut guard = ToolLoopGuard::new(5, 3);
+44
View File
@@ -190,9 +190,53 @@ pub struct ToolResult {
pub call_id: String, pub call_id: String,
pub content: String, pub content: String,
pub status: ToolResultStatus, 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<ToolResultImage>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolResultImage {
pub data: Vec<u8>,
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 { impl ToolResult {
pub fn into_batch_content_parts(results: Vec<Self>) -> Vec<ContentPart> {
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 { pub fn is_error(&self) -> bool {
matches!( matches!(
self.status, self.status,
@@ -169,6 +169,7 @@ fn tool_events_keep_one_call_id_across_permission_and_execution() {
}, },
ToolEvent::Completed { ToolEvent::Completed {
result: ToolResult { result: ToolResult {
images: Vec::new(),
call_id: "call-1".to_string(), call_id: "call-1".to_string(),
content: "ok".to_string(), content: "ok".to_string(),
status: ToolResultStatus::Success, status: ToolResultStatus::Success,
@@ -182,11 +183,13 @@ fn tool_events_keep_one_call_id_across_permission_and_execution() {
#[test] #[test]
fn denied_results_are_errors_but_cancelled_results_are_distinct() { fn denied_results_are_errors_but_cancelled_results_are_distinct() {
let denied = ToolResult { let denied = ToolResult {
images: Vec::new(),
call_id: "denied".to_string(), call_id: "denied".to_string(),
content: "permission denied".to_string(), content: "permission denied".to_string(),
status: ToolResultStatus::Denied, status: ToolResultStatus::Denied,
}; };
let cancelled = ToolResult { let cancelled = ToolResult {
images: Vec::new(),
call_id: "cancelled".to_string(), call_id: "cancelled".to_string(),
content: "cancelled".to_string(), content: "cancelled".to_string(),
status: ToolResultStatus::Cancelled, status: ToolResultStatus::Cancelled,