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

Fix LiteLLM context budgets and direct computer use

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
crates/galaxy_agent_core/src/provider_run.rs
+43 -20
View File
@@ -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<ToolResult>) -> 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,
crates/galaxy_agent_core/src/provider_run_tests.rs
+124 -1
View File
@@ -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<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 {
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,
crates/galaxy_agent_core/src/tool_policy.rs
+82 -14
View File
@@ -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<ToolResultImage>) {
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<ToolHistoryEntry>) {
fn collect_tool_entries<'a>(
messages: &'a [ConversationMessage],
entries: &mut Vec<ToolHistoryEntry<'a>>,
) {
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<Tool
..
} => 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<Tool
| 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(_) => {}
}
@@ -285,7 +352,7 @@ fn pair_result(
tool_use_id: &str,
content: &str,
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 {
return;
@@ -296,6 +363,7 @@ fn pair_result(
name,
input,
result: content.to_string(),
images: Vec::new(),
});
}
crates/galaxy_agent_core/src/tool_policy_tests.rs
+114
View File
@@ -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]
fn loop_guard_detects_repeated_failures_and_resets_after_detection() {
let mut guard = ToolLoopGuard::new(5, 3);
crates/galaxy_agent_core/src/types.rs
+44
View File
@@ -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<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 {
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 {
matches!(
self.status,
crates/galaxy_agent_core/src/types_tests.rs
+3
View File
@@ -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,