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

Improve agent provider resilience

Browse Source
This commit is contained in:
Ryan Ward
2026-09-02 17:09:08 -05:00
parent b115946534
commit 7a33cc7e56
72 changed files with 997 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/runtime/event_translator.rs
+202 -17
View File
@@ -10,11 +10,11 @@ use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
use super::provider_run_coordinator::ProviderRunProjection;
use crate::ai::agent::runtime_activity;
use crate::ai::bedrock::response_translator::{
use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage};
use crate::ai::provider::response_translator::{
build_add_agent_output_message, build_append_text, build_create_task, build_stream_init,
build_user_query_message,
};
use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage};
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub(crate) struct RuntimeResponseConfig {
@@ -26,6 +26,9 @@ pub(crate) struct RuntimeResponseConfig {
pub(crate) max_context_tokens: Option<u32>,
pub(crate) capabilities: RuntimeCapabilities,
pub(crate) empty_output_message: Option<String>,
/// Todo items supplied by the existing task transcript, when one is available.
#[serde(skip)]
pub(crate) todo_items: Option<Vec<api::TodoItem>>,
}
/// Converts the provider-neutral runtime lifecycle into Galaxy's existing
@@ -51,6 +54,7 @@ pub(crate) struct RuntimeResponseTranslator {
pub(crate) struct ProviderRunResponseProjector {
translator: RuntimeResponseTranslator,
has_started_model_turn: bool,
todo_phase: usize,
finished: bool,
}
@@ -59,6 +63,7 @@ impl ProviderRunResponseProjector {
Self {
translator: RuntimeResponseTranslator::new(config),
has_started_model_turn: false,
todo_phase: 0,
finished: false,
}
}
@@ -70,6 +75,8 @@ impl ProviderRunResponseProjector {
Self {
translator: RuntimeResponseTranslator::restored(config, projection_was_initialized),
has_started_model_turn: false,
// Task-list events are part of the already persisted projection.
todo_phase: usize::MAX,
finished: false,
}
}
@@ -87,17 +94,46 @@ impl ProviderRunResponseProjector {
self.translator.begin_followup_turn();
}
self.has_started_model_turn = true;
self.translator.translate(AgentEvent::TurnStarted {
let mut events = self.translator.translate(AgentEvent::TurnStarted {
runtime_request_id: String::new(),
})
})?;
events.extend(self.todo_phase_events());
Ok(events)
}
ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event),
ProviderRunProjection::ModelRetry { .. } => {
Ok(self.translator.discard_failed_turn_output())
}
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()),
| ProviderRunProjection::ModelTurnFinished { .. } => Ok(Vec::new()),
ProviderRunProjection::ToolBatchReady { .. } => {
let todo_index = self.todo_phase.saturating_sub(1);
let Some(todo) = self.todo_items().get(todo_index).cloned() else {
return Ok(Vec::new());
};
self.todo_phase += 1;
let mut events = vec![build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::MarkTodosCompleted(
api::MarkTodosCompleted {
todo_ids: vec![todo.id],
},
),
)];
events.push(build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::UpdatePendingTodos(
api::UpdatePendingTodos {
updated_pending_todos: self
.todo_items()
.into_iter()
.skip(self.todo_phase)
.collect(),
},
),
));
Ok(events)
}
}
}
@@ -115,17 +151,154 @@ impl ProviderRunResponseProjector {
}
self.finished = true;
match outcome {
ProviderRunOutcome::Completed(completion) => Ok(self
.translator
.finish_provider_run(completion.stop_reason.clone(), aggregate_usage)),
ProviderRunOutcome::Failed(failure) => Ok(self
.translator
.provider_failure(&failure.message, aggregate_usage)),
ProviderRunOutcome::Completed(completion) => {
let mut events = self.todo_completion_events();
events.extend(
self.translator
.finish_provider_run(completion.stop_reason.clone(), aggregate_usage),
);
Ok(events)
}
ProviderRunOutcome::Failed(failure) => Ok(self.translator.provider_failure(
&failure.message,
failure.source.as_ref(),
aggregate_usage,
)),
ProviderRunOutcome::Cancelled { .. } => Ok(self
.translator
.finish_provider_run(StopReason::Cancelled, aggregate_usage)),
}
}
// Keep the direct-provider workflow visible in the existing task list protocol. These are
// response events, so the normal history model remains the sole owner of task-list state.
fn todo_phase_events(&mut self) -> Vec<ResponseEvent> {
let events = match self.todo_phase {
0 => vec![build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::CreateTodoList(api::CreateTodoList {
initial_todos: self.todo_items(),
}),
)],
_ => Vec::new(),
};
// The first model turn owns the first phase. Tool batches advance it; this keeps
// arbitrary plan lengths aligned with the UpdateTodos protocol.
if self.todo_phase == 0 {
self.todo_phase = 1;
}
events
}
fn todo_completion_events(&self) -> Vec<ResponseEvent> {
if self.todo_phase == 0 || self.todo_phase == usize::MAX {
return Vec::new();
}
let todos = self.todo_items();
if todos.is_empty() {
return Vec::new();
}
vec![
build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::MarkTodosCompleted(
api::MarkTodosCompleted {
todo_ids: todos.iter().map(|todo| todo.id.clone()).collect(),
},
),
),
build_todo_update(
&self.translator.config.task_id,
api::message::update_todos::Operation::UpdatePendingTodos(
api::UpdatePendingTodos {
updated_pending_todos: Vec::new(),
},
),
),
]
}
fn todo_items(&self) -> Vec<api::TodoItem> {
self.translator
.config
.todo_items
.clone()
.unwrap_or_else(default_workflow_todos)
}
}
fn default_workflow_todos() -> Vec<api::TodoItem> {
[
api::TodoItem {
id: "direct-provider-research".to_owned(),
title: "Research the request".to_owned(),
description: "Inspect the repository and gather relevant evidence".to_owned(),
},
api::TodoItem {
id: "direct-provider-plan".to_owned(),
title: "Create an implementation plan".to_owned(),
description: "Choose an approach grounded in the repository".to_owned(),
},
api::TodoItem {
id: "direct-provider-critique".to_owned(),
title: "Critique the approach".to_owned(),
description: "Check assumptions, risks, and missing cases".to_owned(),
},
api::TodoItem {
id: "direct-provider-revise".to_owned(),
title: "Revise the plan".to_owned(),
description: "Incorporate findings before editing".to_owned(),
},
api::TodoItem {
id: "direct-provider-implement".to_owned(),
title: "Implement the change".to_owned(),
description: "Make the requested edits".to_owned(),
},
api::TodoItem {
id: "direct-provider-verify".to_owned(),
title: "Verify the result".to_owned(),
description: "Run proportionate checks".to_owned(),
},
api::TodoItem {
id: "direct-provider-repair".to_owned(),
title: "Repair validation issues".to_owned(),
description: "Fix failures found during verification".to_owned(),
},
]
.to_vec()
}
fn build_todo_update(
task_id: &str,
operation: api::message::update_todos::Operation,
) -> ResponseEvent {
let message = api::Message {
id: Uuid::new_v4().to_string(),
task_id: task_id.to_owned(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: Vec::new(),
fetched_memories: Vec::new(),
message: Some(api::message::Message::UpdateTodos(
api::message::UpdateTodos {
operation: Some(operation),
},
)),
};
ResponseEvent {
r#type: Some(api::response_event::Type::ClientActions(
api::response_event::ClientActions {
actions: vec![api::ClientAction {
action: Some(api::client_action::Action::AddMessagesToTask(
api::client_action::AddMessagesToTask {
task_id: task_id.to_owned(),
messages: vec![message],
},
)),
}],
},
)),
}
}
impl RuntimeResponseTranslator {
@@ -421,15 +594,27 @@ impl RuntimeResponseTranslator {
self.finished_with_reason(map_stop_reason(reason))
}
fn provider_failure(&mut self, message: &str, aggregate_usage: &Usage) -> Vec<ResponseEvent> {
fn provider_failure(
&mut self,
message: &str,
source: Option<&galaxy_agent_core::AgentError>,
aggregate_usage: &Usage,
) -> Vec<ResponseEvent> {
let mut events = Vec::new();
self.initialize(&mut events);
events.push(self.finished_with_usage(
let reason = if source
.is_some_and(|error| error.kind == galaxy_agent_core::AgentErrorKind::Authentication)
{
stream_finished::Reason::InvalidApiKey(stream_finished::InvalidApiKey {
provider: warp_multi_agent_api::LlmProvider::AwsBedrock as i32,
model_name: self.config.model_id.clone(),
})
} else {
stream_finished::Reason::InternalError(stream_finished::InternalError {
message: message.to_owned(),
}),
aggregate_usage,
));
})
};
events.push(self.finished_with_usage(reason, aggregate_usage));
events
}
app/src/ai/runtime/event_translator_tests.rs
+117 -1
View File
@@ -17,6 +17,7 @@ fn provider_translator() -> RuntimeResponseTranslator {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
})
}
@@ -30,6 +31,7 @@ fn session_translator() -> RuntimeResponseTranslator {
max_context_tokens: None,
capabilities: RuntimeCapabilities::session_runtime(),
empty_output_message: Some("> runtime completed without text".to_owned()),
todo_items: None,
})
}
@@ -44,6 +46,7 @@ fn restored_provider_projection_skips_stream_initialization() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
};
let mut projector = ProviderRunResponseProjector::restored(config, true);
let work_id = galaxy_agent_core::ExternalWorkId {
@@ -82,6 +85,109 @@ fn restored_provider_projection_skips_stream_initialization() {
));
}
#[test]
fn provider_projection_uses_supplied_todos_and_advances_each_id() {
let todos = vec![
warp_multi_agent_api::TodoItem {
id: "research".to_owned(),
title: "Research".to_owned(),
description: "Inspect".to_owned(),
},
warp_multi_agent_api::TodoItem {
id: "implement".to_owned(),
title: "Implement".to_owned(),
description: "Edit".to_owned(),
},
warp_multi_agent_api::TodoItem {
id: "verify".to_owned(),
title: "Verify".to_owned(),
description: "Check".to_owned(),
},
];
let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig {
task_id: "task".to_owned(),
conversation_id: "conversation".to_owned(),
needs_create_task: false,
user_query: None,
model_id: "model".to_owned(),
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: Some(todos),
});
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
};
let initial = projector
.project(ProviderRunProjection::ModelTurnStarted {
work_id: work_id.clone(),
profile: galaxy_agent_core::ProviderRequestProfile::new("base"),
runtime_id: "runtime".to_owned(),
model_id: "model".to_owned(),
runtime_request_id: "request".to_owned(),
retry_attempt: 0,
elapsed_ms: 1,
})
.unwrap();
assert_eq!(initial.len(), 2);
let first = projector
.project(ProviderRunProjection::ToolBatchReady {
batch: galaxy_agent_core::PendingToolBatch {
work_id: galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
},
calls: Vec::new(),
},
})
.unwrap();
let second = projector
.project(ProviderRunProjection::ToolBatchReady {
batch: galaxy_agent_core::PendingToolBatch {
work_id: galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1),
},
calls: Vec::new(),
},
})
.unwrap();
let ids = |events: &[warp_multi_agent_api::ResponseEvent]| {
events
.iter()
.flat_map(|event| match &event.r#type {
Some(response_event::Type::ClientActions(actions)) => actions
.actions
.iter()
.filter_map(|action| match &action.action {
Some(client_action::Action::AddMessagesToTask(add)) => add
.messages
.iter()
.filter_map(|message| match &message.message {
Some(message::Message::UpdateTodos(update)) => match update
.operation
.as_ref()
{
Some(message::update_todos::Operation::MarkTodosCompleted(
mark,
)) => Some(mark.todo_ids[0].clone()),
_ => None,
},
_ => None,
})
.next(),
_ => None,
})
.collect::<Vec<_>>(),
_ => Vec::new(),
})
.collect::<Vec<_>>()
};
assert_eq!(ids(&first), vec!["research"]);
assert_eq!(ids(&second), vec!["implement"]);
}
#[test]
fn restored_uninitialized_projection_replays_init_before_live_delta() {
let config = RuntimeResponseConfig {
@@ -93,6 +199,7 @@ fn restored_uninitialized_projection_replays_init_before_live_delta() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
};
let mut projector = ProviderRunResponseProjector::restored(config, false);
let work_id = galaxy_agent_core::ExternalWorkId {
@@ -120,11 +227,18 @@ fn restored_uninitialized_projection_replays_init_before_live_delta() {
})
.unwrap();
assert_eq!(started.len(), 1);
assert_eq!(started.len(), 2);
assert!(matches!(
started[0].r#type,
Some(response_event::Type::Init(_))
));
let Some(response_event::Type::ClientActions(actions)) = &started[1].r#type else {
panic!("initial provider turn should publish its task list");
};
assert!(matches!(
actions.actions[0].action,
Some(client_action::Action::AddMessagesToTask(_))
));
assert_eq!(delta.len(), 1);
assert!(matches!(
delta[0].r#type,
@@ -143,6 +257,7 @@ fn provider_followup_turn_starts_a_distinct_text_message() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let first_work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
@@ -366,6 +481,7 @@ fn provider_retry_clears_failed_attempt_output_before_new_messages() {
max_context_tokens: Some(1_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let work_id = galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"),
app/src/ai/runtime/provider_run_coordinator.rs
+76
View File
@@ -82,6 +82,7 @@ pub(crate) enum ProviderRunProjection {
pub(crate) enum ProviderRunBlock {
Tools(PendingToolBatch),
ReadyToCallModel,
ContextWindowExceeded,
AwaitingDriver {
work_id: ExternalWorkId,
stop_reason: StopReason,
@@ -175,6 +176,8 @@ pub(crate) struct ProviderRunCoordinator {
profiles: BTreeMap<String, ProviderRunProfile>,
model_start_timeout: Duration,
model_event_idle_timeout: Duration,
context_compaction_requested: bool,
max_context_tokens: Option<u32>,
}
impl ProviderRunCoordinator {
@@ -218,6 +221,8 @@ impl ProviderRunCoordinator {
profiles,
model_start_timeout: PROVIDER_MODEL_START_TIMEOUT,
model_event_idle_timeout: PROVIDER_MODEL_EVENT_IDLE_TIMEOUT,
context_compaction_requested: false,
max_context_tokens: None,
})
}
@@ -239,6 +244,10 @@ impl ProviderRunCoordinator {
self.profiles.get(profile).map(|profile| &profile.request)
}
pub(crate) fn set_max_context_tokens(&mut self, max_context_tokens: Option<u32>) {
self.max_context_tokens = max_context_tokens;
}
pub(crate) fn insert_profile(
&mut self,
profile: impl Into<String>,
@@ -364,8 +373,16 @@ impl ProviderRunCoordinator {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
return self.terminal_block();
}
if self.request_needs_context_compaction(&call) {
self.run.prepare_context_compaction(&call.work_id)?;
return Ok(ProviderRunBlock::ContextWindowExceeded);
}
self.drive_model_call_acknowledged(call, control.clone(), &mut project)
.await?;
if self.context_compaction_requested {
self.context_compaction_requested = false;
return Ok(ProviderRunBlock::ContextWindowExceeded);
}
}
Some(ProviderRunStep::DispatchTools(batch)) => {
if !self.checkpoint_or_fail(&mut checkpoint).await? {
@@ -428,6 +445,17 @@ impl ProviderRunCoordinator {
}
}
fn request_needs_context_compaction(&self, call: &ProviderModelCall) -> bool {
let Some(max_context_tokens) = self.max_context_tokens else {
return false;
};
let Some(profile) = self.profiles.get(call.profile.as_str()) else {
return false;
};
let request = request_for_model_call(profile.request.clone(), call);
estimate_turn_request_tokens(&request) >= u64::from(max_context_tokens)
}
async fn checkpoint_or_fail<C>(
&mut self,
checkpoint: &mut C,
@@ -830,6 +858,11 @@ impl ProviderRunCoordinator {
self.run.cancel("provider model call was cancelled")?;
return Ok(());
}
if reason == StopReason::ContextWindowExceeded {
self.run.prepare_context_compaction(&call.work_id)?;
self.context_compaction_requested = true;
return Ok(());
}
let turn = buffer.complete(reason, advertised_tools);
if let Err(error) = self.run.accept_model_turn(&call.work_id, turn) {
self.run.fail(
@@ -925,6 +958,11 @@ impl ProviderRunCoordinator {
"payload": &error,
}));
let error_message = error.message.clone();
if error.kind == AgentErrorKind::ContextWindowExceeded {
self.run.prepare_context_compaction(&call.work_id)?;
self.context_compaction_requested = true;
return Ok(());
}
let disposition = self
.run
.register_model_failure(&call.work_id, error.clone())?;
@@ -1082,6 +1120,44 @@ fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) -
template
}
const ESTIMATED_CHARS_PER_TOKEN: u64 = 4;
/// Deliberately overestimates request size without requiring provider-specific tokenizers.
fn estimate_turn_request_tokens(request: &TurnRequest) -> u64 {
let mut chars = request
.system_prompt
.as_deref()
.map_or(0, |text| text.chars().count());
chars += request.prompt.as_ref().map_or(0, |prompt| {
serde_json::to_string(prompt)
.unwrap_or_default()
.chars()
.count()
});
chars += request
.messages
.iter()
.map(|message| {
serde_json::to_string(message)
.unwrap_or_default()
.chars()
.count()
})
.sum::<usize>();
chars += request
.tools
.iter()
.map(|tool| {
serde_json::to_string(tool)
.unwrap_or_default()
.chars()
.count()
})
.sum::<usize>();
let input_tokens = (chars as u64).div_ceil(ESTIMATED_CHARS_PER_TOKEN);
input_tokens.saturating_add(request.max_output_tokens.unwrap_or_default())
}
fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> {
match event {
ToolEvent::Proposed { call } => Ok(&call.id),
app/src/ai/runtime/provider_run_coordinator_tests.rs
+44
View File
@@ -164,6 +164,48 @@ fn request() -> TurnRequest {
request
}
#[test]
fn estimates_turn_request_with_system_tools_messages_and_output() {
let mut request = TurnRequest::new(
"test-model",
vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("12345678".to_string()),
}],
);
request.system_prompt = Some("1234".to_string());
request.tools = vec![ToolDefinition {
name: "tool".to_string(),
description: "description".to_string(),
input_schema: serde_json::json!({"type": "object"}),
}];
request.max_output_tokens = Some(10);
let expected_input = (request.system_prompt.as_deref().unwrap().len()
+ serde_json::to_string(&request.messages[0]).unwrap().len()
+ serde_json::to_string(&request.tools[0]).unwrap().len()) as u64;
assert_eq!(
estimate_turn_request_tokens(&request),
expected_input.div_ceil(ESTIMATED_CHARS_PER_TOKEN) + 10
);
}
#[tokio::test]
async fn proactively_compacts_before_starting_an_oversized_request() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
let mut coordinator = coordinator(runtime.clone());
coordinator.set_max_context_tokens(Some(1));
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, |_| Ok(()))
.await
.unwrap();
assert_eq!(block, ProviderRunBlock::ContextWindowExceeded);
assert!(runtime.requests().is_empty());
}
fn started(id: &str) -> ScriptedEvent {
Ok(AgentEvent::TurnStarted {
runtime_request_id: id.to_string(),
@@ -1166,6 +1208,7 @@ async fn transcript_projector_emits_one_ui_stream_for_the_whole_run() {
max_context_tokens: Some(100_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let mut ui_events = Vec::new();
let (_sender, control) = turn_control();
@@ -1224,6 +1267,7 @@ fn transcript_projector_preserves_provider_failure_message() {
max_context_tokens: Some(100_000),
capabilities: RuntimeCapabilities::provider(),
empty_output_message: None,
todo_items: None,
});
let events = projector
.finish(
app/src/ai/runtime/rig.rs
+5 -3
View File
@@ -21,9 +21,9 @@ use super::rig_tool::action_from_tool_call;
use super::ProviderRunProfile;
use crate::ai::agent::api::RequestParams;
use crate::ai::agent::AIAgentAction;
use crate::ai::bedrock::client::BedrockClient;
use crate::ai::bedrock::convert::CachingConfig;
use crate::ai::bedrock::external_config::ExternalBedrockConfig;
use crate::ai::provider::client::BedrockClient;
use crate::ai::provider::convert::CachingConfig;
use crate::ai::provider::external_config::ExternalBedrockConfig;
use crate::ai::provider::types::ConversationMessage;
use crate::ai::runtime::RuntimeResponseConfig;
use crate::settings::OpenAIProviderKind;
@@ -137,6 +137,7 @@ pub(crate) async fn prepare_provider_run(
task_id,
needs_create_task,
user_query,
todo_items,
request,
persistent_messages,
tool_result_archive,
@@ -159,6 +160,7 @@ pub(crate) async fn prepare_provider_run(
max_context_tokens,
capabilities: base_runtime.descriptor().capabilities.clone(),
empty_output_message: None,
todo_items,
};
Ok(PreparedProviderRun {
base_profile: ProviderRunProfile::new(base_runtime, request),
app/src/ai/runtime/rig_request.rs
+31 -3
View File
@@ -11,21 +11,23 @@ use galaxy_agent_core::{
};
use sha2::{Digest as _, Sha256};
use uuid::Uuid;
use warp_multi_agent_api as api;
use warp_multi_agent_api::ToolType;
use crate::ai::agent::api::RequestParams;
use crate::ai::agent::{AIAgentContext, AIAgentInput, MCPContext, UserQueryMode};
use crate::ai::bedrock::request_translator::{
default_tool_definitions, sanitize_messages_for_bedrock, tool_name_is_supported,
};
use crate::ai::openai::client::OpenAIClientConfig;
use crate::ai::openai::request_translator::sanitize_messages_for_openai;
use crate::ai::provider::request_translator::{
default_tool_definitions, sanitize_messages_for_bedrock, tool_name_is_supported,
};
use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn;
pub(crate) struct PreparedRigTurn {
pub task_id: String,
pub needs_create_task: bool,
pub user_query: Option<String>,
pub todo_items: Option<Vec<api::TodoItem>>,
pub request: TurnRequest,
pub persistent_messages: Vec<ConversationMessage>,
pub tool_result_archive: Vec<ConversationMessage>,
@@ -226,6 +228,8 @@ fn prepare_rig_turn_for_provider(
..
} = params;
let todo_items = todo_items_from_tasks(&tasks);
let task_id = root_task_id
.or_else(|| tasks.first().map(|task| task.id.clone()))
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
@@ -299,6 +303,7 @@ fn prepare_rig_turn_for_provider(
task_id,
needs_create_task,
user_query,
todo_items,
request,
persistent_messages,
tool_result_archive,
@@ -307,6 +312,29 @@ fn prepare_rig_turn_for_provider(
}
}
// Reuse a plan emitted by the model when the existing transcript contains one. This keeps
// direct-provider projection aligned with UpdateTodos instead of inventing a second plan.
fn todo_items_from_tasks(tasks: &[api::Task]) -> Option<Vec<api::TodoItem>> {
let mut items = None;
for task in tasks {
for message in &task.messages {
let Some(api::message::Message::UpdateTodos(update)) = &message.message else {
continue;
};
match update.operation.as_ref()? {
api::message::update_todos::Operation::CreateTodoList(create) => {
items = Some(create.initial_todos.clone());
}
api::message::update_todos::Operation::UpdatePendingTodos(update) => {
items = Some(update.updated_pending_todos.clone());
}
api::message::update_todos::Operation::MarkTodosCompleted(_) => {}
}
}
}
items.filter(|items| !items.is_empty())
}
fn input_messages(
inputs: Vec<AIAgentInput>,
tool_results: Vec<ToolResult>,
app/src/ai/runtime/rig_request_tests.rs
-1
View File
@@ -32,7 +32,6 @@ fn config() -> OpenAIClientConfig {
reasoning_effort: None,
max_input_tokens: Some(128_000),
max_output_tokens: Some(8_192),
use_rig: true,
supports_system_messages: true,
}
}
app/src/ai/runtime/rig_tests.rs
-1
View File
@@ -16,7 +16,6 @@ fn openai_config(model: &str) -> OpenAIClientConfig {
reasoning_effort: None,
max_input_tokens: Some(128_000),
max_output_tokens: Some(8_192),
use_rig: true,
supports_system_messages: true,
}
}
app/src/ai/runtime/rig_tool.rs
+1 -1
View File
@@ -102,7 +102,7 @@ pub(crate) fn action_from_tool_call(
optional_bounded_u64(
input,
"wait_seconds",
crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS,
crate::ai::provider::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS,
)?
.unwrap_or(2),
))),