Add tenant_id support for invoke operations (#195)

* Add tenant_id support for invoke operations

Author: Astraea Quinn S

src/aws_durable_execution_sdk_python/config.py

@@ -378,12 +378,34 @@ class MapConfig:
     summary_generator: SummaryGenerator | None = None
 
 
-@dataclass
+@dataclass(frozen=True)
 class InvokeConfig(Generic[P, R]):
+    """
+    Configuration for invoke operations.
+
+    This class configures how function invocations are executed, including
+    timeout behavior, serialization, and tenant isolation.
+
+    Args:
+        timeout: Maximum duration to wait for the invoked function to complete.
+            Default is no timeout. Use this to prevent long-running invocations
+            from blocking execution indefinitely.
+
+        serdes_payload: Custom serialization/deserialization for the payload
+            sent to the invoked function. If None, uses default JSON serialization.
+
+        serdes_result: Custom serialization/deserialization for the result
+            returned from the invoked function. If None, uses default JSON serialization.
+
+        tenant_id: Optional tenant identifier for multi-tenant isolation.
+            If provided, the invocation will be scoped to this tenant.
+    """
+
     # retry_strategy: Callable[[Exception, int], RetryDecision] | None = None
     timeout: Duration = field(default_factory=Duration)
     serdes_payload: SerDes[P] | None = None
     serdes_result: SerDes[R] | None = None
+    tenant_id: str | None = None
 
     @property
     def timeout_seconds(self) -> int:

src/aws_durable_execution_sdk_python/lambda_service.py

@@ -303,17 +303,22 @@ class ChainedInvokeOptions:
     """
 
     function_name: str
+    tenant_id: str | None = None
 
     @classmethod
     def from_dict(cls, data: MutableMapping[str, Any]) -> ChainedInvokeOptions:
         return cls(
             function_name=data["FunctionName"],
+            tenant_id=data.get("TenantId"),
         )
 
     def to_dict(self) -> MutableMapping[str, Any]:
         result: MutableMapping[str, Any] = {
             "FunctionName": self.function_name,
         }
+        if self.tenant_id is not None:
+            result["TenantId"] = self.tenant_id
+
         return result
 
 

src/aws_durable_execution_sdk_python/operation/invoke.py

@@ -40,6 +40,7 @@ def invoke_handler(
 
     if not config:
         config = InvokeConfig[P, R]()
+    tenant_id = config.tenant_id
 
     # Check if we have existing step data
     checkpointed_result = state.get_checkpoint_result(operation_identifier.operation_id)
@@ -87,7 +88,10 @@ def invoke_handler(
     start_operation: OperationUpdate = OperationUpdate.create_invoke_start(
         identifier=operation_identifier,
         payload=serialized_payload,
-        chained_invoke_options=ChainedInvokeOptions(function_name=function_name),
+        chained_invoke_options=ChainedInvokeOptions(
+            function_name=function_name,
+            tenant_id=tenant_id,
+        ),
     )
 
     # Checkpoint invoke START with blocking (is_sync=True, default).

tests/config_test.py

@@ -10,6 +10,7 @@
     ChildConfig,
     CompletionConfig,
     Duration,
+    InvokeConfig,
     ItemBatcher,
     ItemsPerBatchUnit,
     MapConfig,
@@ -275,3 +276,16 @@ def test_step_future_without_name():
 
         result = step_future.result()
         assert result == 42
+
+
+def test_invoke_config_defaults():
+    """Test InvokeConfig defaults."""
+    config = InvokeConfig()
+    assert config.tenant_id is None
+    assert config.timeout_seconds == 0
+
+
+def test_invoke_config_with_tenant_id():
+    """Test InvokeConfig with explicit tenant_id."""
+    config = InvokeConfig(tenant_id="test-tenant")
+    assert config.tenant_id == "test-tenant"

tests/context_test.py

@@ -662,9 +662,11 @@ def test_invoke_with_custom_serdes(mock_handler):
         "arn:aws:durable:us-east-1:123456789012:execution/test"
     )
 
+    payload_serdes = CustomDictSerDes()
+    result_serdes = CustomDictSerDes()
     config = InvokeConfig[dict, dict](
-        serdes_payload=CustomDictSerDes(),
-        serdes_result=CustomDictSerDes(),
+        serdes_payload=payload_serdes,
+        serdes_result=result_serdes,
         timeout=Duration.from_minutes(1),
     )
 
@@ -1680,3 +1682,41 @@ def test_operation_id_generation_unique():
 
     for i in range(len(ids) - 1):
         assert ids[i] != ids[i + 1]
+
+
+@patch("aws_durable_execution_sdk_python.context.invoke_handler")
+def test_invoke_with_explicit_tenant_id(mock_handler):
+    """Test invoke with explicit tenant_id in config."""
+    mock_handler.return_value = "result"
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = (
+        "arn:aws:durable:us-east-1:123456789012:execution/test"
+    )
+
+    config = InvokeConfig(tenant_id="explicit-tenant")
+    context = DurableContext(state=mock_state)
+
+    result = context.invoke("test_function", "payload", config=config)
+
+    assert result == "result"
+    call_args = mock_handler.call_args[1]
+    assert call_args["config"].tenant_id == "explicit-tenant"
+
+
+@patch("aws_durable_execution_sdk_python.context.invoke_handler")
+def test_invoke_without_tenant_id_defaults_to_none(mock_handler):
+    """Test invoke without tenant_id defaults to None."""
+    mock_handler.return_value = "result"
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = (
+        "arn:aws:durable:us-east-1:123456789012:execution/test"
+    )
+
+    context = DurableContext(state=mock_state)
+
+    result = context.invoke("test_function", "payload")
+
+    assert result == "result"
+    # Config should be None when not provided
+    call_args = mock_handler.call_args[1]
+    assert call_args["config"] is None

tests/lambda_service_test.py

@@ -400,16 +400,26 @@ def test_callback_options_from_dict_partial():
 
 def test_invoke_options_from_dict():
     """Test ChainedInvokeOptions.from_dict method."""
-    data = {"FunctionName": "test-function", "TimeoutSeconds": 120}
+    data = {"FunctionName": "test-function", "TenantId": "test-tenant"}
     options = ChainedInvokeOptions.from_dict(data)
     assert options.function_name == "test-function"
+    assert options.tenant_id == "test-tenant"
 
 
 def test_invoke_options_from_dict_required_only():
     """Test ChainedInvokeOptions.from_dict with only required field."""
     data = {"FunctionName": "test-function"}
     options = ChainedInvokeOptions.from_dict(data)
     assert options.function_name == "test-function"
+    assert options.tenant_id is None
+
+
+def test_invoke_options_from_dict_with_none_tenant():
+    """Test ChainedInvokeOptions.from_dict with explicit None tenant_id."""
+    data = {"FunctionName": "test-function", "TenantId": None}
+    options = ChainedInvokeOptions.from_dict(data)
+    assert options.function_name == "test-function"
+    assert options.tenant_id is None
 
 
 def test_context_options_from_dict():

tests/operation/invoke_test.py

@@ -308,10 +308,10 @@ def test_invoke_handler_no_config():
 
     # Verify default config was used
     operation_update = mock_state.create_checkpoint.call_args[1]["operation_update"]
-    assert (
-        operation_update.to_dict()["ChainedInvokeOptions"]["FunctionName"]
-        == "test_function"
-    )
+    chained_invoke_options = operation_update.to_dict()["ChainedInvokeOptions"]
+    assert chained_invoke_options["FunctionName"] == "test_function"
+    # tenant_id should be None when not specified
+    assert "TenantId" not in chained_invoke_options
 
 
 def test_invoke_handler_custom_serdes():
@@ -533,3 +533,82 @@ def test_invoke_handler_suspend_does_not_raise(mock_suspend):
         )
 
     mock_suspend.assert_called_once()
+
+
+def test_invoke_handler_with_tenant_id():
+    """Test invoke_handler passes tenant_id to checkpoint."""
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "test_arn"
+    mock_state.get_checkpoint_result.return_value = (
+        CheckpointedResult.create_not_found()
+    )
+
+    config = InvokeConfig(tenant_id="test-tenant-123")
+
+    with pytest.raises(SuspendExecution):
+        invoke_handler(
+            function_name="test_function",
+            payload="test_input",
+            state=mock_state,
+            operation_identifier=OperationIdentifier("invoke1", None, None),
+            config=config,
+        )
+
+    # Verify checkpoint was called with tenant_id
+    mock_state.create_checkpoint.assert_called_once()
+    operation_update = mock_state.create_checkpoint.call_args[1]["operation_update"]
+    chained_invoke_options = operation_update.to_dict()["ChainedInvokeOptions"]
+    assert chained_invoke_options["FunctionName"] == "test_function"
+    assert chained_invoke_options["TenantId"] == "test-tenant-123"
+
+
+def test_invoke_handler_without_tenant_id():
+    """Test invoke_handler without tenant_id doesn't include it in checkpoint."""
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "test_arn"
+    mock_state.get_checkpoint_result.return_value = (
+        CheckpointedResult.create_not_found()
+    )
+
+    config = InvokeConfig(tenant_id=None)
+
+    with pytest.raises(SuspendExecution):
+        invoke_handler(
+            function_name="test_function",
+            payload="test_input",
+            state=mock_state,
+            operation_identifier=OperationIdentifier("invoke1", None, None),
+            config=config,
+        )
+
+    # Verify checkpoint was called without tenant_id
+    mock_state.create_checkpoint.assert_called_once()
+    operation_update = mock_state.create_checkpoint.call_args[1]["operation_update"]
+    chained_invoke_options = operation_update.to_dict()["ChainedInvokeOptions"]
+    assert chained_invoke_options["FunctionName"] == "test_function"
+    assert "TenantId" not in chained_invoke_options
+
+
+def test_invoke_handler_default_config_no_tenant_id():
+    """Test invoke_handler with default config has no tenant_id."""
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "test_arn"
+    mock_state.get_checkpoint_result.return_value = (
+        CheckpointedResult.create_not_found()
+    )
+
+    with pytest.raises(SuspendExecution):
+        invoke_handler(
+            function_name="test_function",
+            payload="test_input",
+            state=mock_state,
+            operation_identifier=OperationIdentifier("invoke1", None, None),
+            config=None,
+        )
+
+    # Verify checkpoint was called without tenant_id
+    mock_state.create_checkpoint.assert_called_once()
+    operation_update = mock_state.create_checkpoint.call_args[1]["operation_update"]
+    chained_invoke_options = operation_update.to_dict()["ChainedInvokeOptions"]
+    assert chained_invoke_options["FunctionName"] == "test_function"
+    assert "TenantId" not in chained_invoke_options