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

Add unified models UI and Rig Bedrock runtime

Browse Source
This commit is contained in:
rkw6086
2026-08-04 17:25:19 -05:00
parent a3c68e9c30
commit b0ad07f6f2
41 changed files with 2122 additions and 564 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/agent/api/impl.rs
+26 -2
View File
@@ -28,8 +28,8 @@ pub async fn generate_multi_agent_output(
redaction::redact_inputs(&mut params.input);
}
if let ProviderConfig::OpenAI(config) = &provider_config {
if config.use_rig {
match &provider_config {
ProviderConfig::OpenAI(config) if config.use_rig => {
return Ok(crate::ai::runtime::rig_openai_response_stream(
config.clone(),
params,
@@ -38,6 +38,30 @@ pub async fn generate_multi_agent_output(
cancellation_rx,
));
}
ProviderConfig::Bedrock(config) if config.use_rig => {
return match crate::ai::runtime::rig_bedrock_response_stream(
config.clone(),
params,
supported_tools,
supported_cli_agent_tools,
cancellation_rx,
)
.await
{
Ok(stream) => Ok(stream),
Err(error) => {
log::error!("[rig/bedrock] Runtime error: {error}");
let error = Arc::new(crate::server::server_api::AIApiError::Stream {
stream_type: "rig_bedrock",
source: error,
});
let (sender, receiver) = async_channel::unbounded();
let _ = sender.send(Err(error)).await;
Ok(Box::pin(receiver))
}
};
}
ProviderConfig::OpenAI(_) | ProviderConfig::Bedrock(_) | ProviderConfig::None => {}
}
let mut logging_metadata = HashMap::new();
app/src/ai/bedrock/client.rs
+27 -2
View File
@@ -5,6 +5,8 @@ use aws_config::BehaviorVersion;
use aws_credential_types::provider::ProvideCredentials;
use aws_sdk_bedrockruntime::config::Region;
use aws_sdk_bedrockruntime::Client as BedrockRuntimeClient;
use galaxy_agent_core::AgentError;
use galaxy_agent_rig::{BedrockRigConfig, BedrockRuntime};
use super::convert::{build_converse_request, CachingConfig, ConversationMessage, ToolDefinition};
use super::diagnostic::BedrockDiagnosticLogger;
@@ -38,6 +40,7 @@ pub struct BedrockClientConfig {
pub secret_access_key: String,
pub session_token: Option<String>,
pub cross_region_inference: bool,
pub use_rig: bool,
}
impl BedrockClientConfig {
@@ -141,8 +144,8 @@ impl BedrockClient {
match provider.provide_credentials().await {
Ok(creds) => {
log::info!(
"[bedrock] Resolved AWS credentials successfully: access_key_id={:?}, has_session_token={}, expiry={:?}",
creds.access_key_id(),
"[bedrock] Resolved AWS credentials successfully: has_access_key_id={}, has_session_token={}, expiry={:?}",
!creds.access_key_id().is_empty(),
creds.session_token().is_some(),
creds.expiry(),
);
@@ -168,6 +171,28 @@ impl BedrockClient {
})
}
/// Builds the Phase 4 Rig runtime from the AWS SDK client whose region and
/// credentials Galaxy already resolved. This does not change production
/// routing; callers opt in only after the Bedrock parity suite passes.
pub fn rig_runtime(
&self,
model: String,
cross_region_inference: bool,
prompt_caching: bool,
max_output_tokens: Option<u64>,
) -> Result<BedrockRuntime, AgentError> {
BedrockRuntime::from_aws_client(
self.runtime_client.clone(),
BedrockRigConfig {
model,
region: self.region.clone(),
cross_region_inference,
prompt_caching,
max_output_tokens,
},
)
}
#[allow(clippy::too_many_arguments)]
pub async fn converse_stream(
&self,
app/src/ai/bedrock/convert.rs
+12 -2
View File
@@ -3,8 +3,9 @@ use std::collections::HashMap;
use aws_sdk_bedrockruntime::types::{
CachePointBlock, CachePointType, CacheTtl, ContentBlock, ConversationRole, ImageBlock,
ImageFormat, ImageSource, InferenceConfiguration, Message as BedrockMessage,
SystemContentBlock, Tool, ToolConfiguration, ToolInputSchema, ToolResultBlock,
ToolResultContentBlock, ToolResultStatus, ToolSpecification, ToolUseBlock,
ReasoningContentBlock, ReasoningTextBlock, SystemContentBlock, Tool, ToolConfiguration,
ToolInputSchema, ToolResultBlock, ToolResultContentBlock, ToolResultStatus, ToolSpecification,
ToolUseBlock,
};
use aws_smithy_types::{Blob, Document};
use serde_json::Value as JsonValue;
@@ -148,6 +149,15 @@ fn convert_messages(
.into_iter()
.map(|part| match part {
ContentPart::Text(text) => ContentBlock::Text(text),
ContentPart::Reasoning { text, signature } => {
ContentBlock::ReasoningContent(ReasoningContentBlock::ReasoningText(
ReasoningTextBlock::builder()
.text(text)
.set_signature(signature)
.build()
.expect("valid reasoning text block"),
))
}
ContentPart::Image { data, mime_type } => image_content_block(data, &mime_type),
ContentPart::ToolUse {
tool_use_id,
app/src/ai/bedrock/diagnostic.rs
+8
View File
@@ -492,6 +492,14 @@ fn serialize_messages(messages: &[ConversationMessage]) -> JsonValue {
super::convert::ContentPart::Text(t) => {
serde_json::json!({"type": "text", "text": t})
}
super::convert::ContentPart::Reasoning { text, signature } => {
serde_json::json!({
"type": "reasoning",
"char_length": text.len(),
"has_signature": signature.is_some(),
"text": "REDACTED",
})
}
super::convert::ContentPart::Image { data, mime_type } => {
serde_json::json!({
"type": "image",
app/src/ai/bedrock/e2e_tests.rs
+1
View File
@@ -269,6 +269,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(),
session_token: None,
cross_region_inference: false,
use_rig: false,
})
}
app/src/ai/bedrock/external_config.rs
+1
View File
@@ -144,6 +144,7 @@ fn parse_claude_code_model_map(
model_id: arn,
display_name,
vision_supported: true,
use_rig: false,
}
})
.collect()
app/src/ai/bedrock/integration_tests.rs
+1
View File
@@ -24,6 +24,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(),
session_token: None,
cross_region_inference: false,
use_rig: false,
})
}
app/src/ai/bedrock/models.rs
+29
View File
@@ -129,6 +129,7 @@ pub fn get_effective_models(user_models: &[BedrockModelConfig]) -> Vec<BedrockMo
model_id: m.model_id.to_string(),
display_name: m.display_name.to_string(),
vision_supported: m.vision_supported,
use_rig: false,
})
.collect();
for default in defaults {
@@ -145,10 +146,38 @@ pub fn get_effective_models(user_models: &[BedrockModelConfig]) -> Vec<BedrockMo
model_id: m.model_id.to_string(),
display_name: m.display_name.to_string(),
vision_supported: m.vision_supported,
use_rig: false,
})
.collect()
}
pub fn configured_model_uses_rig(
selected_model_id: &str,
configured_models: &[BedrockModelConfig],
region: &str,
cross_region_inference: bool,
) -> bool {
let selected_model_id = strip_context_marker(selected_model_id);
configured_models.iter().any(|model| {
if !model.use_rig {
return false;
}
let configured_model_id = strip_context_marker(&model.model_id);
if configured_model_id == selected_model_id {
return true;
}
galaxy_agent_rig::resolve_bedrock_model_id(&model.model_id, region, cross_region_inference)
.is_ok_and(|resolved| strip_context_marker(&resolved) == selected_model_id)
})
}
fn strip_context_marker(model_id: &str) -> &str {
model_id
.strip_suffix("[1m]")
.or_else(|| model_id.strip_suffix("[1M]"))
.unwrap_or(model_id)
}
pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
if model_id.starts_with("arn:") {
return model_id.to_string();
app/src/ai/bedrock/models_tests.rs
+32 -2
View File
@@ -82,8 +82,8 @@ fn test_cross_region_prefix_unknown_region() {
fn test_get_effective_models_empty_returns_defaults() {
let models = get_effective_models(&[]);
assert_eq!(models.len(), DEFAULT_BEDROCK_MODELS.len());
assert_eq!(models[0].model_id, "anthropic.claude-opus-4-6[1m]");
assert_eq!(models[0].display_name, "Claude Opus 4.6");
assert_eq!(models[0].model_id, "us.anthropic.claude-opus-4-6-v1[1m]");
assert_eq!(models[0].display_name, "Claude Opus 4.6 (1M)");
}
#[test]
@@ -92,6 +92,7 @@ fn test_get_effective_models_custom_overrides() {
model_id: "custom.model-v1:0".to_string(),
display_name: "Custom Model".to_string(),
vision_supported: false,
use_rig: true,
}];
let models = get_effective_models(&custom);
assert_eq!(models.len(), 1);
@@ -103,3 +104,32 @@ fn test_cross_region_prefix_skips_arn() {
let arn = "arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy";
assert_eq!(apply_cross_region_prefix(arn, "us-east-1"), arn);
}
#[test]
fn rig_opt_in_matches_context_markers_and_resolved_inference_profiles() {
let configured = vec![BedrockModelConfig {
model_id: "anthropic.claude-test[1m]".to_string(),
display_name: "Claude Test".to_string(),
vision_supported: false,
use_rig: true,
}];
assert!(configured_model_uses_rig(
"us.anthropic.claude-test",
&configured,
"us-east-1",
true,
));
assert!(configured_model_uses_rig(
"anthropic.claude-test[1M]",
&configured,
"us-east-1",
false,
));
assert!(!configured_model_uses_rig(
"anthropic.other-model",
&configured,
"us-east-1",
false,
));
}
app/src/ai/bedrock/request_translator.rs
+15 -2
View File
@@ -750,6 +750,7 @@ fn persist_input_images_on_latest_user_message(
Some(api::input_context::Image { data, mime_type })
}
Some(ContentPart::Text(_))
| Some(ContentPart::Reasoning { .. })
| Some(ContentPart::ToolUse { .. })
| Some(ContentPart::ToolResult { .. })
| None => None,
@@ -873,11 +874,17 @@ fn is_pure_tool_result(content: &MessageContent) -> bool {
fn strip_tool_result_parts(content: &mut MessageContent) {
if let MessageContent::MultiPart(parts) = content {
parts.retain(|p| !matches!(p, ContentPart::ToolResult { .. }));
if parts.len() == 1 && !matches!(parts.first(), Some(ContentPart::Image { .. })) {
if parts.len() == 1
&& !matches!(
parts.first(),
Some(ContentPart::Image { .. } | ContentPart::Reasoning { .. })
)
{
let part = parts.remove(0);
*content = match part {
ContentPart::Text(t) => MessageContent::Text(t),
ContentPart::Image { .. } => unreachable!(),
ContentPart::Reasoning { .. } => unreachable!(),
ContentPart::ToolUse {
tool_use_id,
name,
@@ -916,11 +923,17 @@ fn strip_orphaned_tool_result_parts(
ContentPart::ToolResult { tool_use_id, .. } => valid_ids.contains(tool_use_id),
_ => true,
});
if parts.len() == 1 && !matches!(parts.first(), Some(ContentPart::Image { .. })) {
if parts.len() == 1
&& !matches!(
parts.first(),
Some(ContentPart::Image { .. } | ContentPart::Reasoning { .. })
)
{
let part = parts.remove(0);
*content = match part {
ContentPart::Text(t) => MessageContent::Text(t),
ContentPart::Image { .. } => unreachable!(),
ContentPart::Reasoning { .. } => unreachable!(),
ContentPart::ToolUse {
tool_use_id,
name,
app/src/ai/bedrock/translator.rs
+7
View File
@@ -180,6 +180,13 @@ fn describe_message_content(content: &crate::ai::bedrock::convert::MessageConten
.iter()
.map(|p| match p {
ContentPart::Text(t) => format!("Text({})", t.len()),
ContentPart::Reasoning { text, signature } => {
format!(
"Reasoning({}chars,signed={})",
text.len(),
signature.is_some()
)
}
ContentPart::Image { data, mime_type } => {
format!("Image({mime_type},{}bytes)", data.len())
}
app/src/ai/blocklist/controller.rs
+7
View File
@@ -4346,6 +4346,7 @@ impl BlocklistAIController {
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
use_rig: false,
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
@@ -4431,6 +4432,9 @@ impl BlocklistAIController {
.iter()
.map(|p| match p {
ContentPart::Text(t) => t.clone(),
ContentPart::Reasoning { text, .. } => {
format!("[Reasoning] {text}")
}
ContentPart::Image { .. } => "[Image attachment]".to_string(),
ContentPart::ToolUse { name, input, .. } => {
format!("[Tool: {}] {}", name, input)
@@ -4540,6 +4544,9 @@ impl BlocklistAIController {
.iter()
.map(|p| match p {
ContentPart::Text(t) => (t.len() / 4) as u32,
ContentPart::Reasoning { text, .. } => {
(text.len() / 4) as u32
}
ContentPart::Image { .. } => 1_600,
ContentPart::ToolUse { input, .. } => {
(input.to_string().len() / 4) as u32
app/src/ai/blocklist/controller/response_stream.rs
+20 -2
View File
@@ -243,15 +243,33 @@ impl ResponseStream {
// Fall back to Bedrock
if *settings.bedrock_enabled.value() {
let auth_method = *settings.bedrock_auth_method.value();
let region = settings.bedrock_region.value().clone();
let cross_region_inference = *settings.bedrock_cross_region_inference.value();
let mut use_rig = crate::ai::bedrock::models::configured_model_uses_rig(
model_id,
settings.bedrock_models.value(),
&region,
cross_region_inference,
);
if use_rig
&& crate::ai::bedrock::external_config::ExternalBedrockConfig::load()
.enable_prompt_caching_1h
{
log::warn!(
"[rig/bedrock] Using the compatibility runtime because Rig does not yet expose Bedrock's one-hour cache TTL"
);
use_rig = false;
}
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let mut config = BedrockClientConfig {
auth_method,
profile: settings.bedrock_profile.value().clone(),
region: settings.bedrock_region.value().clone(),
region,
access_key_id: settings.bedrock_access_key_id.value().clone(),
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
cross_region_inference,
use_rig,
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
app/src/ai/blocklist/usage/context_window_view.rs
+9
View File
@@ -45,6 +45,7 @@ impl View for ContextWindowView {
.iter()
.map(|p| match p {
ContentPart::Text(t) => t.len(),
ContentPart::Reasoning { text, .. } => text.len(),
ContentPart::Image { .. } => 6_400,
ContentPart::ToolUse { input, .. } => input.to_string().len(),
ContentPart::ToolResult { content, .. } => content.len(),
@@ -112,6 +113,14 @@ impl View for ContextWindowView {
ContentPart::Text(t) => {
out.push_str(&format!("[Part {} Text] {}\n", pi, t));
}
ContentPart::Reasoning { text, signature } => {
out.push_str(&format!(
"[Part {} Reasoning] signed={}\n{}\n",
pi,
signature.is_some(),
text
));
}
ContentPart::Image { data, mime_type } => {
out.push_str(&format!(
"[Part {} Image] mime_type={}, bytes={}\n",
app/src/ai/crosscheck/reviewer.rs
+1
View File
@@ -186,6 +186,7 @@ impl CrosscheckReviewer {
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
use_rig: false,
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
app/src/ai/llms.rs
+131 -2
View File
@@ -721,6 +721,7 @@ impl LLMPreferences {
model_id: default.model_id.to_string(),
display_name: default.display_name.to_string(),
vision_supported: default.vision_supported,
use_rig: false,
});
added = true;
}
@@ -1234,6 +1235,80 @@ impl LLMPreferences {
);
}
/// Explicitly refreshes the models for one entry in the OpenAI-compatible
/// provider registry. Unlike the legacy endpoint refresh, this is only
/// called from a user action so configured remote endpoints are never
/// contacted merely because Galaxy started.
#[cfg(not(target_family = "wasm"))]
pub fn fetch_openai_provider_models(
&mut self,
provider_index: usize,
ctx: &mut ModelContext<Self>,
) {
let settings = AISettings::as_ref(ctx);
if !*settings.openai_enabled.value() {
return;
}
let Some(provider) = settings
.openai_providers
.value()
.get(provider_index)
.cloned()
else {
return;
};
if provider.base_url.trim().is_empty() {
return;
}
let requested_base_url = provider.base_url;
let api_key = provider.api_key.filter(|key| !key.is_empty());
let request_base_url = requested_base_url.clone();
let _ = ctx.spawn(
async move {
let base = request_base_url.trim_end_matches('/');
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()
.unwrap_or_default();
if let Some(models) =
fetch_from_litellm_model_info(base, api_key.as_deref(), &client).await
{
return models;
}
fetch_from_openai_models(base, api_key.as_deref(), &client).await
},
move |_, discovered_models, ctx| {
if discovered_models.is_empty() {
return;
}
AISettings::handle(ctx).update(ctx, |settings, ctx| {
let mut providers = settings.openai_providers.value().clone();
let Some(provider) = providers.get_mut(provider_index) else {
return;
};
// Do not apply a response to an entry that was edited or
// reordered while its discovery request was in flight.
if provider.base_url != requested_base_url {
return;
}
provider.models =
merge_discovered_provider_models(&provider.models, discovered_models);
if let Err(err) = settings.openai_providers.set_value(providers, ctx) {
report_error!(err.context("Failed to persist discovered provider models"));
}
});
},
);
}
/// Returns the `LLMInfo` for the base LLM to be used for an Agent Mode request.
pub fn get_active_base_model<'a>(
&'a self,
@@ -1994,6 +2069,52 @@ fn openai_model_context_size(model: &OpenAIModelConfig) -> u32 {
model.max_input_tokens.unwrap_or(model.context_size)
}
/// Merges endpoint metadata into a provider's configured models without
/// discarding local routing choices or manually configured models.
#[cfg(not(target_family = "wasm"))]
fn merge_discovered_provider_models(
existing_models: &[OpenAIModelConfig],
discovered_models: Vec<OpenAIModelConfig>,
) -> Vec<OpenAIModelConfig> {
let mut merged = Vec::with_capacity(discovered_models.len() + existing_models.len());
let mut discovered_ids = HashSet::new();
for mut discovered in discovered_models {
if !discovered_ids.insert(discovered.model_id.clone()) {
continue;
}
if let Some(existing) = existing_models
.iter()
.find(|model| model.model_id == discovered.model_id)
{
discovered.display_name = existing.display_name.clone();
discovered.use_rig = existing.use_rig;
if existing.supports_system_messages.is_some() {
discovered.supports_system_messages = existing.supports_system_messages;
}
if discovered.provider.is_none() {
discovered.provider = existing.provider.clone();
}
} else {
discovered.use_rig = true;
}
if discovered.model_id.starts_with("codex-gpt-") {
discovered.supports_system_messages = Some(false);
}
merged.push(discovered);
}
merged.extend(
existing_models
.iter()
.filter(|model| !discovered_ids.contains(&model.model_id))
.cloned(),
);
merged
}
#[cfg(not(target_family = "wasm"))]
fn openai_model_context_window(model: &OpenAIModelConfig) -> LLMContextWindow {
let context_size = openai_model_context_size(model);
@@ -2118,7 +2239,11 @@ async fn fetch_from_litellm_model_info(
max_output_tokens,
provider,
use_rig: false,
supports_system_messages: model_info["supports_system_messages"].as_bool(),
supports_system_messages: if model_name.starts_with("codex-gpt-") {
Some(false)
} else {
model_info["supports_system_messages"].as_bool()
},
})
})
.collect();
@@ -2241,7 +2366,11 @@ async fn fetch_from_openai_models(
max_output_tokens,
provider,
use_rig: false,
supports_system_messages: m["supports_system_messages"].as_bool(),
supports_system_messages: if id.starts_with("codex-gpt-") {
Some(false)
} else {
m["supports_system_messages"].as_bool()
},
})
})
.collect();
app/src/ai/llms_tests.rs
+65 -1
View File
@@ -10,7 +10,7 @@ use crate::network::NetworkStatus;
use crate::server::cloud_objects::update_manager::UpdateManager;
use crate::server::server_api::ServerApiProvider;
use crate::server::sync_queue::SyncQueue;
use crate::settings::{OpenAIModelConfig, OpenAIProviderConfig};
use crate::settings::OpenAIModelConfig;
use crate::test_util::settings::initialize_settings_for_tests;
use crate::workspaces::team_tester::TeamTesterStatus;
use crate::workspaces::user_workspaces::UserWorkspaces;
@@ -138,3 +138,67 @@ fn llm_info_round_trip_serializes_and_deserializes() {
assert_eq!(info, round_tripped);
}
fn openai_model(model_id: &str) -> OpenAIModelConfig {
OpenAIModelConfig {
model_id: model_id.to_string(),
display_name: model_id.to_string(),
vision_supported: false,
context_size: 200_000,
max_input_tokens: None,
max_output_tokens: None,
provider: None,
use_rig: false,
supports_system_messages: None,
}
}
#[test]
fn provider_discovery_preserves_local_model_overrides() {
let mut existing = openai_model("codex-gpt-5.6-sol-xhigh");
existing.display_name = "My Codex".to_string();
existing.context_size = 100_000;
existing.provider = Some("openai".to_string());
existing.use_rig = true;
// Even stale or incorrect endpoint metadata must not opt ChatGPT-backed
// Codex models back into the system role.
existing.supports_system_messages = Some(true);
let mut discovered = openai_model("codex-gpt-5.6-sol-xhigh");
discovered.display_name = "Codex from endpoint".to_string();
discovered.context_size = 400_000;
discovered.max_output_tokens = Some(32_000);
discovered.supports_system_messages = Some(true);
let merged = merge_discovered_provider_models(&[existing], vec![discovered]);
assert_eq!(merged.len(), 1);
assert_eq!(merged[0].display_name, "My Codex");
assert_eq!(merged[0].context_size, 400_000);
assert_eq!(merged[0].max_output_tokens, Some(32_000));
assert!(merged[0].use_rig);
assert_eq!(merged[0].supports_system_messages, Some(false));
assert_eq!(merged[0].provider.as_deref(), Some("openai"));
}
#[test]
fn codex_models_reject_system_messages_even_with_stale_true_metadata() {
let mut model = openai_model("codex-gpt-5.6-sol-xhigh");
model.supports_system_messages = Some(true);
assert!(!model.supports_system_messages());
}
#[test]
fn provider_discovery_enables_rig_for_new_models_and_keeps_manual_models() {
let manual = openai_model("manual-model");
let discovered = openai_model("codex-gpt-new");
let merged = merge_discovered_provider_models(&[manual], vec![discovered]);
assert_eq!(merged.len(), 2);
assert_eq!(merged[0].model_id, "codex-gpt-new");
assert!(merged[0].use_rig);
assert_eq!(merged[0].supports_system_messages, Some(false));
assert_eq!(merged[1].model_id, "manual-model");
}
app/src/ai/openai/convert.rs
+9
View File
@@ -121,6 +121,9 @@ fn convert_user_message(content: MessageContent) -> ConvertedMessages {
ContentPart::Text(text) => {
user_content_parts.push(UserContentPart::Text(text));
}
ContentPart::Reasoning { text, .. } => {
user_content_parts.push(UserContentPart::Text(text));
}
ContentPart::Image { data, mime_type } => {
user_content_parts.push(UserContentPart::Image { data, mime_type });
}
@@ -201,6 +204,12 @@ fn convert_assistant_message(content: MessageContent) -> ConvertedMessages {
}
text_content.push_str(&text);
}
ContentPart::Reasoning { text, .. } => {
if !text_content.is_empty() {
text_content.push('\n');
}
text_content.push_str(&text);
}
ContentPart::ToolUse {
tool_use_id,
name,
app/src/ai/runtime/mod.rs
+1 -1
View File
@@ -4,4 +4,4 @@ mod rig_request;
mod rig_tool;
pub(crate) use provider::ProviderRuntime;
pub(crate) use rig::rig_openai_response_stream;
pub(crate) use rig::{rig_bedrock_response_stream, rig_openai_response_stream};
app/src/ai/runtime/rig.rs
+111 -19
View File
@@ -11,10 +11,12 @@ use uuid::Uuid;
use warp_multi_agent_api::response_event::stream_finished;
use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent, ToolType};
use super::rig_request::{prepare_rig_turn, PreparedRigTurn};
use super::rig_request::{prepare_bedrock_rig_turn, prepare_rig_turn, PreparedRigTurn};
use super::rig_tool::action_from_tool_call;
use crate::ai::agent::api::{Event, RequestParams, ResponseStream, StreamEvent};
use crate::ai::agent::AIAgentAction;
use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig};
use crate::ai::bedrock::external_config::ExternalBedrockConfig;
use crate::ai::bedrock::response_translator::{
build_add_agent_output_message, build_append_text, build_create_task, build_stream_init,
build_user_query_message,
@@ -32,6 +34,75 @@ pub(crate) fn rig_openai_response_stream(
cancellation_rx: oneshot::Receiver<()>,
) -> ResponseStream {
let skill_path_origin = params.session_context.skill_path_origin();
let prepared = prepare_rig_turn(&config, params, supported_tools, supported_cli_agent_tools);
let model_id = prepared.request.model.as_str().to_string();
let runtime = OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig {
base_url: config.base_url,
api_key: config.api_key,
model: model_id.clone(),
max_output_tokens: config.max_output_tokens.map(u64::from),
supports_system_messages: config.supports_system_messages,
});
rig_response_stream(
runtime,
prepared,
skill_path_origin,
config.max_input_tokens,
"rig_openai_compatible",
cancellation_rx,
)
}
pub(crate) async fn rig_bedrock_response_stream(
config: BedrockClientConfig,
params: RequestParams,
supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
cancellation_rx: oneshot::Receiver<()>,
) -> anyhow::Result<ResponseStream> {
let skill_path_origin = params.session_context.skill_path_origin();
let max_context_tokens = params.context_window_limit;
let model = params.model.as_str().to_string();
let max_output_tokens = Some(64_000);
let cross_region_inference = config.cross_region_inference;
let external_config = ExternalBedrockConfig::load();
let prompt_caching = !external_config.disable_prompt_caching;
let client = BedrockClient::from_config(config).await?;
let runtime = client.rig_runtime(
model.clone(),
cross_region_inference,
prompt_caching,
max_output_tokens,
)?;
let prepared = prepare_bedrock_rig_turn(
model,
max_output_tokens,
params,
supported_tools,
supported_cli_agent_tools,
);
Ok(rig_response_stream(
runtime,
prepared,
skill_path_origin,
max_context_tokens,
"rig_bedrock",
cancellation_rx,
))
}
fn rig_response_stream<R>(
runtime: R,
prepared: PreparedRigTurn,
skill_path_origin: ai::skills::SkillPathOrigin,
max_context_tokens: Option<u32>,
stream_type: &'static str,
cancellation_rx: oneshot::Receiver<()>,
) -> ResponseStream
where
R: AgentRuntime + Send + Sync + 'static,
{
let PreparedRigTurn {
task_id,
needs_create_task,
@@ -40,21 +111,12 @@ pub(crate) fn rig_openai_response_stream(
persistent_messages,
tool_result_archive,
messages_sent,
} = prepare_rig_turn(&config, params, supported_tools, supported_cli_agent_tools);
} = prepared;
store_messages_sent(&messages_sent, &persistent_messages);
let conversation_id = turn_request.conversation_id.clone();
let model_id = turn_request.model.as_str().to_string();
let tool_policy = ToolPolicy::new(&turn_request.tools);
let runtime = OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig {
base_url: config.base_url,
api_key: config.api_key,
model: model_id.clone(),
max_output_tokens: config.max_output_tokens.map(u64::from),
supports_system_messages: config.supports_system_messages,
});
let max_context_tokens = config.max_input_tokens;
let stream = async_stream::stream! {
let (control_sender, control) = turn_control();
let start_future = runtime.start_turn(turn_request, control).fuse();
@@ -67,7 +129,7 @@ pub(crate) fn rig_openai_response_stream(
match start_future.await {
Ok(stream) => stream,
Err(error) => {
yield Err(agent_error(error));
yield Err(agent_error(error, stream_type));
return;
}
}
@@ -75,7 +137,7 @@ pub(crate) fn rig_openai_response_stream(
result = start_future => match result {
Ok(stream) => stream,
Err(error) => {
yield Err(agent_error(error));
yield Err(agent_error(error, stream_type));
return;
}
},
@@ -87,6 +149,8 @@ pub(crate) fn rig_openai_response_stream(
let mut current_text_message_id: Option<String> = None;
let mut current_reasoning_message_id: Option<String> = None;
let mut full_text = String::new();
let mut full_reasoning = String::new();
let mut reasoning_signature = None;
let mut proposed_tools = Vec::new();
let mut assistant_history_index = None;
let mut usage = Usage::default();
@@ -106,7 +170,7 @@ pub(crate) fn rig_openai_response_stream(
let event = match event {
Ok(event) => event,
Err(error) => {
yield Err(agent_error(error));
yield Err(agent_error(error, stream_type));
return;
}
};
@@ -133,6 +197,7 @@ pub(crate) fn rig_openai_response_stream(
}
}
AgentEvent::ReasoningDelta { text } => {
full_reasoning.push_str(&text);
if let Some(message_id) = &current_reasoning_message_id {
yield Ok(StreamEvent::Response(build_append_reasoning(&task_id, message_id, &text)));
} else {
@@ -141,6 +206,17 @@ pub(crate) fn rig_openai_response_stream(
current_reasoning_message_id = Some(message_id);
}
}
AgentEvent::ReasoningCompleted { text, signature } => {
if current_reasoning_message_id.is_none() && !text.is_empty() {
let message_id = Uuid::new_v4().to_string();
yield Ok(StreamEvent::Response(build_add_reasoning(&task_id, &message_id, &text)));
current_reasoning_message_id = Some(message_id);
}
if !text.is_empty() {
full_reasoning = text;
}
reasoning_signature = signature;
}
AgentEvent::UsageUpdated { usage: updated } => usage = updated,
AgentEvent::Tool {
event: ToolEvent::Proposed { call },
@@ -148,6 +224,8 @@ pub(crate) fn rig_openai_response_stream(
proposed_tools.push(call.clone());
sync_assistant_turn(
&messages_sent,
&full_reasoning,
reasoning_signature.as_deref(),
&full_text,
&proposed_tools,
&mut assistant_history_index,
@@ -164,7 +242,7 @@ pub(crate) fn rig_openai_response_stream(
yield Err(agent_error(AgentError::new(
galaxy_agent_core::AgentErrorKind::Protocol,
message,
)));
), stream_type));
return;
}
}
@@ -198,6 +276,8 @@ pub(crate) fn rig_openai_response_stream(
}
sync_assistant_turn(
&messages_sent,
&full_reasoning,
reasoning_signature.as_deref(),
&full_text,
&proposed_tools,
&mut assistant_history_index,
@@ -222,7 +302,7 @@ pub(crate) fn rig_openai_response_stream(
yield Err(agent_error(AgentError::new(
galaxy_agent_core::AgentErrorKind::Protocol,
"the provider runtime attempted to execute a tool outside Galaxy's permission boundary",
)));
), stream_type));
return;
}
}
@@ -264,11 +344,22 @@ fn append_tool_result(
fn sync_assistant_turn(
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
reasoning_text: &str,
reasoning_signature: Option<&str>,
text: &str,
tool_calls: &[ToolCall],
history_index: &mut Option<usize>,
) {
let mut parts = Vec::with_capacity(usize::from(!text.is_empty()) + tool_calls.len());
let has_reasoning = !reasoning_text.is_empty() || reasoning_signature.is_some();
let mut parts = Vec::with_capacity(
usize::from(has_reasoning) + usize::from(!text.is_empty()) + tool_calls.len(),
);
if has_reasoning {
parts.push(ContentPart::Reasoning {
text: reasoning_text.to_string(),
signature: reasoning_signature.map(str::to_string),
});
}
if !text.is_empty() {
parts.push(ContentPart::Text(text.to_string()));
}
@@ -293,6 +384,7 @@ fn sync_assistant_turn(
name,
input,
},
reasoning @ ContentPart::Reasoning { .. } => MessageContent::MultiPart(vec![reasoning]),
ContentPart::Image { .. } | ContentPart::ToolResult { .. } => unreachable!(),
}
} else {
@@ -395,10 +487,10 @@ fn saturating_i32(value: u64) -> i32 {
i32::try_from(value).unwrap_or(i32::MAX)
}
fn agent_error(error: AgentError) -> Arc<AIApiError> {
fn agent_error(error: AgentError, stream_type: &'static str) -> Arc<AIApiError> {
Arc::new(
AIApiError::Stream {
stream_type: "rig_openai_compatible",
stream_type,
source: anyhow::anyhow!(error),
}
.into_quota_limit_if_provider_budget_exhausted(),
app/src/ai/runtime/rig_request.rs
+52 -6
View File
@@ -13,7 +13,9 @@ 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, tool_name_is_supported};
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;
@@ -32,6 +34,47 @@ pub(crate) fn prepare_rig_turn(
params: RequestParams,
supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
) -> PreparedRigTurn {
prepare_rig_turn_for_provider(
config.model.clone(),
config.max_output_tokens.map(u64::from),
RigRequestSanitizer::OpenAICompatible,
params,
supported_tools,
supported_cli_agent_tools,
)
}
pub(crate) fn prepare_bedrock_rig_turn(
model: String,
max_output_tokens: Option<u64>,
params: RequestParams,
supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
) -> PreparedRigTurn {
prepare_rig_turn_for_provider(
Some(model),
max_output_tokens,
RigRequestSanitizer::Bedrock,
params,
supported_tools,
supported_cli_agent_tools,
)
}
#[derive(Clone, Copy)]
enum RigRequestSanitizer {
OpenAICompatible,
Bedrock,
}
fn prepare_rig_turn_for_provider(
model_override: Option<String>,
max_output_tokens: Option<u64>,
sanitizer: RigRequestSanitizer,
params: RequestParams,
supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
) -> PreparedRigTurn {
let RequestParams {
input,
@@ -70,7 +113,12 @@ pub(crate) fn prepare_rig_turn(
for message in &mut persistent_messages {
message.truncate_tool_results_for_provider_request();
}
sanitize_messages_for_openai(&mut persistent_messages);
match sanitizer {
RigRequestSanitizer::OpenAICompatible => {
sanitize_messages_for_openai(&mut persistent_messages)
}
RigRequestSanitizer::Bedrock => sanitize_messages_for_bedrock(&mut persistent_messages),
}
let mut turn_messages = Vec::new();
if let Some(summary) = progressive_summary {
@@ -91,16 +139,14 @@ pub(crate) fn prepare_rig_turn(
}
turn_messages.extend(persistent_messages.clone());
let model_id = config
.model
.clone()
let model_id = model_override
.filter(|model| !model.is_empty() && model != "auto")
.unwrap_or_else(|| model.as_str().to_string());
let mut request = TurnRequest::new(model_id, turn_messages);
request.conversation_id = conversation_token.map(|token| token.as_str().to_string());
request.system_prompt = Some(system_prompt);
request.tools = tools;
request.max_output_tokens = config.max_output_tokens.map(u64::from);
request.max_output_tokens = max_output_tokens;
PreparedRigTurn {
task_id,
app/src/ai/runtime/rig_request_tests.rs
+35 -1
View File
@@ -4,7 +4,7 @@ use std::sync::Arc;
use galaxy_agent_core::{ContentPart, MessageContent, MessageRole, ToolResult, ToolResultStatus};
use warp_multi_agent_api::ToolType;
use super::{input_messages, prepare_rig_turn, tool_definitions};
use super::{input_messages, prepare_bedrock_rig_turn, prepare_rig_turn, tool_definitions};
use crate::ai::agent::api::RequestParams;
use crate::ai::agent::{
AIAgentContext, AIAgentInput, AnyFileContent, FileContext, MCPContext, MCPServer, UserQueryMode,
@@ -114,6 +114,40 @@ fn builds_a_rig_turn_directly_from_galaxy_request_state() {
));
}
#[test]
fn bedrock_rig_turn_uses_bedrock_history_invariants_without_a_proto_round_trip() {
let mut params = RequestParams::new_for_test();
params.message_history = vec![galaxy_agent_core::ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::Text("Prior assistant message".to_string()),
}];
params.input = vec![user_query("Continue safely")];
let prepared = prepare_bedrock_rig_turn(
"anthropic.claude-test".to_string(),
Some(64_000),
params,
Vec::new(),
Vec::new(),
);
assert_eq!(prepared.request.model.as_str(), "anthropic.claude-test");
assert_eq!(prepared.request.max_output_tokens, Some(64_000));
assert_eq!(
prepared
.request
.messages
.first()
.map(|message| message.role),
Some(MessageRole::User)
);
assert_eq!(
prepared.request.messages.last().map(|message| message.role),
Some(MessageRole::User)
);
assert_eq!(prepared.request.messages, prepared.persistent_messages);
}
#[test]
#[allow(deprecated)]
fn grouped_mcp_tool_names_use_the_installation_id_not_the_display_name() {
app/src/ai/runtime/rig_tests.rs
+44 -1
View File
@@ -2,7 +2,7 @@ use std::sync::{Arc, Mutex};
use ai::skills::SkillPathOrigin;
use galaxy_agent_core::{
MessageContent, MessageRole, StopReason, ToolCall, ToolResult, ToolResultStatus,
ContentPart, MessageContent, MessageRole, StopReason, ToolCall, ToolResult, ToolResultStatus,
};
use warp_multi_agent_api::response_event::stream_finished;
@@ -134,12 +134,16 @@ fn assistant_history_is_updated_before_fast_tool_execution_can_continue() {
sync_assistant_turn(
&messages,
"",
None,
"I'll inspect both.",
std::slice::from_ref(&first_call),
&mut history_index,
);
sync_assistant_turn(
&messages,
"",
None,
"I'll inspect both.",
&[first_call, second_call],
&mut history_index,
@@ -162,6 +166,43 @@ fn assistant_history_is_updated_before_fast_tool_execution_can_continue() {
);
}
#[test]
fn signed_reasoning_is_persisted_before_the_tool_call() {
let messages = Arc::new(Mutex::new(Vec::new()));
let mut history_index = None;
let call = ToolCall {
id: "call-1".to_string(),
name: "read_files".to_string(),
arguments: serde_json::json!({"files": ["Cargo.toml"]}),
};
sync_assistant_turn(
&messages,
"I should inspect the manifest.",
Some("signed-reasoning"),
"",
std::slice::from_ref(&call),
&mut history_index,
);
let messages = messages.lock().unwrap();
let MessageContent::MultiPart(parts) = &messages[0].content else {
panic!("expected reasoning and tool call parts");
};
assert!(matches!(
parts.as_slice(),
[
ContentPart::Reasoning {
text,
signature: Some(signature),
},
ContentPart::ToolUse { tool_use_id, .. },
] if text == "I should inspect the manifest."
&& signature == "signed-reasoning"
&& tool_use_id == "call-1"
));
}
#[test]
fn rejected_tool_result_is_paired_with_the_assistant_call_in_history() {
let messages = Arc::new(Mutex::new(Vec::new()));
@@ -174,6 +215,8 @@ fn rejected_tool_result_is_paired_with_the_assistant_call_in_history() {
sync_assistant_turn(
&messages,
"",
None,
"",
std::slice::from_ref(&call),
&mut history_index,
);