From af942749234487fe028484d8768c83332722a9bf Mon Sep 17 00:00:00 2001 From: Lucas Vieira Date: Tue, 29 Sep 2026 18:42:23 -0300 Subject: [PATCH 1/3] fix(bedrock-agent): make flow and prompt responses match the model The flow and prompt handlers returned shapes the Smithy model doesn't have, and most flow operations could not be reached at all: - Every flow handler read the flow from a flowId body member, but the router supplies the @httpLabel as flowIdentifier (aliases as aliasIdentifier), so GetFlow, UpdateFlow, DeleteFlow, PrepareFlow and every version and alias operation failed with a missing-field error over the wire. They now read flowIdentifier/aliasIdentifier and accept the bare ID or the (URL-encoded) ARN. - GetFlow/UpdateFlow wrapped the flow in a flow member keyed flowId; GetPrompt/UpdatePrompt wrapped it in prompt keyed promptId. All now return the top-level output members (id, arn from the stored ARN, name, status, version, timestamps, plus customerEncryptionKeyArn, defaultVariant and definition, now persisted). - Flow versions and aliases return GetFlowVersion/FlowAlias shapes (the alias ARN is flow//alias/); deletes return their modeled id members; PrepareFlow returns {id, status} with the FlowStatus wire value Prepared (not PREPARED) and 202, creates return 201, and UpdateFlow drops the flow back to NotPrepared. - GetPrompt and DeletePrompt honor the promptVersion query member (and a version pinned in the ARN). - POST /flows/validate-definition was routed as PrepareFlow on a flow named validate-definition. It now reaches ValidateFlowDefinition, which reads the definition and reports structural validations (missing starting/ending nodes, unknown connection endpoints and inputs/outputs, duplicate connections, unfulfilled or over-fed node inputs, cycles, unreachable nodes) in the modeled validations list instead of the unmodeled isValid/validationDetails. --- .../src/service/flows.rs | 489 +++++++++++++----- .../src/service/mod.rs | 456 +++++++++++++++- .../src/service/prompts.rs | 186 ++++--- crates/fakecloud-bedrock-agent/src/state.rs | 36 +- 4 files changed, 950 insertions(+), 217 deletions(-) diff --git a/crates/fakecloud-bedrock-agent/src/service/flows.rs b/crates/fakecloud-bedrock-agent/src/service/flows.rs index 2b9f36b50..d7c8de4ac 100644 --- a/crates/fakecloud-bedrock-agent/src/service/flows.rs +++ b/crates/fakecloud-bedrock-agent/src/service/flows.rs @@ -2,6 +2,236 @@ use super::*; +/// The flow ID a `flowIdentifier` names: the identifier is either the bare +/// flow ID or the (URL-encoded) flow ARN (`arn:...:flow/`). +fn flow_id_of(identifier: &str) -> String { + let identifier = decode_label(identifier); + identifier + .rsplit_once(":flow/") + .map_or(identifier.as_str(), |(_, id)| id) + .to_string() +} + +/// The alias ID an `aliasIdentifier` names: the bare alias ID or the alias +/// ARN (`arn:...:flow//alias/`). +fn alias_id_of(identifier: &str) -> String { + let identifier = decode_label(identifier); + identifier + .rsplit_once("/alias/") + .map_or(identifier.as_str(), |(_, id)| id) + .to_string() +} + +fn flow_not_found(id: &str) -> AwsServiceError { + not_found(format!("Flow {id} not found")) +} + +fn flow_validation(severity: &str, kind: &str, message: String, details: Value) -> Value { + json!({ + "severity": severity, + "type": kind, + "message": message, + "details": details, + }) +} + +/// Structural validation of a flow definition, the checks `ValidateFlowDefinition` +/// reports: the flow needs a starting (`Input`) and an ending (`Output`) node; +/// every connection must join existing nodes (and, for data connections, an +/// existing output of the source to an existing input of the target); no two +/// connections may join the same pair of nodes; no node input may be fed by +/// more than one data connection or by none; connections may not form a cycle; +/// and every node should be reachable from a starting node. +fn validate_definition(definition: &Value) -> Vec { + use std::collections::{BTreeMap, BTreeSet}; + + let empty = Vec::new(); + let nodes = definition["nodes"].as_array().unwrap_or(&empty); + let connections = definition["connections"].as_array().unwrap_or(&empty); + let str_of = |v: &Value, k: &str| v[k].as_str().unwrap_or_default().to_string(); + let names_of = |v: &Value, k: &str| -> BTreeSet { + v[k].as_array() + .map(|a| a.iter().map(|x| str_of(x, "name")).collect()) + .unwrap_or_default() + }; + + // node name -> (type, input names, output names) + let by_name: BTreeMap, BTreeSet)> = nodes + .iter() + .map(|n| { + ( + str_of(n, "name"), + ( + str_of(n, "type"), + names_of(n, "inputs"), + names_of(n, "outputs"), + ), + ) + }) + .collect(); + + let mut out = Vec::new(); + if !by_name.values().any(|(t, _, _)| t == "Input") { + out.push(flow_validation( + "Error", + "MissingStartingNodes", + "The flow has no starting node. Add an Input node.".to_string(), + json!({ "missingStartingNodes": {} }), + )); + } + if !by_name.values().any(|(t, _, _)| t == "Output") { + out.push(flow_validation( + "Error", + "MissingEndingNodes", + "The flow has no ending node. Add an Output node.".to_string(), + json!({ "missingEndingNodes": {} }), + )); + } + + let mut seen_pairs = BTreeSet::new(); + let mut input_feeds: BTreeMap<(String, String), usize> = BTreeMap::new(); + let mut edges: BTreeMap> = BTreeMap::new(); + for c in connections { + let name = str_of(c, "name"); + let source = str_of(c, "source"); + let target = str_of(c, "target"); + let source_node = by_name.get(&source); + let target_node = by_name.get(&target); + if source_node.is_none() { + out.push(flow_validation( + "Error", + "UnknownConnectionSource", + format!("Connection {name} has an unknown source node {source}."), + json!({ "unknownConnectionSource": { "connection": name } }), + )); + } + if target_node.is_none() { + out.push(flow_validation( + "Error", + "UnknownConnectionTarget", + format!("Connection {name} has an unknown target node {target}."), + json!({ "unknownConnectionTarget": { "connection": name } }), + )); + } + if !seen_pairs.insert((source.clone(), target.clone())) { + out.push(flow_validation( + "Error", + "DuplicateConnections", + format!("Nodes {source} and {target} are joined by more than one connection."), + json!({ "duplicateConnections": { "source": source, "target": target } }), + )); + } + if let Some(data) = c["configuration"].get("data") { + let source_output = str_of(data, "sourceOutput"); + let target_input = str_of(data, "targetInput"); + if source_node.is_some_and(|(_, _, outputs)| !outputs.contains(&source_output)) { + out.push(flow_validation( + "Error", + "UnknownConnectionSourceOutput", + format!( + "Connection {name} references unknown output {source_output} of node {source}." + ), + json!({ "unknownConnectionSourceOutput": { "connection": name } }), + )); + } + if target_node.is_some_and(|(_, inputs, _)| !inputs.contains(&target_input)) { + out.push(flow_validation( + "Error", + "UnknownConnectionTargetInput", + format!( + "Connection {name} references unknown input {target_input} of node {target}." + ), + json!({ "unknownConnectionTargetInput": { "connection": name } }), + )); + } else if target_node.is_some() { + *input_feeds + .entry((target.clone(), target_input)) + .or_default() += 1; + } + } + if source_node.is_some() && target_node.is_some() { + edges.entry(source).or_default().push((target, name)); + } + } + + for (node, (_, inputs, _)) in &by_name { + for input in inputs { + match input_feeds.get(&(node.clone(), input.clone())).copied() { + None | Some(0) => out.push(flow_validation( + "Error", + "UnfulfilledNodeInput", + format!("Input {input} of node {node} is not connected."), + json!({ "unfulfilledNodeInput": { "node": node, "input": input } }), + )), + Some(1) => {} + Some(_) => out.push(flow_validation( + "Error", + "MultipleNodeInputConnections", + format!("Input {input} of node {node} has more than one connection."), + json!({ "multipleNodeInputConnections": { "node": node, "input": input } }), + )), + } + } + } + + // A connection that leads back to a node still on the DFS stack closes a + // cycle. + fn visit<'a>( + node: &'a str, + edges: &'a BTreeMap>, + state: &mut BTreeMap<&'a str, bool>, + cyclic: &mut Vec, + ) { + state.insert(node, true); + for (target, connection) in edges.get(node).into_iter().flatten() { + match state.get(target.as_str()) { + Some(true) => cyclic.push(connection.clone()), + Some(false) => {} + None => visit(target, edges, state, cyclic), + } + } + state.insert(node, false); + } + let mut dfs_state = BTreeMap::new(); + let mut cyclic = Vec::new(); + for node in by_name.keys() { + if !dfs_state.contains_key(node.as_str()) { + visit(node, &edges, &mut dfs_state, &mut cyclic); + } + } + for connection in cyclic { + out.push(flow_validation( + "Error", + "CyclicConnection", + format!("Connection {connection} creates a cycle."), + json!({ "cyclicConnection": { "connection": connection } }), + )); + } + + let mut reached: BTreeSet<&str> = by_name + .iter() + .filter(|(_, (t, _, _))| t == "Input") + .map(|(n, _)| n.as_str()) + .collect(); + let mut frontier: Vec<&str> = reached.iter().copied().collect(); + while let Some(node) = frontier.pop() { + for (target, _) in edges.get(node).into_iter().flatten() { + if reached.insert(target.as_str()) { + frontier.push(target.as_str()); + } + } + } + for node in by_name.keys().filter(|n| !reached.contains(n.as_str())) { + out.push(flow_validation( + "Warning", + "UnreachableNode", + format!("Node {node} cannot be reached from a starting node."), + json!({ "unreachableNode": { "node": node } }), + )); + } + out +} + impl BedrockAgentService { pub(super) fn create_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); @@ -14,57 +244,35 @@ impl BedrockAgentService { let role_arn = opt_str(&body, "executionRoleArn").unwrap_or_else(|| { crate::arns::default_flow_execution_role_arn(&req.region, &req.account_id, &id) }); - let arn = flow_arn(&req.region, &req.account_id, &id); - let definition = opt_json(&body, "definition"); let flow = Flow { flow_id: id.clone(), - name: name.clone(), + name, description: opt_str(&body, "description"), - execution_role_arn: Some(role_arn.clone()), + execution_role_arn: Some(role_arn), status: "NotPrepared".to_string(), created_at: now_dt, updated_at: now_dt, version: "DRAFT".to_string(), - definition: definition.clone(), - arn: arn.clone(), + definition: opt_json(&body, "definition"), + arn: flow_arn(&req.region, &req.account_id, &id), + customer_encryption_key_arn: opt_str(&body, "customerEncryptionKeyArn"), }; + let out = flow_json(&flow); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); - state.flows.insert(id.clone(), flow); - let mut out = json!({ - "name": name, - "executionRoleArn": role_arn, - "id": id, - "arn": arn, - "status": "NotPrepared", - "createdAt": now_dt.to_rfc3339(), - "updatedAt": now_dt.to_rfc3339(), - "version": "DRAFT", - }); - if let Some(d) = opt_str(&body, "description") { - out["description"] = json!(d); - } - if let Some(k) = opt_str(&body, "customerEncryptionKeyArn") { - out["customerEncryptionKeyArn"] = json!(k); - } - if let Some(def) = definition { - out["definition"] = def; - } - Ok(AwsResponse::ok_json(out)) + state.flows.insert(id, flow); + Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } pub(super) fn get_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "flowId")?; + let id = flow_id_of(&req_str(&body, "flowIdentifier")?); let accts = self.state.read(); - let state = accts + let f = accts .get(&req.account_id) - .ok_or_else(|| not_found(format!("Flow {id} not found")))?; - let f = state - .flows - .get(&id) - .ok_or_else(|| not_found(format!("Flow {id} not found")))?; - Ok(AwsResponse::ok_json(json!({ "flow": flow_json(f) }))) + .and_then(|s| s.flows.get(&id)) + .ok_or_else(|| flow_not_found(&id))?; + Ok(AwsResponse::ok_json(flow_json(f))) } pub(super) fn list_flows(&self, req: &AwsRequest) -> Result { @@ -78,13 +286,13 @@ impl BedrockAgentService { pub(super) fn update_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "flowId")?; + let id = flow_id_of(&req_str(&body, "flowIdentifier")?); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let f = state .flows .get_mut(&id) - .ok_or_else(|| not_found(format!("Flow {id} not found")))?; + .ok_or_else(|| flow_not_found(&id))?; f.updated_at = now(); if let Some(n) = opt_str(&body, "name") { f.name = n; @@ -95,41 +303,46 @@ impl BedrockAgentService { if let Some(r) = opt_str(&body, "executionRoleArn") { f.execution_role_arn = Some(r); } + if let Some(k) = opt_str(&body, "customerEncryptionKeyArn") { + f.customer_encryption_key_arn = Some(k); + } if body.get("definition").is_some() { f.definition = opt_json(&body, "definition"); } - Ok(AwsResponse::ok_json(json!({ "flow": flow_json(f) }))) + // An updated draft has to be prepared again before it can run. + f.status = "NotPrepared".to_string(); + Ok(AwsResponse::ok_json(flow_json(f))) } pub(super) fn delete_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "flowId")?; + let id = flow_id_of(&req_str(&body, "flowIdentifier")?); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); - state - .flows - .remove(&id) - .ok_or_else(|| not_found(format!("Flow {id} not found")))?; + state.flows.remove(&id).ok_or_else(|| flow_not_found(&id))?; state.flow_versions.remove(&id); state.flow_aliases.retain(|_, a| a.flow_id != id); - Ok(AwsResponse::ok_json(json!({}))) + Ok(AwsResponse::ok_json(json!({ "id": id }))) } pub(super) fn prepare_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "flowId")?; + let id = flow_id_of(&req_str(&body, "flowIdentifier")?); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let f = state .flows .get_mut(&id) - .ok_or_else(|| not_found(format!("Flow {id} not found")))?; - f.status = "PREPARED".to_string(); + .ok_or_else(|| flow_not_found(&id))?; + f.status = "Prepared".to_string(); f.updated_at = now(); - Ok(AwsResponse::ok_json(json!({ - "flowId": id, - "status": "PREPARED", - }))) + Ok(AwsResponse::json_value( + StatusCode::ACCEPTED, + json!({ + "id": id, + "status": f.status, + }), + )) } pub(super) fn create_flow_version( @@ -137,32 +350,32 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); let now_dt = now(); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let flow = state .flows .get(&flow_id) - .ok_or_else(|| not_found(format!("Flow {flow_id} not found")))?; + .ok_or_else(|| flow_not_found(&flow_id))? + .clone(); let versions = state.flow_versions.entry(flow_id.clone()).or_default(); let version_num = (versions.len() as u64 + 1).to_string(); let fv = FlowVersion { - flow_version: version_num.clone(), - flow_id: flow_id.clone(), + flow_version: version_num, + flow_id, description: opt_str(&body, "description"), created_at: now_dt, updated_at: now_dt, definition: flow.definition.clone(), + name: Some(flow.name.clone()), + execution_role_arn: flow.execution_role_arn.clone(), + customer_encryption_key_arn: flow.customer_encryption_key_arn.clone(), + status: Some(flow.status.clone()), }; + let out = flow_version_json(&flow, &fv); versions.push(fv); - Ok(AwsResponse::ok_json(json!({ - "flowVersion": { - "flowVersion": version_num, - "flowId": flow_id, - "createdAt": now_dt.to_rfc3339(), - } - }))) + Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } pub(super) fn get_flow_version( @@ -170,20 +383,18 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); let version = req_str(&body, "flowVersion")?; let accts = self.state.read(); - let state = accts - .get(&req.account_id) - .ok_or_else(|| not_found(format!("Flow version {version} not found")))?; + let version_not_found = || not_found(format!("Flow version {version} not found")); + let state = accts.get(&req.account_id).ok_or_else(version_not_found)?; + let flow = state.flows.get(&flow_id).ok_or_else(version_not_found)?; let v = state .flow_versions .get(&flow_id) .and_then(|vec| vec.iter().find(|v| v.flow_version == version)) - .ok_or_else(|| not_found(format!("Flow version {version} not found")))?; - Ok(AwsResponse::ok_json( - json!({ "flowVersion": flow_version_json(v) }), - )) + .ok_or_else(version_not_found)?; + Ok(AwsResponse::ok_json(flow_version_json(flow, v))) } pub(super) fn list_flow_versions( @@ -191,12 +402,31 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); let accts = self.state.read(); - let list: Vec = accts + let state = accts .get(&req.account_id) - .and_then(|s| s.flow_versions.get(&flow_id)) - .map(|vec| vec.iter().map(flow_version_json).collect()) + .ok_or_else(|| flow_not_found(&flow_id))?; + let flow = state + .flows + .get(&flow_id) + .ok_or_else(|| flow_not_found(&flow_id))?; + let list: Vec = state + .flow_versions + .get(&flow_id) + .map(|vec| { + vec.iter() + .map(|v| { + json!({ + "id": flow.flow_id, + "arn": flow.arn, + "status": v.status.as_deref().unwrap_or(&flow.status), + "createdAt": v.created_at.to_rfc3339(), + "version": v.flow_version, + }) + }) + .collect() + }) .unwrap_or_default(); Ok(AwsResponse::ok_json( json!({ "flowVersionSummaries": list }), @@ -208,7 +438,7 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); let version = req_str(&body, "flowVersion")?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); @@ -221,7 +451,10 @@ impl BedrockAgentService { .position(|v| v.flow_version == version) .ok_or_else(|| not_found(format!("Flow version {version} not found")))?; vec.remove(pos); - Ok(AwsResponse::ok_json(json!({}))) + Ok(AwsResponse::ok_json(json!({ + "id": flow_id, + "version": version, + }))) } pub(super) fn validate_flow_definition( @@ -229,10 +462,12 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let _flow_id = req_str(&body, "flowId")?; + let definition = body + .get("definition") + .filter(|d| d.is_object()) + .ok_or_else(|| missing("definition"))?; Ok(AwsResponse::ok_json(json!({ - "isValid": true, - "validationDetails": [], + "validations": validate_definition(definition), }))) } @@ -241,51 +476,47 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); let name = req_str(&body, "name")?; let alias_id = short_id(); let now_dt = now(); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); - if !state.flows.contains_key(&flow_id) { - return Err(not_found(format!("Flow {flow_id} not found"))); - } + let flow_arn = state + .flows + .get(&flow_id) + .ok_or_else(|| flow_not_found(&flow_id))? + .arn + .clone(); let alias = FlowAlias { alias_id: alias_id.clone(), - alias_name: name.clone(), - flow_id: flow_id.clone(), + alias_name: name, + flow_id, routing_configuration: opt_array(&body, "routingConfiguration"), description: opt_str(&body, "description"), created_at: now_dt, updated_at: now_dt, + concurrency_configuration: opt_json(&body, "concurrencyConfiguration"), }; - state.flow_aliases.insert(alias_id.clone(), alias); - Ok(AwsResponse::ok_json(json!({ - "flowAlias": { - "aliasId": alias_id, - "aliasName": name, - "flowId": flow_id, - "createdAt": now_dt.to_rfc3339(), - } - }))) + let out = flow_alias_json(&flow_arn, &alias); + state.flow_aliases.insert(alias_id, alias); + Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } pub(super) fn get_flow_alias(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; - let alias_id = req_str(&body, "aliasId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let alias_id = alias_id_of(&req_str(&body, "aliasIdentifier")?); let accts = self.state.read(); - let state = accts - .get(&req.account_id) - .ok_or_else(|| not_found(format!("Flow alias {alias_id} not found")))?; + let alias_not_found = || not_found(format!("Flow alias {alias_id} not found")); + let state = accts.get(&req.account_id).ok_or_else(alias_not_found)?; + let flow = state.flows.get(&flow_id).ok_or_else(alias_not_found)?; let a = state .flow_aliases .get(&alias_id) .filter(|a| a.flow_id == flow_id) - .ok_or_else(|| not_found(format!("Flow alias {alias_id} not found")))?; - Ok(AwsResponse::ok_json( - json!({ "flowAlias": flow_alias_json(a) }), - )) + .ok_or_else(alias_not_found)?; + Ok(AwsResponse::ok_json(flow_alias_json(&flow.arn, a))) } pub(super) fn list_flow_aliases( @@ -293,18 +524,21 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); let accts = self.state.read(); - let list: Vec = accts + let state = accts .get(&req.account_id) - .map(|s| { - s.flow_aliases - .values() - .filter(|a| a.flow_id == flow_id) - .map(flow_alias_json) - .collect() - }) - .unwrap_or_default(); + .ok_or_else(|| flow_not_found(&flow_id))?; + let flow = state + .flows + .get(&flow_id) + .ok_or_else(|| flow_not_found(&flow_id))?; + let list: Vec = state + .flow_aliases + .values() + .filter(|a| a.flow_id == flow_id) + .map(|a| flow_alias_json(&flow.arn, a)) + .collect(); Ok(AwsResponse::ok_json(json!({ "flowAliasSummaries": list }))) } @@ -313,15 +547,22 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; - let alias_id = req_str(&body, "aliasId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let alias_id = alias_id_of(&req_str(&body, "aliasIdentifier")?); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); + let alias_not_found = || not_found(format!("Flow alias {alias_id} not found")); + let flow_arn = state + .flows + .get(&flow_id) + .ok_or_else(alias_not_found)? + .arn + .clone(); let a = state .flow_aliases .get_mut(&alias_id) .filter(|a| a.flow_id == flow_id) - .ok_or_else(|| not_found(format!("Flow alias {alias_id} not found")))?; + .ok_or_else(alias_not_found)?; a.updated_at = now(); if let Some(n) = opt_str(&body, "name") { a.alias_name = n; @@ -332,9 +573,10 @@ impl BedrockAgentService { if body.get("routingConfiguration").is_some() { a.routing_configuration = opt_array(&body, "routingConfiguration"); } - Ok(AwsResponse::ok_json( - json!({ "flowAlias": flow_alias_json(a) }), - )) + if body.get("concurrencyConfiguration").is_some() { + a.concurrency_configuration = opt_json(&body, "concurrencyConfiguration"); + } + Ok(AwsResponse::ok_json(flow_alias_json(&flow_arn, a))) } pub(super) fn delete_flow_alias( @@ -342,8 +584,8 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = req_str(&body, "flowId")?; - let alias_id = req_str(&body, "aliasId")?; + let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let alias_id = alias_id_of(&req_str(&body, "aliasIdentifier")?); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); match state.flow_aliases.get(&alias_id) { @@ -352,6 +594,9 @@ impl BedrockAgentService { } _ => return Err(not_found(format!("Flow alias {alias_id} not found"))), } - Ok(AwsResponse::ok_json(json!({}))) + Ok(AwsResponse::ok_json(json!({ + "flowId": flow_id, + "id": alias_id, + }))) } } diff --git a/crates/fakecloud-bedrock-agent/src/service/mod.rs b/crates/fakecloud-bedrock-agent/src/service/mod.rs index e5477528b..9e890d18f 100644 --- a/crates/fakecloud-bedrock-agent/src/service/mod.rs +++ b/crates/fakecloud-bedrock-agent/src/service/mod.rs @@ -478,6 +478,15 @@ impl BedrockAgentService { return Some(("ListFlows", params)); } } + // Checked before `/flows/{flowIdentifier}`, whose POST (PrepareFlow) + // would otherwise swallow it as a flow named `validate-definition`. + if segs.len() == 2 + && segs[0] == "flows" + && segs[1] == "validate-definition" + && *m == Method::POST + { + return Some(("ValidateFlowDefinition", params)); + } if segs.len() == 2 && segs[0] == "flows" { params.push(("flowIdentifier".to_string(), segs[1].clone())); if *m == Method::GET { @@ -534,13 +543,6 @@ impl BedrockAgentService { return Some(("DeleteFlowVersion", params)); } } - if segs.len() == 2 - && segs[0] == "flows" - && segs[1] == "validate-definition" - && *m == Method::POST - { - return Some(("ValidateFlowDefinition", params)); - } // Prompts if segs.len() == 1 && segs[0] == "prompts" { @@ -1021,6 +1023,14 @@ fn now() -> DateTime { Utc::now() } +/// Percent-decode an `@httpLabel` path segment. An ARN identifier arrives +/// URL-encoded (`arn%3Aaws%3A...%2F...`); a bare ID is unchanged. +fn decode_label(segment: &str) -> String { + percent_encoding::percent_decode_str(segment) + .decode_utf8_lossy() + .into_owned() +} + fn short_id() -> String { // Smithy ResourceIdentifier shape is @pattern("^[0-9a-zA-Z]{10}$"): the // first 10 hex chars of a v4 UUID. @@ -1247,8 +1257,9 @@ fn agent_summary_json(a: &Agent) -> Value { } /// `FlowSummary` shape: requires `arn`, `id`, `name`, `status`, `createdAt`, -/// `updatedAt`, and `version`. The full `flow_json` exposes `flowId`, -/// `executionRoleArn`, and `definition`, none of which appear on the summary. +/// `updatedAt`, and `version`. The full `flow_json` adds `executionRoleArn`, +/// `customerEncryptionKeyArn`, and `definition`, none of which appear on the +/// summary. fn flow_summary_json(f: &Flow) -> Value { let mut o = json!({ "arn": f.arn, @@ -1298,10 +1309,14 @@ fn prompt_summary_json(p: &Prompt) -> Value { o } +/// `GetFlowResponse` / `CreateFlowResponse` / `UpdateFlowResponse` members: +/// `name`, `description`, `executionRoleArn`, `customerEncryptionKeyArn`, +/// `id`, `arn`, `status`, `createdAt`, `updatedAt`, `version`, `definition`. fn flow_json(f: &Flow) -> Value { let mut o = json!({ - "flowId": f.flow_id, "name": f.name, + "id": f.flow_id, + "arn": f.arn, "status": f.status, "createdAt": f.created_at.to_rfc3339(), "updatedAt": f.updated_at.to_rfc3339(), @@ -1313,33 +1328,60 @@ fn flow_json(f: &Flow) -> Value { if let Some(ref r) = f.execution_role_arn { o["executionRoleArn"] = json!(r); } + if let Some(ref k) = f.customer_encryption_key_arn { + o["customerEncryptionKeyArn"] = json!(k); + } if let Some(ref def) = f.definition { o["definition"] = def.clone(); } o } -fn flow_version_json(v: &FlowVersion) -> Value { +/// `GetFlowVersionResponse` / `CreateFlowVersionResponse` members. `id` and +/// `arn` are the flow's (a flow version has no ARN of its own); the +/// flow-level fields come from the version snapshot, falling back to the flow +/// for versions persisted before they were captured. +fn flow_version_json(f: &Flow, v: &FlowVersion) -> Value { let mut o = json!({ - "flowVersion": v.flow_version, - "flowId": v.flow_id, + "name": v.name.as_deref().unwrap_or(&f.name), + "id": f.flow_id, + "arn": f.arn, + "status": v.status.as_deref().unwrap_or(&f.status), "createdAt": v.created_at.to_rfc3339(), - "updatedAt": v.updated_at.to_rfc3339(), + "version": v.flow_version, }); if let Some(ref d) = v.description { o["description"] = json!(d); } + if let Some(r) = v + .execution_role_arn + .as_ref() + .or(f.execution_role_arn.as_ref()) + { + o["executionRoleArn"] = json!(r); + } + if let Some(k) = v + .customer_encryption_key_arn + .as_ref() + .or(f.customer_encryption_key_arn.as_ref()) + { + o["customerEncryptionKeyArn"] = json!(k); + } if let Some(ref def) = v.definition { o["definition"] = def.clone(); } o } -fn flow_alias_json(a: &FlowAlias) -> Value { +/// `GetFlowAliasResponse` / `FlowAliasSummary` members: `name`, +/// `description`, `routingConfiguration`, `concurrencyConfiguration`, +/// `flowId`, `id`, `arn`, `createdAt`, `updatedAt`. +fn flow_alias_json(flow_arn: &str, a: &FlowAlias) -> Value { let mut o = json!({ - "aliasId": a.alias_id, - "aliasName": a.alias_name, + "name": a.alias_name, "flowId": a.flow_id, + "id": a.alias_id, + "arn": format!("{flow_arn}/alias/{}", a.alias_id), "routingConfiguration": a.routing_configuration, "createdAt": a.created_at.to_rfc3339(), "updatedAt": a.updated_at.to_rfc3339(), @@ -1347,13 +1389,21 @@ fn flow_alias_json(a: &FlowAlias) -> Value { if let Some(ref d) = a.description { o["description"] = json!(d); } + if let Some(ref c) = a.concurrency_configuration { + o["concurrencyConfiguration"] = c.clone(); + } o } +/// `GetPromptResponse` / `CreatePromptResponse` / `UpdatePromptResponse` +/// members for the working draft: `name`, `description`, +/// `customerEncryptionKeyArn`, `defaultVariant`, `variants`, `id`, `arn`, +/// `version`, `createdAt`, `updatedAt`. fn prompt_json(p: &Prompt) -> Value { let mut o = json!({ - "promptId": p.prompt_id, "name": p.name, + "id": p.prompt_id, + "arn": p.arn, "variants": p.variants, "version": p.version, "createdAt": p.created_at.to_rfc3339(), @@ -1362,6 +1412,36 @@ fn prompt_json(p: &Prompt) -> Value { if let Some(ref d) = p.description { o["description"] = json!(d); } + if let Some(ref k) = p.customer_encryption_key_arn { + o["customerEncryptionKeyArn"] = json!(k); + } + if let Some(ref dv) = p.default_variant { + o["defaultVariant"] = json!(dv); + } + o +} + +/// The same members for a numbered prompt version: its ARN is the prompt ARN +/// with a `:` suffix. +fn prompt_version_json(p: &Prompt, v: &PromptVersion) -> Value { + let mut o = json!({ + "name": v.name.as_deref().unwrap_or(&p.name), + "id": p.prompt_id, + "arn": format!("{}:{}", p.arn, v.prompt_version), + "variants": v.variants, + "version": v.prompt_version, + "createdAt": v.created_at.to_rfc3339(), + "updatedAt": v.updated_at.to_rfc3339(), + }); + if let Some(ref d) = v.description { + o["description"] = json!(d); + } + if let Some(ref k) = p.customer_encryption_key_arn { + o["customerEncryptionKeyArn"] = json!(k); + } + if let Some(dv) = v.default_variant.as_ref().or(p.default_variant.as_ref()) { + o["defaultVariant"] = json!(dv); + } o } @@ -1445,6 +1525,346 @@ mod tests { serde_json::from_slice(resp.body.expect_bytes()).unwrap() } + /// Send `method path[?query]` with a JSON body through the service's own + /// REST routing, the way an SDK request arrives. + async fn call( + svc: &BedrockAgentService, + method: Method, + path: &str, + query: &[(&str, &str)], + payload: Value, + ) -> (StatusCode, Value) { + let mut req = cn_request(payload); + req.region = "us-east-1".to_string(); + req.method = method; + req.raw_path = path.to_string(); + req.query_params = query + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + let resp = svc.handle(req).await.unwrap(); + let status = resp.status; + (status, body(resp)) + } + + fn keys(v: &Value) -> Vec<&str> { + let mut k: Vec<&str> = v.as_object().unwrap().keys().map(|k| k.as_str()).collect(); + k.sort_unstable(); + k + } + + #[tokio::test] + async fn flow_responses_carry_the_model_output_members() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let (status, created) = call( + &svc, + Method::POST, + "/flows/", + &[], + json!({ + "name": "f", + "description": "d", + "executionRoleArn": "arn:aws:iam::123456789012:role/r", + "customerEncryptionKeyArn": "arn:aws:kms:us-east-1:123456789012:key/k", + "definition": {"nodes": [], "connections": []}, + }), + ) + .await; + assert_eq!(status, StatusCode::CREATED); + let id = created["id"].as_str().unwrap().to_string(); + let arn = format!("arn:aws:bedrock:us-east-1:123456789012:flow/{id}"); + assert_eq!(created["arn"], arn); + let flow_members = vec![ + "arn", + "createdAt", + "customerEncryptionKeyArn", + "definition", + "description", + "executionRoleArn", + "id", + "name", + "status", + "updatedAt", + "version", + ]; + assert_eq!(keys(&created), flow_members); + + // GetFlow by ID and by ARN returns the same top-level shape. + let (_, got) = call(&svc, Method::GET, &format!("/flows/{id}/"), &[], json!({})).await; + assert_eq!(keys(&got), flow_members); + assert_eq!(got["id"], id); + assert_eq!(got["arn"], arn); + assert_eq!(got["status"], "NotPrepared"); + assert!(got.get("flowId").is_none() && got.get("flow").is_none()); + let encoded_arn = arn.replace(':', "%3A").replace('/', "%2F"); + let (_, by_arn) = call( + &svc, + Method::GET, + &format!("/flows/{encoded_arn}/"), + &[], + json!({}), + ) + .await; + assert_eq!(by_arn["id"], id); + + let (status, prepared) = + call(&svc, Method::POST, &format!("/flows/{id}/"), &[], json!({})).await; + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!(prepared, json!({"id": id, "status": "Prepared"})); + + let (_, updated) = call( + &svc, + Method::PUT, + &format!("/flows/{id}/"), + &[], + json!({"name": "f2", "executionRoleArn": "arn:aws:iam::123456789012:role/r"}), + ) + .await; + assert_eq!(keys(&updated), flow_members); + assert_eq!(updated["name"], "f2"); + assert_eq!(updated["arn"], arn); + assert_eq!(updated["status"], "NotPrepared"); + + let (status, version) = call( + &svc, + Method::POST, + &format!("/flows/{id}/versions"), + &[], + json!({}), + ) + .await; + assert_eq!(status, StatusCode::CREATED); + assert_eq!(version["id"], id); + assert_eq!(version["arn"], arn); + assert_eq!(version["version"], "1"); + assert_eq!(version["name"], "f2"); + let (_, got_version) = call( + &svc, + Method::GET, + &format!("/flows/{id}/versions/1/"), + &[], + json!({}), + ) + .await; + assert_eq!(got_version, version); + + let (status, alias) = call( + &svc, + Method::POST, + &format!("/flows/{id}/aliases"), + &[], + json!({"name": "live", "routingConfiguration": [{"flowVersion": "1"}]}), + ) + .await; + assert_eq!(status, StatusCode::CREATED); + let alias_id = alias["id"].as_str().unwrap().to_string(); + assert_eq!(alias["flowId"], id); + assert_eq!(alias["arn"], format!("{arn}/alias/{alias_id}")); + assert_eq!( + keys(&alias), + vec![ + "arn", + "createdAt", + "flowId", + "id", + "name", + "routingConfiguration", + "updatedAt" + ] + ); + let (_, got_alias) = call( + &svc, + Method::GET, + &format!("/flows/{id}/aliases/{alias_id}"), + &[], + json!({}), + ) + .await; + assert_eq!(got_alias, alias); + + let (_, deleted_alias) = call( + &svc, + Method::DELETE, + &format!("/flows/{id}/aliases/{alias_id}"), + &[], + json!({}), + ) + .await; + assert_eq!(deleted_alias, json!({"flowId": id, "id": alias_id})); + let (_, deleted_version) = call( + &svc, + Method::DELETE, + &format!("/flows/{id}/versions/1/"), + &[], + json!({}), + ) + .await; + assert_eq!(deleted_version, json!({"id": id, "version": "1"})); + let (_, deleted) = call( + &svc, + Method::DELETE, + &format!("/flows/{id}/"), + &[], + json!({}), + ) + .await; + assert_eq!(deleted, json!({"id": id})); + } + + #[tokio::test] + async fn validate_flow_definition_reports_structural_problems() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let input = json!({"name": "In", "type": "Input", + "outputs": [{"name": "document", "type": "String"}]}); + let output = json!({"name": "Out", "type": "Output", + "inputs": [{"name": "document", "type": "String", "expression": "$.data"}]}); + let link = |name: &str, source: &str, target: &str| { + json!({"name": name, "source": source, "target": target, "type": "Data", + "configuration": {"data": {"sourceOutput": "document", "targetInput": "document"}}}) + }; + + // A well-formed Input -> Output flow validates clean. The route must + // not be swallowed by PrepareFlow's POST /flows/{flowIdentifier}. + let (status, ok) = call( + &svc, + Method::POST, + "/flows/validate-definition", + &[], + json!({"definition": {"nodes": [input, output], + "connections": [link("c1", "In", "Out")]}}), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(ok, json!({"validations": []})); + + // No Output node, a connection to a missing node, an unreachable node. + let stray = json!({"name": "Stray", "type": "Prompt", "inputs": [], "outputs": []}); + let (_, bad) = call( + &svc, + Method::POST, + "/flows/validate-definition", + &[], + json!({"definition": {"nodes": [input, stray], + "connections": [link("c1", "In", "Ghost")]}}), + ) + .await; + let types: Vec<&str> = bad["validations"] + .as_array() + .unwrap() + .iter() + .map(|v| v["type"].as_str().unwrap()) + .collect(); + assert_eq!( + types, + vec![ + "MissingEndingNodes", + "UnknownConnectionTarget", + "UnreachableNode" + ] + ); + assert_eq!( + bad["validations"][1]["details"], + json!({"unknownConnectionTarget": {"connection": "c1"}}) + ); + assert_eq!(bad["validations"][2]["severity"], "Warning"); + } + + #[tokio::test] + async fn prompt_responses_carry_the_model_output_members() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let variants = json!([{"name": "v1", "templateType": "TEXT", + "templateConfiguration": {"text": {"text": "hi"}}}]); + let (status, created) = call( + &svc, + Method::POST, + "/prompts/", + &[], + json!({"name": "p", "defaultVariant": "v1", "variants": variants}), + ) + .await; + assert_eq!(status, StatusCode::CREATED); + let id = created["id"].as_str().unwrap().to_string(); + let arn = format!("arn:aws:bedrock:us-east-1:123456789012:prompt/{id}"); + let prompt_members = vec![ + "arn", + "createdAt", + "defaultVariant", + "id", + "name", + "updatedAt", + "variants", + "version", + ]; + + let (_, got) = call( + &svc, + Method::GET, + &format!("/prompts/{id}/"), + &[], + json!({}), + ) + .await; + assert_eq!(keys(&got), prompt_members); + assert_eq!(got["id"], id); + assert_eq!(got["arn"], arn); + assert_eq!(got["version"], "DRAFT"); + assert_eq!(got["defaultVariant"], "v1"); + assert!(got.get("promptId").is_none() && got.get("prompt").is_none()); + + let (_, updated) = call( + &svc, + Method::PUT, + &format!("/prompts/{id}/"), + &[], + json!({"name": "p2", "variants": variants}), + ) + .await; + assert_eq!(keys(&updated), prompt_members); + assert_eq!(updated["name"], "p2"); + assert_eq!(updated["arn"], arn); + + let (_, version) = call( + &svc, + Method::POST, + &format!("/prompts/{id}/versions"), + &[], + json!({}), + ) + .await; + assert_eq!(version["arn"], format!("{arn}:1")); + // GetPrompt with ?promptVersion= returns that version. + let (_, got_version) = call( + &svc, + Method::GET, + &format!("/prompts/{id}/"), + &[("promptVersion", "1")], + json!({}), + ) + .await; + assert_eq!(got_version, version); + assert_eq!(got_version["version"], "1"); + assert_eq!(got_version["name"], "p2"); + + let (_, deleted_version) = call( + &svc, + Method::DELETE, + &format!("/prompts/{id}/"), + &[("promptVersion", "1")], + json!({}), + ) + .await; + assert_eq!(deleted_version, json!({"id": id, "version": "1"})); + let (_, deleted) = call( + &svc, + Method::DELETE, + &format!("/prompts/{id}/"), + &[], + json!({}), + ) + .await; + assert_eq!(deleted, json!({"id": id})); + } + #[test] fn listed_flow_and_prompt_arns_are_the_ones_they_were_created_with() { let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); diff --git a/crates/fakecloud-bedrock-agent/src/service/prompts.rs b/crates/fakecloud-bedrock-agent/src/service/prompts.rs index 2b320990c..4a1065cf4 100644 --- a/crates/fakecloud-bedrock-agent/src/service/prompts.rs +++ b/crates/fakecloud-bedrock-agent/src/service/prompts.rs @@ -2,60 +2,91 @@ use super::*; +/// Split a `promptIdentifier` into the prompt ID and the version it pins, if +/// any. The identifier is the bare prompt ID or the prompt ARN, optionally +/// suffixed `:` (`arn:...:prompt/[:]`). +fn prompt_id_of(identifier: &str) -> (String, Option) { + let identifier = decode_label(identifier); + match identifier.rsplit_once(":prompt/") { + Some((_, rest)) => match rest.split_once(':') { + Some((id, version)) => (id.to_string(), Some(version.to_string())), + None => (rest.to_string(), None), + }, + None => (identifier, None), + } +} + +/// The prompt ID and the version the request targets: `promptVersion` (an +/// HTTP query member) wins over a version pinned in the identifier ARN. +fn prompt_target( + req: &AwsRequest, + body: &Value, +) -> Result<(String, Option), AwsServiceError> { + let (id, pinned) = prompt_id_of(&req_str(body, "promptIdentifier")?); + let version = opt_str(body, "promptVersion") + .or_else(|| req.query_params.get("promptVersion").cloned()) + .or(pinned) + // DRAFT names the working draft, the same as naming no version. + .filter(|v| v != "DRAFT"); + Ok((id, version)) +} + +fn prompt_not_found(id: &str) -> AwsServiceError { + not_found(format!("Prompt {id} not found")) +} + +fn prompt_version_not_found(version: &str) -> AwsServiceError { + not_found(format!("Prompt version {version} not found")) +} + impl BedrockAgentService { pub(super) fn create_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); let name = req_str(&body, "name")?; let id = short_id(); let now_dt = now(); - let variants = opt_array(&body, "variants"); - let arn = prompt_arn(&req.region, &req.account_id, &id); let prompt = Prompt { prompt_id: id.clone(), - name: name.clone(), + name, description: opt_str(&body, "description"), - variants: variants.clone(), + variants: opt_array(&body, "variants"), version: "DRAFT".to_string(), created_at: now_dt, updated_at: now_dt, - arn: arn.clone(), + arn: prompt_arn(&req.region, &req.account_id, &id), + customer_encryption_key_arn: opt_str(&body, "customerEncryptionKeyArn"), + default_variant: opt_str(&body, "defaultVariant"), }; + let out = prompt_json(&prompt); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); - state.prompts.insert(id.clone(), prompt); - let mut out = json!({ - "name": name, - "id": id, - "arn": arn, - "version": "DRAFT", - "createdAt": now_dt.to_rfc3339(), - "updatedAt": now_dt.to_rfc3339(), - "variants": variants, - }); - if let Some(d) = opt_str(&body, "description") { - out["description"] = json!(d); - } - if let Some(k) = opt_str(&body, "customerEncryptionKeyArn") { - out["customerEncryptionKeyArn"] = json!(k); - } - if let Some(dv) = opt_str(&body, "defaultVariant") { - out["defaultVariant"] = json!(dv); - } - Ok(AwsResponse::ok_json(out)) + state.prompts.insert(id, prompt); + Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } pub(super) fn get_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "promptIdentifier")?; + let (id, version) = prompt_target(req, &body)?; let accts = self.state.read(); let state = accts .get(&req.account_id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; + .ok_or_else(|| prompt_not_found(&id))?; let p = state .prompts .get(&id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; - Ok(AwsResponse::ok_json(json!({ "prompt": prompt_json(p) }))) + .ok_or_else(|| prompt_not_found(&id))?; + let out = match version { + None => prompt_json(p), + Some(version) => { + let v = state + .prompt_versions + .get(&id) + .and_then(|vs| vs.iter().find(|v| v.prompt_version == version)) + .ok_or_else(|| prompt_version_not_found(&version))?; + prompt_version_json(p, v) + } + }; + Ok(AwsResponse::ok_json(out)) } pub(super) fn create_prompt_version( @@ -66,39 +97,29 @@ impl BedrockAgentService { // the body under `promptIdentifier`, so we read it back here. The // resulting version is numbered incrementally per the Smithy contract. let body = req.json_body(); - let id = req_str(&body, "promptIdentifier")?; + let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); let now_dt = now(); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let prompt = state .prompts .get(&id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))? + .ok_or_else(|| prompt_not_found(&id))? .clone(); let versions = state.prompt_versions.entry(id.clone()).or_default(); let version_num = (versions.len() as u64 + 1).to_string(); let pv = PromptVersion { - prompt_version: version_num.clone(), - prompt_id: id.clone(), - description: opt_str(&body, "description"), + prompt_version: version_num, + prompt_id: id, + description: opt_str(&body, "description").or(prompt.description.clone()), created_at: now_dt, updated_at: now_dt, variants: prompt.variants.clone(), + name: Some(prompt.name.clone()), + default_variant: prompt.default_variant.clone(), }; + let out = prompt_version_json(&prompt, &pv); versions.push(pv); - let arn = format!("{}:{version_num}", prompt.arn); - let mut out = json!({ - "name": prompt.name, - "id": id, - "arn": arn, - "version": version_num, - "createdAt": now_dt.to_rfc3339(), - "updatedAt": now_dt.to_rfc3339(), - "variants": prompt.variants, - }); - if let Some(d) = opt_str(&body, "description").or(prompt.description) { - out["description"] = json!(d); - } Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } @@ -113,13 +134,13 @@ impl BedrockAgentService { pub(super) fn update_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "promptIdentifier")?; + let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let p = state .prompts .get_mut(&id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; + .ok_or_else(|| prompt_not_found(&id))?; p.updated_at = now(); if let Some(n) = opt_str(&body, "name") { p.name = n; @@ -130,20 +151,46 @@ impl BedrockAgentService { if body.get("variants").is_some() { p.variants = opt_array(&body, "variants"); } - Ok(AwsResponse::ok_json(json!({ "prompt": prompt_json(p) }))) + if let Some(k) = opt_str(&body, "customerEncryptionKeyArn") { + p.customer_encryption_key_arn = Some(k); + } + if let Some(dv) = opt_str(&body, "defaultVariant") { + p.default_variant = Some(dv); + } + Ok(AwsResponse::ok_json(prompt_json(p))) } pub(super) fn delete_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = req_str(&body, "promptIdentifier")?; + let (id, version) = prompt_target(req, &body)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); - state - .prompts - .remove(&id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; - state.prompt_versions.remove(&id); - Ok(AwsResponse::ok_json(json!({}))) + if !state.prompts.contains_key(&id) { + return Err(prompt_not_found(&id)); + } + match version { + // Naming a version deletes just that version. + Some(version) => { + let versions = state + .prompt_versions + .get_mut(&id) + .ok_or_else(|| prompt_version_not_found(&version))?; + let pos = versions + .iter() + .position(|v| v.prompt_version == version) + .ok_or_else(|| prompt_version_not_found(&version))?; + versions.remove(pos); + Ok(AwsResponse::ok_json(json!({ + "id": id, + "version": version, + }))) + } + None => { + state.prompts.remove(&id); + state.prompt_versions.remove(&id); + Ok(AwsResponse::ok_json(json!({ "id": id }))) + } + } } pub(super) fn list_prompt_versions( @@ -151,11 +198,11 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let id = req_str(&body, "promptIdentifier")?; + let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); let accts = self.state.read(); let state = accts .get(&req.account_id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; + .ok_or_else(|| prompt_not_found(&id))?; let prompt_arn = state.prompts.get(&id).map(|p| p.arn.as_str()); let versions: Vec = prompt_arn .zip(state.prompt_versions.get(&id)) @@ -185,32 +232,21 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let id = req_str(&body, "promptIdentifier")?; + let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); let version = req_str(&body, "promptVersion")?; let accts = self.state.read(); let state = accts .get(&req.account_id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; + .ok_or_else(|| prompt_not_found(&id))?; let prompt = state .prompts .get(&id) - .ok_or_else(|| not_found(format!("Prompt {id} not found")))?; + .ok_or_else(|| prompt_not_found(&id))?; let v = state .prompt_versions .get(&id) .and_then(|vs| vs.iter().find(|v| v.prompt_version == version)) - .ok_or_else(|| not_found(format!("Prompt version {version} not found")))?; - let mut out = json!({ - "id": v.prompt_id, - "version": v.prompt_version, - "arn": format!("{}:{}", prompt.arn, v.prompt_version), - "createdAt": v.created_at.to_rfc3339(), - "updatedAt": v.updated_at.to_rfc3339(), - "variants": v.variants, - }); - if let Some(ref d) = v.description { - out["description"] = json!(d); - } - Ok(AwsResponse::ok_json(out)) + .ok_or_else(|| prompt_version_not_found(&version))?; + Ok(AwsResponse::ok_json(prompt_version_json(prompt, v))) } } diff --git a/crates/fakecloud-bedrock-agent/src/state.rs b/crates/fakecloud-bedrock-agent/src/state.rs index 797867dcb..4a5e334ec 100644 --- a/crates/fakecloud-bedrock-agent/src/state.rs +++ b/crates/fakecloud-bedrock-agent/src/state.rs @@ -70,8 +70,15 @@ impl BedrockAgentAccounts { } else { state.region.as_str() }; - for flow in state.flows.values_mut().filter(|f| f.arn.is_empty()) { - flow.arn = crate::arns::flow_arn(region, account_id, &flow.flow_id); + for flow in state.flows.values_mut() { + if flow.arn.is_empty() { + flow.arn = crate::arns::flow_arn(region, account_id, &flow.flow_id); + } + // PrepareFlow used to store the enum name instead of its wire + // value; FlowStatus serializes as `Prepared`. + if flow.status == "PREPARED" { + flow.status = "Prepared".to_string(); + } } for prompt in state.prompts.values_mut().filter(|p| p.arn.is_empty()) { prompt.arn = crate::arns::prompt_arn(region, account_id, &prompt.prompt_id); @@ -267,6 +274,8 @@ pub struct Flow { /// it backfilled on load (see [`BedrockAgentAccounts::backfill_arns`]). #[serde(default)] pub arn: String, + #[serde(default)] + pub customer_encryption_key_arn: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -278,8 +287,13 @@ pub struct FlowAlias { pub description: Option, pub created_at: DateTime, pub updated_at: DateTime, + #[serde(default)] + pub concurrency_configuration: Option, } +/// A numbered, immutable snapshot of a flow. The flow-level fields are +/// captured at CreateFlowVersion time; snapshots written before they were +/// captured leave them `None`, and readers fall back to the parent flow. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FlowVersion { pub flow_version: String, @@ -288,6 +302,14 @@ pub struct FlowVersion { pub created_at: DateTime, pub updated_at: DateTime, pub definition: Option, + #[serde(default)] + pub name: Option, + #[serde(default)] + pub execution_role_arn: Option, + #[serde(default)] + pub customer_encryption_key_arn: Option, + #[serde(default)] + pub status: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -302,8 +324,14 @@ pub struct Prompt { /// The ARN minted at creation; backfilled on load like [`Flow::arn`]. #[serde(default)] pub arn: String, + #[serde(default)] + pub customer_encryption_key_arn: Option, + #[serde(default)] + pub default_variant: Option, } +/// A numbered snapshot of a prompt. `name` and `default_variant` are captured +/// at CreatePromptVersion time; older snapshots fall back to the prompt. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PromptVersion { pub prompt_version: String, @@ -312,6 +340,10 @@ pub struct PromptVersion { pub created_at: DateTime, pub updated_at: DateTime, pub variants: Vec, + #[serde(default)] + pub name: Option, + #[serde(default)] + pub default_variant: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] From 683eeba40192222cf3259c58087218a90e24b443 Mon Sep 17 00:00:00 2001 From: Lucas Vieira Date: Tue, 29 Sep 2026 19:04:40 -0300 Subject: [PATCH 2/3] fix(bedrock-agent): harden flow/prompt identifiers, versions and validation - Flow and prompt version numbers come from a persisted per-resource counter (latest_version, defaulting from the highest live version for older snapshots), so a deleted version's number is never minted again. - DeletePrompt's promptVersion is a NumericalVersion: DRAFT or any other non-numeric value is a ValidationException instead of deleting the whole prompt. GetPrompt still reads DRAFT as the working draft. - ListPrompts honors its promptIdentifier query member, listing the draft and every numbered version of that prompt as PromptSummary. The GET /prompts/{id}/versions[/{v}] routes (ListPromptVersions/GetPromptVersion), which are not in the model, are removed. - Prompt versions capture customerEncryptionKeyArn at creation. Version snapshots are now authoritative, and only versions persisted before the snapshot existed fall back to the live resource. - Path labels are percent-decoded once in handle() instead of per handler. - Identifier ARNs go through arn_resource and must name the caller's account, region and partition (and an alias ARN its own flow). A foreign ARN is ResourceNotFound instead of resolving to a local resource with the same ID. - ValidateFlowDefinition keys duplicate connections on their ends (output and input, or condition), so one node feeding two inputs of another is valid. Nodes with an empty or already-used name are reported as Unspecified validations instead of being silently merged. - create_flow_version/create_prompt_version borrow the resource instead of cloning it. Fixed the stale prompt_summary_json doc comment. --- .../src/service/flows.rs | 165 +++-- .../src/service/mod.rs | 589 +++++++++++++++++- .../src/service/prompts.rs | 202 +++--- crates/fakecloud-bedrock-agent/src/state.rs | 15 +- 4 files changed, 788 insertions(+), 183 deletions(-) diff --git a/crates/fakecloud-bedrock-agent/src/service/flows.rs b/crates/fakecloud-bedrock-agent/src/service/flows.rs index d7c8de4ac..6479b9fd0 100644 --- a/crates/fakecloud-bedrock-agent/src/service/flows.rs +++ b/crates/fakecloud-bedrock-agent/src/service/flows.rs @@ -2,24 +2,40 @@ use super::*; -/// The flow ID a `flowIdentifier` names: the identifier is either the bare -/// flow ID or the (URL-encoded) flow ARN (`arn:...:flow/`). -fn flow_id_of(identifier: &str) -> String { - let identifier = decode_label(identifier); - identifier - .rsplit_once(":flow/") - .map_or(identifier.as_str(), |(_, id)| id) - .to_string() +/// The flow ID a `flowIdentifier` names: the bare flow ID, or the ARN of a +/// flow in the caller's account and region (`arn:...:flow/`). An ARN of +/// anything else names no flow here. +fn flow_id_of(req: &AwsRequest, identifier: &str) -> Result { + match parse_identifier(req, identifier) { + Some(Identifier::Id(id)) => Ok(id.to_string()), + Some(Identifier::Resource(resource)) => resource + .strip_prefix("flow/") + .filter(|id| !id.is_empty() && !id.contains('/')) + .map(str::to_string) + .ok_or_else(|| flow_not_found(identifier)), + None => Err(flow_not_found(identifier)), + } } -/// The alias ID an `aliasIdentifier` names: the bare alias ID or the alias -/// ARN (`arn:...:flow//alias/`). -fn alias_id_of(identifier: &str) -> String { - let identifier = decode_label(identifier); - identifier - .rsplit_once("/alias/") - .map_or(identifier.as_str(), |(_, id)| id) - .to_string() +/// The alias ID an `aliasIdentifier` names: the bare alias ID, or the ARN of +/// an alias of `flow_id` in the caller's account and region +/// (`arn:...:flow//alias/`). +fn alias_id_of( + req: &AwsRequest, + flow_id: &str, + identifier: &str, +) -> Result { + let not_found_alias = || not_found(format!("Flow alias {identifier} not found")); + match parse_identifier(req, identifier) { + Some(Identifier::Id(id)) => Ok(id.to_string()), + Some(Identifier::Resource(resource)) => resource + .strip_prefix("flow/") + .and_then(|rest| rest.split_once("/alias/")) + .filter(|(flow, alias)| *flow == flow_id && !alias.is_empty()) + .map(|(_, alias)| alias.to_string()) + .ok_or_else(not_found_alias), + None => Err(not_found_alias()), + } } fn flow_not_found(id: &str) -> AwsServiceError { @@ -39,7 +55,8 @@ fn flow_validation(severity: &str, kind: &str, message: String, details: Value) /// reports: the flow needs a starting (`Input`) and an ending (`Output`) node; /// every connection must join existing nodes (and, for data connections, an /// existing output of the source to an existing input of the target); no two -/// connections may join the same pair of nodes; no node input may be fed by +/// connections may join the same ends; node names must be present and unique; +/// no node input may be fed by /// more than one data connection or by none; connections may not form a cycle; /// and every node should be reachable from a starting node. fn validate_definition(definition: &Value) -> Vec { @@ -55,22 +72,43 @@ fn validate_definition(definition: &Value) -> Vec { .unwrap_or_default() }; - // node name -> (type, input names, output names) - let by_name: BTreeMap, BTreeSet)> = nodes - .iter() - .map(|n| { + // node name -> (type, input names, output names). A node with no name, + // or a name another node already took, can't be addressed by a + // connection; it is reported (the model has no dedicated validation type + // for either, so as `Unspecified`) rather than silently merged away. + let mut out = Vec::new(); + let mut by_name: BTreeMap, BTreeSet)> = + BTreeMap::new(); + for (index, n) in nodes.iter().enumerate() { + let name = str_of(n, "name"); + if name.is_empty() { + out.push(flow_validation( + "Error", + "Unspecified", + format!("Node at index {index} has no name."), + json!({ "unspecified": {} }), + )); + continue; + } + if by_name.contains_key(&name) { + out.push(flow_validation( + "Error", + "Unspecified", + format!("Node name {name} is used by more than one node."), + json!({ "unspecified": {} }), + )); + continue; + } + by_name.insert( + name, ( - str_of(n, "name"), - ( - str_of(n, "type"), - names_of(n, "inputs"), - names_of(n, "outputs"), - ), - ) - }) - .collect(); + str_of(n, "type"), + names_of(n, "inputs"), + names_of(n, "outputs"), + ), + ); + } - let mut out = Vec::new(); if !by_name.values().any(|(t, _, _)| t == "Input") { out.push(flow_validation( "Error", @@ -113,7 +151,23 @@ fn validate_definition(definition: &Value) -> Vec { json!({ "unknownConnectionTarget": { "connection": name } }), )); } - if !seen_pairs.insert((source.clone(), target.clone())) { + // Two connections duplicate each other when they join the same ends: + // the same output to the same input for a data connection, the same + // condition for a conditional one. One node feeding two different + // inputs of another is two distinct connections. + let ends = match ( + c["configuration"].get("data"), + c["configuration"].get("conditional"), + ) { + (Some(data), _) => format!( + "data:{}:{}", + str_of(data, "sourceOutput"), + str_of(data, "targetInput") + ), + (None, Some(cond)) => format!("conditional:{}", str_of(cond, "condition")), + (None, None) => String::new(), + }; + if !seen_pairs.insert((source.clone(), target.clone(), ends)) { out.push(flow_validation( "Error", "DuplicateConnections", @@ -256,6 +310,7 @@ impl BedrockAgentService { definition: opt_json(&body, "definition"), arn: flow_arn(&req.region, &req.account_id, &id), customer_encryption_key_arn: opt_str(&body, "customerEncryptionKeyArn"), + latest_version: 0, }; let out = flow_json(&flow); let mut accts = self.state.write(); @@ -266,7 +321,7 @@ impl BedrockAgentService { pub(super) fn get_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let accts = self.state.read(); let f = accts .get(&req.account_id) @@ -286,7 +341,7 @@ impl BedrockAgentService { pub(super) fn update_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let f = state @@ -316,7 +371,7 @@ impl BedrockAgentService { pub(super) fn delete_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); state.flows.remove(&id).ok_or_else(|| flow_not_found(&id))?; @@ -327,7 +382,7 @@ impl BedrockAgentService { pub(super) fn prepare_flow(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let f = state @@ -350,17 +405,21 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let now_dt = now(); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); + // `flows` and `flow_versions` are disjoint fields, so the flow can be + // borrowed while its version list is extended. let flow = state .flows - .get(&flow_id) - .ok_or_else(|| flow_not_found(&flow_id))? - .clone(); + .get_mut(&flow_id) + .ok_or_else(|| flow_not_found(&flow_id))?; let versions = state.flow_versions.entry(flow_id.clone()).or_default(); - let version_num = (versions.len() as u64 + 1).to_string(); + let version_num = next_version( + &mut flow.latest_version, + versions.iter().map(|v| v.flow_version.as_str()), + ); let fv = FlowVersion { flow_version: version_num, flow_id, @@ -373,7 +432,7 @@ impl BedrockAgentService { customer_encryption_key_arn: flow.customer_encryption_key_arn.clone(), status: Some(flow.status.clone()), }; - let out = flow_version_json(&flow, &fv); + let out = flow_version_json(flow, &fv); versions.push(fv); Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } @@ -383,7 +442,7 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let version = req_str(&body, "flowVersion")?; let accts = self.state.read(); let version_not_found = || not_found(format!("Flow version {version} not found")); @@ -402,7 +461,7 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let accts = self.state.read(); let state = accts .get(&req.account_id) @@ -438,7 +497,7 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let version = req_str(&body, "flowVersion")?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); @@ -476,7 +535,7 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let name = req_str(&body, "name")?; let alias_id = short_id(); let now_dt = now(); @@ -505,8 +564,8 @@ impl BedrockAgentService { pub(super) fn get_flow_alias(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); - let alias_id = alias_id_of(&req_str(&body, "aliasIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; + let alias_id = alias_id_of(req, &flow_id, &req_str(&body, "aliasIdentifier")?)?; let accts = self.state.read(); let alias_not_found = || not_found(format!("Flow alias {alias_id} not found")); let state = accts.get(&req.account_id).ok_or_else(alias_not_found)?; @@ -524,7 +583,7 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; let accts = self.state.read(); let state = accts .get(&req.account_id) @@ -547,8 +606,8 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); - let alias_id = alias_id_of(&req_str(&body, "aliasIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; + let alias_id = alias_id_of(req, &flow_id, &req_str(&body, "aliasIdentifier")?)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let alias_not_found = || not_found(format!("Flow alias {alias_id} not found")); @@ -584,8 +643,8 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { let body = req.json_body(); - let flow_id = flow_id_of(&req_str(&body, "flowIdentifier")?); - let alias_id = alias_id_of(&req_str(&body, "aliasIdentifier")?); + let flow_id = flow_id_of(req, &req_str(&body, "flowIdentifier")?)?; + let alias_id = alias_id_of(req, &flow_id, &req_str(&body, "aliasIdentifier")?)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); match state.flow_aliases.get(&alias_id) { diff --git a/crates/fakecloud-bedrock-agent/src/service/mod.rs b/crates/fakecloud-bedrock-agent/src/service/mod.rs index 9e890d18f..af682bb33 100644 --- a/crates/fakecloud-bedrock-agent/src/service/mod.rs +++ b/crates/fakecloud-bedrock-agent/src/service/mod.rs @@ -570,16 +570,6 @@ impl BedrockAgentService { if *m == Method::POST { return Some(("CreatePromptVersion", params)); } - if *m == Method::GET { - return Some(("ListPromptVersions", params)); - } - } - if segs.len() == 4 && segs[0] == "prompts" && segs[2] == "versions" { - params.push(("promptIdentifier".to_string(), segs[1].clone())); - params.push(("promptVersion".to_string(), segs[3].clone())); - if *m == Method::GET { - return Some(("GetPromptVersion", params)); - } } // Tags @@ -630,8 +620,11 @@ impl AwsService for BedrockAgentService { if body.is_null() { body = serde_json::Value::Object(serde_json::Map::new()); } + // Path labels arrive percent-encoded (an ARN identifier is + // `arn%3Aaws%3A...`); decode each once here so every handler sees + // the identifier the caller sent. for (k, v) in path_params { - body[k] = serde_json::Value::String(v); + body[k] = serde_json::Value::String(decode_label(&v)); } req.body = serde_json::to_vec(&body).unwrap_or_default().into(); } @@ -677,7 +670,6 @@ impl AwsService for BedrockAgentService { "GetKnowledgeBase" => self.get_knowledge_base(&req), "GetKnowledgeBaseDocuments" => self.get_knowledge_base_documents(&req), "GetPrompt" => self.get_prompt(&req), - "GetPromptVersion" => self.get_prompt_version(&req), "IngestKnowledgeBaseDocuments" => self.ingest_knowledge_base_documents(&req), "ListAgentActionGroups" => self.list_agent_action_groups(&req), "ListAgentAliases" => self.list_agent_aliases(&req), @@ -693,7 +685,6 @@ impl AwsService for BedrockAgentService { "ListKnowledgeBaseDocuments" => self.list_knowledge_base_documents(&req), "ListKnowledgeBases" => self.list_knowledge_bases(&req), "ListPrompts" => self.list_prompts(&req), - "ListPromptVersions" => self.list_prompt_versions(&req), "ListTagsForResource" => self.list_tags_for_resource(&req), "PrepareAgent" => self.prepare_agent(&req), "PrepareFlow" => self.prepare_flow(&req), @@ -1031,6 +1022,42 @@ fn decode_label(segment: &str) -> String { .into_owned() } +/// A resource identifier as a request names it. +enum Identifier<'a> { + /// A bare resource ID. + Id(&'a str), + /// The resource part of a Bedrock ARN in the caller's account and region + /// (`flow/`, `prompt/:`, ...). + Resource(&'a str), +} + +/// Parse an identifier that is either a bare ID or a Bedrock ARN. An ARN that +/// is not Bedrock's, or that names another account, region or partition, +/// yields `None`: it can never resolve to a resource in this account's state. +fn parse_identifier<'a>(req: &AwsRequest, identifier: &'a str) -> Option> { + if !identifier.starts_with("arn:") { + return Some(Identifier::Id(identifier)); + } + let rest = fakecloud_aws::arn::arn_resource(identifier, "bedrock")?; + let mut parts = rest.splitn(3, ':'); + let (region, account, resource) = (parts.next()?, parts.next()?, parts.next()?); + let local = region == req.region + && account == req.account_id + && fakecloud_aws::arn::partition_of(identifier) + == fakecloud_aws::arn::partition_for(region); + local.then_some(Identifier::Resource(resource)) +} + +/// Mint the next version number for a resource whose highest minted number is +/// `latest` and whose live versions are `existing`. Numbers only grow, so a +/// deleted version's number is never reused; `existing` covers state persisted +/// before `latest` was tracked. +fn next_version<'a>(latest: &mut u64, existing: impl Iterator) -> String { + let highest_live = existing.filter_map(|v| v.parse::().ok()).max(); + *latest = (*latest).max(highest_live.unwrap_or(0)) + 1; + latest.to_string() +} + fn short_id() -> String { // Smithy ResourceIdentifier shape is @pattern("^[0-9a-zA-Z]{10}$"): the // first 10 hex chars of a v4 UUID. @@ -1292,8 +1319,10 @@ fn knowledge_base_summary_json(k: &KnowledgeBase) -> Value { o } -/// `PromptSummary`: `arn`, `id`, `name`, `version`, `createdAt`, `updatedAt`. -/// The full prompt JSON keys `promptId` (not `id`) and surfaces `variants`. +/// `PromptSummary` for a prompt's working draft: `arn`, `id`, `name`, +/// `version`, `createdAt`, `updatedAt`, and optional `description`. The full +/// `prompt_json` adds `variants`, `defaultVariant` and +/// `customerEncryptionKeyArn`, which the summary omits. fn prompt_summary_json(p: &Prompt) -> Value { let mut o = json!({ "arn": p.arn, @@ -1342,6 +1371,16 @@ fn flow_json(f: &Flow) -> Value { /// flow-level fields come from the version snapshot, falling back to the flow /// for versions persisted before they were captured. fn flow_version_json(f: &Flow, v: &FlowVersion) -> Value { + // A version's own snapshot is authoritative; only one persisted before the + // snapshot was captured (`name` unset) falls back to the live flow. + let legacy = v.name.is_none(); + let captured = |own: &Option, live: &Option| { + if legacy { + live.clone() + } else { + own.clone() + } + }; let mut o = json!({ "name": v.name.as_deref().unwrap_or(&f.name), "id": f.flow_id, @@ -1353,18 +1392,13 @@ fn flow_version_json(f: &Flow, v: &FlowVersion) -> Value { if let Some(ref d) = v.description { o["description"] = json!(d); } - if let Some(r) = v - .execution_role_arn - .as_ref() - .or(f.execution_role_arn.as_ref()) - { + if let Some(r) = captured(&v.execution_role_arn, &f.execution_role_arn) { o["executionRoleArn"] = json!(r); } - if let Some(k) = v - .customer_encryption_key_arn - .as_ref() - .or(f.customer_encryption_key_arn.as_ref()) - { + if let Some(k) = captured( + &v.customer_encryption_key_arn, + &f.customer_encryption_key_arn, + ) { o["customerEncryptionKeyArn"] = json!(k); } if let Some(ref def) = v.definition { @@ -1422,8 +1456,18 @@ fn prompt_json(p: &Prompt) -> Value { } /// The same members for a numbered prompt version: its ARN is the prompt ARN -/// with a `:` suffix. +/// with a `:` suffix. The version's own snapshot is authoritative; a +/// version persisted before the snapshot was captured (`name` unset) falls +/// back to the prompt. fn prompt_version_json(p: &Prompt, v: &PromptVersion) -> Value { + let legacy = v.name.is_none(); + let captured = |own: &Option, live: &Option| { + if legacy { + live.clone() + } else { + own.clone() + } + }; let mut o = json!({ "name": v.name.as_deref().unwrap_or(&p.name), "id": p.prompt_id, @@ -1436,10 +1480,13 @@ fn prompt_version_json(p: &Prompt, v: &PromptVersion) -> Value { if let Some(ref d) = v.description { o["description"] = json!(d); } - if let Some(ref k) = p.customer_encryption_key_arn { + if let Some(k) = captured( + &v.customer_encryption_key_arn, + &p.customer_encryption_key_arn, + ) { o["customerEncryptionKeyArn"] = json!(k); } - if let Some(dv) = v.default_variant.as_ref().or(p.default_variant.as_ref()) { + if let Some(dv) = captured(&v.default_variant, &p.default_variant) { o["defaultVariant"] = json!(dv); } o @@ -1534,6 +1581,18 @@ mod tests { query: &[(&str, &str)], payload: Value, ) -> (StatusCode, Value) { + let resp = try_call(svc, method, path, query, payload).await.unwrap(); + let status = resp.status; + (status, body(resp)) + } + + async fn try_call( + svc: &BedrockAgentService, + method: Method, + path: &str, + query: &[(&str, &str)], + payload: Value, + ) -> Result { let mut req = cn_request(payload); req.region = "us-east-1".to_string(); req.method = method; @@ -1542,9 +1601,7 @@ mod tests { .iter() .map(|(k, v)| (k.to_string(), v.to_string())) .collect(); - let resp = svc.handle(req).await.unwrap(); - let status = resp.status; - (status, body(resp)) + svc.handle(req).await } fn keys(v: &Value) -> Vec<&str> { @@ -1553,6 +1610,466 @@ mod tests { k } + fn encode(arn: &str) -> String { + arn.replace(':', "%3A").replace('/', "%2F") + } + + fn validation_types(v: &Value) -> Vec<&str> { + v["validations"] + .as_array() + .unwrap() + .iter() + .map(|v| v["type"].as_str().unwrap()) + .collect() + } + + #[tokio::test] + async fn version_numbers_are_never_reused_after_a_delete() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let (_, flow) = call( + &svc, + Method::POST, + "/flows/", + &[], + json!({"name": "f", "executionRoleArn": "arn:aws:iam::123456789012:role/r"}), + ) + .await; + let fid = flow["id"].as_str().unwrap().to_string(); + let (_, prompt) = call(&svc, Method::POST, "/prompts/", &[], json!({"name": "p"})).await; + let pid = prompt["id"].as_str().unwrap().to_string(); + let flow_versions = format!("/flows/{fid}/versions"); + let prompt_versions = format!("/prompts/{pid}/versions"); + let mint = |path: String| { + let svc = &svc; + async move { + let (_, v) = call(svc, Method::POST, &path, &[], json!({})).await; + v["version"].as_str().unwrap().to_string() + } + }; + + assert_eq!(mint(flow_versions.clone()).await, "1"); + assert_eq!(mint(flow_versions.clone()).await, "2"); + call( + &svc, + Method::DELETE, + &format!("/flows/{fid}/versions/1/"), + &[], + json!({}), + ) + .await; + assert_eq!(mint(flow_versions.clone()).await, "3"); + // Deleting the highest version doesn't free its number either. + call( + &svc, + Method::DELETE, + &format!("/flows/{fid}/versions/3/"), + &[], + json!({}), + ) + .await; + assert_eq!(mint(flow_versions).await, "4"); + + assert_eq!(mint(prompt_versions.clone()).await, "1"); + assert_eq!(mint(prompt_versions.clone()).await, "2"); + let prompt_path = format!("/prompts/{pid}/"); + call( + &svc, + Method::DELETE, + &prompt_path, + &[("promptVersion", "1")], + json!({}), + ) + .await; + assert_eq!(mint(prompt_versions.clone()).await, "3"); + call( + &svc, + Method::DELETE, + &prompt_path, + &[("promptVersion", "3")], + json!({}), + ) + .await; + assert_eq!(mint(prompt_versions).await, "4"); + } + + #[tokio::test] + async fn delete_prompt_rejects_a_non_numeric_version() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let (_, prompt) = call(&svc, Method::POST, "/prompts/", &[], json!({"name": "p"})).await; + let path = format!("/prompts/{}/", prompt["id"].as_str().unwrap()); + for bad in ["DRAFT", "x1", "123456"] { + let err = try_call( + &svc, + Method::DELETE, + &path, + &[("promptVersion", bad)], + json!({}), + ) + .await + .err() + .expect("non-numeric promptVersion must be rejected"); + assert_eq!(err.code(), "ValidationException", "{bad}"); + } + // The prompt survived every rejected delete. + let (status, got) = call(&svc, Method::GET, &path, &[], json!({})).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(got["id"], prompt["id"]); + // GetPrompt's promptVersion does accept DRAFT: the working draft. + let (_, draft) = call( + &svc, + Method::GET, + &path, + &[("promptVersion", "DRAFT")], + json!({}), + ) + .await; + assert_eq!(draft, got); + } + + #[tokio::test] + async fn list_prompts_with_an_identifier_lists_that_prompts_versions() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let (_, p) = call(&svc, Method::POST, "/prompts/", &[], json!({"name": "p"})).await; + call( + &svc, + Method::POST, + "/prompts/", + &[], + json!({"name": "other"}), + ) + .await; + let id = p["id"].as_str().unwrap().to_string(); + let arn = p["arn"].as_str().unwrap().to_string(); + call( + &svc, + Method::POST, + &format!("/prompts/{id}/versions"), + &[], + json!({}), + ) + .await; + + let (_, all) = call(&svc, Method::GET, "/prompts/", &[], json!({})).await; + assert_eq!(all["promptSummaries"].as_array().unwrap().len(), 2); + + for identifier in [id.as_str(), arn.as_str()] { + let (_, versions) = call( + &svc, + Method::GET, + "/prompts/", + &[("promptIdentifier", identifier)], + json!({}), + ) + .await; + let summaries = versions["promptSummaries"].as_array().unwrap(); + let listed: Vec<(&str, &str)> = summaries + .iter() + .map(|s| (s["version"].as_str().unwrap(), s["arn"].as_str().unwrap())) + .collect(); + let v1_arn = format!("{arn}:1"); + assert_eq!( + listed, + vec![("DRAFT", arn.as_str()), ("1", v1_arn.as_str())] + ); + assert_eq!( + keys(&summaries[1]), + vec!["arn", "createdAt", "id", "name", "updatedAt", "version"] + ); + } + + let err = try_call( + &svc, + Method::GET, + "/prompts/", + &[("promptIdentifier", "NOSUCHID01")], + json!({}), + ) + .await + .err() + .unwrap(); + assert_eq!(err.code(), "ResourceNotFoundException"); + + // The version routes the model doesn't have are gone. + for (method, path) in [ + (Method::GET, format!("/prompts/{id}/versions")), + (Method::GET, format!("/prompts/{id}/versions/1")), + ] { + let err = try_call(&svc, method, &path, &[], json!({})) + .await + .err() + .unwrap(); + assert!( + matches!(err, AwsServiceError::ActionNotImplemented { .. }), + "{path}" + ); + } + } + + #[tokio::test] + async fn prompt_version_keeps_the_key_it_was_created_with() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let k1 = "arn:aws:kms:us-east-1:123456789012:key/one"; + let k2 = "arn:aws:kms:us-east-1:123456789012:key/two"; + let (_, p) = call( + &svc, + Method::POST, + "/prompts/", + &[], + json!({"name": "p", "customerEncryptionKeyArn": k1}), + ) + .await; + let id = p["id"].as_str().unwrap().to_string(); + call( + &svc, + Method::POST, + &format!("/prompts/{id}/versions"), + &[], + json!({}), + ) + .await; + call( + &svc, + Method::PUT, + &format!("/prompts/{id}/"), + &[], + json!({"name": "p", "customerEncryptionKeyArn": k2}), + ) + .await; + let path = format!("/prompts/{id}/"); + let (_, v1) = call( + &svc, + Method::GET, + &path, + &[("promptVersion", "1")], + json!({}), + ) + .await; + assert_eq!(v1["customerEncryptionKeyArn"], k1); + let (_, draft) = call(&svc, Method::GET, &path, &[], json!({})).await; + assert_eq!(draft["customerEncryptionKeyArn"], k2); + } + + #[tokio::test] + async fn arn_identifiers_must_name_this_account_and_region() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let (_, flow) = call( + &svc, + Method::POST, + "/flows/", + &[], + json!({"name": "f", "executionRoleArn": "arn:aws:iam::123456789012:role/r"}), + ) + .await; + let (_, prompt) = call(&svc, Method::POST, "/prompts/", &[], json!({"name": "p"})).await; + let fid = flow["id"].as_str().unwrap(); + let pid = prompt["id"].as_str().unwrap(); + call( + &svc, + Method::POST, + &format!("/prompts/{pid}/versions"), + &[], + json!({}), + ) + .await; + + // The local ARNs (URL-encoded, as an SDK sends them) resolve. + let (_, by_arn) = call( + &svc, + Method::GET, + &format!("/flows/{}/", encode(flow["arn"].as_str().unwrap())), + &[], + json!({}), + ) + .await; + assert_eq!(by_arn["id"], fid); + let pinned = format!("{}:1", prompt["arn"].as_str().unwrap()); + let (_, v1) = call( + &svc, + Method::GET, + &format!("/prompts/{}/", encode(&pinned)), + &[], + json!({}), + ) + .await; + assert_eq!(v1["version"], "1"); + assert_eq!(v1["arn"], pinned); + + // The same IDs under another account, region or partition do not. + for foreign in [ + format!("arn:aws:bedrock:us-east-1:999999999999:flow/{fid}"), + format!("arn:aws:bedrock:eu-west-1:123456789012:flow/{fid}"), + format!("arn:aws-cn:bedrock:us-east-1:123456789012:flow/{fid}"), + format!("arn:aws:lambda:us-east-1:123456789012:flow/{fid}"), + ] { + let err = try_call( + &svc, + Method::GET, + &format!("/flows/{}/", encode(&foreign)), + &[], + json!({}), + ) + .await + .err() + .unwrap(); + assert_eq!(err.code(), "ResourceNotFoundException", "{foreign}"); + } + let foreign_prompt = format!("arn:aws:bedrock:us-east-1:999999999999:prompt/{pid}"); + let err = try_call( + &svc, + Method::DELETE, + &format!("/prompts/{}/", encode(&foreign_prompt)), + &[], + json!({}), + ) + .await + .err() + .unwrap(); + assert_eq!(err.code(), "ResourceNotFoundException"); + let (status, _) = call( + &svc, + Method::GET, + &format!("/prompts/{pid}/"), + &[], + json!({}), + ) + .await; + assert_eq!( + status, + StatusCode::OK, + "the foreign-ARN delete must not touch it" + ); + + // An alias ARN must belong to the flow it is addressed under. + let (_, alias) = call( + &svc, + Method::POST, + &format!("/flows/{fid}/aliases"), + &[], + json!({"name": "a", "routingConfiguration": []}), + ) + .await; + let alias_arn = alias["arn"].as_str().unwrap(); + let (_, got) = call( + &svc, + Method::GET, + &format!("/flows/{fid}/aliases/{}", encode(alias_arn)), + &[], + json!({}), + ) + .await; + assert_eq!(got["id"], alias["id"]); + let other_flow = alias_arn.replace(fid, "OTHERFLOW1"); + let err = try_call( + &svc, + Method::GET, + &format!("/flows/{fid}/aliases/{}", encode(&other_flow)), + &[], + json!({}), + ) + .await + .err() + .unwrap(); + assert_eq!(err.code(), "ResourceNotFoundException"); + } + + #[tokio::test] + async fn validate_flow_definition_accepts_one_node_feeding_two_inputs() { + // The AWS sample: the flow input's document feeds both the `genre` + // and the `number` input of one prompt node. + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let data = |name: &str, source: &str, target: &str, out: &str, input: &str| { + json!({"name": name, "source": source, "target": target, "type": "Data", + "configuration": {"data": {"sourceOutput": out, "targetInput": input}}}) + }; + let nodes = json!([ + {"name": "FlowInput", "type": "Input", + "outputs": [{"name": "document", "type": "Object"}]}, + {"name": "MakePlaylist", "type": "Prompt", + "inputs": [ + {"name": "genre", "type": "String", "expression": "$.data.genre"}, + {"name": "number", "type": "Number", "expression": "$.data.number"} + ], + "outputs": [{"name": "modelCompletion", "type": "String"}]}, + {"name": "FlowOutput", "type": "Output", + "inputs": [{"name": "document", "type": "String", "expression": "$.data"}]} + ]); + let mut connections = vec![ + data("c1", "FlowInput", "MakePlaylist", "document", "genre"), + data("c2", "FlowInput", "MakePlaylist", "document", "number"), + data( + "c3", + "MakePlaylist", + "FlowOutput", + "modelCompletion", + "document", + ), + ]; + let (_, ok) = call( + &svc, + Method::POST, + "/flows/validate-definition", + &[], + json!({"definition": {"nodes": nodes, "connections": connections}}), + ) + .await; + assert_eq!(ok, json!({"validations": []})); + + // A true duplicate (same output into the same input) is still flagged. + connections.push(data("c4", "FlowInput", "MakePlaylist", "document", "genre")); + let (_, dup) = call( + &svc, + Method::POST, + "/flows/validate-definition", + &[], + json!({"definition": {"nodes": nodes, "connections": connections}}), + ) + .await; + assert_eq!( + validation_types(&dup), + vec!["DuplicateConnections", "MultipleNodeInputConnections"] + ); + } + + #[tokio::test] + async fn validate_flow_definition_reports_duplicate_and_missing_node_names() { + let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); + let (_, got) = call( + &svc, + Method::POST, + "/flows/validate-definition", + &[], + json!({"definition": { + "nodes": [ + {"name": "In", "type": "Input", "outputs": [{"name": "document", "type": "String"}]}, + {"name": "In", "type": "Output", "inputs": []}, + {"type": "Output", "inputs": []} + ], + "connections": [] + }}), + ) + .await; + let validations = got["validations"].as_array().unwrap(); + let unspecified: Vec<&str> = validations + .iter() + .filter(|v| v["type"] == "Unspecified") + .map(|v| v["message"].as_str().unwrap()) + .collect(); + assert_eq!( + unspecified, + vec![ + "Node name In is used by more than one node.", + "Node at index 2 has no name." + ] + ); + assert!(validations + .iter() + .filter(|v| v["type"] == "Unspecified") + .all(|v| v["details"] == json!({"unspecified": {}}) && v["severity"] == "Error")); + // The second `In` (an Output node) was not merged into the first, so + // the flow still has no ending node. + assert!(validation_types(&got).contains(&"MissingEndingNodes")); + } + #[tokio::test] async fn flow_responses_carry_the_model_output_members() { let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); @@ -1884,14 +2401,18 @@ mod tests { format!("{}:1", prompt["arn"].as_str().unwrap()) ); - let mut elsewhere = cn_request(json!({"promptIdentifier": prompt["id"].clone()})); + let mut elsewhere = cn_request(json!({})); elsewhere.region = "us-east-1".to_string(); let flows = body(svc.list_flows(&elsewhere).unwrap()); assert_eq!(flows["flowSummaries"][0]["arn"], flow["arn"]); let prompts = body(svc.list_prompts(&elsewhere).unwrap()); assert_eq!(prompts["promptSummaries"][0]["arn"], prompt["arn"]); - let versions = body(svc.list_prompt_versions(&elsewhere).unwrap()); - assert_eq!(versions["promptSummaries"][0]["arn"], version["arn"]); + elsewhere.query_params.insert( + "promptIdentifier".to_string(), + prompt["id"].as_str().unwrap().to_string(), + ); + let versions = body(svc.list_prompts(&elsewhere).unwrap()); + assert_eq!(versions["promptSummaries"][1]["arn"], version["arn"]); } #[test] diff --git a/crates/fakecloud-bedrock-agent/src/service/prompts.rs b/crates/fakecloud-bedrock-agent/src/service/prompts.rs index 4a1065cf4..5c2eccb24 100644 --- a/crates/fakecloud-bedrock-agent/src/service/prompts.rs +++ b/crates/fakecloud-bedrock-agent/src/service/prompts.rs @@ -3,32 +3,36 @@ use super::*; /// Split a `promptIdentifier` into the prompt ID and the version it pins, if -/// any. The identifier is the bare prompt ID or the prompt ARN, optionally -/// suffixed `:` (`arn:...:prompt/[:]`). -fn prompt_id_of(identifier: &str) -> (String, Option) { - let identifier = decode_label(identifier); - match identifier.rsplit_once(":prompt/") { - Some((_, rest)) => match rest.split_once(':') { - Some((id, version)) => (id.to_string(), Some(version.to_string())), - None => (rest.to_string(), None), - }, - None => (identifier, None), +/// any. The identifier is the bare prompt ID, or the ARN of a prompt in the +/// caller's account and region, optionally suffixed `:` +/// (`arn:...:prompt/[:]`). An ARN of anything else names no +/// prompt here. +fn prompt_id_of( + req: &AwsRequest, + identifier: &str, +) -> Result<(String, Option), AwsServiceError> { + match parse_identifier(req, identifier) { + Some(Identifier::Id(id)) => Ok((id.to_string(), None)), + Some(Identifier::Resource(resource)) => { + let rest = resource + .strip_prefix("prompt/") + .filter(|rest| !rest.is_empty() && !rest.contains('/')) + .ok_or_else(|| prompt_not_found(identifier))?; + Ok(match rest.split_once(':') { + Some((id, version)) => (id.to_string(), Some(version.to_string())), + None => (rest.to_string(), None), + }) + } + None => Err(prompt_not_found(identifier)), } } -/// The prompt ID and the version the request targets: `promptVersion` (an -/// HTTP query member) wins over a version pinned in the identifier ARN. -fn prompt_target( - req: &AwsRequest, - body: &Value, -) -> Result<(String, Option), AwsServiceError> { - let (id, pinned) = prompt_id_of(&req_str(body, "promptIdentifier")?); - let version = opt_str(body, "promptVersion") +/// The version a request's `promptVersion` (an `@httpQuery` member) names, or +/// failing that the version pinned in the identifier ARN. +fn requested_version(req: &AwsRequest, body: &Value, pinned: Option) -> Option { + opt_str(body, "promptVersion") .or_else(|| req.query_params.get("promptVersion").cloned()) .or(pinned) - // DRAFT names the working draft, the same as naming no version. - .filter(|v| v != "DRAFT"); - Ok((id, version)) } fn prompt_not_found(id: &str) -> AwsServiceError { @@ -39,6 +43,22 @@ fn prompt_version_not_found(version: &str) -> AwsServiceError { not_found(format!("Prompt version {version} not found")) } +/// `PromptSummary` for numbered version `v` of prompt `p`. +fn prompt_version_summary_json(p: &Prompt, v: &PromptVersion) -> Value { + let mut o = json!({ + "name": v.name.as_deref().unwrap_or(&p.name), + "id": p.prompt_id, + "arn": format!("{}:{}", p.arn, v.prompt_version), + "version": v.prompt_version, + "createdAt": v.created_at.to_rfc3339(), + "updatedAt": v.updated_at.to_rfc3339(), + }); + if let Some(ref d) = v.description { + o["description"] = json!(d); + } + o +} + impl BedrockAgentService { pub(super) fn create_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); @@ -56,6 +76,7 @@ impl BedrockAgentService { arn: prompt_arn(&req.region, &req.account_id, &id), customer_encryption_key_arn: opt_str(&body, "customerEncryptionKeyArn"), default_variant: opt_str(&body, "defaultVariant"), + latest_version: 0, }; let out = prompt_json(&prompt); let mut accts = self.state.write(); @@ -66,7 +87,10 @@ impl BedrockAgentService { pub(super) fn get_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let (id, version) = prompt_target(req, &body)?; + let (id, pinned) = prompt_id_of(req, &req_str(&body, "promptIdentifier")?)?; + // GetPrompt's `promptVersion` is a `Version`: `DRAFT` names the + // working draft, the same as naming no version. + let version = requested_version(req, &body, pinned).filter(|v| v != "DRAFT"); let accts = self.state.read(); let state = accts .get(&req.account_id) @@ -94,47 +118,77 @@ impl BedrockAgentService { req: &AwsRequest, ) -> Result { // The routing layer surfaces the prompt identifier (path segment) into - // the body under `promptIdentifier`, so we read it back here. The - // resulting version is numbered incrementally per the Smithy contract. + // the body under `promptIdentifier`, so we read it back here. let body = req.json_body(); - let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); + let (id, _) = prompt_id_of(req, &req_str(&body, "promptIdentifier")?)?; let now_dt = now(); let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); + // `prompts` and `prompt_versions` are disjoint fields, so the prompt + // can be borrowed while its version list is extended. let prompt = state .prompts - .get(&id) - .ok_or_else(|| prompt_not_found(&id))? - .clone(); + .get_mut(&id) + .ok_or_else(|| prompt_not_found(&id))?; let versions = state.prompt_versions.entry(id.clone()).or_default(); - let version_num = (versions.len() as u64 + 1).to_string(); + let version_num = next_version( + &mut prompt.latest_version, + versions.iter().map(|v| v.prompt_version.as_str()), + ); let pv = PromptVersion { prompt_version: version_num, prompt_id: id, - description: opt_str(&body, "description").or(prompt.description.clone()), + description: opt_str(&body, "description").or_else(|| prompt.description.clone()), created_at: now_dt, updated_at: now_dt, variants: prompt.variants.clone(), name: Some(prompt.name.clone()), default_variant: prompt.default_variant.clone(), + customer_encryption_key_arn: prompt.customer_encryption_key_arn.clone(), }; - let out = prompt_version_json(&prompt, &pv); + let out = prompt_version_json(prompt, &pv); versions.push(pv); Ok(AwsResponse::json_value(StatusCode::CREATED, out)) } + /// Without `promptIdentifier`, the working draft of every prompt; with + /// it, every version of that one prompt (the draft, then each numbered + /// version). pub(super) fn list_prompts(&self, req: &AwsRequest) -> Result { + let body = req.json_body(); + let identifier = opt_str(&body, "promptIdentifier") + .or_else(|| req.query_params.get("promptIdentifier").cloned()); let accts = self.state.read(); - let list: Vec = accts - .get(&req.account_id) - .map(|s| s.prompts.values().map(prompt_summary_json).collect()) - .unwrap_or_default(); + let state = accts.get(&req.account_id); + let list: Vec = match identifier { + None => state + .map(|s| s.prompts.values().map(prompt_summary_json).collect()) + .unwrap_or_default(), + Some(identifier) => { + let (id, _) = prompt_id_of(req, &identifier)?; + let state = state.ok_or_else(|| prompt_not_found(&id))?; + let p = state + .prompts + .get(&id) + .ok_or_else(|| prompt_not_found(&id))?; + std::iter::once(prompt_summary_json(p)) + .chain( + state + .prompt_versions + .get(&id) + .into_iter() + .flatten() + .map(|v| prompt_version_summary_json(p, v)), + ) + .collect() + } + }; Ok(AwsResponse::ok_json(json!({ "promptSummaries": list }))) } pub(super) fn update_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); + let (id, _) = prompt_id_of(req, &req_str(&body, "promptIdentifier")?)?; let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); let p = state @@ -162,7 +216,24 @@ impl BedrockAgentService { pub(super) fn delete_prompt(&self, req: &AwsRequest) -> Result { let body = req.json_body(); - let (id, version) = prompt_target(req, &body)?; + let (id, pinned) = prompt_id_of(req, &req_str(&body, "promptIdentifier")?)?; + let version = requested_version(req, &body, pinned); + // DeletePrompt's `promptVersion` is a `NumericalVersion` + // (`^[0-9]{1,5}$`): only a numbered version can be deleted on its own, + // and `DRAFT` is not a way to name the whole prompt. + if let Some(ref v) = version { + if v.is_empty() || v.len() > 5 || !v.bytes().all(|b| b.is_ascii_digit()) { + return Err(AwsServiceError::aws_error( + StatusCode::BAD_REQUEST, + "ValidationException", + format!( + "1 validation error detected: Value '{v}' at 'promptVersion' failed to \ + satisfy constraint: Member must satisfy regular expression pattern: \ + ^[0-9]{{1,5}}$" + ), + )); + } + } let mut accts = self.state.write(); let state = accts.get_or_create(&req.account_id, &req.region); if !state.prompts.contains_key(&id) { @@ -192,61 +263,4 @@ impl BedrockAgentService { } } } - - pub(super) fn list_prompt_versions( - &self, - req: &AwsRequest, - ) -> Result { - let body = req.json_body(); - let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); - let accts = self.state.read(); - let state = accts - .get(&req.account_id) - .ok_or_else(|| prompt_not_found(&id))?; - let prompt_arn = state.prompts.get(&id).map(|p| p.arn.as_str()); - let versions: Vec = prompt_arn - .zip(state.prompt_versions.get(&id)) - .map(|(prompt_arn, vs)| { - vs.iter() - .map(|v| { - let mut o = json!({ - "id": v.prompt_id, - "version": v.prompt_version, - "arn": format!("{prompt_arn}:{}", v.prompt_version), - "createdAt": v.created_at.to_rfc3339(), - "updatedAt": v.updated_at.to_rfc3339(), - }); - if let Some(ref d) = v.description { - o["description"] = json!(d); - } - o - }) - .collect() - }) - .unwrap_or_default(); - Ok(AwsResponse::ok_json(json!({ "promptSummaries": versions }))) - } - - pub(super) fn get_prompt_version( - &self, - req: &AwsRequest, - ) -> Result { - let body = req.json_body(); - let (id, _) = prompt_id_of(&req_str(&body, "promptIdentifier")?); - let version = req_str(&body, "promptVersion")?; - let accts = self.state.read(); - let state = accts - .get(&req.account_id) - .ok_or_else(|| prompt_not_found(&id))?; - let prompt = state - .prompts - .get(&id) - .ok_or_else(|| prompt_not_found(&id))?; - let v = state - .prompt_versions - .get(&id) - .and_then(|vs| vs.iter().find(|v| v.prompt_version == version)) - .ok_or_else(|| prompt_version_not_found(&version))?; - Ok(AwsResponse::ok_json(prompt_version_json(prompt, v))) - } } diff --git a/crates/fakecloud-bedrock-agent/src/state.rs b/crates/fakecloud-bedrock-agent/src/state.rs index 4a5e334ec..9b85ea7b7 100644 --- a/crates/fakecloud-bedrock-agent/src/state.rs +++ b/crates/fakecloud-bedrock-agent/src/state.rs @@ -276,6 +276,10 @@ pub struct Flow { pub arn: String, #[serde(default)] pub customer_encryption_key_arn: Option, + /// The highest version number ever minted for this flow, so a deleted + /// version's number is never handed out again. + #[serde(default)] + pub latest_version: u64, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -328,10 +332,15 @@ pub struct Prompt { pub customer_encryption_key_arn: Option, #[serde(default)] pub default_variant: Option, + /// The highest version number ever minted for this prompt; see + /// [`Flow::latest_version`]. + #[serde(default)] + pub latest_version: u64, } -/// A numbered snapshot of a prompt. `name` and `default_variant` are captured -/// at CreatePromptVersion time; older snapshots fall back to the prompt. +/// A numbered snapshot of a prompt. `name`, `default_variant` and +/// `customer_encryption_key_arn` are captured at CreatePromptVersion time; +/// older snapshots fall back to the prompt. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PromptVersion { pub prompt_version: String, @@ -344,6 +353,8 @@ pub struct PromptVersion { pub name: Option, #[serde(default)] pub default_variant: Option, + #[serde(default)] + pub customer_encryption_key_arn: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] From 6621fb8953556e285816f15555c71907f051c3d7 Mon Sep 17 00:00:00 2001 From: Lucas Vieira Date: Tue, 29 Sep 2026 19:08:05 -0300 Subject: [PATCH 3/3] fix(bedrock-agent): re-key percent-encoded tag ARNs on load Tags are now keyed by the decoded resource ARN, but state persisted by builds that did not decode the {resourceArn} path label holds TagResource entries under the percent-encoded ARN, which the decoded lookups in ListTagsForResource/UntagResource never find. The load-time backfill now percent-decodes each tags-map key and merges it into any existing decoded entry. The decoded entry was written later, so its values win. --- .../src/service/mod.rs | 46 +++++++++++++++++++ crates/fakecloud-bedrock-agent/src/state.rs | 33 +++++++++++++ 2 files changed, 79 insertions(+) diff --git a/crates/fakecloud-bedrock-agent/src/service/mod.rs b/crates/fakecloud-bedrock-agent/src/service/mod.rs index af682bb33..32fdb2309 100644 --- a/crates/fakecloud-bedrock-agent/src/service/mod.rs +++ b/crates/fakecloud-bedrock-agent/src/service/mod.rs @@ -2415,6 +2415,52 @@ mod tests { assert_eq!(versions["promptSummaries"][1]["arn"], version["arn"]); } + #[tokio::test] + async fn loading_a_snapshot_rekeys_percent_encoded_tag_arns() { + let arn = "arn:aws:bedrock:us-east-1:123456789012:flow/ABCDEFGHIJ"; + let other = "arn:aws:bedrock:us-east-1:123456789012:prompt/KLMNOPQRST"; + // An older build stored TagResource-over-the-wire tags under the + // still-encoded path label. One ARN has only the encoded entry; the + // other has both, the decoded one written later by a newer build. + let raw = json!({"accounts": {"123456789012": { + "account_id": "123456789012", + "region": "us-east-1", + "agents": {}, "agent_aliases": {}, "agent_versions": {}, + "knowledge_bases": {}, "data_sources": {}, "agent_knowledge_bases": {}, + "agent_collaborators": {}, "flows": {}, "flow_aliases": {}, + "flow_versions": {}, "prompts": {}, "prompt_versions": {}, + "ingestion_jobs": {}, + "tags": { + encode(arn): {"team": "a", "env": "dev"}, + encode(other): {"team": "old", "cost": "1"}, + other: {"team": "new"} + } + }}}); + let loaded: BedrockAgentAccounts = serde_json::from_value(raw).unwrap(); + let tags = &loaded.get("123456789012").unwrap().tags; + assert_eq!( + tags.keys().map(String::as_str).collect::>(), + vec![arn, other], + "every encoded key is re-keyed, none is left behind" + ); + assert_eq!(tags[other]["team"], "new", "the later decoded value wins"); + assert_eq!(tags[other]["cost"], "1", "encoded-only keys are merged in"); + + let svc = BedrockAgentService::new(Arc::new(RwLock::new(loaded))); + let path = format!("/tags/{}", encode(arn)); + let (_, listed) = call(&svc, Method::GET, &path, &[], json!({})).await; + assert_eq!(listed, json!({"tags": {"team": "a", "env": "dev"}})); + + let mut untag = cn_request(json!({})); + untag.region = "us-east-1".to_string(); + untag.method = Method::DELETE; + untag.raw_path = path.clone(); + untag.raw_query = "tagKeys=team&tagKeys=env".to_string(); + svc.handle(untag).await.unwrap(); + let (_, listed) = call(&svc, Method::GET, &path, &[], json!({})).await; + assert_eq!(listed, json!({"tags": {}})); + } + #[test] fn loading_a_snapshot_without_arns_backfills_them_from_the_state_region() { let svc = BedrockAgentService::new(Arc::new(RwLock::new(BedrockAgentAccounts::new()))); diff --git a/crates/fakecloud-bedrock-agent/src/state.rs b/crates/fakecloud-bedrock-agent/src/state.rs index 9b85ea7b7..4ff4e6c00 100644 --- a/crates/fakecloud-bedrock-agent/src/state.rs +++ b/crates/fakecloud-bedrock-agent/src/state.rs @@ -83,6 +83,39 @@ impl BedrockAgentAccounts { for prompt in state.prompts.values_mut().filter(|p| p.arn.is_empty()) { prompt.arn = crate::arns::prompt_arn(region, account_id, &prompt.prompt_id); } + state.rekey_encoded_tags(); + } + } +} + +impl BedrockAgentState { + /// Tags are keyed by the resource ARN as the caller wrote it. Builds that + /// did not decode the `{resourceArn}` path label stored TagResource calls + /// under the percent-encoded ARN (`arn%3Aaws%3Abedrock...`); move those + /// entries to the decoded ARN. When both forms exist, the decoded entry + /// was written later, so its values win and the encoded entry only + /// contributes keys the decoded one lacks. + fn rekey_encoded_tags(&mut self) { + let encoded: Vec = self + .tags + .keys() + .filter(|k| k.contains('%')) + .cloned() + .collect(); + for key in encoded { + let decoded = percent_encoding::percent_decode_str(&key) + .decode_utf8_lossy() + .into_owned(); + if decoded == key { + continue; + } + let Some(old) = self.tags.remove(&key) else { + continue; + }; + let entry = self.tags.entry(decoded).or_default(); + for (k, v) in old { + entry.entry(k).or_insert(v); + } } } }