aws-durable-execution-sdk-python-testing 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

aws_durable_execution_sdk_python_testing/scheduler.py

@@ -3,9 +3,12 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import itertools
 import logging
 import threading
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
 from typing import TYPE_CHECKING, Any
 
 
@@ -64,6 +67,10 @@ class Scheduler:
 
     def __init__(self) -> None:
         self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
+        self._thread_executor = ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="durable-scheduler"
+        )
+        self._loop.set_default_executor(self._thread_executor)
         self._ready_event: threading.Event = threading.Event()
         self._thread: threading.Thread = threading.Thread(
             target=self._start_loop, daemon=True
@@ -99,7 +106,8 @@ class Scheduler:
             return
 
         self._running = False
-        self._loop.call_soon_threadsafe(self._cleanup_and_stop)
+        future = asyncio.run_coroutine_threadsafe(self._cleanup_and_stop(), self._loop)
+        future.result()
         self._thread.join()
 
     def is_started(self) -> bool:
@@ -116,16 +124,24 @@ class Scheduler:
             return 0
         return len(asyncio.all_tasks(self._loop))
 
-    def _cleanup_and_stop(self):
-        """Cancel all tasks and clear all events. Stop the event-loop."""
-        # Cancel all tasks
-        for task in asyncio.all_tasks(self._loop):
+    async def _cleanup_and_stop(self):
+        """Cancel all tasks, clear events, and stop the event loop."""
+        current_task = asyncio.current_task(self._loop)
+        tasks = [
+            task for task in asyncio.all_tasks(self._loop) if task is not current_task
+        ]
+
+        for task in tasks:
             task.cancel()
 
-        # Clear events (don't set them)
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        await self._loop.shutdown_default_executor()
+
         self._events.clear()
 
-        self._loop.stop()
+        self._loop.call_soon(self._loop.stop)
 
     def _start_loop(self):
         """Initialize the event-loop. The ready event notifies that the loop is started."""
@@ -146,8 +162,8 @@ class Scheduler:
     ) -> Future[Any]:
         """Call func after the delay.
 
-        If func is async it runs inside a thread-safe coroutine. If func is sync it runs in its own
-        threadpool, so it won't block the event loop.
+        If func is async it runs inside a thread-safe coroutine. If func is sync it runs in the
+        scheduler worker thread, so it won't block the event loop.
 
         Args:
             func (Callable[[], Any]): The function to call later. This can be an async or a standard
@@ -170,7 +186,7 @@ class Scheduler:
                     await asyncio.sleep(delay)
 
                     try:
-                        if asyncio.iscoroutinefunction(func):
+                        if inspect.iscoroutinefunction(func):
                             result = await func()
                         else:
                             result = await asyncio.to_thread(func)
@@ -215,13 +231,22 @@ class Scheduler:
         if event not in self._events:
             return False
 
+        async def wait_for_event_with_timeout() -> bool:
+            return await asyncio.wait_for(event.wait(), timeout)
+
         future: Future[bool] = asyncio.run_coroutine_threadsafe(
-            asyncio.wait_for(event.wait(), timeout), self._loop
+            wait_for_event_with_timeout(), self._loop
         )
 
         try:
-            return future.result()
-        except TimeoutError:
+            if timeout is None:
+                return future.result()
+            # Enforce the timeout from the waiting thread too. If the scheduler
+            # loop is blocked, the inner asyncio.wait_for timeout cannot fire.
+            margin = min(1.0, max(0.01, timeout * 0.1))
+            return future.result(timeout=timeout + margin)
+        except (TimeoutError, FutureTimeoutError):
+            future.cancel()
             return False
 
     def set_event(self, event: asyncio.Event):

aws_durable_execution_sdk_python_testing/stores/filesystem.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import json
 import logging
+import os
 from pathlib import Path
 
 from aws_durable_execution_sdk_python_testing.exceptions import (
@@ -46,14 +47,21 @@ class FileSystemExecutionStore(BaseExecutionStore):
         file_path = self._get_file_path(execution.durable_execution_arn)
         data = execution.to_json_dict()
 
-        with open(file_path, "w", encoding="utf-8") as f:
+        # Write to a temporary file in the same directory, then replace
+        # the target in one atomic step. A concurrent reader then sees
+        # either the old file or the new file, never a partial or empty
+        # one. Saves for one execution run on its worker lane, so two
+        # writes never target the same temp path at once.
+        tmp_path = file_path.with_name(f"{file_path.name}.tmp")
+        with open(tmp_path, "w", encoding="utf-8") as f:
             json.dump(data, f, indent=2)
+        os.replace(tmp_path, file_path)
 
     def load(self, execution_arn: str) -> Execution:
         """Load execution from file system."""
         file_path = self._get_file_path(execution_arn)
         if not file_path.exists():
-            msg = f"Execution {execution_arn} not found"
+            msg = "Durable Execution does not exist"
             raise ResourceNotFoundException(msg)
 
         with open(file_path, encoding="utf-8") as f:

aws_durable_execution_sdk_python_testing/stores/sqlite.py

@@ -4,6 +4,8 @@ from __future__ import annotations
 
 import json
 import sqlite3
+from collections.abc import Iterator
+from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
 from typing import Any, cast
@@ -43,10 +45,20 @@ class SQLiteExecutionStore(ExecutionStore):
         conn.execute("PRAGMA synchronous=NORMAL;")
         return conn
 
+    @contextmanager
+    def _connection(self) -> Iterator[sqlite3.Connection]:
+        """Open a connection, preserve transaction handling, and always close it."""
+        conn = self._get_connection()
+        try:
+            with conn:
+                yield conn
+        finally:
+            conn.close()
+
     def _init_db(self) -> None:
         """Initialize database schema."""
         try:
-            with self._get_connection() as conn:
+            with self._connection() as conn:
                 conn.execute("""
                     CREATE TABLE IF NOT EXISTS executions (
                         durable_execution_arn TEXT PRIMARY KEY,
@@ -80,7 +92,7 @@ class SQLiteExecutionStore(ExecutionStore):
             execution_op = execution.get_operation_execution_started()
             status: str = execution.current_status().value
 
-            with self._get_connection() as conn:
+            with self._connection() as conn:
                 conn.execute(
                     """
                     INSERT OR REPLACE INTO executions
@@ -111,7 +123,7 @@ class SQLiteExecutionStore(ExecutionStore):
     def load(self, execution_arn: str) -> Execution:
         """Load execution from SQLite."""
         try:
-            with self._get_connection() as conn:
+            with self._connection() as conn:
                 cursor: sqlite3.Cursor = conn.execute(
                     "SELECT data FROM executions WHERE durable_execution_arn = ?",
                     (execution_arn,),
@@ -119,7 +131,7 @@ class SQLiteExecutionStore(ExecutionStore):
                 row: tuple[str] | None = cursor.fetchone()
 
             if not row:
-                raise ResourceNotFoundException(f"Execution {execution_arn} not found")
+                raise ResourceNotFoundException("Durable Execution does not exist")
 
             return Execution.from_json_dict(json.loads(row[0]))
         except sqlite3.Error as e:
@@ -204,7 +216,7 @@ class SQLiteExecutionStore(ExecutionStore):
                 )
                 params_with_limit = params
 
-            with self._get_connection() as conn:
+            with self._connection() as conn:
                 # Get total count for pagination
                 total_count: int = int(conn.execute(count_query, params).fetchone()[0])
 
@@ -247,7 +259,7 @@ class SQLiteExecutionStore(ExecutionStore):
     def get_execution_metadata(self, execution_arn: str) -> dict[str, Any] | None:
         """Get just the metadata without full deserialization for performance."""
         try:
-            with self._get_connection() as conn:
+            with self._connection() as conn:
                 cursor: sqlite3.Cursor = conn.execute(
                     "SELECT function_name, execution_name, status, start_timestamp, end_timestamp FROM executions WHERE durable_execution_arn = ?",
                     (execution_arn,),

aws_durable_execution_sdk_python_testing/threads.py

@@ -0,0 +1,157 @@
+"""A bounded pool of daemon threads for blocking calls.
+
+The runner runs two kinds of blocking call off the scheduler loop: a
+handler invocation, and the synchronous Invoke of a plain chained-invoke
+target. Each blocks until the endpoint answers or the read timeout
+expires. Python cannot interrupt a thread blocked in a socket read. So
+the runner cannot stop such a call. It can stop the call from holding
+the process.
+
+``concurrent.futures.ThreadPoolExecutor`` does not allow that. Its
+workers are not daemon threads, and the interpreter joins them at exit.
+So a runner closed while one Invoke is blocked keeps the process alive
+until that read timeout. This pool has the same shape (a worker cap, a
+FIFO queue, lazily started workers) and differs in two ways. Its workers
+are daemon threads, which the interpreter does not wait for. Its
+``shutdown`` drops queued work, refuses new work, and never joins.
+"""
+
+from __future__ import annotations
+
+import os
+import queue
+import threading
+from concurrent.futures import Executor, Future
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+def default_max_workers() -> int:
+    """The worker cap ``ThreadPoolExecutor`` uses when none is given."""
+    return min(32, (os.cpu_count() or 1) + 4)
+
+
+class _WorkItem:
+    def __init__(
+        self,
+        future: Future[Any],
+        fn: Callable[..., Any],
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> None:
+        self.future = future
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+
+    def run(self) -> None:
+        if not self.future.set_running_or_notify_cancel():
+            return
+        try:
+            result = self.fn(*self.args, **self.kwargs)
+        except BaseException as exc:  # noqa: BLE001 — the future carries the error
+            self.future.set_exception(exc)
+        else:
+            self.future.set_result(result)
+
+
+class _Stop:
+    """Sentinel that makes an idle worker return."""
+
+
+class DaemonThreadPool(Executor):
+    """Run callables on at most ``max_workers`` daemon threads, in FIFO order.
+
+    A worker is started when work is submitted and no worker is idle,
+    until the cap is reached. Beyond the cap, work waits in the queue in
+    submission order, as in ``ThreadPoolExecutor``.
+
+    ``shutdown`` marks the pool closed, cancels every queued future, and
+    tells idle workers to return. A worker busy in a blocking call is
+    left alone: it finishes when its call returns, and its result lands
+    on a future nobody reads. ``wait`` is accepted for API compatibility
+    and ignored, because joining a blocked worker would wait for the read
+    timeout this pool exists to avoid. Workers are daemon threads, so
+    process exit does not wait for them either.
+    """
+
+    def __init__(
+        self, max_workers: int | None = None, thread_name_prefix: str = ""
+    ) -> None:
+        if max_workers is None:
+            max_workers = default_max_workers()
+        if max_workers <= 0:
+            msg = "max_workers must be greater than 0"
+            raise ValueError(msg)
+        self._max_workers = max_workers
+        self._thread_name_prefix = thread_name_prefix or "daemon-pool"
+        self._queue: queue.SimpleQueue[_WorkItem | _Stop] = queue.SimpleQueue()
+        self._idle = threading.Semaphore(0)
+        self._threads: list[threading.Thread] = []
+        self._closed = False
+        self._lock = threading.Lock()
+
+    @property
+    def max_workers(self) -> int:
+        return self._max_workers
+
+    def submit(
+        self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any
+    ) -> Future[Any]:
+        with self._lock:
+            if self._closed:
+                msg = "cannot schedule new futures after shutdown"
+                raise RuntimeError(msg)
+            future: Future[Any] = Future()
+            self._queue.put(_WorkItem(future, fn, args, kwargs))
+            self._adjust_thread_count()
+            return future
+
+    def _adjust_thread_count(self) -> None:
+        # An idle worker will take the item; do not start another.
+        if self._idle.acquire(blocking=False):
+            return
+        if len(self._threads) >= self._max_workers:
+            return
+        thread = threading.Thread(
+            target=self._worker,
+            name=f"{self._thread_name_prefix}_{len(self._threads)}",
+            daemon=True,
+        )
+        self._threads.append(thread)
+        thread.start()
+
+    def _worker(self) -> None:
+        while True:
+            item: _WorkItem | _Stop = self._queue.get()
+            if isinstance(item, _Stop):
+                return
+            item.run()
+            del item
+            self._idle.release()
+
+    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:  # noqa: ARG002, FBT001, FBT002
+        """Close the pool without waiting for busy workers.
+
+        Queued work is cancelled whatever ``cancel_futures`` says: work
+        that has not started must not start after shutdown. ``wait`` is
+        ignored, see the class docstring.
+        """
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            drained: list[_WorkItem] = []
+            while True:
+                try:
+                    item = self._queue.get_nowait()
+                except queue.Empty:
+                    break
+                if isinstance(item, _WorkItem):
+                    drained.append(item)
+            for _ in self._threads:
+                self._queue.put(_Stop())
+        for item in drained:
+            item.future.cancel()

aws_durable_execution_sdk_python_testing/token.py

@@ -13,9 +13,14 @@ class CheckpointToken:
 
     execution_arn: str
     token_sequence: int
+    invocation_id: str = ""
 
     def to_str(self) -> str:
-        data = {"arn": self.execution_arn, "seq": self.token_sequence}
+        data = {
+            "arn": self.execution_arn,
+            "seq": self.token_sequence,
+            "inv": self.invocation_id,
+        }
         json_str = json.dumps(data, separators=(",", ":"))
         # str -> bytes -> base64 bytes -> str
         return base64.b64encode(json_str.encode()).decode()
@@ -25,7 +30,11 @@ class CheckpointToken:
         # str -> base64 bytes -> str
         decoded = base64.b64decode(token).decode()
         data = json.loads(decoded)
-        return cls(execution_arn=data["arn"], token_sequence=data["seq"])
+        return cls(
+            execution_arn=data["arn"],
+            token_sequence=data["seq"],
+            invocation_id=data.get("inv", ""),
+        )
 
 
 @dataclass(frozen=True)

aws_durable_execution_sdk_python_testing/web/handlers.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import base64
 import json
 import logging
 from abc import ABC, abstractmethod
@@ -26,7 +25,6 @@ from aws_durable_execution_sdk_python_testing.model import (
     ListDurableExecutionsResponse,
     SendDurableExecutionCallbackFailureRequest,
     SendDurableExecutionCallbackFailureResponse,
-    SendDurableExecutionCallbackHeartbeatRequest,
     SendDurableExecutionCallbackHeartbeatResponse,
     SendDurableExecutionCallbackSuccessResponse,
     StartDurableExecutionInput,
@@ -438,7 +436,7 @@ class StopDurableExecutionHandler(EndpointHandler):
 class GetDurableExecutionStateHandler(EndpointHandler):
     """Handler for GET /2025-12-01/durable-executions/{arn}/state."""
 
-    def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:  # noqa: ARG002
+    def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:
         """Handle get durable execution state request.
 
         Args:
@@ -452,8 +450,19 @@ class GetDurableExecutionStateHandler(EndpointHandler):
             state_route = cast(GetDurableExecutionStateRoute, parsed_route)
             execution_arn: str = state_route.arn
 
+            checkpoint_token: str | None = self._parse_query_param(
+                request, "CheckpointToken"
+            )
+            marker: str | None = self._parse_query_param(request, "Marker") or None
+            max_items: str | None = self._parse_query_param(request, "MaxItems")
+
             state_response: GetDurableExecutionStateResponse = (
-                self.executor.get_execution_state(execution_arn)
+                self.executor.get_execution_state(
+                    execution_arn,
+                    checkpoint_token=checkpoint_token,
+                    marker=marker,
+                    max_items=int(max_items) if max_items else None,
+                )
             )
 
             response_data: dict[str, Any] = state_response.to_dict()
@@ -493,12 +502,18 @@ class GetDurableExecutionHistoryHandler(EndpointHandler):
                 if include_execution_data_str
                 else False
             )
+            reverse_order_str: str | None = self._parse_query_param(
+                request, "ReverseOrder"
+            )
+            reverse_order: bool = (
+                reverse_order_str == "true" if reverse_order_str else False
+            )
 
             history_response: GetDurableExecutionHistoryResponse = (
                 self.executor.get_execution_history(
                     execution_arn,
                     include_execution_data=include_execution_data,
-                    reverse_order=False,
+                    reverse_order=reverse_order,
                     marker=marker,
                     max_items=int(max_items) if max_items else None,
                 )
@@ -794,17 +809,28 @@ class UpdateLambdaEndpointHandler(EndpointHandler):
         try:
             body = self._parse_json_body(request)
             endpoint_url = body.get("EndpointUrl")
-            region_name = body.get("RegionName", "us-east-1")
+            region_name = body.get("RegionName")
 
             if not endpoint_url:
                 return self._handle_aws_exception(
                     InvalidParameterValueException("EndpointUrl is required")
                 )
+            # The runner emulates one region, fixed at startup, and every
+            # execution is created in it. A request that names another
+            # region would give the runner two regions at once, so it is
+            # rejected instead of applied to the invoker alone.
+            if region_name is not None and region_name != self.executor.region:
+                return self._handle_aws_exception(
+                    InvalidParameterValueException(
+                        f"RegionName must be {self.executor.region}, "
+                        f"the region this runner emulates; got {region_name}"
+                    )
+                )
 
             # Update the invoker's Lambda endpoint
             invoker = self.executor._invoker  # noqa: SLF001
             logger.info("Updating lambda endpoint to %s", endpoint_url)
-            invoker.update_endpoint(endpoint_url, region_name)
+            invoker.update_endpoint(endpoint_url, self.executor.region)
             return self._success_response(
                 {"message": "Lambda endpoint updated successfully"}
             )

aws_durable_execution_sdk_python_testing/web/routes.py

@@ -39,14 +39,22 @@ class Route:
     def from_string(cls, path: str) -> Route:
         """Create a Route from a string.
 
+        Each segment is URL-decoded; ``raw_path`` is preserved as the
+        original wire path. Splitting on ``/`` happens before decoding so
+        that an encoded ``%2F`` inside a captured value (e.g. an ARN that
+        contains ``/``) stays inside its segment instead of being treated
+        as a path separator.
+
         Args:
             path: The raw path string
 
         Returns:
-            Route instance with parsed segments
+            Route instance with parsed, URL-decoded segments
         """
-        # Remove leading/trailing slashes and split into segments
-        segments = [s for s in path.strip("/").split("/") if s]
+        # Remove leading/trailing slashes, split on '/', then URL-decode each
+        # segment. Order matters: split on the literal '/' first so '%2F'-
+        # encoded slashes inside values don't act as separators.
+        segments = [unquote(s) for s in path.strip("/").split("/") if s]
         return cls(raw_path=path, segments=segments)
 
     def matches_pattern(self, pattern: list[str]) -> bool:
@@ -445,7 +453,7 @@ class CallbackSuccessRoute(BytesPayloadRoute):
         return cls(
             raw_path=route.raw_path,
             segments=route.segments,
-            callback_id=unquote(route.segments[2]),
+            callback_id=route.segments[2],
         )
 
 
@@ -488,7 +496,7 @@ class CallbackFailureRoute(BytesPayloadRoute):
         return cls(
             raw_path=route.raw_path,
             segments=route.segments,
-            callback_id=unquote(route.segments[2]),
+            callback_id=route.segments[2],
         )
 
 
@@ -531,7 +539,7 @@ class CallbackHeartbeatRoute(Route):
         return cls(
             raw_path=route.raw_path,
             segments=route.segments,
-            callback_id=unquote(route.segments[2]),
+            callback_id=route.segments[2],
         )
 
 

aws_durable_execution_sdk_python_testing/web/serialization.py

@@ -7,11 +7,10 @@ along with AWS-compatible implementations using boto's rest-json serializers.
 from __future__ import annotations
 
 import json
-import os
+import os  # noqa: F401 - accessed via module namespace in tests (mock.patch)
 from typing import Any, Protocol
 from datetime import datetime
 
-import aws_durable_execution_sdk_python
 import botocore.loaders  # type: ignore
 from botocore.model import ServiceModel  # type: ignore
 from botocore.parsers import create_parser  # type: ignore

aws_durable_execution_sdk_python_testing/web/server.py

@@ -1,4 +1,4 @@
-"""Local testing web server for AWS Lambda Durable Functions that mimics the actual Lambda backend services."""
+"""Local testing web server for AWS Lambda Durable Functions for local testing."""
 
 from __future__ import annotations
 

aws_durable_execution_sdk_python_testing/worker/__init__.py

@@ -0,0 +1,9 @@
+"""Per-execution serial actor model for the local runner.
+
+Each execution is served by a single worker that runs all of its
+operations (checkpoint, get-state, invoke, timer, callback, completion)
+one at a time on its own lane. Serializing the operations of one
+execution keeps its state consistent when several callers act on it at
+once; different executions are served by different workers and run in
+parallel.
+"""