fix: add InvocationCompleted event support

Author: bchampp

src/aws_durable_execution_sdk_python_testing/execution.py

@@ -60,6 +60,7 @@ def __init__(
         self.start_input: StartDurableExecutionInput = start_input
         self.operations: list[Operation] = operations
         self.updates: list[OperationUpdate] = []
+        self.invocation_completions: list[tuple[float, float, str]] = []  # (start_ts, end_ts, request_id)
         self.used_tokens: set[str] = set()
         # TODO: this will need to persist/rehydrate depending on inmemory vs sqllite store
         self._token_sequence: int = 0
@@ -101,6 +102,7 @@ def to_dict(self) -> dict[str, Any]:
             "StartInput": self.start_input.to_dict(),
             "Operations": [op.to_dict() for op in self.operations],
             "Updates": [update.to_dict() for update in self.updates],
+            "InvocationCompletions": [[start, end, req_id] for start, end, req_id in self.invocation_completions],
             "UsedTokens": list(self.used_tokens),
             "TokenSequence": self._token_sequence,
             "IsComplete": self.is_complete,
@@ -129,6 +131,9 @@ def from_dict(cls, data: dict[str, Any]) -> Execution:
         execution.updates = [
             OperationUpdate.from_dict(update_data) for update_data in data["Updates"]
         ]
+        execution.invocation_completions = [
+            tuple(item) for item in data.get("InvocationCompletions", [])
+        ]
         execution.used_tokens = set(data["UsedTokens"])
         execution._token_sequence = data["TokenSequence"]  # noqa: SLF001
         execution.is_complete = data["IsComplete"]
@@ -215,6 +220,12 @@ def has_pending_operations(self, execution: Execution) -> bool:
                 return True
         return False
 
+    def record_invocation_completion(
+        self, start_timestamp: float, end_timestamp: float, request_id: str
+    ) -> None:
+        """Record an invocation completion event."""
+        self.invocation_completions.append((start_timestamp, end_timestamp, request_id))
+
     def complete_success(self, result: str | None) -> None:
         """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION)."""
         self.result = DurableExecutionInvocationOutput(

src/aws_durable_execution_sdk_python_testing/executor.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import time
 import uuid
 from datetime import UTC, datetime
 from typing import TYPE_CHECKING
@@ -413,6 +414,20 @@ def get_execution_history(
         updates_dict: dict[str, OperationUpdate] = {u.operation_id: u for u in updates}
         durable_execution_arn: str = execution.durable_execution_arn
 
+        # Add InvocationCompleted events
+        for start_ts, end_ts, request_id in execution.invocation_completions:
+            invocation_event = HistoryEvent(
+                event_id=0,  # Temporary, will be reassigned
+                event_type="InvocationCompleted",
+                event_timestamp=datetime.fromtimestamp(end_ts, tz=UTC),
+                invocation_completed_details={
+                    "StartTimestamp": start_ts,
+                    "EndTimestamp": end_ts,
+                    "RequestId": request_id,
+                },
+            )
+            all_events.append(invocation_event)
+
         # Generate all events first (without final event IDs)
         for op in ops:
             operation_update: OperationUpdate | None = updates_dict.get(
@@ -769,14 +784,21 @@ async def invoke() -> None:
 
                 self._store.save(execution)
 
-                response: DurableExecutionInvocationOutput = self._invoker.invoke(
+                invocation_start = time.time()
+                response, request_id = self._invoker.invoke(
                     execution.start_input.function_name,
                     invocation_input,
                     execution.start_input.lambda_endpoint,
                 )
+                invocation_end = time.time()
 
                 # Reload execution after invocation in case it was completed via checkpoint
                 execution = self._store.load(execution_arn)
+                
+                # Record invocation completion and save immediately
+                execution.record_invocation_completion(invocation_start, invocation_end, request_id)
+                self._store.save(execution)
+                
                 if execution.is_complete:
                     logger.info(
                         "[%s] Execution completed during invocation, ignoring result",

src/aws_durable_execution_sdk_python_testing/invoker.py

@@ -3,6 +3,7 @@
 import json
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Protocol
+from uuid import uuid4
 
 import boto3  # type: ignore
 from aws_durable_execution_sdk_python.execution import (
@@ -65,7 +66,7 @@ def invoke(
         function_name: str,
         input: DurableExecutionInvocationInput,
         endpoint_url: str | None = None,
-    ) -> DurableExecutionInvocationOutput: ...  # pragma: no cover
+    ) -> tuple[DurableExecutionInvocationOutput, str]: ...  # pragma: no cover
 
     def update_endpoint(
         self, endpoint_url: str, region_name: str
@@ -96,14 +97,15 @@ def invoke(
         function_name: str,  # noqa: ARG002
         input: DurableExecutionInvocationInput,
         endpoint_url: str | None = None,  # noqa: ARG002
-    ) -> DurableExecutionInvocationOutput:
+    ) -> tuple[DurableExecutionInvocationOutput, str]:
         # TODO: reasses if function_name will be used in future
         input_with_client = DurableExecutionInvocationInputWithClient.from_durable_execution_invocation_input(
             input, self.service_client
         )
         context = create_test_lambda_context()
         response_dict = self.handler(input_with_client, context)
-        return DurableExecutionInvocationOutput.from_dict(response_dict)
+        output = DurableExecutionInvocationOutput.from_dict(response_dict)
+        return output, context.aws_request_id
 
     def update_endpoint(self, endpoint_url: str, region_name: str) -> None:
         """No-op for in-process invoker."""
@@ -192,7 +194,7 @@ def invoke(
         function_name: str,
         input: DurableExecutionInvocationInput,
         endpoint_url: str | None = None,
-    ) -> DurableExecutionInvocationOutput:
+    ) -> tuple[DurableExecutionInvocationOutput, str]:
         """Invoke AWS Lambda function and return durable execution result.
 
         Args:
@@ -201,7 +203,7 @@ def invoke(
             endpoint_url: Lambda endpoint url
 
         Returns:
-            DurableExecutionInvocationOutput: Result of the function execution
+            tuple: (DurableExecutionInvocationOutput, request_id)
 
         Raises:
             ResourceNotFoundException: If function does not exist
@@ -247,8 +249,13 @@ def invoke(
             response_payload = response["Payload"].read().decode("utf-8")
             response_dict = json.loads(response_payload)
 
+            # Extract request ID from response headers (x-amzn-RequestId or x-amzn-request-id)
+            headers = response.get("ResponseMetadata", {}).get("HTTPHeaders", {})
+            request_id = headers.get("x-amzn-RequestId") or headers.get("x-amzn-request-id") or f"local-{uuid4()}"
+
             # Convert to DurableExecutionInvocationOutput
-            return DurableExecutionInvocationOutput.from_dict(response_dict)
+            output = DurableExecutionInvocationOutput.from_dict(response_dict)
+            return output, request_id
 
         except client.exceptions.ResourceNotFoundException as e:
             msg = f"Function not found: {function_name}"

src/aws_durable_execution_sdk_python_testing/model.py

@@ -1329,6 +1329,7 @@ class Event:
     callback_succeeded_details: CallbackSucceededDetails | None = None
     callback_failed_details: CallbackFailedDetails | None = None
     callback_timed_out_details: CallbackTimedOutDetails | None = None
+    invocation_completed_details: dict[str, Any] | None = None
 
     @classmethod
     def from_dict(cls, data: dict) -> Event:
@@ -1447,6 +1448,8 @@ def from_dict(cls, data: dict) -> Event:
         if details_data := data.get("CallbackTimedOutDetails"):
             callback_timed_out_details = CallbackTimedOutDetails.from_dict(details_data)
 
+        invocation_completed_details = data.get("InvocationCompletedDetails")
+
         return cls(
             event_type=data["EventType"],
             event_timestamp=data["EventTimestamp"],
@@ -1479,6 +1482,7 @@ def from_dict(cls, data: dict) -> Event:
             callback_succeeded_details=callback_succeeded_details,
             callback_failed_details=callback_failed_details,
             callback_timed_out_details=callback_timed_out_details,
+            invocation_completed_details=invocation_completed_details,
         )
 
     def to_dict(self) -> dict[str, Any]:
@@ -1563,6 +1567,8 @@ def to_dict(self) -> dict[str, Any]:
             result["CallbackTimedOutDetails"] = (
                 self.callback_timed_out_details.to_dict()
             )
+        if self.invocation_completed_details is not None:
+            result["InvocationCompletedDetails"] = self.invocation_completed_details
         return result
 
     # region execution