This is an automated email from the ASF dual-hosted git repository. gnodet pushed a commit to branch backport/29157-to-camel-4.22.x in repository https://gitbox.apache.org/repos/asf/camel.git
commit db51aae83c73902c993e18f9eea60cd631786118 Author: Guillaume Nodet <[email protected]> AuthorDate: Wed Nov 4 15:36:48 2026 +0000 [backport camel-4.22.x] CAMEL-26584: camel-aws-bedrock - clear the retrieve token on the last page --- .../agentruntime/BedrockAgentRuntimeProducer.java | 4 ++ .../BedrockAgentRuntimeProducerRetrieveTest.java | 58 ++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/components/camel-aws/camel-aws-bedrock/src/main/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducer.java b/components/camel-aws/camel-aws-bedrock/src/main/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducer.java index a68b54bd764b..cea3699066e9 100644 --- a/components/camel-aws/camel-aws-bedrock/src/main/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducer.java +++ b/components/camel-aws/camel-aws-bedrock/src/main/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducer.java @@ -265,9 +265,13 @@ public class BedrockAgentRuntimeProducer extends DefaultProducer { message.setHeader(BedrockAgentRuntimeConstants.RETRIEVED_RESULTS, results); if (ObjectHelper.isNotEmpty(response.nextToken())) { message.setHeader(BedrockAgentRuntimeConstants.NEXT_TOKEN, response.nextToken()); + } else { + message.removeHeader(BedrockAgentRuntimeConstants.NEXT_TOKEN); } if (ObjectHelper.isNotEmpty(response.guardrailActionAsString())) { message.setHeader(BedrockAgentRuntimeConstants.RETRIEVE_GUARDRAIL_ACTION, response.guardrailActionAsString()); + } else { + message.removeHeader(BedrockAgentRuntimeConstants.RETRIEVE_GUARDRAIL_ACTION); } message.setBody(results); } diff --git a/components/camel-aws/camel-aws-bedrock/src/test/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducerRetrieveTest.java b/components/camel-aws/camel-aws-bedrock/src/test/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducerRetrieveTest.java index fe6a98447484..76467f4c90ce 100644 --- a/components/camel-aws/camel-aws-bedrock/src/test/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducerRetrieveTest.java +++ b/components/camel-aws/camel-aws-bedrock/src/test/java/org/apache/camel/component/aws2/bedrock/agentruntime/BedrockAgentRuntimeProducerRetrieveTest.java @@ -40,6 +40,7 @@ import software.amazon.awssdk.services.bedrockagentruntime.model.RetrieveRespons import software.amazon.awssdk.services.bedrockagentruntime.model.SearchType; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; @@ -47,6 +48,7 @@ import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) public class BedrockAgentRuntimeProducerRetrieveTest { @@ -161,6 +163,62 @@ public class BedrockAgentRuntimeProducerRetrieveTest { assertEquals("page-2-token", result.getMessage().getHeader(BedrockAgentRuntimeConstants.NEXT_TOKEN)); } + @Test + public void retrieveClearsNextTokenOnLastPage() { + assertLastPageClearsNextToken(null); + } + + @Test + public void retrieveClearsEmptyNextTokenOnLastPage() { + assertLastPageClearsNextToken(""); + } + + private void assertLastPageClearsNextToken(String nextToken) { + KnowledgeBaseRetrievalResult chunk = KnowledgeBaseRetrievalResult.builder() + .content(RetrievalResultContent.builder().text("last chunk").build()) + .score(0.87) + .build(); + when(syncClient.retrieve(any(RetrieveRequest.class))).thenAnswer(invocation -> { + capturedRequest.set(invocation.getArgument(0, RetrieveRequest.class)); + return RetrieveResponse.builder().retrievalResults(chunk).nextToken(nextToken).guardrailAction("NONE").build(); + }); + + Exchange result = template.send("direct:retrieve", exchange -> { + exchange.getMessage().setBody("what is camel?"); + exchange.getMessage().setHeader(BedrockAgentRuntimeConstants.NEXT_TOKEN, "last-page-token"); + exchange.getMessage().setHeader("correlationId", "request-1"); + }); + + assertNull(result.getException()); + assertEquals("last-page-token", capturedRequest.get().nextToken()); + assertFalse(result.getMessage().getHeaders().containsKey(BedrockAgentRuntimeConstants.NEXT_TOKEN)); + assertEquals(List.of(chunk), result.getMessage().getBody()); + assertSame(result.getMessage().getBody(), + result.getMessage().getHeader(BedrockAgentRuntimeConstants.RETRIEVED_RESULTS)); + assertEquals("NONE", result.getMessage().getHeader(BedrockAgentRuntimeConstants.RETRIEVE_GUARDRAIL_ACTION)); + assertEquals("request-1", result.getMessage().getHeader("correlationId")); + } + + @Test + public void retrieveClearsAbsentGuardrailAction() { + when(syncClient.retrieve(any(RetrieveRequest.class))).thenReturn( + RetrieveResponse.builder().nextToken("next-page-token").build()); + + Exchange result = template.send("direct:retrieve", exchange -> { + exchange.getMessage().setBody("what is camel?"); + exchange.getMessage().setHeader(BedrockAgentRuntimeConstants.RETRIEVE_GUARDRAIL_ACTION, "INTERVENED"); + exchange.getMessage().setHeader("correlationId", "request-1"); + }); + + assertNull(result.getException()); + assertFalse(result.getMessage().getHeaders().containsKey(BedrockAgentRuntimeConstants.RETRIEVE_GUARDRAIL_ACTION)); + assertEquals("next-page-token", result.getMessage().getHeader(BedrockAgentRuntimeConstants.NEXT_TOKEN)); + assertEquals(List.of(), result.getMessage().getBody()); + assertSame(result.getMessage().getBody(), + result.getMessage().getHeader(BedrockAgentRuntimeConstants.RETRIEVED_RESULTS)); + assertEquals("request-1", result.getMessage().getHeader("correlationId")); + } + @Test public void retrieveAppliesRetrievalConfigHeaders() { template.send("direct:retrieve", exchange -> {
