Fix provider tool history handling

This commit is contained in:
2026-08-12 14:19:51 -05:00
parent c79634e76f
commit b3f3a72435
18 changed files with 838 additions and 34 deletions
+15 -1
View File
@@ -1912,7 +1912,7 @@ pub fn default_tool_definitions() -> Vec<ToolDefinition> {
},
ToolDefinition {
name: "run_agents".to_string(),
description: "Start one or more child agents that share run-wide configuration. Use independent, uniquely named tasks and do not launch duplicate agents for follow-up work. Omitted run-wide fields use the embedded local child runtime and inherit the parent model.".to_string(),
description: "Start one or more child agents that share run-wide configuration. Use independent, uniquely named tasks and do not launch duplicate agents for follow-up work. Omitted run-wide fields use the embedded local child runtime and inherit the parent model. After launch, call wait_for_events when you need child-agent results instead of repeating their work yourself.".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
@@ -1958,6 +1958,20 @@ pub fn default_tool_definitions() -> Vec<ToolDefinition> {
"required": ["summary", "agent_run_configs"]
}),
},
ToolDefinition {
name: "wait_for_events".to_string(),
description: "Yield after starting child agents or other asynchronous work. Use this when you are waiting for child-agent updates instead of repeating the same investigation yourself.".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"idle_timeout_seconds": {
"type": "integer",
"default": 0,
"description": "Optional idle timeout. 0 lets Galaxy choose the default."
}
}
}),
},
ToolDefinition {
name: "start_agent".to_string(),
description: "Start a sub-agent to handle a specific task autonomously. Use for delegating independent work that can run in parallel. The agent gets its own conversation context and tool access. IMPORTANT: Only use this for the initial investigation or when genuinely new research is needed. Do NOT re-spawn agents for follow-up questions if you already have their output in context — just answer from the information you already have.".to_string(),
+127 -1
View File
@@ -322,7 +322,7 @@ pub fn bedrock_stream_to_response_events(
current_tool_name, current_tool_use_id, current_tool_input_json
));
}
if current_tool_name == "start_agent" {
if matches!(current_tool_name.as_str(), "start_agent" | "run_agents") {
has_start_agent_calls = true;
}
let tool_msg = build_tool_call_message(
@@ -1244,6 +1244,78 @@ pub fn build_tool_call_message(
api::message::tool_call::EditDocuments { diffs },
))
}
"run_agents" => {
let agent_run_configs = input
.get("agent_run_configs")
.and_then(|value| value.as_array())
.map(|configs| {
configs
.iter()
.map(|config| api::run_agents::AgentRunConfig {
name: config
.get("name")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
prompt: config
.get("prompt")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
title: config
.get("title")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
})
.collect()
})
.unwrap_or_default();
let execution_mode = input
.get("execution_mode")
.and_then(run_agents_execution_mode_from_json);
Some(api::message::tool_call::Tool::RunAgents(api::RunAgents {
summary: input
.get("summary")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
base_prompt: input
.get("base_prompt")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
skills: Vec::new(),
model_id: input
.get("model_id")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
harness: input
.get("harness_type")
.and_then(|value| value.as_str())
.and_then(run_agents_harness_from_str),
agent_run_configs,
execution_mode,
plan_id: input
.get("plan_id")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
}))
}
"wait_for_events" => {
let idle_timeout_seconds = input
.get("idle_timeout_seconds")
.and_then(|value| value.as_i64())
.and_then(|value| value.try_into().ok())
.unwrap_or(0);
Some(api::message::tool_call::Tool::WaitForEvents(
api::message::tool_call::WaitForEvents {
idle_timeout_seconds,
},
))
}
"start_agent" => {
let name = input
.get("name")
@@ -1470,6 +1542,58 @@ pub fn build_tool_call_message(
}
}
fn run_agents_execution_mode_from_json(
execution_mode: &serde_json::Value,
) -> Option<api::run_agents::ExecutionMode> {
let mode_type = execution_mode
.get("type")
.and_then(|value| value.as_str())
.or_else(|| execution_mode.as_str());
match mode_type {
Some("remote") => Some(api::run_agents::ExecutionMode::Remote(
api::run_agents::Remote {
environment_id: execution_mode
.get("environment_id")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
worker_host: execution_mode
.get("worker_host")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string(),
computer_use_enabled: execution_mode
.get("computer_use_enabled")
.and_then(|value| value.as_bool())
.unwrap_or(false),
},
)),
Some("local") | Some(_) | None => Some(api::run_agents::ExecutionMode::Local(
api::run_agents::Local {},
)),
}
}
fn run_agents_harness_from_str(harness_type: &str) -> Option<api::Harness> {
let variant = match harness_type
.trim()
.to_ascii_lowercase()
.replace('_', "-")
.as_str()
{
"oz" => api::harness::Variant::Oz(api::harness::Oz {}),
"claude" | "claude-code" => api::harness::Variant::ClaudeCode(api::harness::ClaudeCode {}),
"opencode" | "open-code" => api::harness::Variant::OpenCode(api::harness::OpenCode {}),
"gemini" => api::harness::Variant::Gemini(api::harness::Gemini {}),
"codex" => api::harness::Variant::Codex(api::harness::Codex {}),
"" | "unknown" => return None,
_ => return None,
};
Some(api::Harness {
variant: Some(variant),
})
}
/// Built-in tools that Galaxy knows how to execute directly.
const KNOWN_TOOLS: &[&str] = &[
"run_shell_command",
@@ -1493,6 +1617,8 @@ const KNOWN_TOOLS: &[&str] = &[
"read_documents",
"create_documents",
"edit_documents",
"run_agents",
"wait_for_events",
"start_agent",
"ask_user_question",
"read_skill",
@@ -350,6 +350,56 @@ fn development_tool_calls_preserve_focused_reads_file_lifecycle_and_search_filte
assert_eq!(search.path_filters, vec!["app/src/ai", "crates/ai"]);
}
#[test]
fn orchestration_tool_calls_build_run_agents_and_wait_for_events() {
let run_tool = tool_from_event(build_tool_call_message(
"task-1",
"tool-run-agents",
"run_agents",
r#"{
"summary": "Investigate in parallel",
"base_prompt": "Shared instructions",
"model_id": "coding-assistant-max",
"harness_type": "codex",
"execution_mode": {
"type": "local"
},
"agent_run_configs": [
{
"name": "code",
"prompt": "Inspect code",
"title": "Code inspection"
}
],
"plan_id": "plan-1"
}"#,
));
let api::message::tool_call::Tool::RunAgents(run_agents) = run_tool else {
panic!("expected run_agents");
};
assert_eq!(run_agents.summary, "Investigate in parallel");
assert_eq!(run_agents.base_prompt, "Shared instructions");
assert_eq!(run_agents.model_id, "coding-assistant-max");
assert!(matches!(
run_agents.execution_mode,
Some(api::run_agents::ExecutionMode::Local(_))
));
assert_eq!(run_agents.agent_run_configs.len(), 1);
assert_eq!(run_agents.agent_run_configs[0].name, "code");
assert_eq!(run_agents.agent_run_configs[0].prompt, "Inspect code");
let wait_tool = tool_from_event(build_tool_call_message(
"task-1",
"tool-wait",
"wait_for_events",
r#"{"idle_timeout_seconds": 120}"#,
));
let api::message::tool_call::Tool::WaitForEvents(wait) = wait_tool else {
panic!("expected wait_for_events");
};
assert_eq!(wait.idle_timeout_seconds, 120);
}
#[test]
fn test_context_window_for_model_1m_marker() {
assert_eq!(
@@ -467,6 +517,8 @@ fn test_cost_zero_for_zero_tokens() {
fn direct_provider_known_tools_exclude_hosted_only_tools() {
assert!(!is_known_tool("send_message_to_agent"));
assert!(!is_known_tool("suggest_next_prompt"));
assert!(is_known_tool("run_agents"));
assert!(is_known_tool("wait_for_events"));
assert!(is_known_tool("recall_tool_history"));
assert!(is_known_tool("interrupt_shell_command"));
}
+8
View File
@@ -9,6 +9,7 @@ use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig, BedrockErro
use crate::ai::bedrock::convert::ConversationMessage;
use crate::ai::bedrock::diagnostic::BedrockDiagnosticLogger;
use crate::ai::bedrock::request_translator;
use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn;
pub struct TranslatorRequest {
pub config: BedrockClientConfig,
@@ -108,6 +109,9 @@ pub async fn execute(
let system_prompt = request_translator::extract_system_prompt(request, &params.global_rules);
let tools = request_translator::extract_tools(request);
if tools_are_inline_only(&tools) {
flatten_tool_history_for_no_tools_turn(&mut messages);
}
log::info!(
"[bedrock] Sending {} messages, system_prompt={}, progressive_summary={}, tools={}",
@@ -163,6 +167,10 @@ pub async fn execute(
Ok(stream)
}
fn tools_are_inline_only(tools: &[crate::ai::bedrock::convert::ToolDefinition]) -> bool {
tools.iter().all(|tool| tool.name == "recall_tool_history")
}
fn describe_message_content(content: &crate::ai::bedrock::convert::MessageContent) -> String {
use crate::ai::bedrock::convert::{ContentPart, MessageContent};
match content {
@@ -37,6 +37,8 @@ use crate::ai::document::plan_publication::{
prepare_plan_publications, wait_for_plan_publications,
};
use crate::ai::local_harness_setup::local_harness_product_disabled_message;
#[cfg(not(target_family = "wasm"))]
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
/// Per-child spawn timeout. If a child agent doesn't report back within
/// this window (e.g. binary not found, server error), the slot is failed
@@ -169,12 +171,36 @@ impl RunAgentsExecutor {
if self.pending.contains_key(&action_id) {
log::warn!("RunAgentsExecutor: dispatch reentered for {action_id:?}; rejecting");
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents dispatch rejected",
serde_json::json!({
"event": "run_agents_dispatch_rejected",
"reason": "reentered_pending_action",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
}),
);
let _ = sender.try_send(RunAgentsResult::Cancelled);
return receiver;
}
if let Err(error) = validate_request(&request) {
log::warn!("RunAgentsExecutor: validation failure: {error}");
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents validation failed",
serde_json::json!({
"event": "run_agents_validation_failed",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"error": remote_logging::sanitize_error(&error),
}),
);
let _ = sender.try_send(RunAgentsResult::Failure { error });
return receiver;
}
@@ -185,6 +211,19 @@ impl RunAgentsExecutor {
};
self.pending
.insert(action_id.clone(), PendingRunAgents::Publishing);
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Info,
"RunAgents plan publication wait started",
serde_json::json!({
"event": "run_agents_plan_publication_wait_started",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"agent_count": snapshot.agent_count,
"plan_id_present": !request.plan_id.trim().is_empty(),
}),
);
ctx.emit(RunAgentsExecutorEvent::SpawningStarted {
action_id: action_id.clone(),
snapshot,
@@ -241,6 +280,23 @@ impl RunAgentsExecutor {
..
} = request;
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Info,
"RunAgents child dispatch started",
serde_json::json!({
"event": "run_agents_child_dispatch_started",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"agent_count": agent_run_configs.len(),
"execution_mode": run_agents_execution_mode_label(&run_execution_mode),
"harness_type": harness_type.as_str(),
"model_id_present": !model_id.trim().is_empty(),
"parent_run_id_present": parent_run_id.is_some(),
}),
);
let mut slots: Vec<ChildSlot> = Vec::with_capacity(agent_run_configs.len());
for cfg in &agent_run_configs {
let prompt = compose_run_agents_child_prompt(&base_prompt, &cfg.prompt);
@@ -254,6 +310,19 @@ impl RunAgentsExecutor {
) {
Ok(mode) => mode,
Err(err) => {
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents child dispatch failed before launch",
serde_json::json!({
"event": "run_agents_child_dispatch_prelaunch_failed",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"agent_name": cfg.name.as_str(),
"error": remote_logging::sanitize_error(&err),
}),
);
slots.push(ChildSlot::Failed(err));
continue;
}
@@ -261,11 +330,37 @@ impl RunAgentsExecutor {
if matches!(run_execution_mode, RunAgentsExecutionMode::Remote { .. })
&& parent_run_id.is_none()
{
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents remote child dispatch missing parent run_id",
serde_json::json!({
"event": "run_agents_child_dispatch_prelaunch_failed",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"agent_name": cfg.name.as_str(),
"error": "Remote child agents require the parent run_id to be available.",
}),
);
slots.push(ChildSlot::Failed(
"Remote child agents require the parent run_id to be available.".to_string(),
));
continue;
}
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Info,
"RunAgents child dispatch queued",
serde_json::json!({
"event": "run_agents_child_dispatch_queued",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"agent_name": cfg.name.as_str(),
"execution_mode": start_agent_execution_mode_label(&mode),
}),
);
let recv = self.start_agent_executor.update(ctx, |executor, exec_ctx| {
executor.dispatch(
cfg.name.clone(),
@@ -286,11 +381,20 @@ impl RunAgentsExecutor {
let run_harness_type = harness_type.clone();
let run_execution_mode_for_aggr = run_execution_mode.clone();
let parent_conversation_id_for_result = parent_conversation_id;
#[cfg(not(target_family = "wasm"))]
let action_id_for_async_log = action_id.clone();
#[cfg(not(target_family = "wasm"))]
let parent_conversation_id_for_async_log = parent_conversation_id;
#[cfg(not(target_family = "wasm"))]
let agent_names_for_async_log = agent_run_configs
.iter()
.map(|cfg| cfg.name.clone())
.collect::<Vec<_>>();
ctx.spawn(
async move {
let mut outcomes: Vec<RunAgentsAgentOutcomeKind> = Vec::with_capacity(slots.len());
for slot in slots {
for (slot_index, slot) in slots.into_iter().enumerate() {
let kind = match slot {
ChildSlot::Failed(error) => RunAgentsAgentOutcomeKind::Failed { error },
ChildSlot::Pending(recv) => {
@@ -331,6 +435,19 @@ impl RunAgentsExecutor {
}
}
};
#[cfg(not(target_family = "wasm"))]
log::info!(
"RunAgents child launch outcome action_id={} parent_conversation_id={} \
agent_name={} slot_index={} outcome={}",
action_id_for_async_log,
parent_conversation_id_for_async_log,
agent_names_for_async_log
.get(slot_index)
.map(String::as_str)
.unwrap_or("<unknown>"),
slot_index,
run_agents_agent_outcome_kind_label(&kind)
);
outcomes.push(kind);
}
outcomes
@@ -345,6 +462,35 @@ impl RunAgentsExecutor {
})
.collect();
me.record_launched_agents(parent_conversation_id_for_result, &agents);
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Info,
"RunAgents launch outcomes resolved",
serde_json::json!({
"event": "run_agents_launch_outcomes_resolved",
"action_id": action_id_for_aggr.to_string(),
"parent_conversation_id": parent_conversation_id_for_result.to_string(),
"agent_count": agents.len(),
"launched_count": agents.iter().filter(|agent| matches!(agent.kind, RunAgentsAgentOutcomeKind::Launched { .. })).count(),
"failed_count": agents.iter().filter(|agent| matches!(agent.kind, RunAgentsAgentOutcomeKind::Failed { .. })).count(),
"agents": agents
.iter()
.map(|agent| match &agent.kind {
RunAgentsAgentOutcomeKind::Launched { agent_id } => serde_json::json!({
"name": agent.name.as_str(),
"status": "launched",
"agent_id": agent_id.as_str(),
}),
RunAgentsAgentOutcomeKind::Failed { error } => serde_json::json!({
"name": agent.name.as_str(),
"status": "failed",
"error": remote_logging::sanitize_error(error),
}),
})
.collect::<Vec<_>>(),
}),
);
let launched_mode = match &run_execution_mode_for_aggr {
RunAgentsExecutionMode::Local => RunAgentsLaunchedExecutionMode::Local,
RunAgentsExecutionMode::Remote {
@@ -391,6 +537,18 @@ impl RunAgentsExecutor {
&self.launched_agents,
ctx,
) {
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents execution denied",
serde_json::json!({
"event": "run_agents_execution_denied",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"reason": remote_logging::sanitize_error(&reason),
}),
);
return ActionExecution::Sync(AIAgentActionResultType::RunAgents(
RunAgentsResult::Denied { reason },
));
@@ -450,6 +608,47 @@ impl RunAgentsExecutor {
#[path = "run_agents_tests.rs"]
mod tests;
#[cfg(not(target_family = "wasm"))]
fn log_run_agents_event(
ctx: &mut ModelContext<RunAgentsExecutor>,
level: RemoteLogLevel,
message: impl Into<String>,
context: serde_json::Value,
) {
remote_logging::log_model_event(
ctx,
RemoteLogRecord {
level,
message: message.into(),
context,
},
);
}
#[cfg(not(target_family = "wasm"))]
fn run_agents_execution_mode_label(mode: &RunAgentsExecutionMode) -> &'static str {
match mode {
RunAgentsExecutionMode::Local => "local",
RunAgentsExecutionMode::Remote { .. } => "remote",
}
}
#[cfg(not(target_family = "wasm"))]
fn start_agent_execution_mode_label(mode: &StartAgentExecutionMode) -> &'static str {
match mode {
StartAgentExecutionMode::Local { .. } => "local",
StartAgentExecutionMode::Remote { .. } => "remote",
}
}
#[cfg(not(target_family = "wasm"))]
fn run_agents_agent_outcome_kind_label(kind: &RunAgentsAgentOutcomeKind) -> &'static str {
match kind {
RunAgentsAgentOutcomeKind::Launched { .. } => "launched",
RunAgentsAgentOutcomeKind::Failed { .. } => "failed",
}
}
enum ChildSlot {
Failed(String),
Pending(async_channel::Receiver<StartAgentOutcome>),
+56
View File
@@ -41,6 +41,8 @@ use crate::ai::artifacts::Artifact;
use crate::ai::document::ai_document_model::AIDocumentModel;
#[cfg(not(target_family = "wasm"))]
use crate::ai::llms::LLMPreferences;
#[cfg(not(target_family = "wasm"))]
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
use crate::input_suggestions::HistoryOrder;
use crate::persistence::model::{
AcpConversationData, AgentBackend, AgentConversation, AgentConversationData,
@@ -1360,7 +1362,21 @@ impl BlocklistAIHistoryModel {
ctx: &mut ModelContext<Self>,
) {
if let Some(conversation) = self.conversations_by_id.get_mut(&conversation_id) {
#[cfg(not(target_family = "wasm"))]
let remote_log_context =
remote_status_log_context(conversation, &status, error.as_ref());
conversation.update_status_with_error(status, error, terminal_surface_id, ctx);
#[cfg(not(target_family = "wasm"))]
if let Some(context) = remote_log_context {
remote_logging::log_model_event(
ctx,
RemoteLogRecord {
level: RemoteLogLevel::Info,
message: "Agent conversation status changed".to_string(),
context,
},
);
}
}
}
@@ -2870,6 +2886,46 @@ fn agent_id_key_from_persisted_data(conversation_data: &AgentConversationData) -
conversation_data.run_id.as_deref()
}
#[cfg(not(target_family = "wasm"))]
fn remote_status_log_context(
conversation: &AIConversation,
new_status: &ConversationStatus,
error: Option<&RenderableAIError>,
) -> Option<serde_json::Value> {
let prev_status = conversation.status();
if prev_status == new_status {
return None;
}
Some(serde_json::json!({
"event": "agent_conversation_status_changed",
"conversation_id": conversation.id().to_string(),
"parent_conversation_id": conversation.parent_conversation_id().map(|id| id.to_string()),
"agent_id": conversation.orchestration_agent_id(),
"agent_name": conversation.agent_name(),
"harness_type": conversation.orchestration_harness_type(),
"is_child": conversation.parent_conversation_id().is_some(),
"is_remote_child": conversation.is_remote_child(),
"previous_status": conversation_status_label(prev_status),
"new_status": conversation_status_label(new_status),
"new_status_is_terminal": new_status.is_done(),
"error": error.map(remote_logging::sanitize_error),
}))
}
#[cfg(not(target_family = "wasm"))]
fn conversation_status_label(status: &ConversationStatus) -> &'static str {
match status {
ConversationStatus::InProgress => "in_progress",
ConversationStatus::Success => "success",
ConversationStatus::Error => "error",
ConversationStatus::TransientError => "transient_error",
ConversationStatus::Cancelled => "cancelled",
ConversationStatus::Blocked { .. } => "blocked",
ConversationStatus::WaitingForEvents => "waiting_for_events",
}
}
/// Whether an `UpdatedConversationStatus` event represents a restoration
/// (the conversation was re-loaded for a terminal surface; the underlying
/// `ConversationStatus` did not change) or a real status set, in which case
+58 -5
View File
@@ -1,3 +1,5 @@
use std::collections::HashSet;
use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent, MessageRole};
/// Sanitizes messages for OpenAI API compatibility.
@@ -9,6 +11,7 @@ use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageConten
/// - When routed to Bedrock via LiteLLM, the conversation must end with a user message
pub fn sanitize_messages_for_openai(messages: &mut Vec<ConversationMessage>) {
remove_orphaned_tool_results(messages);
remove_misplaced_tool_results(messages);
synthesize_missing_tool_results(messages);
ensure_ends_with_user_message(messages);
}
@@ -16,8 +19,7 @@ pub fn sanitize_messages_for_openai(messages: &mut Vec<ConversationMessage>) {
/// Removes tool_result messages that reference tool_use_ids not found in any
/// preceding assistant message.
fn remove_orphaned_tool_results(messages: &mut Vec<ConversationMessage>) {
let mut known_tool_use_ids: std::collections::HashSet<String> =
std::collections::HashSet::new();
let mut known_tool_use_ids = HashSet::new();
// First pass: collect all tool_use_ids from assistant messages
for msg in messages.iter() {
@@ -50,12 +52,63 @@ fn remove_orphaned_tool_results(messages: &mut Vec<ConversationMessage>) {
});
}
/// LiteLLM may route OpenAI-compatible requests to Bedrock, which requires a
/// user turn containing tool_result blocks to directly answer the tool_use
/// blocks from the immediately previous assistant turn. Late results from
/// cancelled or superseded actions are valid history globally, but invalid in
/// that later user turn, so drop them before request conversion.
fn remove_misplaced_tool_results(messages: &mut Vec<ConversationMessage>) {
let mut i = 0;
while i < messages.len() {
if messages[i].role != MessageRole::User {
i += 1;
continue;
}
let mut allowed_tool_use_ids = if i > 0 && messages[i - 1].role == MessageRole::Assistant {
let mut ids = HashSet::new();
collect_tool_use_ids(&messages[i - 1].content, &mut ids);
ids
} else {
HashSet::new()
};
if retain_allowed_tool_results(&mut messages[i].content, &mut allowed_tool_use_ids) {
i += 1;
} else {
messages.remove(i);
}
}
}
fn retain_allowed_tool_results(
content: &mut MessageContent,
allowed_tool_use_ids: &mut HashSet<String>,
) -> bool {
match content {
MessageContent::Text(_) | MessageContent::ToolUse { .. } => true,
MessageContent::ToolResult { tool_use_id, .. } => allowed_tool_use_ids.remove(tool_use_id),
MessageContent::MultiPart(parts) => {
parts.retain(|part| match part {
ContentPart::ToolResult { tool_use_id, .. } => {
allowed_tool_use_ids.remove(tool_use_id)
}
ContentPart::Text(_)
| ContentPart::Reasoning { .. }
| ContentPart::Image { .. }
| ContentPart::ToolUse { .. } => true,
});
!parts.is_empty()
}
}
}
/// For any assistant tool_use that doesn't have a matching tool_result in a
/// subsequent user message, synthesize a result immediately after the tool_use.
/// This satisfies Bedrock's requirement (via LiteLLM) that tool_result blocks
/// appear immediately after the corresponding tool_use message.
fn synthesize_missing_tool_results(messages: &mut Vec<ConversationMessage>) {
let mut answered_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
let mut answered_ids = HashSet::new();
// First pass: collect all existing tool_result IDs
for msg in messages.iter() {
@@ -173,7 +226,7 @@ fn synthesize_missing_tool_results(messages: &mut Vec<ConversationMessage>) {
}
}
fn collect_tool_use_ids(content: &MessageContent, ids: &mut std::collections::HashSet<String>) {
fn collect_tool_use_ids(content: &MessageContent, ids: &mut HashSet<String>) {
match content {
MessageContent::ToolUse { tool_use_id, .. } => {
ids.insert(tool_use_id.clone());
@@ -205,7 +258,7 @@ fn collect_tool_use_ids_vec(content: &MessageContent, ids: &mut Vec<String>) {
}
}
fn collect_tool_result_ids(content: &MessageContent, ids: &mut std::collections::HashSet<String>) {
fn collect_tool_result_ids(content: &MessageContent, ids: &mut HashSet<String>) {
match content {
MessageContent::ToolResult { tool_use_id, .. } => {
ids.insert(tool_use_id.clone());
@@ -168,6 +168,106 @@ fn test_multipart_tool_uses_all_get_results() {
}
}
#[test]
fn sanitizer_drops_stale_tool_result_from_current_user_turn() {
let mut messages = vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "old_tool".to_string(),
name: "read_shell_command_output".to_string(),
input: json!({"block_id": "block-1"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "old_tool".to_string(),
content: "cancelled".to_string(),
is_error: true,
},
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "current_tool".to_string(),
name: "read_notebook".to_string(),
input: json!({"document_id": "doc-1"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::MultiPart(vec![
ContentPart::ToolResult {
tool_use_id: "old_tool".to_string(),
content: "late cancellation".to_string(),
is_error: true,
},
ContentPart::ToolResult {
tool_use_id: "current_tool".to_string(),
content: "notebook contents".to_string(),
is_error: false,
},
]),
},
];
sanitize_messages_for_openai(&mut messages);
assert_eq!(messages.len(), 4);
let MessageContent::MultiPart(parts) = &messages[3].content else {
panic!("expected current user message to remain multipart");
};
assert_eq!(parts.len(), 1);
assert!(matches!(
&parts[0],
ContentPart::ToolResult { tool_use_id, content, is_error }
if tool_use_id == "current_tool" && content == "notebook contents" && !is_error
));
}
#[test]
fn sanitizer_drops_tool_result_message_not_following_its_tool_use() {
let mut messages = vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "old_tool".to_string(),
name: "read_shell_command_output".to_string(),
input: json!({"block_id": "block-1"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "old_tool".to_string(),
content: "cancelled".to_string(),
is_error: true,
},
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::Text("Continuing.".to_string()),
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "old_tool".to_string(),
content: "late cancellation".to_string(),
is_error: true,
},
},
];
sanitize_messages_for_openai(&mut messages);
assert_eq!(messages.len(), 4);
let MessageContent::Text(text) = &messages[3].content else {
panic!("expected appended continuation message after stale tool result was dropped");
};
assert_eq!(text, "Continue.");
}
#[test]
fn test_ensure_ends_with_user_message_no_op_when_already_user() {
let mut messages = vec![
+2
View File
@@ -679,6 +679,8 @@ const KNOWN_TOOLS: &[&str] = &[
"read_documents",
"create_documents",
"edit_documents",
"run_agents",
"wait_for_events",
"start_agent",
"ask_user_question",
"read_skill",
@@ -136,6 +136,8 @@ async fn recall_tool_history_does_not_emit_a_client_tool_call() {
fn direct_provider_known_tools_exclude_hosted_only_tools() {
assert!(!is_known_tool("send_message_to_agent"));
assert!(!is_known_tool("suggest_next_prompt"));
assert!(is_known_tool("run_agents"));
assert!(is_known_tool("wait_for_events"));
assert!(is_known_tool("recall_tool_history"));
assert!(is_known_tool("interrupt_shell_command"));
}
+12 -2
View File
@@ -8,7 +8,9 @@ use super::request_translator::sanitize_messages_for_openai;
use super::response_translator::{openai_stream_to_response_events, OpenAIStreamContext};
use crate::ai::agent::api::LegacyResponseStream;
use crate::ai::bedrock::request_translator;
use crate::ai::provider::types::{ConversationMessage, MessageContent, MessageRole};
use crate::ai::provider::types::{
flatten_tool_history_for_no_tools_turn, ConversationMessage, MessageContent, MessageRole,
};
const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 64_000;
@@ -92,6 +94,10 @@ pub(crate) fn prepare_turn(params: &TranslatorRequest, request: &mut api::Reques
message.truncate_tool_results_for_provider_request();
}
sanitize_messages_for_openai(&mut messages);
let tools = request_translator::extract_tools(request);
if tools_are_inline_only(&tools) {
flatten_tool_history_for_no_tools_turn(&mut messages);
}
PreparedTurn {
task_id,
@@ -99,12 +105,16 @@ pub(crate) fn prepare_turn(params: &TranslatorRequest, request: &mut api::Reques
user_query: request_translator::extract_user_query_text(request),
messages,
system_prompt: request_translator::extract_system_prompt(request, &params.global_rules),
tools: request_translator::extract_tools(request),
tools,
model_id,
persistent_message_count,
}
}
fn tools_are_inline_only(tools: &[crate::ai::provider::types::ToolDefinition]) -> bool {
tools.iter().all(|tool| tool.name == "recall_tool_history")
}
pub async fn execute(
params: TranslatorRequest,
request: &mut api::Request,
+58
View File
@@ -4,3 +4,61 @@ pub use galaxy_agent_core::{
ContentPart, ConversationMessage, MessageContent, MessageRole, ToolDefinition,
MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST,
};
pub(crate) fn flatten_tool_history_for_no_tools_turn(messages: &mut [ConversationMessage]) {
for message in messages {
message.content = flatten_tool_history_content(std::mem::replace(
&mut message.content,
MessageContent::Text(String::new()),
));
}
}
fn flatten_tool_history_content(content: MessageContent) -> MessageContent {
match content {
MessageContent::Text(text) => MessageContent::Text(text),
MessageContent::ToolUse {
tool_use_id,
name,
input,
} => MessageContent::Text(flattened_tool_use_text(&tool_use_id, &name, &input)),
MessageContent::ToolResult {
tool_use_id,
content,
is_error,
} => MessageContent::Text(flattened_tool_result_text(&tool_use_id, &content, is_error)),
MessageContent::MultiPart(parts) => MessageContent::MultiPart(
parts
.into_iter()
.map(flatten_tool_history_content_part)
.collect(),
),
}
}
fn flatten_tool_history_content_part(part: ContentPart) -> ContentPart {
match part {
ContentPart::Text(text) => ContentPart::Text(text),
ContentPart::Reasoning { text, signature } => ContentPart::Reasoning { text, signature },
ContentPart::Image { data, mime_type } => ContentPart::Image { data, mime_type },
ContentPart::ToolUse {
tool_use_id,
name,
input,
} => ContentPart::Text(flattened_tool_use_text(&tool_use_id, &name, &input)),
ContentPart::ToolResult {
tool_use_id,
content,
is_error,
} => ContentPart::Text(flattened_tool_result_text(&tool_use_id, &content, is_error)),
}
}
fn flattened_tool_use_text(tool_use_id: &str, name: &str, input: &serde_json::Value) -> String {
format!("Previous tool call `{name}` (id: {tool_use_id}) with input:\n{input}")
}
fn flattened_tool_result_text(tool_use_id: &str, content: &str, is_error: bool) -> String {
let status = if is_error { "error" } else { "success" };
format!("Previous tool result for id `{tool_use_id}` ({status}):\n{content}")
}
+8
View File
@@ -20,6 +20,7 @@ use crate::ai::bedrock::request_translator::{
};
use crate::ai::openai::client::OpenAIClientConfig;
use crate::ai::openai::request_translator::sanitize_messages_for_openai;
use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn;
pub(crate) struct PreparedRigTurn {
pub task_id: String,
@@ -155,6 +156,9 @@ fn prepare_rig_turn_for_provider(
});
}
turn_messages.extend(persistent_messages.clone());
if tools_are_inline_only(&tools) {
flatten_tool_history_for_no_tools_turn(&mut turn_messages);
}
let model_id = model_override
.filter(|model| !model.is_empty() && model != "auto")
@@ -516,6 +520,10 @@ fn tool_definitions(
(tools, mcp_tool_aliases)
}
fn tools_are_inline_only(tools: &[ToolDefinition]) -> bool {
tools.iter().all(|tool| tool.name == "recall_tool_history")
}
const MAX_PROVIDER_TOOL_NAME_BYTES: usize = 64;
const MCP_TOOL_HASH_BYTES: usize = 8;
+111
View File
@@ -124,6 +124,117 @@ fn builds_a_rig_turn_directly_from_galaxy_request_state() {
));
}
#[test]
fn no_tools_turn_flattens_historical_tool_protocol_messages() {
let mut params = RequestParams::new_for_test();
params.message_history = vec![
galaxy_agent_core::ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "call-1".to_string(),
name: "run_shell_command".to_string(),
input: serde_json::json!({
"command": "find . -name package.json",
"wait_until_complete": true,
}),
},
},
galaxy_agent_core::ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "call-1".to_string(),
content: "command exited with code 1".to_string(),
is_error: true,
},
},
];
params.input = vec![user_query("Summarize what happened")];
let prepared = prepare_rig_turn(&config(), params, Vec::new(), Vec::new());
assert_eq!(prepared.request.tools.len(), 1);
assert_eq!(prepared.request.tools[0].name, "recall_tool_history");
assert!(prepared
.request
.messages
.iter()
.all(|message| !message.content.contains_tool_protocol_blocks()));
assert!(prepared.request.messages.iter().any(|message| matches!(
&message.content,
MessageContent::Text(text)
if text.contains("Previous tool call `run_shell_command`")
&& text.contains("call-1")
)));
assert!(prepared.request.messages.iter().any(|message| matches!(
&message.content,
MessageContent::Text(text)
if text.contains("Previous tool result for id `call-1` (error)")
)));
assert!(prepared.persistent_messages.iter().any(|message| matches!(
&message.content,
MessageContent::ToolUse { tool_use_id, .. } if tool_use_id == "call-1"
)));
}
#[test]
fn tool_enabled_turn_preserves_structured_tool_history() {
let mut params = RequestParams::new_for_test();
params.message_history = vec![
galaxy_agent_core::ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "call-1".to_string(),
name: "read_files".to_string(),
input: serde_json::json!({"files": ["Cargo.toml"]}),
},
},
galaxy_agent_core::ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "call-1".to_string(),
content: "[package]\nname = \"galaxy\"".to_string(),
is_error: false,
},
},
];
params.input = vec![user_query("Keep inspecting")];
let prepared = prepare_rig_turn(&config(), params, vec![ToolType::ReadFiles], Vec::new());
assert!(prepared
.request
.tools
.iter()
.any(|tool| tool.name == "read_files"));
assert!(prepared.request.messages.iter().any(|message| matches!(
&message.content,
MessageContent::ToolUse { tool_use_id, .. } if tool_use_id == "call-1"
)));
assert!(prepared.request.messages.iter().any(|message| matches!(
&message.content,
MessageContent::ToolResult { tool_use_id, .. } if tool_use_id == "call-1"
)));
}
trait MessageContentTestExt {
fn contains_tool_protocol_blocks(&self) -> bool;
}
impl MessageContentTestExt for MessageContent {
fn contains_tool_protocol_blocks(&self) -> bool {
match self {
MessageContent::ToolUse { .. } | MessageContent::ToolResult { .. } => true,
MessageContent::MultiPart(parts) => parts.iter().any(|part| {
matches!(
part,
ContentPart::ToolUse { .. } | ContentPart::ToolResult { .. }
)
}),
MessageContent::Text(_) => false,
}
}
}
#[test]
fn rig_prompt_requires_follow_through_without_manual_continue_prompts() {
let mut params = RequestParams::new_for_test();