fix: parse bytes for callback operations (#59)

Author: bchampp

src/aws_durable_execution_sdk_python_testing/model.py

@@ -1411,13 +1411,13 @@ class SendDurableExecutionCallbackFailureRequest:
     error: ErrorObject | None = None
 
     @classmethod
-    def from_dict(cls, data: dict) -> SendDurableExecutionCallbackFailureRequest:
-        error = None
-        if error_data := data.get("Error"):
-            error = ErrorObject.from_dict(error_data)
+    def from_dict(
+        cls, data: dict, callback_id: str
+    ) -> SendDurableExecutionCallbackFailureRequest:
+        error = ErrorObject.from_dict(data) if data else None
 
         return cls(
-            callback_id=data["CallbackId"],
+            callback_id=callback_id,
             error=error,
         )
 

src/aws_durable_execution_sdk_python_testing/web/handlers.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, cast
@@ -27,7 +28,6 @@
     SendDurableExecutionCallbackFailureResponse,
     SendDurableExecutionCallbackHeartbeatRequest,
     SendDurableExecutionCallbackHeartbeatResponse,
-    SendDurableExecutionCallbackSuccessRequest,
     SendDurableExecutionCallbackSuccessResponse,
     StartDurableExecutionInput,
     StartDurableExecutionOutput,
@@ -37,7 +37,6 @@
 from aws_durable_execution_sdk_python_testing.web.models import (
     HTTPRequest,
     HTTPResponse,
-    parse_json_body,
 )
 from aws_durable_execution_sdk_python_testing.web.routes import (
     CallbackFailureRoute,
@@ -92,9 +91,21 @@ def _parse_json_body(self, request: HTTPRequest) -> dict[str, Any]:
             dict: The parsed JSON data
 
         Raises:
-            ValueError: If the request body is empty or invalid JSON
+            InvalidParameterValueException: If the request body is empty or invalid JSON
         """
-        return parse_json_body(request)
+        if not request.body:
+            msg = "Request body is required"
+            raise InvalidParameterValueException(msg)
+
+        # Handle both dict and bytes body types
+        if isinstance(request.body, dict):
+            return request.body
+
+        try:
+            return json.loads(request.body.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            msg = f"Invalid JSON in request body: {e}"
+            raise InvalidParameterValueException(msg) from e
 
     def _json_response(
         self,
@@ -631,20 +642,24 @@ def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:
             HTTPResponse: The HTTP response to send to the client
         """
         try:
-            body_data: dict[str, Any] = self._parse_json_body(request)
-            callback_request: SendDurableExecutionCallbackSuccessRequest = (
-                SendDurableExecutionCallbackSuccessRequest.from_dict(body_data)
-            )
-
             callback_route = cast(CallbackSuccessRoute, parsed_route)
             callback_id: str = callback_route.callback_id
 
+            # For binary payload operations, body is raw bytes
+            result_bytes = request.body if isinstance(request.body, bytes) else b""
+
             callback_response: SendDurableExecutionCallbackSuccessResponse = (  # noqa: F841
                 self.executor.send_callback_success(
-                    callback_id=callback_id, result=callback_request.result
+                    callback_id=callback_id, result=result_bytes
                 )
             )
 
+            logger.debug(
+                "Callback %s succeeded with result: %s",
+                callback_id,
+                result_bytes.decode("utf-8", errors="replace"),
+            )
+
             # Callback success response is empty
             return self._success_response({})
 
@@ -672,20 +687,26 @@ def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:
             HTTPResponse: The HTTP response to send to the client
         """
         try:
+            callback_route = cast(CallbackFailureRoute, parsed_route)
+            callback_id: str = callback_route.callback_id
+
             body_data: dict[str, Any] = self._parse_json_body(request)
             callback_request: SendDurableExecutionCallbackFailureRequest = (
-                SendDurableExecutionCallbackFailureRequest.from_dict(body_data)
+                SendDurableExecutionCallbackFailureRequest.from_dict(
+                    body_data, callback_id
+                )
             )
 
-            callback_route = cast(CallbackFailureRoute, parsed_route)
-            callback_id: str = callback_route.callback_id
-
             callback_response: SendDurableExecutionCallbackFailureResponse = (  # noqa: F841
                 self.executor.send_callback_failure(
                     callback_id=callback_id, error=callback_request.error
                 )
             )
 
+            logger.debug(
+                "Callback %s failed with error: %s", callback_id, callback_request.error
+            )
+
             # Callback failure response is empty
             return self._success_response({})
 

src/aws_durable_execution_sdk_python_testing/web/models.py

@@ -25,13 +25,38 @@
 
 @dataclass(frozen=True)
 class HTTPRequest:
-    """HTTP request data model with dict body for handler logic."""
+    """HTTP request data model with dict or bytes body for handler logic."""
 
     method: str
     path: Route
     headers: dict[str, str]
     query_params: dict[str, list[str]]
-    body: dict[str, Any]
+    body: dict[str, Any] | bytes
+
+    @classmethod
+    def from_raw_bytes(
+        cls,
+        body_bytes: bytes,
+        method: str = "POST",
+        path: Route | None = None,
+        headers: dict[str, str] | None = None,
+        query_params: dict[str, list[str]] | None = None,
+    ) -> HTTPRequest:
+        """Create HTTPRequest with raw bytes body (no parsing)."""
+        if headers is None:
+            headers = {}
+        if query_params is None:
+            query_params = {}
+        if path is None:
+            path = Route.from_string("")
+
+        return cls(
+            method=method,
+            path=path,
+            headers=headers,
+            query_params=query_params,
+            body=body_bytes,
+        )
 
     @classmethod
     def from_bytes(
@@ -269,22 +294,3 @@ def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse:
             HTTPResponse: The HTTP response to send to the client
         """
         ...  # pragma: no cover
-
-
-def parse_json_body(request: HTTPRequest) -> dict[str, Any]:
-    """Parse JSON body from HTTP request.
-
-    Args:
-        request: The HTTP request containing the dict body
-
-    Returns:
-        dict: The parsed JSON data (now just returns the body directly)
-
-    Raises:
-        ValueError: If the request body is empty
-    """
-    if not request.body:
-        msg = "Request body is required"
-        raise InvalidParameterValueException(msg)
-
-    return request.body

src/aws_durable_execution_sdk_python_testing/web/routes.py

@@ -401,7 +401,12 @@ def from_route(cls, route: Route) -> ListDurableExecutionsByFunctionRoute:
 
 
 @dataclass(frozen=True)
-class CallbackSuccessRoute(Route):
+class BytesPayloadRoute(Route):
+    """Base class for routes that handle raw bytes payloads instead of JSON."""
+
+
+@dataclass(frozen=True)
+class CallbackSuccessRoute(BytesPayloadRoute):
     """Route: POST /2025-12-01/durable-execution-callbacks/{callback_id}/succeed"""
 
     callback_id: str
@@ -444,7 +449,7 @@ def from_route(cls, route: Route) -> CallbackSuccessRoute:
 
 
 @dataclass(frozen=True)
-class CallbackFailureRoute(Route):
+class CallbackFailureRoute(BytesPayloadRoute):
     """Route: POST /2025-12-01/durable-execution-callbacks/{callback_id}/fail"""
 
     callback_id: str

src/aws_durable_execution_sdk_python_testing/web/server.py

@@ -42,6 +42,7 @@
     HTTPResponse,
 )
 from aws_durable_execution_sdk_python_testing.web.routes import (
+    BytesPayloadRoute,
     CallbackFailureRoute,
     CallbackHeartbeatRoute,
     CallbackSuccessRoute,
@@ -120,15 +121,25 @@ def _handle_request(self, method: str) -> None:
                 self.rfile.read(content_length) if content_length > 0 else b""
             )
 
-            # Create strongly-typed HTTP request object with pre-parsed body
-            request: HTTPRequest = HTTPRequest.from_bytes(
-                body_bytes=body_bytes,
-                operation_name=None,  # Could be enhanced to map routes to AWS operation names
-                method=method,
-                path=parsed_route,
-                headers=dict(self.headers),
-                query_params=query_params,
-            )
+            # For callback operations, use raw bytes directly
+            if isinstance(parsed_route, BytesPayloadRoute):
+                request = HTTPRequest.from_raw_bytes(
+                    body_bytes=body_bytes,
+                    method=method,
+                    path=parsed_route,
+                    headers=dict(self.headers),
+                    query_params=query_params,
+                )
+            else:
+                # Create strongly-typed HTTP request object with pre-parsed body
+                request = HTTPRequest.from_bytes(
+                    body_bytes=body_bytes,
+                    operation_name=None,
+                    method=method,
+                    path=parsed_route,
+                    headers=dict(self.headers),
+                    query_params=query_params,
+                )
 
             # Handle request with appropriate handler
             response: HTTPResponse = handler.handle(parsed_route, request)

tests/model_test.py

@@ -798,32 +798,38 @@ def test_send_durable_execution_callback_success_response_creation():
 
 def test_send_durable_execution_callback_failure_request_serialization():
     """Test SendDurableExecutionCallbackFailureRequest from_dict/to_dict round-trip."""
-    data = {
-        "CallbackId": "callback-123",
-        "Error": {"ErrorMessage": "callback failed"},
-    }
+    data = {"ErrorMessage": "callback failed"}
 
-    request_obj = SendDurableExecutionCallbackFailureRequest.from_dict(data)
+    request_obj = SendDurableExecutionCallbackFailureRequest.from_dict(
+        data, "callback-123"
+    )
     assert request_obj.callback_id == "callback-123"
     assert request_obj.error.message == "callback failed"
 
     result_data = request_obj.to_dict()
-    assert result_data == data
+    expected_data = {
+        "CallbackId": "callback-123",
+        "Error": {"ErrorMessage": "callback failed"},
+    }
+    assert result_data == expected_data
 
     # Test round-trip
-    round_trip = SendDurableExecutionCallbackFailureRequest.from_dict(result_data)
+    round_trip = SendDurableExecutionCallbackFailureRequest.from_dict(
+        result_data.get("Error", {}), result_data["CallbackId"]
+    )
     assert round_trip == request_obj
 
 
 def test_send_durable_execution_callback_failure_request_minimal():
     """Test SendDurableExecutionCallbackFailureRequest with only required fields."""
-    data = {"CallbackId": "callback-123"}
 
-    request_obj = SendDurableExecutionCallbackFailureRequest.from_dict(data)
+    request_obj = SendDurableExecutionCallbackFailureRequest.from_dict(
+        {}, "callback-123"
+    )
     assert request_obj.error is None
 
     result_data = request_obj.to_dict()
-    assert result_data == data
+    assert result_data == {"CallbackId": "callback-123"}
 
 
 def test_send_durable_execution_callback_failure_response_creation():