diff --git a/samcli/lib/utils/durable_callback_handler.py b/samcli/lib/utils/durable_callback_handler.py new file mode 100644 index 00000000000..2efc32be5c5 --- /dev/null +++ b/samcli/lib/utils/durable_callback_handler.py @@ -0,0 +1,137 @@ +""" +Interactive callback handler for durable function executions. +""" + +import logging +from typing import Optional + +import click + +from samcli.lib.clients.lambda_client import DurableFunctionsClient + +LOG = logging.getLogger(__name__) + +# Menu choice constants +CHOICE_SUCCESS = 1 +CHOICE_FAILURE = 2 +CHOICE_HEARTBEAT = 3 +CHOICE_STOP = 4 + + +class DurableCallbackHandler: + """ + Handles interactive callback detection and response for durable executions. + """ + + def __init__(self, client: DurableFunctionsClient): + self.client = client + self._prompted_callbacks: set[str] = set() # Track which callbacks we've already prompted for + + def check_for_pending_callbacks(self, execution_arn: str) -> Optional[str]: + """ + Check execution history for pending callbacks. + + Returns: + callback_id if found, None otherwise + """ + try: + LOG.debug("Checking for pending callbacks in execution: %s", execution_arn) + history = self.client.get_durable_execution_history(execution_arn) + events = history.get("Events", []) + + if events: + callback_states = {} + + for event in events: + event_type = event.get("EventType") + event_id = event.get("Id") + + if event_type == "CallbackStarted": + callback_id = event.get("CallbackStartedDetails", {}).get("CallbackId") + callback_states[event_id] = {"callback_id": callback_id, "status": "STARTED", "event": event} + elif event_type in ["CallbackCompleted", "CallbackFailed", "CallbackSucceeded"]: + if event_id in callback_states: + callback_states[event_id]["status"] = "COMPLETED" + + # Find callbacks that are started but not completed + for callback_id, state in callback_states.items(): + if state["status"] == "STARTED" and state["callback_id"]: + return str(state["callback_id"]) + + except Exception as e: + LOG.error("Failed to check callback history: %s", e) + + return None + + def prompt_callback_response(self, execution_arn: str, callback_id: str, execution_complete=None) -> bool: + """ + Prompt user for callback response and send it. + + Args: + execution_arn: The execution ARN for stop execution operation + callback_id: The callback ID to respond to + execution_complete: Optional threading.Event to check if execution finished + + Returns: + True if callback was sent, False if user chose to continue waiting + """ + # Only prompt once per callback ID to avoid blocking on timed-out callbacks + if callback_id in self._prompted_callbacks: + return False + + self._prompted_callbacks.add(callback_id) + + # Check if execution already completed before prompting + if execution_complete and execution_complete.is_set(): + return False + + click.echo(f"\nšŸ”„ Execution is waiting for callback: {callback_id}") + click.echo("Choose an action:") + click.echo(" 1. Send callback success") + click.echo(" 2. Send callback failure") + click.echo(" 3. Send callback heartbeat") + click.echo(" 4. Stop execution") + + choice = click.prompt("Enter choice", type=click.IntRange(1, 4), default=CHOICE_SUCCESS) + + # Check again after user makes selection in case execution completed + if execution_complete and execution_complete.is_set(): + click.echo("āš ļø Execution already completed, callback no longer needed") + return False + + try: + if choice == CHOICE_SUCCESS: + result = click.prompt("Enter success result (optional)", default="", show_default=False) + self.client.send_callback_success(callback_id=callback_id, result=result) + click.echo("āœ… Callback success sent") + return True + + elif choice == CHOICE_FAILURE: + error_message = click.prompt("Enter error message", default="User cancelled") + error_type = click.prompt("Enter error type (optional)", default="", show_default=False) or None + + self.client.send_callback_failure( + callback_id=callback_id, error_message=error_message, error_type=error_type + ) + click.echo("āŒ Callback failure sent") + return True + + elif choice == CHOICE_HEARTBEAT: + self.client.send_callback_heartbeat(callback_id=callback_id) + click.echo("šŸ’“ Callback heartbeat sent") + return False # Continue waiting after heartbeat + + else: # CHOICE_STOP + error_message = click.prompt("Enter error message", default="Execution stopped by user") + error_type = click.prompt("Enter error type (optional)", default="", show_default=False) or None + + self.client.stop_durable_execution( + durable_execution_arn=execution_arn, error_message=error_message, error_type=error_type + ) + click.echo("šŸ›‘ Execution stopped") + return True + + except Exception as e: + LOG.error("Failed to send callback: %s", e) + click.echo(f"āŒ Failed to send callback: {e}") + return False diff --git a/samcli/local/docker/durable_lambda_container.py b/samcli/local/docker/durable_lambda_container.py index 6804706db84..890973fdd0d 100644 --- a/samcli/local/docker/durable_lambda_container.py +++ b/samcli/local/docker/durable_lambda_container.py @@ -9,6 +9,7 @@ import click from flask import has_request_context +from samcli.lib.utils.durable_callback_handler import DurableCallbackHandler from samcli.lib.utils.durable_formatters import format_execution_details, format_next_commands_after_invoke from samcli.local.docker.lambda_container import LambdaContainer @@ -146,8 +147,11 @@ def _write_execution_result_to_stdout(self, execution_details: dict, stdout): def _wait_for_execution(self, execution_arn): """Poll the execution status until completion and return the final result.""" - # TODO - poll until the execution timeout is hit + callback_handler = DurableCallbackHandler(self.emulator_container.lambda_client) execution_details = None + callback_thread = None + stop_callback_prompts = threading.Event() + try: while True: try: @@ -156,13 +160,36 @@ def _wait_for_execution(self, execution_arn): status = execution_details.get("Status") if status != "RUNNING": + stop_callback_prompts.set() # Signal callback thread to stop + if callback_thread and callback_thread.is_alive(): + callback_thread.join(timeout=0.5) # Brief wait for thread cleanup return execution_details + # Check for pending callbacks (only in CLI context) + if self._is_cli_context(): + callback_id = callback_handler.check_for_pending_callbacks(execution_arn) + if callback_id: + + def _prompt_in_thread(): + if not stop_callback_prompts.is_set(): + # give the function logs time to settle after the invocation is suspended + time.sleep(0.5) + callback_sent = callback_handler.prompt_callback_response( + execution_arn, callback_id, stop_callback_prompts + ) + if callback_sent: + click.echo("\n" + "ā”€" * 80) + + # Start callback prompt in separate thread so it doesn't block polling + callback_thread = threading.Thread(target=_prompt_in_thread, daemon=True) + callback_thread.start() + time.sleep(1) # Poll every second except Exception as e: LOG.error("Error polling execution status: %s", e) break finally: + stop_callback_prompts.set() # Ensure callback thread knows to stop self._cleanup_if_needed() return execution_details diff --git a/tests/unit/lib/utils/test_durable_callback_handler.py b/tests/unit/lib/utils/test_durable_callback_handler.py new file mode 100644 index 00000000000..76c5b211eef --- /dev/null +++ b/tests/unit/lib/utils/test_durable_callback_handler.py @@ -0,0 +1,197 @@ +""" +Unit tests for DurableCallbackHandler +""" + +import threading +from unittest import TestCase +from unittest.mock import Mock, patch + +from parameterized import parameterized + +from samcli.lib.utils.durable_callback_handler import ( + DurableCallbackHandler, + CHOICE_SUCCESS, + CHOICE_FAILURE, + CHOICE_HEARTBEAT, + CHOICE_STOP, +) + + +class TestDurableCallbackHandler(TestCase): + def setUp(self): + self.mock_client = Mock() + self.handler = DurableCallbackHandler(self.mock_client) + + def test_init(self): + """Test handler initializes with client and empty prompted callbacks set""" + self.assertEqual(self.handler.client, self.mock_client) + self.assertEqual(self.handler._prompted_callbacks, set()) + + @parameterized.expand( + [ + ( + "with_pending_callback", + { + "Events": [ + { + "Id": 1, + "EventType": "CallbackStarted", + "CallbackStartedDetails": {"CallbackId": "callback-123"}, + }, + {"Id": 2, "EventType": "StepStarted"}, + ] + }, + "callback-123", + ), + ( + "with_completed_callback", + { + "Events": [ + { + "Id": 1, + "EventType": "CallbackStarted", + "CallbackStartedDetails": {"CallbackId": "callback-123"}, + }, + {"Id": 1, "EventType": "CallbackCompleted"}, + ] + }, + None, + ), + ( + "no_callbacks", + {"Events": [{"Id": 1, "EventType": "StepStarted"}]}, + None, + ), + ] + ) + def test_check_for_pending_callbacks(self, name, history_response, expected_callback_id): + """Test checking for pending callbacks""" + self.mock_client.get_durable_execution_history.return_value = history_response + + callback_id = self.handler.check_for_pending_callbacks("test-arn") + + self.assertEqual(callback_id, expected_callback_id) + self.mock_client.get_durable_execution_history.assert_called_once_with("test-arn") + + def test_check_for_pending_callbacks_handles_exception(self): + """Test that exceptions during callback check are handled gracefully""" + self.mock_client.get_durable_execution_history.side_effect = Exception("API error") + + callback_id = self.handler.check_for_pending_callbacks("test-arn") + + self.assertIsNone(callback_id) + + @parameterized.expand( + [ + ( + "success", + CHOICE_SUCCESS, + ["test result"], + "send_callback_success", + {"callback_id": "callback-123", "result": "test result"}, + True, + ), + ( + "failure", + CHOICE_FAILURE, + ["Error occurred", "CustomError"], + "send_callback_failure", + {"callback_id": "callback-123", "error_message": "Error occurred", "error_type": "CustomError"}, + True, + ), + ( + "heartbeat", + CHOICE_HEARTBEAT, + [], + "send_callback_heartbeat", + {"callback_id": "callback-123"}, + False, + ), + ( + "stop_execution", + CHOICE_STOP, + ["Execution stopped by user", "StopError"], + "stop_durable_execution", + { + "durable_execution_arn": "test-arn", + "error_message": "Execution stopped by user", + "error_type": "StopError", + }, + True, + ), + ] + ) + @patch("samcli.lib.utils.durable_callback_handler.click.prompt") + @patch("samcli.lib.utils.durable_callback_handler.click.echo") + def test_prompt_callback_response( + self, + name, + choice, + prompt_responses, + method_name, + expected_call_args, + expected_result, + mock_echo, + mock_prompt, + ): + """Test prompting for different callback response types""" + mock_prompt.side_effect = [choice] + prompt_responses + + result = self.handler.prompt_callback_response("test-arn", "callback-123") + + self.assertEqual(result, expected_result) + lambda_client_api_call = getattr(self.mock_client, method_name) + lambda_client_api_call.assert_called_once_with(**expected_call_args) + self.assertIn("callback-123", self.handler._prompted_callbacks) + + def test_prompt_callback_response_only_prompts_once(self): + """Test that callback is only prompted once per ID""" + self.handler._prompted_callbacks.add("callback-123") + + result = self.handler.prompt_callback_response("test-arn", "callback-123") + + self.assertFalse(result) + self.mock_client.send_callback_success.assert_not_called() + + @parameterized.expand( + [ + ("before_prompt", True, False), + ("after_selection", False, True), + ] + ) + @patch("samcli.lib.utils.durable_callback_handler.click.prompt") + @patch("samcli.lib.utils.durable_callback_handler.click.echo") + def test_prompt_callback_response_checks_execution_complete( + self, name, set_before_prompt, set_during_prompt, mock_echo, mock_prompt + ): + """Test that prompt respects execution_complete event""" + execution_complete = threading.Event() + + if set_before_prompt: + execution_complete.set() + elif set_during_prompt: + + def prompt_side_effect(*args, **kwargs): + execution_complete.set() + return CHOICE_SUCCESS + + mock_prompt.side_effect = prompt_side_effect + + result = self.handler.prompt_callback_response("test-arn", "callback-123", execution_complete) + + self.assertFalse(result) + self.mock_client.send_callback_success.assert_not_called() + if set_before_prompt: + mock_prompt.assert_not_called() + + @patch("samcli.lib.utils.durable_callback_handler.click.prompt") + @patch("samcli.lib.utils.durable_callback_handler.click.echo") + def test_prompt_callback_response_handles_exception(self, mock_echo, mock_prompt): + """Test that exceptions during callback send are handled gracefully""" + mock_prompt.return_value = CHOICE_HEARTBEAT + + self.mock_client.send_callback_heartbeat.side_effect = Exception("API error") + + result = self.handler.prompt_callback_response("test-arn", "callback-123") + + self.assertFalse(result) diff --git a/tests/unit/local/docker/test_durable_lambda_container.py b/tests/unit/local/docker/test_durable_lambda_container.py index 87fe62d7417..ee624567ee6 100644 --- a/tests/unit/local/docker/test_durable_lambda_container.py +++ b/tests/unit/local/docker/test_durable_lambda_container.py @@ -2,7 +2,7 @@ Unit tests for DurableLambdaContainer """ -from unittest import TestCase +from unittest import TestCase, mock from unittest.mock import Mock, patch, MagicMock from parameterized import parameterized @@ -76,6 +76,7 @@ def test_creates_lambda_container_with_emulator(self): (False, "http://host.docker.internal:8080", True), # is_external_emulator=False, HTTP context ] ) + @patch("samcli.local.docker.durable_lambda_container.DurableCallbackHandler") @patch("samcli.local.docker.durable_lambda_container.click.secho") @patch("samcli.local.docker.durable_lambda_container.format_next_commands_after_invoke") @patch("samcli.local.docker.durable_lambda_container.format_execution_details") @@ -91,12 +92,18 @@ def test_sync_wait_for_result( mock_format_execution_details, mock_format_next_commands, mock_secho, + mock_callback_handler_class, ): """Test wait_for_result for sync invocation waits for completion and shows commands based on context""" mock_has_request_context.return_value = has_flask_request_context mock_format_execution_details.return_value = "Execution details" mock_format_next_commands.return_value = "Next commands" + # Mock callback handler to return no pending callbacks + mock_callback_handler = Mock() + mock_callback_handler.check_for_pending_callbacks.return_value = None + mock_callback_handler_class.return_value = mock_callback_handler + mock_emulator = Mock() mock_emulator.start_or_attach = Mock() mock_emulator._is_external_emulator = Mock(return_value=is_external_emulator) @@ -147,6 +154,9 @@ def test_sync_wait_for_result( {"ExecutionTimeout": 900, "RetentionPeriodInDays": 7}, ) + # Verify callback handler was created + mock_callback_handler_class.assert_called_once_with(mock_emulator.lambda_client) + # Verify execution was polled multiple times until completion self.assertEqual(mock_emulator.lambda_client.get_durable_execution.call_count, 3) mock_emulator.lambda_client.get_durable_execution.assert_called_with(mock_execution_arn) @@ -175,16 +185,106 @@ def test_sync_wait_for_result( # Verify headers are returned self.assertEqual(headers["X-Amz-Durable-Execution-Arn"], mock_execution_arn) + @parameterized.expand( + [ + (False,), # has_request_context=False (CLI context) - user is prompted to respond to callbacks + (True,), # has_request_context=True (HTTP context) - user is not prompted + ] + ) + @patch("samcli.local.docker.durable_lambda_container.has_request_context") + @patch("samcli.local.docker.durable_lambda_container.threading.Thread") + @patch("samcli.local.docker.durable_lambda_container.DurableCallbackHandler") + @patch("samcli.local.docker.durable_lambda_container.LambdaContainer.start") + def test_sync_wait_for_result_with_callbacks( + self, has_request_context, mock_start, mock_callback_handler_class, mock_thread_class, mock_has_request_context + ): + """Test callback thread lifecycle: detection, prompt, and cleanup on completion""" + mock_has_request_context.return_value = has_request_context + is_cli_context = not has_request_context + mock_execution_arn = "arn:123" + + mock_emulator = Mock() + mock_emulator._is_external_emulator = Mock(return_value=False) + mock_emulator.start_durable_execution = Mock(return_value={"ExecutionArn": mock_execution_arn}) + # Execution runs for 3 polls, callback detected on 3rd, then completes + mock_emulator.lambda_client.get_durable_execution = Mock( + side_effect=[ + {"Status": "RUNNING"}, + {"Status": "RUNNING"}, + {"Status": "RUNNING"}, + {"Status": "SUCCEEDED", "Output": "result"}, + ] + ) + + # Mock callback handler - no callback for first 2 polls, then callback detected + mock_callback_handler = Mock() + mock_callback_handler.check_for_pending_callbacks.side_effect = [None, None, "callback-123"] + mock_callback_handler.prompt_callback_response.return_value = True + mock_callback_handler_class.return_value = mock_callback_handler + + # Mock callback thread - capture target and execute it to verify prompt is called + mock_callback_thread = Mock() + mock_callback_thread.is_alive.return_value = True # So join() gets called at cleanup + + def capture_and_execute_thread(target, **kwargs): + target() # Execute the thread function immediately + return mock_callback_thread + + mock_thread_class.side_effect = capture_and_execute_thread + + container = self._create_container(mock_emulator) + container.start_logs_thread_if_not_alive = Mock() + container.get_port = Mock(return_value=8080) + container._wait_for_socket_connection = Mock() + + container.wait_for_result( + full_path="test-function", + event={"test": "event"}, + stdout=Mock(), + stderr=Mock(), + durable_execution_name="test-execution", + ) + + if is_cli_context: + # CLI context - validate the user was prompted to respond to the callback + self.assertEqual(mock_callback_handler.check_for_pending_callbacks.call_count, 3) + mock_callback_handler.check_for_pending_callbacks.assert_called_with(mock_execution_arn) + + # Verify callback prompt was called with correct arguments + mock_callback_handler.prompt_callback_response.assert_called_once_with( + mock_execution_arn, "callback-123", mock.ANY + ) + + # Verify callback thread was created with daemon=True and started + mock_thread_class.assert_called_once_with(target=mock.ANY, daemon=True) + mock_callback_thread.start.assert_called_once() + + # Verify thread cleanup: join called when execution completes + mock_callback_thread.join.assert_called_once_with(timeout=0.5) + else: + # HTTP context - the user should not have been prompted + mock_callback_handler.check_for_pending_callbacks.assert_not_called() + mock_callback_handler.prompt_callback_response.assert_not_called() + mock_thread_class.assert_not_called() + @parameterized.expand( [ (False, "http://host.docker.internal:8080"), # internal emulator (True, "http://localhost:8080"), # external emulator ] ) + @patch("samcli.local.docker.durable_lambda_container.DurableCallbackHandler") @patch("samcli.local.docker.durable_lambda_container.LambdaContainer.start") @patch("samcli.local.docker.durable_lambda_container.threading.Thread") - def test_async_wait_for_result(self, is_external_emulator, expected_emulator_endpoint, mock_thread, mock_start): + def test_async_wait_for_result( + self, is_external_emulator, expected_emulator_endpoint, mock_thread, mock_start, mock_callback_handler_class + ): """Test wait_for_result with async invocation returns immediately and polls in background""" + # Mock callback handler to return no pending callbacks + mock_callback_handler = Mock() + mock_callback_handler.check_for_pending_callbacks.return_value = None + mock_callback_handler_class.return_value = mock_callback_handler + mock_emulator = Mock() mock_emulator.start_or_attach = Mock() mock_emulator._is_external_emulator = Mock(return_value=is_external_emulator)