From b1778272fc91d53922ee62fe4be304ff72767f0b Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Thu, 16 Apr 2026 16:30:22 +0200 Subject: [PATCH] fix: clean ai memory and cache bedrock prompts (#8847) * fix: avoid persisting system prompts in ai memory Co-Authored-By: Claude Opus 4.5 * fix: keep ai memory cleanup write-side only Co-Authored-By: Claude Opus 4.5 * feat: add bedrock prompt caching for claude Co-Authored-By: Claude Opus 4.5 * test: add bedrock memory regression Co-Authored-By: Claude Opus 4.5 * fix: gate bedrock prompt caching by model id Co-Authored-By: Claude Opus 4.5 * docs: link bedrock caching allowlist source Co-Authored-By: Claude Opus 4.5 --------- Co-authored-by: Claude Opus 4.5 --- backend/windmill-api/src/bedrock.rs | 19 +- backend/windmill-common/src/ai_bedrock.rs | 203 +++++++++++++++++- .../src/ai/providers/bedrock.rs | 19 +- backend/windmill-worker/src/ai_executor.rs | 139 ++++++++++-- .../ai_agent_tests/test_memory.py | 96 ++++++++- 5 files changed, 447 insertions(+), 29 deletions(-) diff --git a/backend/windmill-api/src/bedrock.rs b/backend/windmill-api/src/bedrock.rs index 3824f8966d..8a5c81bf47 100644 --- a/backend/windmill-api/src/bedrock.rs +++ b/backend/windmill-api/src/bedrock.rs @@ -131,6 +131,7 @@ async fn create_bedrock_client( fn build_tool_config_from_request( tools: Option<&[OpenAIToolDef]>, tool_choice: Option<&serde_json::Value>, + enable_prompt_caching: bool, ) -> Result> { if let Some(tools) = tools { let tool_defs: Vec = tools @@ -163,7 +164,7 @@ fn build_tool_config_from_request( .map(|tc| tc == "required" || tc.as_str() == Some("required")) .unwrap_or(false); - build_tool_config(Some(&tool_defs), force_tool_use) + build_tool_config(Some(&tool_defs), force_tool_use, enable_prompt_caching) } else { Ok(None) } @@ -365,8 +366,13 @@ pub async fn handle_bedrock_sdk_streaming( .await?; // Convert messages using shared conversion + let enable_prompt_caching = + windmill_common::ai_bedrock::bedrock_model_supports_prompt_caching(model); let (bedrock_messages, system_prompts) = - windmill_common::ai_bedrock::openai_messages_to_bedrock(&openai_req.messages)?; + windmill_common::ai_bedrock::openai_messages_to_bedrock( + &openai_req.messages, + enable_prompt_caching, + )?; // Build inference configuration let inference_config = windmill_common::ai_bedrock::create_inference_config( @@ -378,6 +384,7 @@ pub async fn handle_bedrock_sdk_streaming( let tool_config = build_tool_config_from_request( openai_req.tools.as_deref(), openai_req.tool_choice.as_ref(), + enable_prompt_caching, )?; // Build the SDK request @@ -625,8 +632,13 @@ pub async fn handle_bedrock_sdk_non_streaming( .await?; // Convert messages using shared conversion + let enable_prompt_caching = + windmill_common::ai_bedrock::bedrock_model_supports_prompt_caching(model); let (bedrock_messages, system_prompts) = - windmill_common::ai_bedrock::openai_messages_to_bedrock(&openai_req.messages)?; + windmill_common::ai_bedrock::openai_messages_to_bedrock( + &openai_req.messages, + enable_prompt_caching, + )?; // Build inference configuration let inference_config = windmill_common::ai_bedrock::create_inference_config( @@ -638,6 +650,7 @@ pub async fn handle_bedrock_sdk_non_streaming( let tool_config = build_tool_config_from_request( openai_req.tools.as_deref(), openai_req.tool_choice.as_ref(), + enable_prompt_caching, )?; // Build the SDK request (non-streaming) diff --git a/backend/windmill-common/src/ai_bedrock.rs b/backend/windmill-common/src/ai_bedrock.rs index a740b202cb..5d869b8de8 100644 --- a/backend/windmill-common/src/ai_bedrock.rs +++ b/backend/windmill-common/src/ai_bedrock.rs @@ -94,6 +94,86 @@ pub async fn check_env_credentials() -> BedrockCredentialsCheck { /// Constants for commonly used strings to avoid allocations pub const FUNCTION_TYPE: &str = "function"; +// AWS documents Bedrock prompt-caching support as a model allowlist rather than a +// capability exposed by the model metadata APIs: +// https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html +const BEDROCK_PROMPT_CACHING_SUPPORTED_MODEL_IDS: &[&str] = &[ + "anthropic.claude-opus-4-5-20251101-v1:0", + "anthropic.claude-opus-4-1-20250805-v1:0", + "anthropic.claude-opus-4-20250514-v1:0", + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-haiku-4-5-20251001-v1:0", + "anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-3-7-sonnet-20250219-v1:0", + "anthropic.claude-3-5-haiku-20241022-v1:0", + "anthropic.claude-3-5-sonnet-20241022-v2:0", +]; + +fn build_default_cache_point() -> aws_sdk_bedrockruntime::types::CachePointBlock { + aws_sdk_bedrockruntime::types::CachePointBlock::builder() + .r#type(aws_sdk_bedrockruntime::types::CachePointType::Default) + .build() + .expect("cache point type is required") +} + +fn normalize_bedrock_model_id(model: &str) -> String { + let model = model + .rsplit('/') + .next() + .unwrap_or(model) + .to_ascii_lowercase(); + + for prefix in ["global.", "us.", "eu.", "apac."] { + if let Some(normalized_model) = model.strip_prefix(prefix) { + return normalized_model.to_string(); + } + } + + model +} + +pub fn bedrock_model_supports_prompt_caching(model: &str) -> bool { + let normalized_model = normalize_bedrock_model_id(model); + BEDROCK_PROMPT_CACHING_SUPPORTED_MODEL_IDS.contains(&normalized_model.as_str()) +} + +fn append_cache_point_to_system_prompts(system_prompts: &mut Vec) { + if system_prompts.is_empty() + || matches!( + system_prompts.last(), + Some(SystemContentBlock::CachePoint(_)) + ) + { + return; + } + + system_prompts.push(SystemContentBlock::CachePoint(build_default_cache_point())); +} + +fn append_cache_point_to_last_message(messages: &mut [Message]) -> Result<(), Error> { + let Some(last_message) = messages.last_mut() else { + return Ok(()); + }; + + if matches!( + last_message.content().last(), + Some(ContentBlock::CachePoint(_)) + ) { + return Ok(()); + } + + let mut content = last_message.content().to_vec(); + content.push(ContentBlock::CachePoint(build_default_cache_point())); + + *last_message = Message::builder() + .role(last_message.role().clone()) + .set_content(Some(content)) + .build() + .map_err(|e| Error::internal_err(format!("Failed to append cache point: {}", e)))?; + + Ok(()) +} + #[derive(Debug, Clone)] pub struct BearerTokenProvider { token: String, @@ -277,6 +357,7 @@ pub fn json_to_document(value: serde_json::Value) -> aws_smithy_types::Document /// Tuple of (conversation_messages, system_prompts) pub fn openai_messages_to_bedrock( messages: &[OpenAIMessage], + enable_prompt_caching: bool, ) -> Result<(Vec, Vec), Error> { let mut bedrock_messages = Vec::new(); let mut system_prompts = Vec::new(); @@ -334,6 +415,11 @@ pub fn openai_messages_to_bedrock( bedrock_messages.push(tool_result_message); } + if enable_prompt_caching { + append_cache_point_to_system_prompts(&mut system_prompts); + append_cache_point_to_last_message(&mut bedrock_messages)?; + } + Ok((bedrock_messages, system_prompts)) } @@ -703,9 +789,15 @@ pub fn streaming_tool_calls_to_openai(tool_calls: Vec) -> Vec pub fn build_tool_config( tools: Option<&[ToolDef]>, force_tool_use: bool, + enable_prompt_caching: bool, ) -> Result, Error> { if let Some(tools) = tools { - let bedrock_tools = openai_tools_to_bedrock(tools)?; + let mut bedrock_tools = openai_tools_to_bedrock(tools)?; + + if enable_prompt_caching && !matches!(bedrock_tools.last(), Some(Tool::CachePoint(_))) { + bedrock_tools.push(Tool::CachePoint(build_default_cache_point())); + } + let mut tool_config_builder = aws_sdk_bedrockruntime::types::ToolConfiguration::builder() .set_tools(Some(bedrock_tools)); @@ -724,3 +816,112 @@ pub fn build_tool_config( Ok(None) } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::value::RawValue; + + fn text_message(role: &str, content: &str) -> OpenAIMessage { + OpenAIMessage { + role: role.to_string(), + content: Some(OpenAIContent::Text(content.to_string())), + ..Default::default() + } + } + + fn test_tool() -> ToolDef { + ToolDef { + r#type: FUNCTION_TYPE.to_string(), + function: crate::ai_types::ToolDefFunction { + name: "test_tool".to_string(), + description: Some("A test tool".to_string()), + parameters: RawValue::from_string( + r#"{"type":"object","properties":{},"additionalProperties":false}"#.to_string(), + ) + .expect("valid raw json"), + }, + } + } + + #[test] + fn openai_messages_to_bedrock_adds_cache_points_when_enabled() { + let messages = vec![ + text_message("system", "Reply concisely"), + text_message("user", "Tell me a joke"), + ]; + + let (bedrock_messages, system_prompts) = + openai_messages_to_bedrock(&messages, true).expect("bedrock conversion succeeds"); + + assert!(matches!( + system_prompts.last(), + Some(SystemContentBlock::CachePoint(_)) + )); + assert!(matches!( + bedrock_messages + .last() + .and_then(|message| message.content().last()), + Some(ContentBlock::CachePoint(_)) + )); + } + + #[test] + fn openai_messages_to_bedrock_skips_cache_points_when_disabled() { + let messages = vec![ + text_message("system", "Reply concisely"), + text_message("user", "Tell me a joke"), + ]; + + let (bedrock_messages, system_prompts) = + openai_messages_to_bedrock(&messages, false).expect("bedrock conversion succeeds"); + + assert!(!matches!( + system_prompts.last(), + Some(SystemContentBlock::CachePoint(_)) + )); + assert!(!matches!( + bedrock_messages + .last() + .and_then(|message| message.content().last()), + Some(ContentBlock::CachePoint(_)) + )); + } + + #[test] + fn build_tool_config_adds_cache_point_when_enabled() { + let tools = vec![test_tool()]; + let tool_config = + build_tool_config(Some(&tools), false, true).expect("tool config succeeds"); + + assert!(matches!( + tool_config + .as_ref() + .and_then(|config| config.tools().last()), + Some(Tool::CachePoint(_)) + )); + } + + #[test] + fn bedrock_prompt_caching_supports_documented_claude_model_ids() { + assert!(bedrock_model_supports_prompt_caching( + "anthropic.claude-haiku-4-5-20251001-v1:0" + )); + assert!(bedrock_model_supports_prompt_caching( + "global.anthropic.claude-haiku-4-5-20251001-v1:0" + )); + assert!(bedrock_model_supports_prompt_caching( + "arn:aws:bedrock:us-east-1::inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + )); + } + + #[test] + fn bedrock_prompt_caching_rejects_unsupported_or_opaque_model_ids() { + assert!(!bedrock_model_supports_prompt_caching( + "anthropic.claude-3-haiku-20240307-v1:0" + )); + assert!(!bedrock_model_supports_prompt_caching( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-profile" + )); + } +} diff --git a/backend/windmill-worker/src/ai/providers/bedrock.rs b/backend/windmill-worker/src/ai/providers/bedrock.rs index 18f44f0573..809a222526 100644 --- a/backend/windmill-worker/src/ai/providers/bedrock.rs +++ b/backend/windmill-worker/src/ai/providers/bedrock.rs @@ -18,10 +18,11 @@ use windmill_common::{client::AuthedClient, error::Error}; // Re-export from shared module for use by other parts of the worker use windmill_common::ai_bedrock::{ - bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text, - bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, build_tool_config, - create_inference_config, format_bedrock_error, openai_messages_to_bedrock, - streaming_tool_calls_to_openai, StreamingToolCall, + bedrock_model_supports_prompt_caching, bedrock_stream_event_is_block_stop, + bedrock_stream_event_to_text, bedrock_stream_event_to_tool_delta, + bedrock_stream_event_to_tool_start, build_tool_config, create_inference_config, + format_bedrock_error, openai_messages_to_bedrock, streaming_tool_calls_to_openai, + StreamingToolCall, }; pub use windmill_common::ai_bedrock::{check_env_credentials, BedrockClient}; @@ -71,13 +72,19 @@ impl BedrockQueryBuilder { let prepared_messages = prepare_messages_for_api(messages, client, workspace_id).await?; // Convert messages to Bedrock format (separates system prompts) - let (bedrock_messages, system_prompts) = openai_messages_to_bedrock(&prepared_messages)?; + let enable_prompt_caching = bedrock_model_supports_prompt_caching(model); + let (bedrock_messages, system_prompts) = + openai_messages_to_bedrock(&prepared_messages, enable_prompt_caching)?; // Build inference configuration using shared helper let inference_config = create_inference_config(temperature, max_tokens.map(|t| t as i32)); // Build tool configuration with optional ToolChoice - let tool_config = build_tool_config(tools, structured_output_tool_name.is_some())?; + let tool_config = build_tool_config( + tools, + structured_output_tool_name.is_some(), + enable_prompt_caching, + )?; self.execute_converse_stream( &bedrock_client, diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 594830cebd..99197b9e9a 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -92,6 +92,38 @@ lazy_static::lazy_static! { const DEFAULT_MAX_AGENT_ITERATIONS: usize = 10; const HARD_MAX_AGENT_ITERATIONS: usize = 1000; +fn strip_system_messages(messages: &[OpenAIMessage]) -> Vec { + messages + .iter() + .filter(|message| message.role != "system") + .cloned() + .collect() +} + +fn strip_leading_tool_messages(messages: Vec) -> Vec { + match messages.iter().position(|message| message.role != "tool") { + Some(first_non_tool_index) => messages.into_iter().skip(first_non_tool_index).collect(), + None => Vec::new(), + } +} + +fn prepare_auto_memory_messages_for_request( + loaded_messages: &[OpenAIMessage], + context_length: usize, +) -> Vec { + let start_idx = loaded_messages.len().saturating_sub(context_length); + strip_leading_tool_messages(loaded_messages[start_idx..].to_vec()) +} + +fn prepare_auto_memory_messages_for_persistence( + all_messages: &[OpenAIMessage], + context_length: usize, +) -> Vec { + let non_system_messages = strip_system_messages(all_messages); + let start_idx = non_system_messages.len().saturating_sub(context_length); + non_system_messages[start_idx..].to_vec() +} + fn find_module_by_id( modules: &Vec, target_id: &str, @@ -654,18 +686,10 @@ pub async fn run_agent( // Read messages from memory match read_from_memory(db, &job.workspace_id, memory_id, step_id).await { Ok(Some(loaded_messages)) => { - // Take the last n messages - let start_idx = - loaded_messages.len().saturating_sub(*context_length); - let mut messages_to_load = loaded_messages[start_idx..].to_vec(); - let first_non_tool_message_index = - messages_to_load.iter().position(|m| m.role != "tool"); - - // Remove the first messages if their role is "tool" to avoid OpenAI API error - if let Some(index) = first_non_tool_message_index { - messages_to_load = messages_to_load[index..].to_vec(); - } - + let messages_to_load = prepare_auto_memory_messages_for_request( + &loaded_messages, + *context_length, + ); messages.extend(messages_to_load); } Ok(None) => {} @@ -729,9 +753,7 @@ pub async fn run_agent( .unwrap_or(false); if has_message && has_attachments { - let mut parts = vec![ContentPart::Text { - text: args.user_message.clone().unwrap(), - }]; + let mut parts = vec![ContentPart::Text { text: args.user_message.clone().unwrap() }]; for attachment in args.user_attachments.as_ref().unwrap() { if !attachment.s3.is_empty() { parts.push(ContentPart::S3Object { s3_object: attachment.clone() }); @@ -1306,9 +1328,10 @@ pub async fn run_agent( final_messages.iter().map(|m| m.message.clone()).collect(); if !all_messages.is_empty() { - // Keep only the last n messages - let start_idx = all_messages.len().saturating_sub(*context_length); - let messages_to_persist = all_messages[start_idx..].to_vec(); + let messages_to_persist = prepare_auto_memory_messages_for_persistence( + &all_messages, + *context_length, + ); if let Some(memory_id) = memory_id { if let Err(e) = write_to_memory( @@ -1349,6 +1372,86 @@ pub async fn run_agent( })) } +#[cfg(test)] +mod tests { + use super::*; + + fn text_message(role: &str, content: &str) -> OpenAIMessage { + OpenAIMessage { + role: role.to_string(), + content: Some(OpenAIContent::Text(content.to_string())), + ..Default::default() + } + } + + #[test] + fn auto_memory_request_preserves_messages_within_context_window() { + let loaded_messages = vec![ + text_message("system", "instructions-a"), + text_message("user", "first-user"), + text_message("assistant", "first-assistant"), + text_message("system", "instructions-b"), + text_message("user", "second-user"), + text_message("assistant", "second-assistant"), + ]; + + let prepared = prepare_auto_memory_messages_for_request(&loaded_messages, 3); + let roles: Vec<&str> = prepared + .iter() + .map(|message| message.role.as_str()) + .collect(); + let contents: Vec<&str> = prepared + .iter() + .map(|message| match message.content.as_ref() { + Some(OpenAIContent::Text(text)) => text.as_str(), + _ => "", + }) + .collect(); + + assert_eq!(roles, vec!["system", "user", "assistant"]); + assert_eq!( + contents, + vec!["instructions-b", "second-user", "second-assistant"] + ); + } + + #[test] + fn auto_memory_request_drops_leading_tool_messages() { + let loaded_messages = vec![ + text_message("tool", "stale-tool-result"), + text_message("user", "hello"), + text_message("assistant", "hi"), + ]; + + let prepared = prepare_auto_memory_messages_for_request(&loaded_messages, 10); + let roles: Vec<&str> = prepared + .iter() + .map(|message| message.role.as_str()) + .collect(); + + assert_eq!(roles, vec!["user", "assistant"]); + } + + #[test] + fn auto_memory_persistence_excludes_system_messages() { + let all_messages = vec![ + text_message("system", "instructions"), + text_message("user", "hello"), + text_message("assistant", "hi"), + text_message("system", "duplicate-instructions"), + text_message("user", "follow-up"), + ]; + + let persisted = prepare_auto_memory_messages_for_persistence(&all_messages, 10); + let roles: Vec<&str> = persisted + .iter() + .map(|message| message.role.as_str()) + .collect(); + + assert_eq!(roles, vec!["user", "assistant", "user"]); + } +} + /// Handle credentials check mode - check credentials without making API calls async fn handle_credentials_check(provider: &ProviderWithResource) -> Result, Error> { let result = match &provider.kind { diff --git a/integration_tests/ai_agent_tests/test_memory.py b/integration_tests/ai_agent_tests/test_memory.py index cf51c20a1d..4948c57440 100644 --- a/integration_tests/ai_agent_tests/test_memory.py +++ b/integration_tests/ai_agent_tests/test_memory.py @@ -4,17 +4,56 @@ Memory tests for AI agents. Tests that AI agents correctly handle conversation memory/history. """ +import os import pytest import uuid from typing import Any from .conftest import AIAgentTestClient, create_ai_agent_flow -from .providers import ALL_PROVIDERS, get_provider_ids +from .providers import ( + ALL_PROVIDERS, + BEDROCK_API_KEY, + BEDROCK_ENV, + BEDROCK_IAM, + BEDROCK_IAM_SESSION, + get_provider_ids, +) class TestMemory: """Test AI agent memory functionality.""" + @staticmethod + def _pick_bedrock_provider() -> dict[str, Any]: + if os.environ.get("BEDROCK_API_KEY"): + return BEDROCK_API_KEY + + if ( + os.environ.get("BEDROCK_IAM_ACCESS_KEY_ID") + and os.environ.get("BEDROCK_IAM_SECRET_ACCESS_KEY") + ): + return BEDROCK_IAM + + if ( + os.environ.get("BEDROCK_SESSION_ACCESS_KEY_ID") + and os.environ.get("BEDROCK_SESSION_SECRET_ACCESS_KEY") + and os.environ.get("BEDROCK_SESSION_TOKEN") + ): + return BEDROCK_IAM_SESSION + + aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID") + aws_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY") + aws_session_token = os.environ.get("AWS_SESSION_TOKEN") + + if aws_access_key_id and aws_secret_access_key: + if aws_access_key_id.startswith("ASIA") and not aws_session_token: + pytest.skip( + "AWS_SESSION_TOKEN required for Bedrock env fallback when AWS_ACCESS_KEY_ID is temporary (ASIA...)" + ) + return BEDROCK_ENV + + pytest.skip("No Bedrock credentials available for the memory regression test") + @pytest.mark.parametrize( "provider_config", ALL_PROVIDERS, @@ -65,6 +104,61 @@ class TestMemory: print(f"Memory test passed for {provider_config['name']}") + def test_memory_keeps_single_system_prompt_across_turns( + self, + client: AIAgentTestClient, + setup_providers, + ): + """ + Test that auto memory persists conversation turns without persisting + duplicate system prompts across repeated runs. + """ + provider_config = self._pick_bedrock_provider() + + system_prompt = "You are a helpful assistant. Keep responses short." + flow_value = create_ai_agent_flow( + provider_input_transform=provider_config["input_transform"], + system_prompt=system_prompt, + context_length=6, + ) + + memory_id = str(uuid.uuid4()) + turns = [ + "Turn one: say hi.", + "Turn two: say hi again.", + "Turn three: say hi one more time.", + ] + + result = None + for turn in turns: + result = client.run_preview_flow( + flow_value=flow_value, + args={"user_message": turn}, + memory_id=memory_id, + ) + assert result is not None + assert "error" not in result, f"Turn failed for {provider_config['name']}: {result}" + + assert result is not None + messages = result.get("messages", []) + system_messages = [msg for msg in messages if msg.get("role") == "system"] + + assert len(system_messages) == 1, ( + f"Expected exactly one system message for {provider_config['name']}, got: {messages}" + ) + assert system_messages[0].get("content") == system_prompt + + user_messages = [ + msg.get("content") + for msg in messages + if msg.get("role") == "user" + ] + assert user_messages == turns, ( + f"Expected all prior user turns to remain in memory for {provider_config['name']}: {messages}" + ) + + print(f"System prompt dedupe test passed for {provider_config['name']}") + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])