Fix legacy entity protocol: propagate call_entity failures and stop double-encoding results (#167)

* Propagate entity operation failures over legacy entity protocol
* Fix double-encoded call_entity result over legacy entity protocol
* Pop legacy entity map entries after handling response

Author: andystaples

CHANGELOG.md

@@ -19,6 +19,9 @@ CHANGED
 
 FIXED
 
+- Fixed `OrchestrationContext.call_entity` not propagating entity operation failures when running under the legacy entity protocol (used by the Azure Functions Durable extension). A failed entity operation now raises `TaskFailedError` in the calling orchestration instead of silently completing, matching the behavior of the current entity protocol and the .NET SDK.
+- Fixed `OrchestrationContext.call_entity` returning a double-encoded result under the legacy entity protocol. The entity's return value was left as a raw serialized JSON string (for example, a returned string arrived with extra quotes and dicts/lists arrived as strings); it is now fully deserialized and coerced to the requested `return_type`, matching the current entity protocol.
+- Fixed unbounded growth of the internal entity request/lock tracking maps when using entities over the legacy entity protocol. Entries are now released as each response is handled, reducing memory use in long-running orchestrations that call or lock entities.
 - Fixed schedules created with `durabletask.scheduled` not appearing in the Durable Task Scheduler dashboard. The schedule entity state is now persisted in a format compatible with the .NET SDK: the `status` is serialized as its numeric enum value, the `interval` as a .NET `TimeSpan` string, and the `orchestration_input` as a JSON-encoded string, all using .NET property names, so the dashboard can read it without a JSON deserialization error.
 
 ## v1.7.0

durabletask/internal/helpers.py

@@ -3,6 +3,7 @@
 
 import traceback
 from datetime import datetime, timezone
+from typing import Any, cast
 
 from google.protobuf import timestamp_pb2, wrappers_pb2
 
@@ -136,6 +137,59 @@ def new_failure_details(ex: Exception, _visited: set[int] | None = None) -> pb.T
     )
 
 
+def _failure_details_from_core_dict(fd: dict[str, Any]) -> pb.TaskFailureDetails:
+    """Convert a serialized DurableTask.Core ``FailureDetails`` dict to protobuf."""
+    inner = fd.get("InnerFailure")
+    stack_trace = fd.get("StackTrace")
+    return pb.TaskFailureDetails(
+        errorType=str(fd.get("ErrorType") or ""),
+        errorMessage=str(fd.get("ErrorMessage") or ""),
+        stackTrace=get_string_value(str(stack_trace) if stack_trace is not None else None),
+        innerFailure=_failure_details_from_core_dict(cast(dict[str, Any], inner)) if isinstance(inner, dict) else None,
+        isNonRetriable=bool(fd.get("IsNonRetriable", False)),
+    )
+
+
+def entity_response_failure_details(
+        response: dict[str, Any],
+        error_content: Any = None) -> pb.TaskFailureDetails:
+    """Build failure details from a failed legacy-protocol entity ``ResponseMessage``.
+
+    Call this only for responses that :func:`is_entity_error_response` reports as
+    failures. In the WebJobs "old protocol" ``ResponseMessage`` (see
+    ``EntityScheduler/ResponseMessage.cs``), a failed operation serializes the
+    human-readable content into ``result`` while ``exceptionType`` carries only
+    the exception's type name (a presence marker) -- it is *not* the message.
+    This mirrors ``azure-functions-durable-python`` / ``-js``, which read the
+    message from ``result`` and ignore ``exceptionType``'s value.
+
+    Parameters
+    ----------
+    response:
+        The deserialized ``ResponseMessage`` dict.
+    error_content:
+        The already-deserialized ``result`` payload, used as the failure
+        message. A structured ``failureDetails`` object, if present, takes
+        precedence (current-protocol shape).
+    """
+    failure_details = response.get("failureDetails")
+    if isinstance(failure_details, dict):
+        return _failure_details_from_core_dict(cast(dict[str, Any], failure_details))
+    error_type = str(response.get("exceptionType") or "")
+    error_message = "" if error_content is None else str(error_content)
+    return pb.TaskFailureDetails(errorType=error_type, errorMessage=error_message)
+
+
+def is_entity_error_response(response: dict[str, Any]) -> bool:
+    """Return ``True`` if a legacy-protocol entity ``ResponseMessage`` is a failure.
+
+    In the WebJobs "old protocol" a failed operation is marked by the presence of
+    an ``exceptionType`` field (successful responses omit it). Current-protocol
+    payloads may instead carry a structured ``failureDetails`` object.
+    """
+    return "exceptionType" in response or isinstance(response.get("failureDetails"), dict)
+
+
 def new_event_sent_event(event_id: int, instance_id: str, input: str):
     return pb.HistoryEvent(
         eventId=event_id,

durabletask/worker.py

@@ -2588,10 +2588,10 @@ def _cancel_timer() -> None:
                     raise TypeError("Unexpected sub-orchestration task type")
             elif event.HasField("eventRaised"):
                 if event.eventRaised.name in ctx._entity_task_id_map:  # pyright: ignore[reportPrivateUsage]
-                    entity_id, operation, task_id = ctx._entity_task_id_map.get(event.eventRaised.name, (None, None, None))  # pyright: ignore[reportPrivateUsage]
-                    self._handle_entity_event_raised(ctx, event, entity_id, task_id, False)
+                    entity_id, operation, task_id = ctx._entity_task_id_map.pop(event.eventRaised.name, (None, None, None))  # pyright: ignore[reportPrivateUsage]
+                    self._handle_entity_event_raised(ctx, event, entity_id, task_id, False, operation)
                 elif event.eventRaised.name in ctx._entity_lock_task_id_map:  # pyright: ignore[reportPrivateUsage]
-                    entity_id, task_id = ctx._entity_lock_task_id_map.get(event.eventRaised.name, (None, None))  # pyright: ignore[reportPrivateUsage]
+                    entity_id, task_id = ctx._entity_lock_task_id_map.pop(event.eventRaised.name, (None, None))  # pyright: ignore[reportPrivateUsage]
                     self._handle_entity_event_raised(ctx, event, entity_id, task_id, True)
                 else:
                     # event names are case-insensitive
@@ -2818,7 +2818,8 @@ def _handle_entity_event_raised(self,
                                     event: pb.HistoryEvent,
                                     entity_id: EntityInstanceId | None,
                                     task_id: int | None,
-                                    is_lock_event: bool):
+                                    is_lock_event: bool,
+                                    operation: str | None = None):
         # This eventRaised represents the result of an entity operation after being translated to the old
         # entity protocol by the Durable WebJobs extension
         if entity_id is None:
@@ -2828,16 +2829,38 @@ def _handle_entity_event_raised(self,
         entity_task = ctx._pending_tasks.pop(task_id, None)  # pyright: ignore[reportPrivateUsage]
         if not entity_task:
             raise RuntimeError(f"Could not retrieve entity task for entity-related eventRaised with ID '{event.eventId}'")
-        result = None
+        response: dict[str, Any] | None = None
         if not ph.is_empty(event.eventRaised.input):
-            # TODO: Investigate why the event result is wrapped in a dict with "result" key
-            # The expected type applies to the unwrapped result value, not the
-            # transport wrapper. Unwrap first, then coerce the already-parsed
-            # inner value to the expected type via the converter (no redundant
-            # re-serialization round-trip).
-            unwrapped = self._data_converter.deserialize(event.eventRaised.input.value)["result"]
-            result = self._data_converter.coerce(
-                unwrapped,
+            response = self._data_converter.deserialize(event.eventRaised.input.value)
+
+        # For entity operation calls (lock acquisitions never fail this way), the legacy WebJobs
+        # "old protocol" ResponseMessage signals a failed operation via the presence of an
+        # "exceptionType" marker (or a structured "failureDetails" object). The human-readable
+        # message lives in the *serialized* "result" field -- "exceptionType" is only the exception
+        # type name, not the message -- so deserialize "result" to recover it, matching
+        # azure-functions-durable-python / -js. Propagate as a task failure so an awaiting
+        # call_entity raises, like the current entity protocol and the .NET SDK.
+        if not is_lock_event and isinstance(response, dict) and ph.is_entity_error_response(response):
+            raw_result = response.get("result")
+            error_content = (
+                self._data_converter.deserialize(raw_result) if isinstance(raw_result, str) else raw_result
+            )
+            failure_details = ph.entity_response_failure_details(response, error_content)
+            failure = EntityOperationFailedException(entity_id, operation or "", failure_details)
+            ctx._entity_context.recover_lock_after_call(entity_id)  # pyright: ignore[reportPrivateUsage]
+            entity_task.fail(str(failure), failure)
+            ctx.resume()
+            return
+
+        result = None
+        if response is not None:
+            # The legacy protocol wraps the result as {"result": <serialized>},
+            # where the value is a serialized JSON string (like the new protocol's
+            # entityOperationCompleted.output). Deserialize it -- not coerce -- so
+            # the value is fully parsed and the expected type applied; coercing
+            # would skip JSON parsing and leave it double-encoded (e.g. '"done"').
+            result = self._data_converter.deserialize(
+                response["result"],
                 entity_task._expected_type,  # pyright: ignore[reportPrivateUsage]
             )
         if is_lock_event:

tests/durabletask/test_orchestration_executor.py

@@ -2139,6 +2139,176 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     assert actions[0].sendEntityMessage.entityOperationCalled.targetInstanceId.value == str(test_entity_id)
 
 
+def test_entity_call_failure_propagated_over_old_protocol():
+    """A failed entity response carrying a structured ``failureDetails`` object
+    (defensively supported, current-protocol shape) must surface as a
+    TaskFailedError with the structured message."""
+    test_entity_id = entities.EntityInstanceId("Counter", "myCounter")
+
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        try:
+            yield ctx.call_entity(test_entity_id, "set", 1)
+        except task.TaskFailedError as e:
+            return e.details.message
+        return "no error"
+
+    registry = worker._Registry()
+    name = registry.add_orchestrator(orchestrator)
+
+    started_events = [
+        helpers.new_orchestrator_started_event(),
+        helpers.new_execution_started_event(name, TEST_INSTANCE_ID, None),
+    ]
+
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER, JsonDataConverter())
+    result1 = executor.execute(TEST_INSTANCE_ID, [], started_events)
+    actions = result1.actions
+    assert len(actions) == 1
+    assert actions[0].sendEntityMessage.HasField("entityOperationCalled")
+    request_id = actions[0].sendEntityMessage.entityOperationCalled.requestId
+
+    # A structured FailureDetails object, when present, takes precedence over the
+    # WebJobs result/exceptionType fields.
+    response_message = {
+        "result": None,
+        "failureDetails": {
+            "ErrorType": "ValueError",
+            "ErrorMessage": "Something went wrong!",
+            "StackTrace": None,
+            "InnerFailure": None,
+            "IsNonRetriable": False,
+        },
+    }
+    new_events = [
+        helpers.new_event_sent_event(1, str(test_entity_id), json.dumps({"id": request_id})),
+        helpers.new_event_raised_event(request_id, json.dumps(response_message)),
+    ]
+    result2 = executor.execute(TEST_INSTANCE_ID, started_events, new_events)
+    complete_action = get_and_validate_complete_orchestration_action_list(1, result2.actions)
+    assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    output = json.loads(complete_action.result.value)
+    assert output == (
+        "Operation 'set' on entity '@counter@myCounter' failed with error: Something went wrong!"
+    )
+
+
+def test_entity_call_failure_propagated_over_old_protocol_exception_type():
+    """A WebJobs "old protocol" failure carries the message in the serialized
+    ``result`` field, while ``exceptionType`` is only the exception type name.
+    The surfaced message must come from ``result``, not ``exceptionType``."""
+    test_entity_id = entities.EntityInstanceId("Counter", "myCounter")
+
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        try:
+            yield ctx.call_entity(test_entity_id, "set", 1)
+        except task.TaskFailedError as e:
+            return e.details.message
+        return "no error"
+
+    registry = worker._Registry()
+    name = registry.add_orchestrator(orchestrator)
+
+    started_events = [
+        helpers.new_orchestrator_started_event(),
+        helpers.new_execution_started_event(name, TEST_INSTANCE_ID, None),
+    ]
+
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER, JsonDataConverter())
+    result1 = executor.execute(TEST_INSTANCE_ID, [], started_events)
+    actions = result1.actions
+    assert len(actions) == 1
+    request_id = actions[0].sendEntityMessage.entityOperationCalled.requestId
+
+    # Real WebJobs shape: message serialized into "result"; "exceptionType" is the
+    # (assembly-qualified) exception type name and must NOT be used as the message.
+    response_message = {
+        "result": json.dumps("Something went wrong!"),
+        "exceptionType": "System.InvalidOperationException, mscorlib",
+    }
+    new_events = [
+        helpers.new_event_sent_event(1, str(test_entity_id), json.dumps({"id": request_id})),
+        helpers.new_event_raised_event(request_id, json.dumps(response_message)),
+    ]
+    result2 = executor.execute(TEST_INSTANCE_ID, started_events, new_events)
+    complete_action = get_and_validate_complete_orchestration_action_list(1, result2.actions)
+    assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    output = json.loads(complete_action.result.value)
+    assert output == (
+        "Operation 'set' on entity '@counter@myCounter' failed with error: Something went wrong!"
+    )
+    # Guard against the regression of surfacing the type name instead of the message.
+    assert "InvalidOperationException" not in output
+
+
+def test_entity_call_success_over_old_protocol():
+    """A successful entity operation delivered via the legacy protocol must
+    complete the call_entity task with the unwrapped result."""
+    test_entity_id = entities.EntityInstanceId("Counter", "myCounter")
+
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        return (yield ctx.call_entity(test_entity_id, "get", return_type=int))
+
+    registry = worker._Registry()
+    name = registry.add_orchestrator(orchestrator)
+
+    started_events = [
+        helpers.new_orchestrator_started_event(),
+        helpers.new_execution_started_event(name, TEST_INSTANCE_ID, None),
+    ]
+
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER, JsonDataConverter())
+    result1 = executor.execute(TEST_INSTANCE_ID, [], started_events)
+    actions = result1.actions
+    assert len(actions) == 1
+    request_id = actions[0].sendEntityMessage.entityOperationCalled.requestId
+
+    response_message = {"result": json.dumps(42)}
+    new_events = [
+        helpers.new_event_sent_event(1, str(test_entity_id), json.dumps({"id": request_id})),
+        helpers.new_event_raised_event(request_id, json.dumps(response_message)),
+    ]
+    result2 = executor.execute(TEST_INSTANCE_ID, started_events, new_events)
+    complete_action = get_and_validate_complete_orchestration_action_list(1, result2.actions)
+    assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert json.loads(complete_action.result.value) == 42
+
+
+@pytest.mark.parametrize("entity_result", ["done", 42, 3.5, True, {"a": 1}, [1, 2, 3]])
+def test_entity_call_success_over_old_protocol_round_trips_result(entity_result):
+    """The legacy-protocol ``result`` field holds a *serialized* JSON string, so
+    the caller must receive the fully deserialized value regardless of type and
+    without needing an explicit ``return_type`` (i.e. no double-encoding)."""
+    test_entity_id = entities.EntityInstanceId("Counter", "myCounter")
+
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        return (yield ctx.call_entity(test_entity_id, "get"))
+
+    registry = worker._Registry()
+    name = registry.add_orchestrator(orchestrator)
+
+    started_events = [
+        helpers.new_orchestrator_started_event(),
+        helpers.new_execution_started_event(name, TEST_INSTANCE_ID, None),
+    ]
+
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER, JsonDataConverter())
+    result1 = executor.execute(TEST_INSTANCE_ID, [], started_events)
+    actions = result1.actions
+    assert len(actions) == 1
+    request_id = actions[0].sendEntityMessage.entityOperationCalled.requestId
+
+    # The entity worker places the *serialized* return value into the "result" field.
+    response_message = {"result": json.dumps(entity_result)}
+    new_events = [
+        helpers.new_event_sent_event(1, str(test_entity_id), json.dumps({"id": request_id})),
+        helpers.new_event_raised_event(request_id, json.dumps(response_message)),
+    ]
+    result2 = executor.execute(TEST_INSTANCE_ID, started_events, new_events)
+    complete_action = get_and_validate_complete_orchestration_action_list(1, result2.actions)
+    assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert json.loads(complete_action.result.value) == entity_result
+
+
 def get_and_validate_complete_orchestration_action_list(expected_action_count: int, actions: list[pb.OrchestratorAction]) -> pb.CompleteOrchestrationAction:
     assert len(actions) == expected_action_count
     assert type(actions[-1]) is pb.OrchestratorAction