33import json
44from dataclasses import replace
55from datetime import UTC , datetime
6+ from enum import Enum
67from threading import Lock
78from typing import Any
89from uuid import uuid4
2021 OperationUpdate ,
2122)
2223
23- # Import AWS exceptions
2424from aws_durable_execution_sdk_python_testing .exceptions import (
2525 IllegalStateException ,
2626 InvalidParameterValueException ,
2727)
28+
29+ # Import AWS exceptions
2830from aws_durable_execution_sdk_python_testing .model import (
2931 StartDurableExecutionInput ,
3032)
3133from aws_durable_execution_sdk_python_testing .token import CheckpointToken
3234
3335
36+ class CloseStatus (Enum ):
37+ """Close status for completed executions (mimics backend SWF CloseStatus)."""
38+
39+ COMPLETED = "COMPLETED"
40+ FAILED = "FAILED"
41+ TERMINATED = "TERMINATED"
42+ TIMED_OUT = "TIMED_OUT"
43+
44+
45+ class ExecutionStatus (Enum ):
46+ """Execution status for API responses (mimics backend ExecutionStatus)."""
47+
48+ RUNNING = "RUNNING"
49+ SUCCEEDED = "SUCCEEDED"
50+ FAILED = "FAILED"
51+ STOPPED = "STOPPED"
52+ TIMED_OUT = "TIMED_OUT"
53+
54+
3455class Execution :
3556 """Execution state."""
3657
@@ -52,12 +73,37 @@ def __init__(
5273 self .is_complete : bool = False
5374 self .result : DurableExecutionInvocationOutput | None = None
5475 self .consecutive_failed_invocation_attempts : int = 0
76+ self .close_status : CloseStatus | None = (
77+ None # Track close status like backend SWF
78+ )
5579
5680 @property
5781 def token_sequence (self ) -> int :
5882 """Get current token sequence value."""
5983 return self ._token_sequence
6084
85+ @property
86+ def status (self ) -> str :
87+ """Get execution status string (mimics backend ExecutionStatusConverter)."""
88+ if not self .is_complete :
89+ return ExecutionStatus .RUNNING .value
90+
91+ if not self .close_status :
92+ return ExecutionStatus .FAILED .value
93+
94+ # Convert CloseStatus to ExecutionStatus (like backend ExecutionStatusConverter)
95+ match self .close_status :
96+ case CloseStatus .COMPLETED :
97+ return ExecutionStatus .SUCCEEDED .value
98+ case CloseStatus .FAILED :
99+ return ExecutionStatus .FAILED .value
100+ case CloseStatus .TERMINATED :
101+ return ExecutionStatus .STOPPED .value
102+ case CloseStatus .TIMED_OUT :
103+ return ExecutionStatus .TIMED_OUT .value
104+ case _:
105+ return ExecutionStatus .FAILED .value
106+
61107 @staticmethod
62108 def new (input : StartDurableExecutionInput ) -> Execution : # noqa: A002
63109 # make a nicer arn
@@ -79,6 +125,7 @@ def to_dict(self) -> dict[str, Any]:
79125 "IsComplete" : self .is_complete ,
80126 "Result" : self .result .to_dict () if self .result else None ,
81127 "ConsecutiveFailedInvocationAttempts" : self .consecutive_failed_invocation_attempts ,
128+ "CloseStatus" : self .close_status .value if self .close_status else None ,
82129 }
83130
84131 @classmethod
@@ -112,6 +159,10 @@ def from_dict(cls, data: dict[str, Any]) -> Execution:
112159 execution .consecutive_failed_invocation_attempts = data [
113160 "ConsecutiveFailedInvocationAttempts"
114161 ]
162+ close_status_str = data .get ("CloseStatus" )
163+ execution .close_status = (
164+ CloseStatus (close_status_str ) if close_status_str else None
165+ )
115166
116167 return execution
117168
@@ -184,16 +235,36 @@ def has_pending_operations(self, execution: Execution) -> bool:
184235 return False
185236
186237 def complete_success (self , result : str | None ) -> None :
238+ """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION)."""
187239 self .result = DurableExecutionInvocationOutput (
188240 status = InvocationStatus .SUCCEEDED , result = result
189241 )
190242 self .is_complete = True
243+ self .close_status = CloseStatus .COMPLETED
191244
192245 def complete_fail (self , error : ErrorObject ) -> None :
246+ """Complete execution with failure (DecisionType.FAIL_WORKFLOW_EXECUTION)."""
247+ self .result = DurableExecutionInvocationOutput (
248+ status = InvocationStatus .FAILED , error = error
249+ )
250+ self .is_complete = True
251+ self .close_status = CloseStatus .FAILED
252+
253+ def complete_timeout (self , error : ErrorObject ) -> None :
254+ """Complete execution with timeout (SWF workflow timeout)."""
255+ self .result = DurableExecutionInvocationOutput (
256+ status = InvocationStatus .FAILED , error = error
257+ )
258+ self .is_complete = True
259+ self .close_status = CloseStatus .TIMED_OUT
260+
261+ def complete_stopped (self , error : ErrorObject ) -> None :
262+ """Complete execution as terminated (TerminateWorkflowExecutionV2Request)."""
193263 self .result = DurableExecutionInvocationOutput (
194264 status = InvocationStatus .FAILED , error = error
195265 )
196266 self .is_complete = True
267+ self .close_status = CloseStatus .TERMINATED
197268
198269 def find_operation (self , operation_id : str ) -> tuple [int , Operation ]:
199270 """Find operation by ID, return index and operation."""
0 commit comments