Repository: strands-agents/agent-builder

Issue:
Knowledge Base Ingestion Error: ValidationException on dataSourceType mismatch
## Bug Description
When ingesting documents into a knowledge base, the following error occurs:

```
Error ingesting into knowledge base: An error occurred (ValidationException) when calling the IngestKnowledgeBaseDocuments operation: The dataSourceType that you specified in CUSTOM doesn't match the type of the data source specified in the request header. Check that both types match and retry your request.
```

## Related Information
This issue appears to have been solved in PR #8 (https://github.com/strands-agents/agent-builder/pull/8), but the error is still occurring.

## Environment
- GITHUB_TOKEN is available in the environment

## Example Log
```
list your tools
[...tool list output...]
Error ingesting into knowledge base: An error occurred (ValidationException) when calling the IngestKnowledgeBaseDocuments operation: The dataSourceType that you specified in CUSTOM doesn't match the type of the data source specified in the request header. Check that both types match and retry your request.
```

Hints from issue comments:
@jsamuel1: To reproduce, create your first datasource as a S3 ingestion datasource on the knowledgebase.  The second on as Custom (for use by Strands Agents).@cagataycali: Fixed in #8!@rbpotter-aws: I'm still seeing this behavior:

running:
strands-agents 0.2.1
strands-agents-tools 0.1.8

```
{'status': 'error',
 'content': [{'text': "❌ Error during store operation: An error occurred (ValidationException) when calling the IngestKnowledgeBaseDocuments operation: The dataSourceType that you specified in CUSTOM doesn't match the type of the data source specified in the request header. Check that both types match and retry your request."}],
 'toolUseId': 'tooluse_memory_552170041'}
```

<img width="1747" height="546" alt="Image" src="https://github.com/user-attachments/assets/6af4f770-f502-4da5-b83a-6c95eba0e9a1" />

**Question:** Is it assuming that the knowledge base only has 1 data source for this?

Solution patch:
diff --git a/README.md b/README.md
index 5280000..105b15c 100644
--- a/README.md
+++ b/README.md
@@ -177,7 +177,7 @@ strands --model-provider ollama --model-config '{"model_id": <ID>}'
 
 Strands Agent Builder is packaged with `bedrock` and `ollama`.
 
-If you have implemented a custom model provider ([instructions](<LINK>)) and would like to use it with strands, create a python module under the directory "$CWD/.models" and expose an `instance` function that returns an instance of your provider. As an example, assume you have:
+If you have implemented a custom model provider ([instructions](https://strandsagents.com/latest/user-guide/concepts/model-providers/custom_model_provider/)) and would like to use it with strands, create a python module under the directory "$CWD/.models" and expose an `instance` function that returns an instance of your provider. As an example, assume you have:
 
 ```bash
 $ cat ./.models/custom_model.py
diff --git a/pyproject.toml b/pyproject.toml
index 7fb0456..bb3e0f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,7 +49,7 @@ dev = [
     "commitizen>=4.4.0,<5.0.0",
     "hatch>=1.0.0,<2.0.0",
     "mypy>=1.0.0,<2.0.0",
-    "pre-commit>=2.20.0,<3.0.0",
+    "pre-commit>=3.2.0,<4.2.0",
     "ruff>=0.4.4,<1.0.0",
 ]
 docs = [
diff --git a/src/strands_agents_builder/utils/kb_utils.py b/src/strands_agents_builder/utils/kb_utils.py
index 6aae46b..38b1da2 100644
--- a/src/strands_agents_builder/utils/kb_utils.py
+++ b/src/strands_agents_builder/utils/kb_utils.py
@@ -8,7 +8,7 @@ from pathlib import Path
 
 def store_conversation_in_kb(agent, user_input, response=None, knowledge_base_id=None):
     """
-    Store conversation between user and assistant in knowledge base
+    Store conversation between user and assistant in knowledge base asynchronously
 
     Args:
         agent: Strands agent instance
@@ -75,11 +75,12 @@ def store_conversation_in_kb(agent, user_input, response=None, knowledge_base_id
         if len(user_input) > max_title_length:
             conversation_title += "..."
 
-        # Store in knowledge base
+        # Store in knowledge base asynchronously with record_direct_tool_call=False
         agent.tool.store_in_kb(
             content=conversation_content,
             title=conversation_title,
             knowledge_base_id=knowledge_base_id,
+            record_direct_tool_call=False,
         )
     except Exception as e:
         print(f"Error storing conversation in knowledge base: {str(e)}")
diff --git a/src/strands_agents_builder/utils/model_utils.py b/src/strands_agents_builder/utils/model_utils.py
index d1482e1..fd67a7c 100644
--- a/src/strands_agents_builder/utils/model_utils.py
+++ b/src/strands_agents_builder/utils/model_utils.py
@@ -11,7 +11,7 @@ from strands.types.models import Model
 
 # Default model configuration
 DEFAULT_MODEL_CONFIG = {
-    "model_id": "us.anthropic.claude-sonnet-4-20250514-v1:0",
+    "model_id": os.getenv("STRANDS_MODEL_ID", "us.anthropic.claude-sonnet-4-20250514-v1:0"),
     "max_tokens": int(os.getenv("STRANDS_MAX_TOKENS", "32768")),
     "boto_client_config": Config(
         read_timeout=900,
@@ -20,10 +20,12 @@ DEFAULT_MODEL_CONFIG = {
     ),
     "additional_request_fields": {
         "thinking": {
-            "type": "enabled",
+            "type": os.getenv("STRANDS_THINKING_TYPE", "enabled"),
             "budget_tokens": int(os.getenv("STRANDS_BUDGET_TOKENS", "2048")),
         }
     },
+    "cache_tools": os.getenv("STRANDS_CACHE_TOOLS", "default"),
+    "cache_prompt": os.getenv("STRANDS_CACHE_PROMPT", "default"),
 }
 
 
diff --git a/tools/store_in_kb.py b/tools/store_in_kb.py
index e400c88..14a10e6 100644
--- a/tools/store_in_kb.py
+++ b/tools/store_in_kb.py
@@ -1,12 +1,14 @@
 """
-Tool for storing data in Bedrock Knowledge Base
+Tool for storing data in Bedrock Knowledge Base asynchronously
 """
 
 import json
 import logging
 import os
+import threading
 import time
 import uuid
+from typing import Any, Dict
 
 import boto3
 from strands import tool
@@ -15,43 +17,30 @@ from strands import tool
 logger = logging.getLogger(__name__)
 
 
-@tool
-def store_in_kb(content: str, title: str = None, knowledge_base_id: str = None) -> dict:
+def _store_in_kb_background(content: str, title: str, kb_id: str, region_name: str) -> None:
     """
-    Store content in a Bedrock Knowledge Base using real-time ingestion.
+    Background worker function that performs the actual KB storage.
 
-    Args:
-        content: The text content to store in the knowledge base.
-        title: Optional title for the content. If not provided, a timestamp will be used.
-        knowledge_base_id: Optional knowledge base ID. If not provided, will use the KNOWLEDGE_BASE_ID env variable.
+    This runs in a separate thread and handles all the KB operations.
+    All validation is done in the main thread before calling this function.
 
-    Returns:
-        A dictionary containing the result of the operation.
+    Args:
+        content: The text content to store
+        title: The title for the content
+        kb_id: The knowledge base ID
+        region_name: The AWS region to use
     """
-    # Get knowledge base ID from parameter or environment variable
-    kb_id = knowledge_base_id or os.getenv("KNOWLEDGE_BASE_ID")
-
-    # Validate required inputs
-    if not kb_id:
-        return {
-            "status": "error",
-            "content": [
-                {"text": "❌ No knowledge base ID provided or found in environment variables KNOWLEDGE_BASE_ID"}
-            ],
-        }
-
-    if not content or not content.strip():
-        return {"status": "error", "content": [{"text": "❌ Content cannot be empty"}]}
-
     try:
         # Generate document ID with timestamp for traceability
         timestamp = time.strftime("%Y%m%d_%H%M%S")
         doc_id = f"memory_{timestamp}_{str(uuid.uuid4())[:8]}"
 
-        # Create a document title if not provided
-        doc_title = title or f"Strands Memory {timestamp}"
-
-        region_name = os.getenv("AWS_REGION", "us-west-2")
+        # Package content with title for better organization
+        content_with_metadata = {
+            "title": title,
+            "action": "store",
+            "content": content,
+        }
 
         # Initialize Bedrock agent client
         bedrock_agent_client = boto3.client("bedrock-agent", region_name=region_name)
@@ -60,61 +49,121 @@ def store_in_kb(content: str, title: str = None, knowledge_base_id: str = None)
         data_sources = bedrock_agent_client.list_data_sources(knowledgeBaseId=kb_id)
 
         if not data_sources.get("dataSourceSummaries"):
-            return {
-                "status": "error",
-                "content": [{"text": f"❌ No data sources found for knowledge base {kb_id}, region {region_name}."}],
-            }
-
-        data_source_id = data_sources["dataSourceSummaries"][0]["dataSourceId"]
-
-        # Package content with title for better organization
-        content_with_metadata = {
-            "title": doc_title,
-            "action": "store",
-            "content": content,
-        }
-
-        # Prepare document for ingestion
-        ingest_request = {
-            "knowledgeBaseId": kb_id,
-            "dataSourceId": data_source_id,
-            "documents": [
-                {
-                    "content": {
-                        "dataSourceType": "CUSTOM",
-                        "custom": {
-                            "customDocumentIdentifier": {"id": doc_id},
-                            "inlineContent": {
-                                "textContent": {"data": json.dumps(content_with_metadata)},
-                                "type": "TEXT",
+            logger.error(f"No data sources found for knowledge base {kb_id}, region {region_name}.")
+            return
+
+        # Look for a CUSTOM data source type first, as it's required for inline content ingestion
+        data_source_id = None
+        source_type = None
+
+        for ds in data_sources["dataSourceSummaries"]:
+            # Get the data source details to check its type
+            ds_detail = bedrock_agent_client.get_data_source(knowledgeBaseId=kb_id, dataSourceId=ds["dataSourceId"])
+
+            # Check if this is a CUSTOM type data source
+            if ds_detail["dataSource"]["dataSourceConfiguration"]["type"] == "CUSTOM":
+                data_source_id = ds["dataSourceId"]
+                source_type = "CUSTOM"
+                logger.debug(f"Found CUSTOM data source: {data_source_id}")
+                break
+
+        # If no CUSTOM data source found, use the first available one but log a warning
+        if not data_source_id and data_sources["dataSourceSummaries"]:
+            data_source_id = data_sources["dataSourceSummaries"][0]["dataSourceId"]
+            ds_detail = bedrock_agent_client.get_data_source(knowledgeBaseId=kb_id, dataSourceId=data_source_id)
+            source_type = ds_detail["dataSource"]["dataSourceConfiguration"]["type"]
+            logger.debug(f"No CUSTOM data source found. Using {source_type} data source: {data_source_id}")
+
+        if not data_source_id:
+            logger.error(f"No suitable data source found for knowledge base {kb_id}.")
+            return
+
+        # Prepare document for ingestion based on the data source type
+        if source_type == "CUSTOM":
+            ingest_request = {
+                "knowledgeBaseId": kb_id,
+                "dataSourceId": data_source_id,
+                "documents": [
+                    {
+                        "content": {
+                            "dataSourceType": "CUSTOM",
+                            "custom": {
+                                "customDocumentIdentifier": {"id": doc_id},
+                                "inlineContent": {
+                                    "textContent": {"data": json.dumps(content_with_metadata)},
+                                    "type": "TEXT",
+                                },
+                                "sourceType": "IN_LINE",
                             },
-                            "sourceType": "IN_LINE",
-                        },
+                        }
                     }
-                }
-            ],
-        }
+                ],
+            }
+        elif source_type == "S3":
+            # S3 source types need a different ingestion approach
+            logger.error("S3 data source type is not supported for direct ingestion with this tool.")
+            return
+        else:
+            logger.error(f"Unsupported data source type: {source_type}")
+            return
 
         # Ingest document into knowledge base
         _ = bedrock_agent_client.ingest_knowledge_base_documents(**ingest_request)
 
         # Log success
-        logger.debug(f"Successfully ingested document into knowledge base {kb_id}: {doc_id}")
-
-        # Return in standard Strands format
-        return {
-            "status": "success",
-            "content": [
-                {"text": "✅ Successfully stored content in knowledge base:"},
-                {"text": f"📝 Title: {doc_title}"},
-                {"text": f"🔑 Document ID: {doc_id}"},
-                {"text": f"🗄️ Knowledge Base ID: {kb_id}"},
-            ],
-        }
+        logger.info(f"Successfully ingested document into knowledge base {kb_id}: {doc_id}")
 
     except Exception as e:
         logger.error(f"Error ingesting into knowledge base: {str(e)}")
+
+
+@tool
+def store_in_kb(content: str, title: str = None, knowledge_base_id: str = None) -> Dict[str, Any]:
+    """
+    Store content in a Bedrock Knowledge Base using real-time ingestion.
+
+    This version runs asynchronously in a background thread and returns immediately.
+
+    Args:
+        content: The text content to store in the knowledge base.
+        title: Optional title for the content. If not provided, a timestamp will be used.
+        knowledge_base_id: Optional knowledge base ID. If not provided, will use the STRANDS_KNOWLEDGE_BASE_ID env.
+
+    Returns:
+        A dictionary containing the result of the operation.
+    """
+    # All validation done in main thread before spawning background thread
+
+    # Validate content first
+    if not content or not content.strip():
+        return {"status": "error", "content": [{"text": "❌ Content cannot be empty"}]}
+
+    # Resolve and validate knowledge base ID early (addresses environment variable race condition)
+    kb_id = knowledge_base_id or os.getenv("STRANDS_KNOWLEDGE_BASE_ID")
+    if not kb_id:
         return {
             "status": "error",
-            "content": [{"text": f"❌ Failed to store content in knowledge base: {str(e)}"}],
+            "content": [
+                {"text": "❌ No knowledge base ID provided or found in environment variables STRANDS_KNOWLEDGE_BASE_ID"}
+            ],
         }
+
+    region_name = os.getenv("AWS_REGION", "us-west-2")
+
+    doc_title = title or f"Strands Memory {time.strftime('%Y%m%d_%H%M%S')}"
+
+    thread = threading.Thread(
+        target=_store_in_kb_background, args=(content, doc_title, kb_id, region_name), daemon=True
+    )
+    thread.start()
+
+    # Return immediately with status
+    return {
+        "status": "success",
+        "content": [
+            {"text": "✅ Started background task to store content in knowledge base:"},
+            {"text": f"📝 Title: {doc_title}"},
+            {"text": f"🗄️ Knowledge Base ID: {kb_id}"},
+            {"text": "⏱️ Processing in background..."},
+        ],
+    }


Test patch:
diff --git a/tests/tools/test_store_in_kb.py b/tests/tools/test_store_in_kb.py
index 49da9e3..2ba8e47 100644
--- a/tests/tools/test_store_in_kb.py
+++ b/tests/tools/test_store_in_kb.py
@@ -7,284 +7,445 @@ import json
 import os
 from unittest import mock
 
-from tools.store_in_kb import store_in_kb
+import pytest
+
+from tools.store_in_kb import _store_in_kb_background, store_in_kb
+
+
+@pytest.fixture(autouse=True)
+def mock_environment():
+    """Mock environment variables for all tests using proper mocking pattern"""
+    mock_env = {}
+    with mock.patch.object(os, "environ", mock_env):
+        yield mock_env
+
+
+@pytest.fixture
+def mock_boto_client():
+    """Mock boto3.client for tests that need it"""
+    with mock.patch("boto3.client") as mock_client:
+        yield mock_client
+
+
+@pytest.fixture
+def mock_logger():
+    """Mock logger for tests that need it"""
+    with mock.patch("tools.store_in_kb.logger") as mock_log:
+        yield mock_log
+
+
+@pytest.fixture
+def mock_thread():
+    """Mock threading.Thread for tests that need it"""
+    with mock.patch("threading.Thread") as mock_thread_class:
+        yield mock_thread_class
 
 
 class TestStoreInKbTool:
     """Test cases for the store_in_kb tool"""
 
-    def test_missing_kb_id(self):
+    def test_missing_kb_id(self, mock_environment):
         """Test storing without knowledge base ID"""
-        with mock.patch.dict(os.environ, {}), mock.patch("boto3.client") as mock_boto:
-            # Force boto client to raise an exception when no KB ID is provided
-            mock_boto.side_effect = Exception("No knowledge base ID provided")
-
-            result = store_in_kb(content="Test content")
+        # Environment is already clean from fixture
+        # Call the function with no KB ID parameter
+        result = store_in_kb(content="Test content")
 
-            assert result["status"] == "error"
-            assert "No knowledge base ID provided" in result["content"][0]["text"]
+        # The implementation should return an error when KB ID is missing
+        assert result["status"] == "error"
+        assert "No knowledge base ID" in result["content"][0]["text"]
 
-    def test_empty_content(self):
+    def test_empty_content(self, mock_environment):
         """Test storing empty content"""
-        with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}):
-            result = store_in_kb(content="   ")  # Only whitespace
+        # Set a test KB ID in the environment
+        mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "test-kb-id"
 
-            assert result["status"] == "error"
-            assert "Content cannot be empty" in result["content"][0]["text"]
+        # Call with empty content
+        result = store_in_kb(content="   ")  # Only whitespace
 
-    def test_successful_storage_env_kb(self):
+        # Should get an error
+        assert result["status"] == "error"
+        assert "Content cannot be empty" in result["content"][0]["text"]
+
+    def test_successful_storage_env_kb(self, mock_environment, mock_thread):
         """Test successful storage using environment variable KB ID"""
-        with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "env-kb-id"}), mock.patch("boto3.client") as mock_boto:
-            # Mock the list_data_sources response
-            mock_boto.return_value.list_data_sources.return_value = {
-                "dataSourceSummaries": [{"dataSourceId": "test-data-source"}]
-            }
+        # Set the env var
+        mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "env-kb-id"
 
-            # Call the function
-            result = store_in_kb(content="Test content")
+        # Call the function
+        result = store_in_kb(content="Test content")
 
-            # Verify boto3 client was created correctly
-            mock_boto.assert_called_once_with("bedrock-agent", region_name=mock.ANY)
+        # Verify the thread was started
+        mock_thread.assert_called_once()
 
-            # Verify list_data_sources was called with correct KB ID
-            mock_boto.return_value.list_data_sources.assert_called_once_with(knowledgeBaseId="env-kb-id")
+        # Verify thread was initialized with daemon=True and the target function
+        kwargs = mock_thread.call_args.kwargs
+        assert kwargs.get("daemon") is True
+        assert callable(kwargs.get("target"))
 
-            # Verify ingest_knowledge_base_documents was called with correct parameters
-            mock_boto.return_value.ingest_knowledge_base_documents.assert_called_once()
-            ingest_args = mock_boto.return_value.ingest_knowledge_base_documents.call_args.kwargs
+        # Verify the response
+        assert result["status"] == "success"
+        assert "Started background task" in result["content"][0]["text"]
+        assert "Knowledge Base ID: env-kb-id" in result["content"][2]["text"]
 
-            assert ingest_args["knowledgeBaseId"] == "env-kb-id"
-            assert ingest_args["dataSourceId"] == "test-data-source"
-            assert len(ingest_args["documents"]) == 1
+    def test_successful_storage_param_kb(self, mock_environment, mock_thread):
+        """Test successful storage using parameter KB ID"""
+        # Call the function with explicit KB ID
+        result = store_in_kb(content="Test content", title="Test Title", knowledge_base_id="param-kb-id")
 
-            # Check content of document
-            doc_content = ingest_args["documents"][0]["content"]["custom"]["inlineContent"]["textContent"]["data"]
-            # Parse JSON content
-            content_obj = json.loads(doc_content)
-            assert content_obj["content"] == "Test content"
-            assert "title" in content_obj
+        # Verify thread was started with the right parameters
+        mock_thread.assert_called_once()
 
-            # Verify success response
-            assert result["status"] == "success"
-            assert "Successfully stored content" in result["content"][0]["text"]
+        # Check that the response shows the KB ID and title
+        assert result["status"] == "success"
+        assert "Title: Test Title" in result["content"][1]["text"]
+        assert "Knowledge Base ID: param-kb-id" in result["content"][2]["text"]
 
-    def test_successful_storage_param_kb(self):
-        """Test successful storage using parameter KB ID"""
-        with mock.patch("boto3.client") as mock_boto:
-            # Mock the list_data_sources response
-            mock_boto.return_value.list_data_sources.return_value = {
-                "dataSourceSummaries": [{"dataSourceId": "test-data-source"}]
-            }
-
-            # Call the function with explicit KB ID
-            _ = store_in_kb(content="Test content", title="Test Title", knowledge_base_id="param-kb-id")
-
-            # Verify list_data_sources was called with correct KB ID
-            mock_boto.return_value.list_data_sources.assert_called_once_with(knowledgeBaseId="param-kb-id")
-
-            # Verify ingest_knowledge_base_documents was called with correct parameters
-            ingest_args = mock_boto.return_value.ingest_knowledge_base_documents.call_args.kwargs
-            assert ingest_args["knowledgeBaseId"] == "param-kb-id"
-
-            # Check content of document
-            doc_content = ingest_args["documents"][0]["content"]["custom"]["inlineContent"]["textContent"]["data"]
-            content_obj = json.loads(doc_content)
-            assert content_obj["title"] == "Test Title"  # Custom title used
-
-    def test_no_data_sources(self):
-        """Test when no data sources are found"""
-        with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}), mock.patch("boto3.client") as mock_boto:
-            # Mock empty data sources response
-            mock_boto.return_value.list_data_sources.return_value = {"dataSourceSummaries": []}
-
-            # Call the function
-            result = store_in_kb(content="Test content")
-
-            # Verify error response
-            assert result["status"] == "error"
-            assert "No data sources found" in result["content"][0]["text"]
-
-    def test_boto3_exception(self):
-        """Test handling of boto3 exceptions"""
-        with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}), mock.patch("boto3.client") as mock_boto:
-            # Make list_data_sources raise an exception
-            mock_boto.return_value.list_data_sources.side_effect = Exception("AWS API Error")
-
-            # Call the function
-            result = store_in_kb(content="Test content")
-
-            # Verify error response
-            assert result["status"] == "error"
-            assert "Failed to store content" in result["content"][0]["text"]
-            assert "AWS API Error" in result["content"][0]["text"]
-
-    def test_custom_region(self):
+    def test_custom_region(self, mock_environment, mock_thread):
         """Test using a custom AWS region"""
-        with (
-            mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id", "AWS_REGION": "us-east-1"}),
-            mock.patch("boto3.client") as mock_boto,
-        ):
-            # Mock successful data sources response
-            mock_boto.return_value.list_data_sources.return_value = {
-                "dataSourceSummaries": [{"dataSourceId": "test-data-source"}]
-            }
-
-            # Call the function
-            store_in_kb(content="Test content")
-
-            # Verify boto3 client was created with correct region
-            mock_boto.assert_called_once_with("bedrock-agent", region_name="us-east-1")
+        # Setup environment variables
+        mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "test-kb-id"
+        mock_environment["AWS_REGION"] = "us-east-1"
 
+        # Call the function
+        result = store_in_kb(content="Test content")
 
-def test_store_in_kb_custom_region():
-    """Test store_in_kb with custom AWS region from environment variable"""
+        # Verify thread was started
+        mock_thread.assert_called_once()
 
-    # Mock responses
-    mock_list_response = {"dataSourceSummaries": [{"dataSourceId": "test-data-source"}]}
-    mock_ingest_response = {"status": "success"}
+        # We can't directly test the region here, but can verify success response
+        assert result["status"] == "success"
 
-    # Set up mocks
-    mock_client = mock.MagicMock()
-    mock_client.list_data_sources.return_value = mock_list_response
-    mock_client.ingest_knowledge_base_documents.return_value = mock_ingest_response
+    # NEW TESTS FOR BACKGROUND FUNCTION
+    def test_background_function_direct(self, mock_logger, mock_boto_client):
+        """Test direct execution of the background function with all successful paths"""
+        # Setup mock boto3 client and responses
+        mock_bedrock = mock.MagicMock()
+        mock_boto_client.return_value = mock_bedrock
 
-    with (
-        mock.patch.dict(os.environ, {"AWS_REGION": "ap-southeast-2"}),
-        mock.patch("boto3.client", return_value=mock_client) as mock_boto3,
-    ):
-        result = store_in_kb(content="Test content", title="Test title", knowledge_base_id="test-kb")
+        # Setup data source response
+        mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "test-ds-id"}]}
+        mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "CUSTOM"}}}
+        mock_bedrock.ingest_knowledge_base_documents.return_value = {"jobId": "test-job"}
 
-        # Verify boto3 client was created with the right region
-        mock_boto3.assert_called_once_with("bedrock-agent", region_name="ap-southeast-2")
+        # Call background function directly with new signature
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
 
-        # Verify success response
-        assert result["status"] == "success"
+        # Verify correct API calls were made
+        mock_bedrock.list_data_sources.assert_called_once_with(knowledgeBaseId="test-kb-id")
+        mock_bedrock.get_data_source.assert_called_once()
+        mock_bedrock.ingest_knowledge_base_documents.assert_called_once()
 
+        # Verify success was logged
+        mock_logger.info.assert_called_once()
+        assert "Successfully ingested" in mock_logger.info.call_args[0][0]
+        assert "test-kb-id" in mock_logger.info.call_args[0][0]
 
-def test_store_in_kb_exception_path():
-    """Test exception handling in store_in_kb"""
+    def test_background_no_data_sources(self, mock_logger, mock_boto_client):
+        """Test background function without any data sources"""
+        # Setup mock
+        mock_bedrock = mock.MagicMock()
+        mock_boto_client.return_value = mock_bedrock
 
-    # Mock client to raise exception
-    mock_client = mock.MagicMock()
-    mock_client.list_data_sources.side_effect = Exception("Test boto3 error")
+        # Return empty data sources
+        mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": []}
 
-    with mock.patch("boto3.client", return_value=mock_client):
-        result = store_in_kb(content="Test content", title="Test title", knowledge_base_id="test-kb")
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
 
-        # Verify error response
-        assert result["status"] == "error"
-        assert "Test boto3 error" in result["content"][0]["text"]
+        # Should not call get_data_source or ingest
+        mock_bedrock.get_data_source.assert_not_called()
+        mock_bedrock.ingest_knowledge_base_documents.assert_not_called()
 
+        # Should log error
+        mock_logger.error.assert_called_once()
+        error_message = mock_logger.error.call_args[0][0]
+        assert "No data sources found" in error_message
+
+    def test_background_custom_data_source_selection(self, mock_logger, mock_boto_client):
+        """Test background function with multiple data sources and CUSTOM selection"""
+        # Setup mock boto3 client and responses
+        mock_bedrock = mock.MagicMock()
+        mock_boto_client.return_value = mock_bedrock
+
+        # Setup multiple data sources with different types
+        mock_bedrock.list_data_sources.return_value = {
+            "dataSourceSummaries": [{"dataSourceId": "s3-source-id"}, {"dataSourceId": "custom-source-id"}]
+        }
 
-def test_no_data_sources_empty_dict():
-    """Test when data sources response is empty dict (no dataSourceSummaries key)"""
-    with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}), mock.patch("boto3.client") as mock_boto:
-        # Mock empty response (no dataSourceSummaries key)
-        mock_boto.return_value.list_data_sources.return_value = {}
+        # Mock get_data_source to return different types based on dataSourceId
+        def mock_get_data_source(**kwargs):
+            if kwargs["dataSourceId"] == "s3-source-id":
+                return {"dataSource": {"dataSourceConfiguration": {"type": "S3"}}}
+            elif kwargs["dataSourceId"] == "custom-source-id":
+                return {"dataSource": {"dataSourceConfiguration": {"type": "CUSTOM"}}}
+            return {}
+
+        mock_bedrock.get_data_source.side_effect = mock_get_data_source
+
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
+
+        # Verify CUSTOM type was selected and ingest was called
+        mock_bedrock.ingest_knowledge_base_documents.assert_called_once()
+        request = mock_bedrock.ingest_knowledge_base_documents.call_args[1]
+        assert request["dataSourceId"] == "custom-source-id"
+
+        # Verify success was logged
+        mock_logger.info.assert_called_once()
+
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_fallback_to_non_custom_source(self, mock_logger, mock_boto):
+        """Test background function falling back to first data source when no CUSTOM source exists"""
+        # Setup mock
+        mock_bedrock = mock.MagicMock()
+        mock_boto.return_value = mock_bedrock
+
+        # Only provide S3 source
+        mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "s3-source-id"}]}
+        mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "S3"}}}
+
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
+
+        # For S3 source, should log error and not call ingest
+        mock_bedrock.ingest_knowledge_base_documents.assert_not_called()
+        mock_logger.error.assert_called()
+        error_message = mock_logger.error.call_args[0][0]
+        assert "S3 data source type is not supported" in error_message
+
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_unknown_source_type(self, mock_logger, mock_boto):
+        """Test background function with unknown data source type"""
+        # Setup mock
+        mock_bedrock = mock.MagicMock()
+        mock_boto.return_value = mock_bedrock
+
+        # Provide unknown source type
+        mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "unknown-source-id"}]}
+        mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "UNKNOWN"}}}
+
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
+
+        # Should log error about unsupported type and not call ingest
+        mock_bedrock.ingest_knowledge_base_documents.assert_not_called()
+        mock_logger.error.assert_called()
+        error_message = mock_logger.error.call_args[0][0]
+        assert "Unsupported data source type" in error_message
+
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_no_datasourcesummaries_key(self, mock_logger, mock_boto):
+        """Test background function with missing dataSourceSummaries key"""
+        # Setup mock
+        mock_bedrock = mock.MagicMock()
+        mock_boto.return_value = mock_bedrock
+
+        # Return response without dataSourceSummaries key
+        mock_bedrock.list_data_sources.return_value = {}
+
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
+
+        # Should not call get_data_source or ingest
+        mock_bedrock.get_data_source.assert_not_called()
+        mock_bedrock.ingest_knowledge_base_documents.assert_not_called()
+
+        # Should log error
+        mock_logger.error.assert_called_once()
+        error_message = mock_logger.error.call_args[0][0]
+        assert "No data sources found" in error_message
 
-        # Call the function
-        result = store_in_kb(content="Test content")
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_boto3_exception(self, mock_logger, mock_boto):
+        """Test background function exception handling"""
+        # Make boto3.client raise exception
+        mock_boto.side_effect = Exception("AWS service error")
 
-        # Verify error response
-        assert result["status"] == "error"
-        assert "No data sources found" in result["content"][0]["text"]
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
 
+        # Should log error
+        mock_logger.error.assert_called_once()
+        error_message = mock_logger.error.call_args[0][0]
+        assert "Error ingesting into knowledge base" in error_message
+        assert "AWS service error" in error_message
 
-def test_data_source_id_selection():
-    """Test that the first data source ID is selected when multiple are available"""
-    with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}), mock.patch("boto3.client") as mock_boto:
-        # Mock multiple data sources
-        mock_boto.return_value.list_data_sources.return_value = {
-            "dataSourceSummaries": [{"dataSourceId": "data-source-1"}, {"dataSourceId": "data-source-2"}]
-        }
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_list_data_sources_exception(self, mock_logger, mock_boto):
+        """Test background function with exception during list_data_sources call"""
+        # Setup mock boto3 client
+        mock_bedrock = mock.MagicMock()
+        mock_boto.return_value = mock_bedrock
 
-        # Call the function
-        result = store_in_kb(content="Test content")
+        # Make list_data_sources raise exception
+        mock_bedrock.list_data_sources.side_effect = Exception("List data sources failed")
 
-        # Verify the first data source was used
-        ingest_args = mock_boto.return_value.ingest_knowledge_base_documents.call_args.kwargs
-        assert ingest_args["dataSourceId"] == "data-source-1"
-        assert result["status"] == "success"
+        # Call background function
+        _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
 
+        # Should log error
+        mock_logger.error.assert_called_once()
+        error_message = mock_logger.error.call_args[0][0]
+        assert "Error ingesting into knowledge base" in error_message
+        assert "List data sources failed" in error_message
+
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_auto_title(self, mock_logger, mock_boto):
+        """Test background function with auto-generated title"""
+        # Setup mocks for successful path
+        mock_bedrock = mock.MagicMock()
+        mock_boto.return_value = mock_bedrock
+        mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "test-ds-id"}]}
+        mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "CUSTOM"}}}
+
+        # Call background function with auto-generated title (simulating what main function does)
+        _store_in_kb_background("Test content", "Strands Memory 20241225_120000", "test-kb-id", "us-west-2")
+
+        # Verify ingest was called with correct parameters
+        mock_bedrock.ingest_knowledge_base_documents.assert_called_once()
+        request = mock_bedrock.ingest_knowledge_base_documents.call_args[1]
+
+        # Check the JSON data in the request includes an auto-generated title
+        text_content = request["documents"][0]["content"]["custom"]["inlineContent"]["textContent"]["data"]
+        data = json.loads(text_content)
+        assert "title" in data
+        assert "Strands Memory" in data["title"]
+
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_custom_title(self, mock_logger, mock_boto):
+        """Test background function with custom title"""
+        # Setup mocks for successful path
+        mock_bedrock = mock.MagicMock()
+        mock_boto.return_value = mock_bedrock
+        mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "test-ds-id"}]}
+        mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "CUSTOM"}}}
+
+        # Call background function with custom title
+        _store_in_kb_background("Test content", "My Custom Title", "test-kb-id", "us-west-2")
+
+        # Verify ingest was called with correct parameters
+        mock_bedrock.ingest_knowledge_base_documents.assert_called_once()
+        request = mock_bedrock.ingest_knowledge_base_documents.call_args[1]
+
+        # Check the JSON data in the request includes the provided title
+        text_content = request["documents"][0]["content"]["custom"]["inlineContent"]["textContent"]["data"]
+        data = json.loads(text_content)
+        assert data["title"] == "My Custom Title"
+
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_env_region(self, mock_logger, mock_boto):
+        """Test background function with region from environment"""
+        # Set environment region
+        with mock.patch.dict(os.environ, {"AWS_REGION": "us-east-1"}):
+            # Setup successful mocks
+            mock_bedrock = mock.MagicMock()
+            mock_boto.return_value = mock_bedrock
+            mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "test-ds-id"}]}
+            mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "CUSTOM"}}}
+
+            # Call background function
+            _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-east-1")
+
+            # Verify boto3 client was created with the right region
+            mock_boto.assert_called_once_with("bedrock-agent", region_name="us-east-1")
 
-def test_logging_functionality():
-    """Test that logging functionality works properly"""
-    with (
-        mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}),
-        mock.patch("boto3.client") as mock_boto,
-        mock.patch("tools.store_in_kb.logger") as mock_logger,
-    ):
-        # Mock successful data sources response
-        mock_boto.return_value.list_data_sources.return_value = {
-            "dataSourceSummaries": [{"dataSourceId": "test-data-source"}]
-        }
+    @mock.patch("boto3.client")
+    @mock.patch("tools.store_in_kb.logger")
+    def test_background_default_region(self, mock_logger, mock_boto):
+        """Test background function with default region when env not set"""
+        # Remove environment region if present
+        with mock.patch.dict(os.environ, {}, clear=True):
+            # Setup successful mocks
+            mock_bedrock = mock.MagicMock()
+            mock_boto.return_value = mock_bedrock
+            mock_bedrock.list_data_sources.return_value = {"dataSourceSummaries": [{"dataSourceId": "test-ds-id"}]}
+            mock_bedrock.get_data_source.return_value = {"dataSource": {"dataSourceConfiguration": {"type": "CUSTOM"}}}
 
-        # Call the function
-        store_in_kb(content="Test content")
+            # Call background function
+            _store_in_kb_background("Test content", "Test Title", "test-kb-id", "us-west-2")
 
-        # Verify logger was called
-        mock_logger.debug.assert_called_once()
-        # The log message should contain "Successfully ingested document"
-        assert "Successfully ingested document" in mock_logger.debug.call_args.args[0]
+            # Verify boto3 client was created with default region
+            mock_boto.assert_called_once_with("bedrock-agent", region_name="us-west-2")
 
 
-def test_exception_logging():
-    """Test that exceptions are properly logged"""
-    with (
-        mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}),
-        mock.patch("boto3.client") as mock_boto,
-        mock.patch("tools.store_in_kb.logger") as mock_logger,
-    ):
-        # Make the API call raise an exception
-        mock_boto.return_value.list_data_sources.side_effect = Exception("API error")
+def test_store_in_kb_custom_region(mock_environment, mock_thread):
+    """Test store_in_kb with custom AWS region from environment variable"""
+    # Set environment variables
+    mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "test-kb"
+    mock_environment["AWS_REGION"] = "ap-southeast-2"
 
-        # Call the function
-        store_in_kb(content="Test content")
+    # Call function
+    result = store_in_kb(content="Test content", title="Test title")
 
-        # Verify error was logged
-        mock_logger.error.assert_called_once()
-        # The log message should contain "Error ingesting"
-        assert "Error ingesting" in mock_logger.error.call_args.args[0]
+    # Verify thread was started
+    mock_thread.assert_called_once()
 
+    # Verify success response
+    assert result["status"] == "success"
+    # Verify title appears in response
+    assert "Title: Test title" in result["content"][1]["text"]
 
-def test_content_metadata_format():
-    """Test that the content metadata is properly formatted"""
-    with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}), mock.patch("boto3.client") as mock_boto:
-        # Mock successful data sources response
-        mock_boto.return_value.list_data_sources.return_value = {
-            "dataSourceSummaries": [{"dataSourceId": "test-data-source"}]
-        }
 
-        # Call the function with specific title
-        custom_title = "My Custom Title"
-        store_in_kb(content="Test content", title=custom_title)
+def test_store_in_kb_exception_path(mock_environment, mock_thread):
+    """Test exception handling in store_in_kb"""
+    # Set required environment
+    mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "test-kb"
 
-        # Get the ingest request document content
-        ingest_args = mock_boto.return_value.ingest_knowledge_base_documents.call_args.kwargs
-        doc_content_json = ingest_args["documents"][0]["content"]["custom"]["inlineContent"]["textContent"]["data"]
-        doc_content = json.loads(doc_content_json)
+    # Our implementation starts a thread regardless of potential errors
+    result = store_in_kb(content="Test content", title="Test title", knowledge_base_id="test-kb")
 
-        # Verify metadata format
-        assert doc_content["title"] == custom_title
-        assert doc_content["action"] == "store"
-        assert doc_content["content"] == "Test content"
+    # Verify thread was started
+    mock_thread.assert_called_once()
 
+    # Verify success response (main function succeeds even if background will fail)
+    assert result["status"] == "success"
 
-def test_auto_generated_title():
-    """Test that a title is auto-generated when not provided"""
-    with mock.patch.dict(os.environ, {"KNOWLEDGE_BASE_ID": "test-kb-id"}), mock.patch("boto3.client") as mock_boto:
-        # Mock successful data sources response
-        mock_boto.return_value.list_data_sources.return_value = {
-            "dataSourceSummaries": [{"dataSourceId": "test-data-source"}]
-        }
 
-        # Call the function without a title
-        store_in_kb(content="Test content")
+def test_data_source_id_selection(mock_environment, mock_thread):
+    """Test thread creation with the target function and proper args"""
+    mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "test-kb-id"
+
+    # Call the function
+    result = store_in_kb(content="Test content")
+
+    # Verify thread was initialized correctly
+    mock_thread.assert_called_once()
+
+    # Verify that args are passed to target function
+    assert "args" in mock_thread.call_args.kwargs
+    args = mock_thread.call_args.kwargs["args"]
+    assert len(args) == 4  # content, title, kb_id, region_name
+    assert args[0] == "Test content"  # First arg is content
+
+    # Verify thread was started
+    mock_thread.return_value.start.assert_called_once()
+
+    # Verify success response
+    assert result["status"] == "success"
+
+
+def test_auto_generated_title(mock_environment, mock_thread):
+    """Test that title is auto-generated when not provided"""
+    mock_environment["STRANDS_KNOWLEDGE_BASE_ID"] = "test-kb-id"
+
+    # Call the function without a title
+    result = store_in_kb(content="Test content")
 
-        # Get the ingest request document content
-        ingest_args = mock_boto.return_value.ingest_knowledge_base_documents.call_args.kwargs
-        doc_content_json = ingest_args["documents"][0]["content"]["custom"]["inlineContent"]["textContent"]["data"]
-        doc_content = json.loads(doc_content_json)
+    # Verify thread creation
+    mock_thread.assert_called_once()
+    args = mock_thread.call_args.kwargs["args"]
+    assert len(args) >= 2
+    assert args[0] == "Test content"  # First arg is content
+    # Second arg (title) should be auto-generated, not None
+    assert "Strands Memory" in args[1]  # Second arg should contain auto-generated title
 
-        # Verify a title was generated (should start with "Strands Memory")
-        assert doc_content["title"].startswith("Strands Memory")
+    # Title format in response should include "Strands Memory"
+    assert "Title: Strands Memory" in result["content"][1]["text"]
diff --git a/tests/utils/test_kb_utils.py b/tests/utils/test_kb_utils.py
index 8d85625..8960ba9 100644
--- a/tests/utils/test_kb_utils.py
+++ b/tests/utils/test_kb_utils.py
@@ -99,7 +99,10 @@ def test_store_conversation_in_kb():
     # Verify agent called store_in_kb with correct parameters
     expected_content = "User: Test query\n\nAssistant Reasoning: Test reasoning\n\nAssistant Response: Test response"
     mock_agent.tool.store_in_kb.assert_called_with(
-        content=expected_content, title="Conversation: Test query", knowledge_base_id="test-kb-id"
+        content=expected_content,
+        title="Conversation: Test query",
+        knowledge_base_id="test-kb-id",
+        record_direct_tool_call=False,
     )
 
 
@@ -118,7 +121,10 @@ def test_store_conversation_without_reasoning():
     # Verify agent called store_in_kb with correct parameters
     expected_content = "User: Test query\n\nAssistant: Test response"
     mock_agent.tool.store_in_kb.assert_called_with(
-        content=expected_content, title="Conversation: Test query", knowledge_base_id="test-kb-id"
+        content=expected_content,
+        title="Conversation: Test query",
+        knowledge_base_id="test-kb-id",
+        record_direct_tool_call=False,
     )
 
 
@@ -161,7 +167,10 @@ def test_store_conversation_direct_mode():
     # Verify agent called store_in_kb with just the user query
     expected_content = "User: Test query"
     mock_agent.tool.store_in_kb.assert_called_with(
-        content=expected_content, title="Conversation: Test query", knowledge_base_id="test-kb-id"
+        content=expected_content,
+        title="Conversation: Test query",
+        knowledge_base_id="test-kb-id",
+        record_direct_tool_call=False,
     )
 
 
@@ -233,4 +242,5 @@ def test_store_conversation_exception_in_response_str():
             content="User: Test query\n\nAssistant: ",
             title="Conversation: Test query",
             knowledge_base_id="test-kb-id",
+            record_direct_tool_call=False,
         )