feat: implement callback token generation and processing

- Add CallbackToken generation in callback processor with observer integration
- Implement SendCallbackSuccess, SendCallbackFailure, and SendCallbackHeartbeat
- Add callback operation lookup and completion methods to execution
- Ensure unique token generation across executions

Author: Rares Polenciuc

examples/test/test_step.py

@@ -1,8 +1,8 @@
 """Tests for step example."""
 
 import pytest
-from aws_durable_execution_sdk_python.execution import InvocationStatus
 
+from aws_durable_execution_sdk_python.execution import InvocationStatus
 from src import step
 from test.conftest import deserialize_operation_payload
 

src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py

@@ -98,9 +98,7 @@ def _create_callback_details(
     ) -> CallbackDetails | None:
         """Create CallbackDetails from OperationUpdate."""
         return (
-            CallbackDetails(
-                callback_id="placeholder", result=update.payload, error=update.error
-            )
+            CallbackDetails(callback_id="", result=update.payload, error=update.error)
             if update.operation_type == OperationType.CALLBACK
             else None
         )

src/aws_durable_execution_sdk_python_testing/checkpoint/processors/callback.py

@@ -17,6 +17,7 @@
 from aws_durable_execution_sdk_python_testing.exceptions import (
     InvalidParameterValueException,
 )
+from aws_durable_execution_sdk_python_testing.token import CallbackToken
 
 
 if TYPE_CHECKING:
@@ -36,13 +37,39 @@ def process(
         """Process CALLBACK operation update with scheduler integration for activities."""
         match update.action:
             case OperationAction.START:
-                # TODO: create CallbackToken (see token module). Add Observer/Notifier for on_callback_created possibly,
-                # but token might well have enough so don't need to maintain token list on execution itself
-                return self._translate_update_to_operation(
+                callback_token = CallbackToken(
+                    execution_arn=execution_arn,
+                    operation_id=update.operation_id,
+                )
+                callback_id = callback_token.to_str()
+
+                notifier.notify_callback_created(
+                    execution_arn=execution_arn,
+                    operation_id=update.operation_id,
+                    callback_id=callback_id,
+                )
+
+                operation = self._translate_update_to_operation(
                     update=update,
                     current_operation=current_op,
                     status=OperationStatus.STARTED,
                 )
+
+                # Replace callback_details with actual callback_id
+                if operation.callback_details:
+                    from dataclasses import replace
+                    from aws_durable_execution_sdk_python.lambda_service import (
+                        CallbackDetails,
+                    )
+
+                    new_callback_details = replace(
+                        operation.callback_details, callback_id=callback_id
+                    )
+                    operation = replace(
+                        operation, callback_details=new_callback_details
+                    )
+
+                return operation
             case _:
                 msg: str = "Invalid action for CALLBACK operation."
 

src/aws_durable_execution_sdk_python_testing/execution.py

@@ -30,7 +30,10 @@
 from aws_durable_execution_sdk_python_testing.model import (
     StartDurableExecutionInput,
 )
-from aws_durable_execution_sdk_python_testing.token import CheckpointToken
+from aws_durable_execution_sdk_python_testing.token import (
+    CheckpointToken,
+    CallbackToken,
+)
 
 
 class CloseStatus(Enum):
@@ -305,6 +308,18 @@ def find_operation(self, operation_id: str) -> tuple[int, Operation]:
         msg: str = f"Attempting to update state of an Operation [{operation_id}] that doesn't exist"
         raise IllegalStateException(msg)
 
+    def find_callback_operation(self, callback_id: str) -> tuple[int, Operation]:
+        """Find callback operation by callback_id, return index and operation."""
+        for i, operation in enumerate(self.operations):
+            if (
+                operation.operation_type == OperationType.CALLBACK
+                and operation.callback_details
+                and operation.callback_details.callback_id == callback_id
+            ):
+                return i, operation
+        msg: str = f"Callback operation with callback_id [{callback_id}] not found"
+        raise IllegalStateException(msg)
+
     def complete_wait(self, operation_id: str) -> Operation:
         """Complete WAIT operation when timer fires."""
         index, operation = self.find_operation(operation_id)
@@ -362,3 +377,56 @@ def complete_retry(self, operation_id: str) -> Operation:
             # Assign
             self.operations[index] = updated_operation
             return updated_operation
+
+    def complete_callback_success(
+        self, callback_id: str, result: bytes | None = None
+    ) -> Operation:
+        """Complete CALLBACK operation with success."""
+        index, operation = self.find_callback_operation(callback_id)
+
+        if operation.status != OperationStatus.STARTED:
+            msg: str = f"Callback operation [{callback_id}] is not in STARTED state"
+            raise IllegalStateException(msg)
+
+        with self._state_lock:
+            self._token_sequence += 1
+            updated_callback_details = None
+            if operation.callback_details:
+                updated_callback_details = replace(
+                    operation.callback_details,
+                    result=result.decode() if result else None,
+                )
+
+            self.operations[index] = replace(
+                operation,
+                status=OperationStatus.SUCCEEDED,
+                end_timestamp=datetime.now(UTC),
+                callback_details=updated_callback_details,
+            )
+            return self.operations[index]
+
+    def complete_callback_failure(
+        self, callback_id: str, error: ErrorObject
+    ) -> Operation:
+        """Complete CALLBACK operation with failure."""
+        index, operation = self.find_callback_operation(callback_id)
+
+        if operation.status != OperationStatus.STARTED:
+            msg: str = f"Callback operation [{callback_id}] is not in STARTED state"
+            raise IllegalStateException(msg)
+
+        with self._state_lock:
+            self._token_sequence += 1
+            updated_callback_details = None
+            if operation.callback_details:
+                updated_callback_details = replace(
+                    operation.callback_details, error=error
+                )
+
+            self.operations[index] = replace(
+                operation,
+                status=OperationStatus.FAILED,
+                end_timestamp=datetime.now(UTC),
+                callback_details=updated_callback_details,
+            )
+            return self.operations[index]

src/aws_durable_execution_sdk_python_testing/executor.py

@@ -46,6 +46,7 @@
     Execution as ExecutionSummary,
 )
 from aws_durable_execution_sdk_python_testing.observer import ExecutionObserver
+from aws_durable_execution_sdk_python_testing.token import CallbackToken
 
 
 if TYPE_CHECKING:
@@ -509,7 +510,7 @@ def checkpoint_execution(
     def send_callback_success(
         self,
         callback_id: str,
-        result: bytes | None = None,  # noqa: ARG002
+        result: bytes | None = None,
     ) -> SendDurableExecutionCallbackSuccessResponse:
         """Send callback success response.
 
@@ -528,16 +529,23 @@ def send_callback_success(
             msg: str = "callback_id is required"
             raise InvalidParameterValueException(msg)
 
-        # TODO: Implement actual callback success logic
-        # This would involve finding the callback operation and completing it
-        logger.info("Callback success sent for callback_id: %s", callback_id)
+        try:
+            callback_token = CallbackToken.from_str(callback_id)
+            execution = self.get_execution(callback_token.execution_arn)
+            execution.complete_callback_success(callback_id, result)
+            self._store.update(execution)
+            self._invoke_execution(callback_token.execution_arn)
+            logger.info("Callback success completed for callback_id: %s", callback_id)
+        except Exception as e:
+            msg = f"Failed to process callback success: {e}"
+            raise ResourceNotFoundException(msg) from e
 
         return SendDurableExecutionCallbackSuccessResponse()
 
     def send_callback_failure(
         self,
         callback_id: str,
-        error: ErrorObject | None = None,  # noqa: ARG002
+        error: ErrorObject | None = None,
     ) -> SendDurableExecutionCallbackFailureResponse:
         """Send callback failure response.
 
@@ -556,9 +564,18 @@ def send_callback_failure(
             msg: str = "callback_id is required"
             raise InvalidParameterValueException(msg)
 
-        # TODO: Implement actual callback failure logic
-        # This would involve finding the callback operation and failing it
-        logger.info("Callback failure sent for callback_id: %s", callback_id)
+        callback_error = error or ErrorObject.from_message("Callback failed")
+
+        try:
+            callback_token = CallbackToken.from_str(callback_id)
+            execution = self.get_execution(callback_token.execution_arn)
+            execution.complete_callback_failure(callback_id, callback_error)
+            self._store.update(execution)
+            self._invoke_execution(callback_token.execution_arn)
+            logger.info("Callback failure completed for callback_id: %s", callback_id)
+        except Exception as e:
+            msg = f"Failed to process callback failure: {e}"
+            raise ResourceNotFoundException(msg) from e
 
         return SendDurableExecutionCallbackFailureResponse()
 
@@ -581,9 +598,20 @@ def send_callback_heartbeat(
             msg: str = "callback_id is required"
             raise InvalidParameterValueException(msg)
 
-        # TODO: Implement actual callback heartbeat logic
-        # This would involve updating the callback timeout
-        logger.info("Callback heartbeat sent for callback_id: %s", callback_id)
+        try:
+            callback_token = CallbackToken.from_str(callback_id)
+            execution = self.get_execution(callback_token.execution_arn)
+
+            # Find callback operation to verify it exists and is active
+            _, operation = execution.find_callback_operation(callback_id)
+            if operation.status != OperationStatus.STARTED:
+                msg = f"Callback {callback_id} is not active"
+                raise ResourceNotFoundException(msg)
+
+            logger.info("Callback heartbeat processed for callback_id: %s", callback_id)
+        except Exception as e:
+            msg = f"Failed to process callback heartbeat: {e}"
+            raise ResourceNotFoundException(msg) from e
 
         return SendDurableExecutionCallbackHeartbeatResponse()
 
@@ -907,4 +935,15 @@ def retry_handler() -> None:
             retry_handler, delay=delay, completion_event=completion_event
         )
 
+    def on_callback_created(
+        self, execution_arn: str, operation_id: str, callback_id: str
+    ) -> None:
+        """Handle callback creation. Observer method triggered by notifier."""
+        logger.debug(
+            "[%s] Callback created for operation %s with callback_id: %s",
+            execution_arn,
+            operation_id,
+            callback_id,
+        )
+
     # endregion ExecutionObserver

src/aws_durable_execution_sdk_python_testing/observer.py

@@ -44,6 +44,12 @@ def on_step_retry_scheduled(
     ) -> None:
         """Called when step retry scheduled."""
 
+    @abstractmethod
+    def on_callback_created(
+        self, execution_arn: str, operation_id: str, callback_id: str
+    ) -> None:
+        """Called when callback is created."""
+
 
 class ExecutionNotifier:
     """Notifies observers about execution events. Thread-safe."""
@@ -111,4 +117,15 @@ def notify_step_retry_scheduled(
             delay=delay,
         )
 
+    def notify_callback_created(
+        self, execution_arn: str, operation_id: str, callback_id: str
+    ) -> None:
+        """Notify observers about callback creation."""
+        self._notify_observers(
+            ExecutionObserver.on_callback_created,
+            execution_arn=execution_arn,
+            operation_id=operation_id,
+            callback_id=callback_id,
+        )
+
     # endregion event emitters

tests/checkpoint/processors/base_test.py

@@ -249,7 +249,7 @@ def test_create_callback_details():
     result = processor.create_callback_details(update)
 
     assert isinstance(result, CallbackDetails)
-    assert result.callback_id == "placeholder"
+    assert result.callback_id == ""
     assert result.result == "test-payload"
     assert result.error == error
 

tests/execution_edge_cases_test.py

@@ -0,0 +1,77 @@
+"""Tests for edge cases in execution functionality."""
+
+import pytest
+from unittest.mock import Mock, patch
+from aws_durable_execution_sdk_python_testing.execution import Execution, CloseStatus
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+
+def test_status_unexpected_close_status():
+    """Test status property with unexpected close status."""
+    execution = Execution("test-arn", Mock(), [])
+    execution.is_complete = True
+    execution.close_status = "UNKNOWN_STATUS"  # Invalid close status
+
+    with pytest.raises(InvalidParameterValueException, match="Unexpected close status"):
+        _ = execution.status
+
+
+def test_start_no_invocation_id():
+    """Test start method when invocation_id is None."""
+    start_input = Mock()
+    start_input.invocation_id = None
+
+    execution = Execution("test-arn", start_input, [])
+
+    with pytest.raises(
+        InvalidParameterValueException, match="invocation_id is required"
+    ):
+        execution.start()
+
+
+def test_from_dict_with_none_result():
+    """Test from_dict with None result."""
+    data = {
+        "DurableExecutionArn": "test-arn",
+        "StartInput": {"function_name": "test"},
+        "Operations": [],
+        "Updates": [],
+        "UsedTokens": [],
+        "TokenSequence": 0,
+        "IsComplete": False,
+        "Result": None,  # None result
+        "ConsecutiveFailedInvocationAttempts": 0,
+        "CloseStatus": None,
+    }
+
+    with patch(
+        "aws_durable_execution_sdk_python_testing.model.StartDurableExecutionInput.from_dict"
+    ) as mock_from_dict:
+        mock_from_dict.return_value = Mock()
+        execution = Execution.from_dict(data)
+        assert execution.result is None
+
+
+def test_from_dict_with_none_close_status():
+    """Test from_dict with None close status."""
+    data = {
+        "DurableExecutionArn": "test-arn",
+        "StartInput": {"function_name": "test"},
+        "Operations": [],
+        "Updates": [],
+        "UsedTokens": [],
+        "TokenSequence": 0,
+        "IsComplete": False,
+        "Result": None,
+        "ConsecutiveFailedInvocationAttempts": 0,
+        "CloseStatus": None,  # None close status
+    }
+
+    with patch(
+        "aws_durable_execution_sdk_python_testing.model.StartDurableExecutionInput.from_dict"
+    ) as mock_from_dict:
+        mock_from_dict.return_value = Mock()
+        execution = Execution.from_dict(data)
+        assert execution.close_status is None