Convert LambdaContext from class to TypedDict and add proper type annotations

- Convert LambdaContext from class to TypedDict with required fields
- Add fields: invoke_id, client_context, cognito_identity, epoch_deadline_time_in_ms, invoked_function_arn, tenant_id
- Update DurableContext.from_lambda_context() to use LambdaContext type
- Update durable_handler decorator to use proper type annotations
- Add LambdaContext to package exports in __init__.py

Signed-off-by: Astraea Sinclair <quinsclr@amazon.com>

Author: FullyTyped

src/aws_durable_execution_sdk_python/__init__.py

@@ -1 +0,0 @@
-"""AWS Lambda Durable Executions Python SDK."""

src/aws_durable_execution_sdk_python/context.py

@@ -52,6 +52,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
+    from aws_durable_execution_sdk_python.lambda_context import LambdaContext
     from aws_durable_execution_sdk_python.state import CheckpointedResult
 
 P = TypeVar("P")  # Payload type
@@ -149,7 +150,7 @@ class DurableContext(DurableContextProtocol):
     def __init__(
         self,
         state: ExecutionState,
-        lambda_context: Any | None = None,
+        lambda_context: LambdaContext | None = None,
         parent_id: str | None = None,
         logger: Logger | None = None,
     ) -> None:
@@ -175,7 +176,7 @@ def __init__(
     @staticmethod
     def from_lambda_context(
         state: ExecutionState,
-        lambda_context: Any,
+        lambda_context: LambdaContext,
     ):
         return DurableContext(
             state=state,

src/aws_durable_execution_sdk_python/execution.py

@@ -25,6 +25,8 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, MutableMapping
 
+    from aws_durable_execution_sdk_python.lambda_context import LambdaContext
+
 
 logger = logging.getLogger(__name__)
 
@@ -187,10 +189,10 @@ def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput:
 
 def durable_handler(
     func: Callable[[Any, DurableContext], Any],
-) -> Callable[[Any, Any], Any]:
+) -> Callable[[Any, LambdaContext], Any]:
     logger.debug("Starting durable execution handler...")
 
-    def wrapper(event: Any, context: Any) -> MutableMapping[str, Any]:
+    def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
         invocation_input: DurableExecutionInvocationInput
         service_client: DurableServiceClient
 

src/aws_durable_execution_sdk_python/lambda_context.py

@@ -0,0 +1,10 @@
+from typing import Any, TypedDict
+
+
+class LambdaContext(TypedDict):
+    invoke_id: str
+    client_context: dict[str, Any]
+    cognito_identity: dict[str, str]
+    epoch_deadline_time_in_ms: int
+    invoked_function_arn: str
+    tenant_id: str