Filter unsupported reasoning from Bedrock request history

This commit is contained in:
Ryan Ward
2026-10-28 18:26:02 +00:00
parent 3cce041a6d
commit f399ada2a3
7 changed files with 462 additions and 16 deletions
+82 -4
View File
@@ -1,8 +1,9 @@
use async_trait::async_trait;
use aws_sdk_bedrockruntime::Client as AwsBedrockClient;
use galaxy_agent_core::{
AgentError, AgentErrorKind, AgentEventStream, AgentRuntime, RuntimeCapabilities,
RuntimeDescriptor, RuntimeKind, TurnControl, TurnRequest,
AgentError, AgentErrorKind, AgentEventStream, AgentRuntime, ContentPart, ConversationMessage,
MessageContent, MessageRole, RuntimeCapabilities, RuntimeDescriptor, RuntimeKind, TurnControl,
TurnRequest,
};
use rig_bedrock::client::Client as RigBedrockClient;
use rig_bedrock::completion::CompletionModel;
@@ -33,6 +34,7 @@ pub struct BedrockRuntime {
client: RigBedrockClient,
config: BedrockRigConfig,
resolved_model: String,
reasoning_model: String,
descriptor: RuntimeDescriptor,
}
@@ -53,6 +55,7 @@ impl BedrockRuntime {
Ok(Self {
client: RigBedrockClient::from(client),
config,
reasoning_model: resolved_model.clone(),
resolved_model,
descriptor,
})
@@ -62,6 +65,13 @@ impl BedrockRuntime {
&self.resolved_model
}
/// Use the discovered foundation model's capabilities when invocation targets an
/// opaque application inference profile rather than a foundation-model ID.
pub fn with_reasoning_model(mut self, model: impl Into<String>) -> Self {
self.reasoning_model = model.into();
self
}
pub fn completion_model(&self) -> CompletionModel {
let model = self.client.completion_model(&self.resolved_model);
if self.config.prompt_caching {
@@ -84,8 +94,11 @@ impl AgentRuntime for BedrockRuntime {
control: TurnControl,
) -> Result<AgentEventStream, AgentError> {
let max_output_tokens = request.max_output_tokens.or(self.config.max_output_tokens);
let mut completion_request =
build_bedrock_completion_request(request, self.config.max_output_tokens)?;
let mut completion_request = build_bedrock_completion_request_for_model(
request,
self.config.max_output_tokens,
&self.reasoning_model,
)?;
// Context markers and inference-profile expansion are Galaxy model
// configuration, not identifiers that Rig should send unchanged.
completion_request.model = Some(self.resolved_model.clone());
@@ -103,9 +116,74 @@ pub fn build_bedrock_completion_request(
request: TurnRequest,
configured_max_output_tokens: Option<u64>,
) -> Result<CompletionRequest, AgentError> {
let model = request.model.as_str().to_owned();
build_bedrock_completion_request_for_model(request, configured_max_output_tokens, &model)
}
fn build_bedrock_completion_request_for_model(
mut request: TurnRequest,
configured_max_output_tokens: Option<u64>,
model: &str,
) -> Result<CompletionRequest, AgentError> {
request.messages = prepare_bedrock_history(model, request.messages);
build_completion_request(request, configured_max_output_tokens, true, true, None)
}
/// Project durable history into Bedrock's input format without changing the saved
/// transcript. Reasoning output is not universally accepted as reasoning input:
/// DeepSeek, Nova and other providers return unsigned, display-only reasoning.
/// Claude tool continuations need their complete signed blocks unchanged.
pub fn prepare_bedrock_history(
model_id: &str,
messages: Vec<ConversationMessage>,
) -> Vec<ConversationMessage> {
let supports_signed_reasoning = accepts_signed_reasoning(model_id);
messages
.into_iter()
.filter_map(|mut message| {
if let MessageContent::MultiPart(parts) = &mut message.content {
parts.retain(|part| match part {
ContentPart::Reasoning { signature, .. } => {
message.role == MessageRole::Assistant
&& supports_signed_reasoning
&& signature
.as_ref()
.is_some_and(|signature| !signature.trim().is_empty())
}
ContentPart::Text(_)
| ContentPart::Image { .. }
| ContentPart::ToolUse { .. }
| ContentPart::ToolResult { .. } => true,
});
if parts.is_empty() {
return None;
}
}
Some(message)
})
.collect()
}
fn accepts_signed_reasoning(model_id: &str) -> bool {
let model_id = model_id.trim();
// An opaque application profile may point to Claude. Preserve its signed
// blocks when no foundation-model hint is available; unsigned output is
// still omitted. Production discovery supplies that hint to the runtime.
if model_id.starts_with("arn:") && model_id.contains(":application-inference-profile/") {
return true;
}
let model = strip_context_marker(model_id.rsplit('/').next().unwrap_or(model_id));
let model = INFERENCE_PROFILE_PREFIXES
.iter()
.find_map(|prefix| model.strip_prefix(prefix))
.unwrap_or(model);
model.starts_with("anthropic.claude-")
&& !model.starts_with("anthropic.claude-instant-")
&& !model.starts_with("anthropic.claude-v1")
&& !model.starts_with("anthropic.claude-v2")
&& (!model.starts_with("anthropic.claude-3-") || model.starts_with("anthropic.claude-3-7-"))
}
pub fn resolve_bedrock_model_id(
configured_model: &str,
region: &str,
+149 -1
View File
@@ -177,7 +177,7 @@ async fn cancellation_before_bedrock_stream_start_never_contacts_aws() {
#[test]
fn bedrock_request_preserves_system_image_reasoning_tool_and_token_semantics() {
let mut request = TurnRequest::new(
"anthropic.claude-test",
"anthropic.claude-sonnet-4-6",
vec![
ConversationMessage {
role: MessageRole::User,
@@ -278,3 +278,151 @@ fn bedrock_request_preserves_system_image_reasoning_tool_and_token_semantics() {
Some(ToolResultContent::Text(text)) if text.text == "[ERROR] permission denied"
));
}
#[test]
fn bedrock_follow_ups_omit_unsigned_reasoning_without_losing_tool_calls_or_results() {
for model in [
"deepseek.r1-v1:0",
"us.amazon.nova-2-lite-v1:0",
"openai.gpt-oss-120b-1:0",
"qwen.qwen3-32b-v1:0",
"anthropic.claude-sonnet-4-6",
] {
let request = TurnRequest::new(
model,
vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Inspect the project".to_owned()),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::MultiPart(vec![
ContentPart::Reasoning {
text: "display-only reasoning".to_owned(),
signature: None,
},
ContentPart::Text("Reading the manifest".to_owned()),
ContentPart::ToolUse {
tool_use_id: "manifest-call".to_owned(),
name: "read_files".to_owned(),
input: serde_json::json!({"files": ["Cargo.toml"]}),
},
]),
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "manifest-call".to_owned(),
content: "manifest contents".to_owned(),
is_error: false,
},
},
],
);
let converted = build_bedrock_completion_request(request, None).unwrap();
let messages = converted.chat_history.iter().collect::<Vec<_>>();
let [
_,
Message::Assistant { content, .. },
Message::User { content: result },
] = messages.as_slice()
else {
panic!("expected the tool exchange for {model}");
};
assert!(matches!(content.iter().collect::<Vec<_>>().as_slice(),
[AssistantContent::Text(text), AssistantContent::ToolCall(call)]
if text.text == "Reading the manifest" && call.id == "manifest-call"
));
assert!(
matches!(result.iter().next(), Some(UserContent::ToolResult(result))
if result.id == "manifest-call")
);
}
}
#[test]
fn bedrock_model_switch_drops_foreign_signed_reasoning_and_empty_messages() {
for model in [
"amazon.nova-pro-v1:0",
"us.deepseek.r1-v1:0",
"openai.gpt-oss-120b-1:0",
"anthropic.claude-3-5-sonnet-20241022-v2:0",
"anthropic.claude-3-haiku-20240307-v1:0",
"anthropic.claude-instant-v1",
"anthropic.claude-v1",
"anthropic.claude-v2:1",
] {
let messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Question".to_owned()),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::MultiPart(vec![ContentPart::Reasoning {
text: "Earlier thinking".to_owned(),
signature: Some("original-signature".to_owned()),
}]),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::MultiPart(vec![
ContentPart::Reasoning {
text: "Earlier thinking".to_owned(),
signature: Some("original-signature".to_owned()),
},
ContentPart::Text("Visible answer".to_owned()),
]),
},
];
// The runtime's selected model takes precedence over a historical request model.
let request = TurnRequest::new("anthropic.claude-sonnet-4-6", messages);
let converted = build_bedrock_completion_request_for_model(request, None, model).unwrap();
assert_eq!(converted.chat_history.len(), 2);
assert!(
matches!(converted.chat_history.iter().nth(1), Some(Message::Assistant { content, .. })
if matches!(content.iter().collect::<Vec<_>>().as_slice(), [AssistantContent::Text(text)] if text.text == "Visible answer"))
);
}
}
#[test]
fn bedrock_preserves_signed_claude_blocks_exactly_across_identifier_forms() {
let signed = ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::MultiPart(vec![
ContentPart::Reasoning {
text: "Unmodified thinking".to_owned(),
signature: Some("signed-text".to_owned()),
},
ContentPart::Reasoning {
text: String::new(),
signature: Some("signature-only".to_owned()),
},
ContentPart::Text("Answer".to_owned()),
]),
};
for model in [
"anthropic.claude-3-7-sonnet-20250219-v1:0",
"global.anthropic.claude-sonnet-4-6[1m]",
"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-sonnet-4-6",
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-sonnet-4-6",
"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque-profile",
] {
assert_eq!(
prepare_bedrock_history(model, vec![signed.clone()]),
vec![signed.clone()]
);
}
for signature in [None, Some(String::new()), Some(" ".to_owned())] {
let message = ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::MultiPart(vec![ContentPart::Reasoning {
text: "Display only".to_owned(),
signature,
}]),
};
assert!(prepare_bedrock_history("anthropic.claude-sonnet-4-6", vec![message]).is_empty());
}
}