fix(embeddings): honor aws_session_token in AWS Bedrock embeddings (#5566)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

mem0/configs/embeddings/base.py

@@ -39,6 +39,7 @@ class BaseEmbedderConfig(ABC):
         # AWS Bedrock specific
         aws_access_key_id: Optional[str] = None,
         aws_secret_access_key: Optional[str] = None,
+        aws_session_token: Optional[str] = None,
         aws_region: Optional[str] = None,
     ):
         """
@@ -106,5 +107,6 @@ class BaseEmbedderConfig(ABC):
         # AWS Bedrock specific
         self.aws_access_key_id = aws_access_key_id
         self.aws_secret_access_key = aws_secret_access_key
+        self.aws_session_token = aws_session_token
         self.aws_region = aws_region or os.environ.get("AWS_REGION") or "us-west-2"
 

mem0/embeddings/aws_bedrock.py

@@ -34,7 +34,12 @@ class AWSBedrockEmbedding(EmbeddingBase):
             aws_access_key = self.config.aws_access_key_id
         if hasattr(self.config, "aws_secret_access_key"):
             aws_secret_key = self.config.aws_secret_access_key
-        
+        # Honor a session token supplied via config (temporary credentials from
+        # STS / assume-role), falling back to the env var when unset. The LLM
+        # Bedrock provider already supports this; mirror it here.
+        if getattr(self.config, "aws_session_token", None):
+            aws_session_token = self.config.aws_session_token
+
         # AWS region is always set in config - see BaseEmbedderConfig
         aws_region = self.config.aws_region or "us-west-2"
 

tests/embeddings/test_aws_bedrock_embeddings.py

@@ -0,0 +1,60 @@
+from unittest.mock import Mock, patch
+
+import pytest
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.aws_bedrock import AWSBedrockEmbedding
+
+
+@pytest.fixture
+def mock_boto3_client():
+    with patch("mem0.embeddings.aws_bedrock.boto3.client") as mock_client:
+        mock_client.return_value = Mock()
+        yield mock_client
+
+
+def test_session_token_from_config_is_passed_to_client(mock_boto3_client):
+    """A session token supplied via config must reach the bedrock-runtime client.
+
+    Temporary credentials (STS / assume-role) require a session token; the LLM
+    Bedrock provider already honors it, so the embedding provider should too.
+    """
+    with patch("mem0.embeddings.aws_bedrock.os.environ", {}):
+        config = BaseEmbedderConfig(
+            model="amazon.titan-embed-text-v2:0",
+            aws_access_key_id="AKIA_TEST",
+            aws_secret_access_key="SECRET_TEST",
+            aws_session_token="SESSION_TOKEN_TEST",
+            aws_region="eu-central-1",
+        )
+        AWSBedrockEmbedding(config)
+
+    _, kwargs = mock_boto3_client.call_args
+    assert kwargs["aws_session_token"] == "SESSION_TOKEN_TEST"
+    assert kwargs["aws_access_key_id"] == "AKIA_TEST"
+    assert kwargs["aws_secret_access_key"] == "SECRET_TEST"
+    assert kwargs["region_name"] == "eu-central-1"
+
+
+def test_session_token_falls_back_to_env(mock_boto3_client):
+    """When config doesn't set a session token, the env var is still used."""
+    with patch.dict(
+        "mem0.embeddings.aws_bedrock.os.environ",
+        {"AWS_SESSION_TOKEN": "ENV_SESSION_TOKEN"},
+        clear=True,
+    ):
+        config = BaseEmbedderConfig(model="amazon.titan-embed-text-v2:0")
+        AWSBedrockEmbedding(config)
+
+    _, kwargs = mock_boto3_client.call_args
+    assert kwargs["aws_session_token"] == "ENV_SESSION_TOKEN"
+
+
+def test_no_session_token_passes_none(mock_boto3_client):
+    """Without a token in config or env, the client receives None (unchanged)."""
+    with patch("mem0.embeddings.aws_bedrock.os.environ", {}):
+        config = BaseEmbedderConfig(model="amazon.titan-embed-text-v2:0")
+        AWSBedrockEmbedding(config)
+
+    _, kwargs = mock_boto3_client.call_args
+    assert kwargs["aws_session_token"] is None