Python: Fix AG-UI resume rejecting JSON for typed responses

python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py

@@ -10,7 +10,8 @@
 import uuid
 from collections.abc import AsyncGenerator
 from functools import partial
-from typing import Any, cast, get_args, get_origin
+from types import UnionType
+from typing import Any, Union, cast, get_args, get_origin, get_type_hints
 
 from ag_ui.core import (
     ActivitySnapshotEvent,
@@ -34,6 +35,10 @@
     Workflow,
     WorkflowRunState,
 )
+from agent_framework._workflows._typing_utils import (  # pyright: ignore[reportPrivateUsage]
+    is_instance_of,
+    try_coerce_to_type,
+)
 from agent_framework.observability import (
     _use_telemetry_conversation_id,  # pyright: ignore[reportPrivateUsage]
 )
@@ -535,6 +540,48 @@ def _coerce_compact_approval_response(request_data: Content, candidate: dict[str
     return response
 
 
+def _without_optional(annotation: Any) -> Any:
+    """Unwrap ``X | None`` so optional fields normalize like their plain counterpart."""
+    if get_origin(annotation) not in (Union, UnionType):
+        return annotation
+    members = [member for member in get_args(annotation) if member is not type(None)]
+    return members[0] if len(members) == 1 else annotation
+
+
+def _normalize_agui_message_fields(response_type: Any, candidate: Any) -> Any:
+    """Convert AG-UI message payloads for fields the response type declares as Message.
+
+    Core coercion only understands the canonical ``contents`` form, so AG-UI wire shapes
+    (``{"role": ..., "content": ...}`` and bare strings) are translated here.
+    """
+    if not isinstance(candidate, dict):
+        return candidate
+
+    try:
+        field_types = get_type_hints(response_type)
+    except Exception:
+        return candidate
+
+    normalized = dict(cast(dict[str, Any], candidate))
+    for name, annotation in field_types.items():
+        if name not in normalized:
+            continue
+        annotation = _without_optional(annotation)
+        target_type = get_origin(annotation) or annotation
+        if target_type is Message:
+            message = _coerce_message(normalized[name])
+            if message is not None:
+                normalized[name] = message
+        elif target_type is list and get_args(annotation)[:1] == (Message,):
+            items = normalized[name]
+            if not isinstance(items, list):
+                continue
+            messages = [_coerce_message(item) for item in cast(list[Any], items)]
+            if all(message is not None for message in messages):
+                normalized[name] = messages
+    return normalized
+
+
 def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None:
     """Coerce a candidate value into the request's expected response type."""
     response_type = getattr(request_event, "response_type", None)
@@ -598,7 +645,8 @@ def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None:
     if target_type is float:
         return candidate if isinstance(candidate, (int, float)) and not isinstance(candidate, bool) else None
     if isinstance(target_type, type):
-        return candidate if isinstance(candidate, target_type) else None
+        coerced = try_coerce_to_type(_normalize_agui_message_fields(response_type, candidate), response_type)
+        return coerced if is_instance_of(coerced, response_type) else None
 
     # Unknown typing metadata: preserve value as-is.
     return candidate

python/packages/ag-ui/tests/ag_ui/test_endpoint.py

@@ -8,6 +8,7 @@
 import sys
 from collections import Counter
 from collections.abc import AsyncIterator, Callable
+from dataclasses import dataclass
 from inspect import signature
 from typing import Any, cast
 
@@ -3776,6 +3777,63 @@ def _build_workflow_request_info_app(
     return app
 
 
+async def test_endpoint_workflow_request_info_resumes_dataclass_response_from_json():
+    """Dataclass response types resume from plain JSON payloads, as AG-UI clients send them."""
+
+    @dataclass
+    class PlanReview:
+        review: list[Message]
+
+    class PlanReviewExecutor(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="plan_review")
+
+        @handler
+        async def start(self, message: Any, ctx: WorkflowContext[Any, Any]) -> None:
+            del message
+            await ctx.request_info({"plan": "ship it"}, PlanReview, request_id="plan-review")
+
+        @response_handler
+        async def handle_review(
+            self, original_request: dict[str, Any], response: PlanReview, ctx: WorkflowContext[Any, Any]
+        ) -> None:
+            del original_request
+            verdict = "approved" if not response.review else response.review[0].text
+            await ctx.yield_output(f"Plan {verdict}")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]
+
+    app = FastAPI()
+    add_agent_framework_fastapi_endpoint(
+        app, WorkflowBuilder(start_executor=PlanReviewExecutor()).build(), path="/workflow"
+    )
+
+    with TestClient(app) as client:
+        pause_response = client.post(
+            "/workflow",
+            json={
+                "runId": "run-pause",
+                "threadId": "thread-plan",
+                "messages": [{"role": "user", "content": "Draft a plan"}],
+            },
+        )
+        assert pause_response.status_code == 200
+
+        resume_response = client.post(
+            "/workflow",
+            json={
+                "runId": "run-resume",
+                "threadId": "thread-plan",
+                "messages": [],
+                "resume": [{"interruptId": "plan-review", "status": "resolved", "payload": {"review": []}}],
+            },
+        )
+
+        assert resume_response.status_code == 200
+        resume_events = _decode_sse_events(resume_response)
+        assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"]
+        text_deltas = [event["delta"] for event in resume_events if event.get("type") == "TEXT_MESSAGE_CONTENT"]
+        assert "Plan approved" in text_deltas
+
+
 async def test_endpoint_workflow_request_info_emits_canonical_interrupt_and_resumes():
     """Workflow request_info pauses and resumes through canonical AG-UI interrupt payloads."""
     app = _build_workflow_request_info_app()

python/packages/ag-ui/tests/ag_ui/test_workflow_run.py

@@ -4,6 +4,7 @@
 
 import json
 from collections.abc import AsyncIterator
+from dataclasses import dataclass, make_dataclass
 from enum import Enum
 from types import SimpleNamespace
 from typing import Any, cast
@@ -28,7 +29,9 @@
     response_handler,
     tool,
 )
+from agent_framework.orchestrations import MagenticPlanReviewResponse
 from conftest import StreamingChatClientStub  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from pydantic import BaseModel
 
 from agent_framework_ag_ui._workflow_run import (
     _coerce_content,
@@ -1314,6 +1317,93 @@ def test_coerce_response_for_request_bool_int_float_and_mismatch() -> None:
     assert _coerce_response_for_request(dict_request, "[1,2,3]") is None
 
 
+def test_coerce_response_for_request_builds_dataclass_from_json() -> None:
+    """JSON objects should map onto dataclass response types such as plan review."""
+    request = SimpleNamespace(response_type=MagenticPlanReviewResponse)
+
+    approved = _coerce_response_for_request(request, {"review": []})
+    assert isinstance(approved, MagenticPlanReviewResponse)
+    assert approved.review == []
+
+    revised = _coerce_response_for_request(request, '{"review": [{"role": "user", "content": "add tests"}]}')
+    assert isinstance(revised, MagenticPlanReviewResponse)
+    assert len(revised.review) == 1
+    assert revised.review[0].text == "add tests"
+
+    from_strings = _coerce_response_for_request(request, {"review": ["add tests"]})
+    assert isinstance(from_strings, MagenticPlanReviewResponse)
+    assert from_strings.review[0].text == "add tests"
+
+    assert _coerce_response_for_request(request, {"unknown": 1}) is None
+    assert _coerce_response_for_request(request, "approve") is None
+
+
+def test_coerce_response_for_request_leaves_non_message_fields_untouched() -> None:
+    """Message normalization must follow field annotations, not payload shape."""
+
+    @dataclass
+    class Tagged:
+        note: Message
+        revision: Message | None
+        metadata: dict[str, str]
+
+    request = SimpleNamespace(response_type=Tagged)
+    message_shaped = {"role": "admin", "content": "keep raw"}
+
+    tagged = _coerce_response_for_request(
+        request,
+        {
+            "note": {"role": "user", "content": "translate me"},
+            "revision": {"role": "user", "content": "optional fields too"},
+            "metadata": message_shaped,
+        },
+    )
+    assert isinstance(tagged, Tagged)
+    assert tagged.note.text == "translate me"
+    assert tagged.revision is not None
+    assert tagged.revision.text == "optional fields too"
+    assert tagged.metadata == message_shaped
+
+
+def test_coerce_response_for_request_rejects_malformed_message_field_payloads() -> None:
+    """A Message-typed field that is not message-shaped must fail, not crash normalization."""
+
+    @dataclass
+    class Review:
+        notes: list[Message]
+
+    request = SimpleNamespace(response_type=Review)
+
+    assert _coerce_response_for_request(request, {"notes": "not-a-list"}) is None
+    assert _coerce_response_for_request(request, {"notes": [42]}) is None
+
+
+def test_coerce_response_for_request_skips_normalization_without_resolvable_hints() -> None:
+    """Unresolvable annotations skip message normalization instead of failing the resume."""
+
+    mystery_type = make_dataclass("Mystery", [("value", "DoesNotExist")])
+    request = SimpleNamespace(response_type=mystery_type)
+
+    mystery = _coerce_response_for_request(request, {"value": 1})
+    assert type(mystery) is mystery_type
+    assert vars(mystery) == {"value": 1}
+
+
+def test_coerce_response_for_request_builds_pydantic_model_from_json() -> None:
+    """JSON objects should validate into pydantic response types."""
+
+    class ReviewDecision(BaseModel):
+        approved: bool
+
+    request = SimpleNamespace(response_type=ReviewDecision)
+
+    decision = _coerce_response_for_request(request, {"approved": True})
+    assert isinstance(decision, ReviewDecision)
+    assert decision.approved is True
+
+    assert _coerce_response_for_request(request, {"approved": "not-a-bool"}) is None
+
+
 async def test_workflow_run_emits_run_error_when_stream_raises() -> None:
     """Unexpected stream exceptions should be converted into RUN_ERROR events."""