Skip to content

Commit 4ee5685

Browse files
fix: clamp Bedrock tool schema integer bounds (#3308)
Co-authored-by: ForgeCode <noreply@forgecode.dev>
1 parent a1f650e commit 4ee5685

1 file changed

Lines changed: 87 additions & 1 deletion

File tree

crates/forge_repo/src/provider/bedrock.rs

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -901,14 +901,56 @@ impl FromDomain<schemars::Schema> for aws_sdk_bedrockruntime::types::ToolInputSc
901901
use aws_sdk_bedrockruntime::types::ToolInputSchema;
902902

903903
// Serialize Schema to JSON value first
904-
let json_value =
904+
let mut json_value =
905905
serde_json::to_value(&schema).with_context(|| "Failed to serialize Schema")?;
906+
sanitize_bedrock_tool_schema_numbers(&mut json_value);
906907

907908
// Convert JSON value to Document and wrap in ToolInputSchema
908909
Ok(ToolInputSchema::Json(json_value_to_document(json_value)))
909910
}
910911
}
911912

913+
fn sanitize_bedrock_tool_schema_numbers(value: &mut serde_json::Value) {
914+
match value {
915+
serde_json::Value::Object(map) => {
916+
for key in [
917+
"maxLength",
918+
"minLength",
919+
"maximum",
920+
"minimum",
921+
"maxItems",
922+
"minItems",
923+
"maxProperties",
924+
"minProperties",
925+
"multipleOf",
926+
] {
927+
if let Some(number_value) = map.get_mut(key) {
928+
clamp_bedrock_schema_number(number_value);
929+
}
930+
}
931+
932+
for value in map.values_mut() {
933+
sanitize_bedrock_tool_schema_numbers(value);
934+
}
935+
}
936+
serde_json::Value::Array(items) => {
937+
for value in items {
938+
sanitize_bedrock_tool_schema_numbers(value);
939+
}
940+
}
941+
_ => {}
942+
}
943+
}
944+
945+
fn clamp_bedrock_schema_number(value: &mut serde_json::Value) {
946+
let Some(number) = value.as_i64() else { return };
947+
948+
let clamped = number.clamp(i32::MIN as i64, i32::MAX as i64);
949+
if clamped != number {
950+
*value = serde_json::Value::Number(clamped.into());
951+
}
952+
}
953+
912954
/// Converts ToolCallArguments to AWS Smithy Document
913955
impl FromDomain<forge_domain::ToolCallArguments> for aws_smithy_types::Document {
914956
fn from_domain(args: forge_domain::ToolCallArguments) -> anyhow::Result<Self> {
@@ -1226,6 +1268,50 @@ mod tests {
12261268
assert_eq!(actual, expected);
12271269
}
12281270

1271+
#[test]
1272+
fn test_sanitize_bedrock_tool_schema_numbers_clamps_nested_integer_bounds() {
1273+
let mut fixture = serde_json::json!({
1274+
"type": "object",
1275+
"properties": {
1276+
"path": {
1277+
"type": "string",
1278+
"maxLength": 9_007_199_254_740_991_i64
1279+
},
1280+
"items": {
1281+
"type": "array",
1282+
"maxItems": 9_007_199_254_740_991_i64,
1283+
"items": {
1284+
"type": "number",
1285+
"minimum": -9_007_199_254_740_991_i64,
1286+
"maximum": 9_007_199_254_740_991_i64
1287+
}
1288+
}
1289+
}
1290+
});
1291+
1292+
sanitize_bedrock_tool_schema_numbers(&mut fixture);
1293+
let actual = fixture;
1294+
let expected = serde_json::json!({
1295+
"type": "object",
1296+
"properties": {
1297+
"path": {
1298+
"type": "string",
1299+
"maxLength": i32::MAX
1300+
},
1301+
"items": {
1302+
"type": "array",
1303+
"maxItems": i32::MAX,
1304+
"items": {
1305+
"type": "number",
1306+
"minimum": i32::MIN,
1307+
"maximum": i32::MAX
1308+
}
1309+
}
1310+
}
1311+
});
1312+
assert_eq!(actual, expected);
1313+
}
1314+
12291315
#[test]
12301316
fn test_json_value_to_document_null() {
12311317
let fixture = serde_json::Value::Null;

0 commit comments

Comments
 (0)