feat(logging): attach structured JSON logging with sessionId to root logger

Author: luyuzhe111

examples/strands_math_agent/basic_app.py

@@ -1,9 +1,14 @@
+import logging
+
 from bedrock_agentcore.runtime import BedrockAgentCoreApp
 from dotenv import load_dotenv
 from strands import Agent
 from strands.models import BedrockModel
 from strands_tools import calculator
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 app = BedrockAgentCoreApp()
 
 load_dotenv()
@@ -28,7 +33,7 @@ def invoke_agent(payload):
     """
     user_input = payload.get("prompt")
 
-    print("User input:", user_input)
+    logger.info("User input: %s", user_input)
 
     response = agent(user_input)
 

examples/strands_math_agent/rl_app.py

@@ -1,10 +1,15 @@
+import logging
+
 from reward import GSM8KReward
 from strands import Agent
 from strands.models.openai import OpenAIModel
 from strands_tools import calculator
 
 from agentcore_rl_toolkit import AgentCoreRLApp
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 app = AgentCoreRLApp()
 
 system_prompt = (
@@ -51,7 +56,7 @@ def invoke_agent(payload: dict):
     user_input = payload.get("prompt")
     answer = payload.get("answer")  # used for computing reward
 
-    print("User input:", user_input)
+    logger.info("User input: %s", user_input)
 
     response = agent(user_input)
 

examples/strands_migration_agent/reward.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import shutil
 import tempfile
@@ -7,6 +8,8 @@
 
 from agentcore_rl_toolkit import RewardFunction
 
+logger = logging.getLogger(__name__)
+
 
 class MigrationReward(RewardFunction):
     def __call__(
@@ -33,21 +36,21 @@ def __call__(
             shutil.copytree(repo_dir, temp_repo_dir)
 
             if self.eval_build_success(repo_dir=temp_repo_dir, require_maximal_migration=require_maximal_migration):
-                print("build succeeded!")
+                logger.info("build succeeded!")
                 reward += 0.5
             else:
-                print("build failed!")
+                logger.info("build failed!")
                 return reward
 
             if self.eval_test_equivalence(
                 repo_dir=temp_repo_dir,
                 original_num_tests=original_num_tests,
                 original_commit_id=original_commit_id,
             ):
-                print("test equivalence check passed!")
+                logger.info("test equivalence check passed!")
                 reward += 0.5
             else:
-                print("test equivalence check failed!")
+                logger.info("test equivalence check failed!")
 
         return reward
 

src/agentcore_rl_toolkit/app.py

@@ -9,6 +9,10 @@
 import boto3
 from bedrock_agentcore.runtime import BedrockAgentCoreApp
 
+from .logging import configure_logging
+
+logger = logging.getLogger(__name__)
+
 _S3_CONFIG_FIELDS = ("exp_id", "input_id", "s3_bucket")
 
 
@@ -36,6 +40,7 @@ def from_dict(cls, data: dict) -> "RolloutConfig":
 class AgentCoreRLApp(BedrockAgentCoreApp):
     def __init__(self):
         super().__init__()
+        configure_logging()
         self.s3_client = boto3.client("s3")
 
     def save_result(self, result: dict, rollout_config: dict, result_key: str, payload: dict = None):
@@ -67,7 +72,7 @@ def save_result(self, result: dict, rollout_config: dict, result_key: str, paylo
         try:
             config = RolloutConfig.from_dict(rollout_config)
         except ValueError as e:
-            logging.error(f"Invalid rollout configuration: {e}")
+            logger.error(f"Invalid rollout configuration: {e}")
             raise
 
         if "status_code" not in result:
@@ -93,9 +98,9 @@ def save_result(self, result: dict, rollout_config: dict, result_key: str, paylo
                 Body=json.dumps(result, indent=2),
                 ContentType="application/json",
             )
-            logging.info(f"Stored complete results at {result_key}")
+            logger.info(f"Stored complete results at {result_key}")
         except Exception as e:
-            logging.error(f"Failed to store results in S3: {e}")
+            logger.error(f"Failed to store results in S3: {e}")
             raise
 
     def rollout_entrypoint(self, func):
@@ -160,7 +165,7 @@ async def rollout_background_task(payload, context, result_key):
                         payload=payload,
                         result_key=result_key,
                     )
-                    logging.info(f"Result saved for function: {func.__name__}")
+                    logger.info(f"Result saved for function: {func.__name__}")
 
                 return result
 
@@ -181,9 +186,9 @@ async def rollout_background_task(payload, context, result_key):
                             payload=payload,
                             result_key=result_key,
                         )
-                        logging.error(f"Error result saved for function: {func.__name__}: {e}")
+                        logger.error(f"Error result saved for function: {func.__name__}: {e}")
                     except Exception:
-                        logging.error(f"Failed to save error result for function: {func.__name__}", exc_info=True)
+                        logger.error(f"Failed to save error result for function: {func.__name__}", exc_info=True)
                 raise
             finally:
                 # Complete the async task for logging and ping status

src/agentcore_rl_toolkit/logging.py

@@ -0,0 +1,58 @@
+import json
+import logging
+import traceback
+from datetime import datetime, timezone
+
+
+class CorrelatedFormatter(logging.Formatter):
+    """JSON formatter that injects sessionId and requestId from ACR request context."""
+
+    def format(self, record):
+        from bedrock_agentcore.runtime import BedrockAgentCoreContext
+
+        log_entry = {
+            "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z",
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "logger": record.name,
+        }
+
+        request_id = BedrockAgentCoreContext.get_request_id()
+        if request_id:
+            log_entry["requestId"] = request_id
+
+        session_id = BedrockAgentCoreContext.get_session_id()
+        if session_id:
+            log_entry["sessionId"] = session_id
+
+        if record.exc_info and record.exc_info[0]:
+            log_entry["errorType"] = record.exc_info[0].__name__
+            log_entry["errorMessage"] = str(record.exc_info[1])
+            log_entry["stackTrace"] = traceback.format_exception(*record.exc_info)
+
+        return json.dumps(log_entry, ensure_ascii=False)
+
+
+def configure_logging(level=logging.INFO):
+    """Attach CorrelatedFormatter to the root logger for automatic sessionId injection.
+
+    Idempotent — safe to call multiple times.
+    """
+    root = logging.getLogger()
+    if getattr(root, "_art_logging_configured", False):
+        return
+
+    root.handlers.clear()
+
+    handler = logging.StreamHandler()
+    handler.setFormatter(CorrelatedFormatter())
+    root.addHandler(handler)
+    # Use the more verbose level: honors user's basicConfig(level=DEBUG) even though we default to INFO.
+    # Root default is WARNING (30); min(30, 20) = INFO; min(10, 20) = DEBUG.
+    root.setLevel(min(root.level, level))
+
+    # The SDK's bedrock_agentcore.app logger already has its own JSON handler.
+    # Suppress propagation to avoid duplicate output.
+    logging.getLogger("bedrock_agentcore.app").propagate = False
+
+    root._art_logging_configured = True

tests/test_logging.py

@@ -0,0 +1,130 @@
+"""Tests for the structured logging module."""
+
+import json
+import logging
+from unittest.mock import patch
+
+import pytest
+
+from agentcore_rl_toolkit.logging import CorrelatedFormatter, configure_logging
+
+
+@pytest.fixture(autouse=True)
+def reset_root_logger():
+    """Reset root logger state before each test."""
+    root = logging.getLogger()
+    original_handlers = root.handlers[:]
+    original_level = root.level
+    if hasattr(root, "_art_logging_configured"):
+        delattr(root, "_art_logging_configured")
+    yield
+    root.handlers = original_handlers
+    root.setLevel(original_level)
+    if hasattr(root, "_art_logging_configured"):
+        delattr(root, "_art_logging_configured")
+
+
+class TestCorrelatedFormatter:
+    def test_outputs_valid_json(self):
+        formatter = CorrelatedFormatter()
+        record = logging.LogRecord(
+            name="test.logger",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="hello %s",
+            args=("world",),
+            exc_info=None,
+        )
+
+        output = formatter.format(record)
+        parsed = json.loads(output)
+
+        assert parsed["level"] == "INFO"
+        assert parsed["message"] == "hello world"
+        assert parsed["logger"] == "test.logger"
+        assert "timestamp" in parsed
+
+    @patch("bedrock_agentcore.runtime.BedrockAgentCoreContext.get_session_id", return_value="sess-123")
+    @patch("bedrock_agentcore.runtime.BedrockAgentCoreContext.get_request_id", return_value="req-456")
+    def test_includes_session_and_request_id(self, mock_req, mock_sess):
+        formatter = CorrelatedFormatter()
+        record = logging.LogRecord(
+            name="test", level=logging.INFO, pathname="", lineno=0, msg="msg", args=(), exc_info=None
+        )
+
+        parsed = json.loads(formatter.format(record))
+
+        assert parsed["sessionId"] == "sess-123"
+        assert parsed["requestId"] == "req-456"
+
+    @patch("bedrock_agentcore.runtime.BedrockAgentCoreContext.get_session_id", return_value=None)
+    @patch("bedrock_agentcore.runtime.BedrockAgentCoreContext.get_request_id", return_value=None)
+    def test_omits_ids_when_no_context(self, mock_req, mock_sess):
+        formatter = CorrelatedFormatter()
+        record = logging.LogRecord(
+            name="test", level=logging.INFO, pathname="", lineno=0, msg="msg", args=(), exc_info=None
+        )
+
+        parsed = json.loads(formatter.format(record))
+
+        assert "sessionId" not in parsed
+        assert "requestId" not in parsed
+
+    def test_includes_exception_info(self):
+        formatter = CorrelatedFormatter()
+
+        try:
+            raise ValueError("test error")
+        except ValueError:
+            import sys
+
+            exc_info = sys.exc_info()
+
+        record = logging.LogRecord(
+            name="test", level=logging.ERROR, pathname="", lineno=0, msg="failed", args=(), exc_info=exc_info
+        )
+
+        parsed = json.loads(formatter.format(record))
+
+        assert parsed["errorType"] == "ValueError"
+        assert parsed["errorMessage"] == "test error"
+        assert isinstance(parsed["stackTrace"], list)
+        assert len(parsed["stackTrace"]) > 0
+
+
+class TestConfigureLogging:
+    def test_attaches_handler_to_root(self):
+        configure_logging()
+
+        root = logging.getLogger()
+        assert len(root.handlers) == 1
+        assert isinstance(root.handlers[0].formatter, CorrelatedFormatter)
+
+    def test_sets_root_level(self):
+        configure_logging(level=logging.DEBUG)
+
+        assert logging.getLogger().level == logging.DEBUG
+
+    def test_idempotent(self):
+        configure_logging()
+        configure_logging()
+        configure_logging()
+
+        assert len(logging.getLogger().handlers) == 1
+
+    def test_clears_existing_handlers(self):
+        logging.basicConfig(level=logging.INFO)
+        assert len(logging.getLogger().handlers) >= 1
+
+        configure_logging()
+
+        root = logging.getLogger()
+        assert len(root.handlers) == 1
+        assert isinstance(root.handlers[0].formatter, CorrelatedFormatter)
+
+    def test_suppresses_sdk_logger_propagation(self):
+        configure_logging()
+
+        sdk_logger = logging.getLogger("bedrock_agentcore.app")
+        assert sdk_logger.propagate is False