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,39 @@ 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+ msg : str = "close_status cannot be None"
93+ raise IllegalStateException (msg )
94+
95+ # Convert CloseStatus to ExecutionStatus (like backend ExecutionStatusConverter)
96+ match self .close_status :
97+ case CloseStatus .COMPLETED :
98+ return ExecutionStatus .SUCCEEDED .value
99+ case CloseStatus .FAILED :
100+ return ExecutionStatus .FAILED .value
101+ case CloseStatus .TERMINATED :
102+ return ExecutionStatus .STOPPED .value
103+ case CloseStatus .TIMED_OUT :
104+ return ExecutionStatus .TIMED_OUT .value
105+ case _:
106+ msg : str = f"Unexpected close status: { self .close_status } "
107+ raise InvalidParameterValueException (msg )
108+
61109 @staticmethod
62110 def new (input : StartDurableExecutionInput ) -> Execution : # noqa: A002
63111 # make a nicer arn
@@ -79,6 +127,7 @@ def to_dict(self) -> dict[str, Any]:
79127 "IsComplete" : self .is_complete ,
80128 "Result" : self .result .to_dict () if self .result else None ,
81129 "ConsecutiveFailedInvocationAttempts" : self .consecutive_failed_invocation_attempts ,
130+ "CloseStatus" : self .close_status .value if self .close_status else None ,
82131 }
83132
84133 @classmethod
@@ -112,6 +161,10 @@ def from_dict(cls, data: dict[str, Any]) -> Execution:
112161 execution .consecutive_failed_invocation_attempts = data [
113162 "ConsecutiveFailedInvocationAttempts"
114163 ]
164+ close_status_str = data .get ("CloseStatus" )
165+ execution .close_status = (
166+ CloseStatus (close_status_str ) if close_status_str else None
167+ )
115168
116169 return execution
117170
@@ -184,16 +237,36 @@ def has_pending_operations(self, execution: Execution) -> bool:
184237 return False
185238
186239 def complete_success (self , result : str | None ) -> None :
240+ """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION)."""
187241 self .result = DurableExecutionInvocationOutput (
188242 status = InvocationStatus .SUCCEEDED , result = result
189243 )
190244 self .is_complete = True
245+ self .close_status = CloseStatus .COMPLETED
191246
192247 def complete_fail (self , error : ErrorObject ) -> None :
248+ """Complete execution with failure (DecisionType.FAIL_WORKFLOW_EXECUTION)."""
249+ self .result = DurableExecutionInvocationOutput (
250+ status = InvocationStatus .FAILED , error = error
251+ )
252+ self .is_complete = True
253+ self .close_status = CloseStatus .FAILED
254+
255+ def complete_timeout (self , error : ErrorObject ) -> None :
256+ """Complete execution with timeout (SWF workflow timeout)."""
257+ self .result = DurableExecutionInvocationOutput (
258+ status = InvocationStatus .FAILED , error = error
259+ )
260+ self .is_complete = True
261+ self .close_status = CloseStatus .TIMED_OUT
262+
263+ def complete_stopped (self , error : ErrorObject ) -> None :
264+ """Complete execution as terminated (TerminateWorkflowExecutionV2Request)."""
193265 self .result = DurableExecutionInvocationOutput (
194266 status = InvocationStatus .FAILED , error = error
195267 )
196268 self .is_complete = True
269+ self .close_status = CloseStatus .TERMINATED
197270
198271 def find_operation (self , operation_id : str ) -> tuple [int , Operation ]:
199272 """Find operation by ID, return index and operation."""
0 commit comments