feat(client): add input_id attribute to RolloutFuture (#23)

Author: luyuzhe111

src/agentcore_rl_toolkit/client.py

@@ -45,13 +45,15 @@ def __init__(
         max_interval: float = 30.0,
         backoff_factor: float = 1.5,
         session_id: str = None,
+        input_id: str = None,
         agentcore_client=None,
         agent_runtime_arn: str = None,
     ):
         self.s3_client = s3_client
         self.s3_bucket = s3_bucket
         self.result_key = result_key
         self.session_id = session_id
+        self.input_id = input_id
         self.agentcore_client = agentcore_client
         self.agent_runtime_arn = agent_runtime_arn
         self._result = None
@@ -284,6 +286,7 @@ def _rate_limited_invoke(self, payload: dict, session_id: str, input_id: str) ->
             s3_bucket=data["s3_bucket"],
             result_key=data["result_key"],
             session_id=session_id,
+            input_id=input_id,
             agentcore_client=self.agentcore_client,
             agent_runtime_arn=self.agent_runtime_arn,
         )

tests/test_client.py

@@ -443,6 +443,31 @@ def test_invoke_future_has_session_id(self):
             assert future.agentcore_client is mock_acr
             assert future.agent_runtime_arn == "arn:aws:bedrock-agentcore:us-west-2:123:agent/test"
 
+    def test_invoke_future_has_input_id(self):
+        """Test invoke() returns future with input_id set."""
+        with patch("agentcore_rl_toolkit.client.boto3") as mock_boto3:
+            mock_acr = MagicMock()
+            mock_s3 = MagicMock()
+            mock_boto3.client.side_effect = lambda service, **kwargs: (
+                mock_acr if service == "bedrock-agentcore" else mock_s3
+            )
+
+            mock_acr.invoke_agent_runtime.return_value = {
+                "response": mock_streaming_body(
+                    {"status": "processing", "s3_bucket": "test-bucket", "result_key": "exp/key.json"}
+                )
+            }
+
+            client = RolloutClient(
+                agent_runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123:agent/test",
+                s3_bucket="test-bucket",
+                exp_id="exp-001",
+            )
+
+            future = client.invoke({"prompt": "test"}, input_id="my-input")
+
+            assert future.input_id == "my-input"
+
     def test_run_batch_returns_batch_result(self):
         """Test run_batch() returns a BatchResult."""
         with patch("agentcore_rl_toolkit.client.boto3"):