refactor(app): derive session_id from ACR header instead of _rollout payload (#22)

Author: luyuzhe111

AGENTS.md

@@ -204,7 +204,7 @@ Since the client won't get results directly from HTTP:
 **AgentCoreRLApp** (`src/agentcore_rl_toolkit/app.py`)
 - Inherits `BedrockAgentCoreApp` - drop-in replacement
 - Provides `@app.rollout_entrypoint` decorator
-- Expects `_rollout` dict in payload following `RolloutConfig` model (experiment id, session id, input id, base_url, model_id)
+- Expects `_rollout` dict in payload following `RolloutConfig` model (experiment id, input id, s3_bucket, base_url, model_id)
 - Framework-agnostic: works with any agent framework, not just Strands
 
 #### Utilities

src/agentcore_rl_toolkit/app.py

@@ -2,21 +2,21 @@
 import json
 import logging
 import os
+import uuid
 from dataclasses import dataclass
 from functools import wraps
 
 import boto3
 from bedrock_agentcore.runtime import BedrockAgentCoreApp
 
-_S3_CONFIG_FIELDS = ("exp_id", "session_id", "input_id", "s3_bucket")
+_S3_CONFIG_FIELDS = ("exp_id", "input_id", "s3_bucket")
 
 
 @dataclass
 class RolloutConfig:
     """Rollout configuration for rollout collection and storage."""
 
     exp_id: str
-    session_id: str
     input_id: str
     s3_bucket: str
 
@@ -26,7 +26,6 @@ def from_dict(cls, data: dict) -> "RolloutConfig":
         try:
             return cls(
                 exp_id=data["exp_id"],
-                session_id=data["session_id"],
                 input_id=data["input_id"],
                 s3_bucket=data["s3_bucket"],
             )
@@ -115,7 +114,7 @@ def _validate_and_normalize_rollout(self, rollout_dict: dict) -> dict:
         rollout_dict["rewards"] = rewards
         return rollout_dict
 
-    def save_rollout(self, rollout_data: dict, rollout_config: dict, payload: dict = None, result_key: str = None):
+    def save_rollout(self, rollout_data: dict, rollout_config: dict, result_key: str, payload: dict = None):
         """
         Save rollout data to S3.
 
@@ -124,7 +123,6 @@ def save_rollout(self, rollout_data: dict, rollout_config: dict, payload: dict =
             rollout_config: Rollout configuration dict containing:
                 - s3_bucket: S3 bucket name
                 - exp_id: Experiment ID for organizing data
-                - session_id: Session id for the current task
                 - input_id: id for discriminating different input data examples
             payload: Original request payload (included in saved result for debugging)
             result_key: S3 key for the result (computed externally for consistency)
@@ -136,10 +134,6 @@ def save_rollout(self, rollout_data: dict, rollout_config: dict, payload: dict =
             logging.error(f"Invalid rollout configuration: {e}")
             raise
 
-        # Use provided result_key or compute it
-        if result_key is None:
-            result_key = f"{config.exp_id}/{config.input_id}_{config.session_id}.json"
-
         if "status_code" not in rollout_data:
             rollout_data["status_code"] = 200
 
@@ -249,7 +243,10 @@ async def rollout_entrypoint_wrapper(payload, context):
             rollout_config = None
             if rollout_dict is not None and any(f in rollout_dict for f in _S3_CONFIG_FIELDS):
                 rollout_config = RolloutConfig.from_dict(rollout_dict)
-                result_key = f"{rollout_config.exp_id}/{rollout_config.input_id}_{rollout_config.session_id}.json"
+                # session_id comes from ACR's HTTP header (set via runtimeSessionId),
+                # fall back to UUID for local testing.
+                session_id = context.session_id if context.session_id else str(uuid.uuid4())
+                result_key = f"{rollout_config.exp_id}/{rollout_config.input_id}/{session_id}.json"
 
             # Start background task without waiting
             asyncio.create_task(rollout_background_task(payload, context, result_key))

src/agentcore_rl_toolkit/client.py

@@ -257,7 +257,6 @@ def _rate_limited_invoke(self, payload: dict, session_id: str, input_id: str) ->
         # Build rollout config
         rollout_config = {
             "exp_id": self.exp_id,
-            "session_id": session_id,
             "input_id": input_id,
             "s3_bucket": self.s3_bucket,
             **self.extra_config,

tests/test_client.py

@@ -410,7 +410,6 @@ def test_invoke_builds_rollout_config(self):
 
             assert payload["prompt"] == "test"
             assert payload["_rollout"]["exp_id"] == "exp-001"
-            assert payload["_rollout"]["session_id"] == "sess-1"
             assert payload["_rollout"]["input_id"] == "input-1"
             assert payload["_rollout"]["s3_bucket"] == "test-bucket"
             assert payload["_rollout"]["base_url"] == "http://localhost:8000"

tests/test_rollout_entrypoint.py

@@ -100,18 +100,18 @@ async def handler(payload: dict):
             "prompt": "test",
             "_rollout": {
                 "exp_id": "exp-123",
-                "session_id": "sess-456",
                 "input_id": "input-789",
                 "s3_bucket": "my-bucket",
             },
         },
+        headers={"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "sess-456"},
     )
 
     assert response.status_code == 200
     result = response.json()
     assert result["status"] == "processing"
     assert result["s3_bucket"] == "my-bucket"
-    assert result["result_key"] == "exp-123/input-789_sess-456.json"
+    assert result["result_key"] == "exp-123/input-789/sess-456.json"
 
 
 def test_response_without_rollout_config():
@@ -180,7 +180,7 @@ async def handler(payload: dict):
     assert "result_key" not in result
 
 
-@pytest.mark.parametrize("missing_field", ["exp_id", "session_id", "input_id", "s3_bucket"])
+@pytest.mark.parametrize("missing_field", ["exp_id", "input_id", "s3_bucket"])
 def test_entrypoint_rejects_partial_s3_config(missing_field):
     """Test that providing some but not all S3 fields returns HTTP 500."""
     app = AgentCoreRLApp()
@@ -191,7 +191,6 @@ async def handler(payload: dict):
 
     complete_config = {
         "exp_id": "exp-123",
-        "session_id": "sess-456",
         "input_id": "input-789",
         "s3_bucket": "my-bucket",
     }
@@ -201,6 +200,7 @@ async def handler(payload: dict):
     response = client.post(
         "/invocations",
         json={"prompt": "test", "_rollout": incomplete_config},
+        headers={"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "sess-456"},
     )
 
     assert response.status_code == 500