DX-2949: bypass qstash v3 header prefixing when starting workflow runs

tests/asyncio/test_context.py

@@ -132,11 +132,24 @@ async def execute() -> None:
         response_fields=ResponseFields(status=200, body="msgId"),
         receives_request=RequestFields(
             method="POST",
-            url=f"{MOCK_QSTASH_SERVER_URL}/v2/publish/{WORKFLOW_ENDPOINT}",
+            url=f"{MOCK_QSTASH_SERVER_URL}/v2/batch",
             token="mock-token",
-            body="my-payload",
-            headers={
-                "Upstash-Redact-Fields": "body,header[Authorization]",
-            },
+            body=[
+                {
+                    "destination": WORKFLOW_ENDPOINT,
+                    "headers": {
+                        "Content-Type": "application/json",
+                        # workflow control headers must reach QStash as-is,
+                        # not prefixed with `Upstash-Forward-`
+                        "Upstash-Workflow-Init": "true",
+                        "Upstash-Workflow-RunId": "wfr-id",
+                        "Upstash-Workflow-Url": WORKFLOW_ENDPOINT,
+                        "Upstash-Feature-Set": "LazyFetch,InitialBody,WF_DetectTrigger",
+                        "Upstash-Forward-Upstash-Workflow-Sdk-Version": "1",
+                        "Upstash-Redact-Fields": "body,header[Authorization]",
+                    },
+                    "body": "my-payload",
+                }
+            ],
         ),
     )

tests/test_context.py

@@ -130,11 +130,24 @@ def execute() -> None:
         response_fields=ResponseFields(status=200, body="msgId"),
         receives_request=RequestFields(
             method="POST",
-            url=f"{MOCK_QSTASH_SERVER_URL}/v2/publish/{WORKFLOW_ENDPOINT}",
+            url=f"{MOCK_QSTASH_SERVER_URL}/v2/batch",
             token="mock-token",
-            body="my-payload",
-            headers={
-                "Upstash-Redact-Fields": "body,header[Authorization]",
-            },
+            body=[
+                {
+                    "destination": WORKFLOW_ENDPOINT,
+                    "headers": {
+                        "Content-Type": "application/json",
+                        # workflow control headers must reach QStash as-is,
+                        # not prefixed with `Upstash-Forward-`
+                        "Upstash-Workflow-Init": "true",
+                        "Upstash-Workflow-RunId": "wfr-id",
+                        "Upstash-Workflow-Url": WORKFLOW_ENDPOINT,
+                        "Upstash-Feature-Set": "LazyFetch,InitialBody,WF_DetectTrigger",
+                        "Upstash-Forward-Upstash-Workflow-Sdk-Version": "1",
+                        "Upstash-Redact-Fields": "body,header[Authorization]",
+                    },
+                    "body": "my-payload",
+                }
+            ],
         ),
     )

tests/test_redact.py

@@ -1,82 +1,102 @@
-"""Tests for redact parameter being passed to qstash client."""
+"""Tests for redact parameter being converted to the Upstash-Redact-Fields header."""
 
+import json
+from typing import Any, Dict, Optional
 from unittest.mock import MagicMock
+
+import pytest
+
 from upstash_workflow.workflow_requests import _trigger_first_invocation
 from upstash_workflow.types import Redact
 
 
-def _make_workflow_context():
+def _make_workflow_context() -> MagicMock:
     ctx = MagicMock()
     ctx.workflow_run_id = "wfr-test-id"
     ctx.url = "https://example.com"
     ctx.headers = {}
     ctx.request_payload = '{"test": true}'
-    ctx.qstash_client.message.publish_json = MagicMock()
+    ctx.failure_url = None
+    ctx.qstash_client.http.request = MagicMock()
     return ctx
 
 
-def test_trigger_passes_redact_body() -> None:
-    """Test that redact with body is passed to publish_json."""
-    ctx = _make_workflow_context()
-    redact: Redact = {"body": True}
-
-    _trigger_first_invocation(ctx, retries=3, redact=redact)
-
-    ctx.qstash_client.message.publish_json.assert_called_once()
-    call_kwargs = ctx.qstash_client.message.publish_json.call_args
-    assert call_kwargs.kwargs["redact"] == {"body": True}
-
-
-def test_trigger_passes_redact_header_all() -> None:
-    """Test that redact with all headers is passed to publish_json."""
-    ctx = _make_workflow_context()
-    redact: Redact = {"header": True}
-
-    _trigger_first_invocation(ctx, retries=3, redact=redact)
-
-    call_kwargs = ctx.qstash_client.message.publish_json.call_args
-    assert call_kwargs.kwargs["redact"] == {"header": True}
-
-
-def test_trigger_passes_redact_specific_headers() -> None:
-    """Test that redact with specific headers is passed to publish_json."""
-    ctx = _make_workflow_context()
-    redact: Redact = {"header": ["Authorization", "X-API-Key"]}
-
-    _trigger_first_invocation(ctx, retries=3, redact=redact)
-
-    call_kwargs = ctx.qstash_client.message.publish_json.call_args
-    assert call_kwargs.kwargs["redact"] == {"header": ["Authorization", "X-API-Key"]}
-
-
-def test_trigger_passes_redact_body_and_headers() -> None:
-    """Test that redact with body and specific headers is passed to publish_json."""
-    ctx = _make_workflow_context()
-    redact: Redact = {"body": True, "header": ["Authorization"]}
-
-    _trigger_first_invocation(ctx, retries=3, redact=redact)
-
-    call_kwargs = ctx.qstash_client.message.publish_json.call_args
-    assert call_kwargs.kwargs["redact"] == {"body": True, "header": ["Authorization"]}
-
-
-def test_trigger_passes_no_redact() -> None:
-    """Test that redact=None is passed when no redact specified."""
-    ctx = _make_workflow_context()
-
-    _trigger_first_invocation(ctx, retries=3, redact=None)
-
-    call_kwargs = ctx.qstash_client.message.publish_json.call_args
-    assert call_kwargs.kwargs["redact"] is None
-
-
-def test_trigger_no_redact_headers_in_headers() -> None:
-    """Test that Upstash-Redact-Fields is NOT in the headers (qstash client handles it)."""
+def _trigger(
+    redact: Optional[Redact], failure_url: Optional[str] = None
+) -> Dict[str, Any]:
+    """Triggers the first invocation and returns the single batch message sent."""
     ctx = _make_workflow_context()
-    redact: Redact = {"body": True, "header": ["Authorization"]}
+    ctx.failure_url = failure_url
 
     _trigger_first_invocation(ctx, retries=3, redact=redact)
 
-    call_kwargs = ctx.qstash_client.message.publish_json.call_args
-    headers = call_kwargs.kwargs["headers"]
-    assert "Upstash-Redact-Fields" not in headers
+    ctx.qstash_client.http.request.assert_called_once()
+    # publish_json must not be used: qstash-py (>=3) prefixes every header with
+    # `Upstash-Forward-` which corrupts the workflow control headers.
+    ctx.qstash_client.message.publish_json.assert_not_called()
+
+    kwargs = ctx.qstash_client.http.request.call_args.kwargs
+    assert kwargs["path"] == "/v2/batch"
+    assert kwargs["method"] == "POST"
+
+    batch_body = json.loads(kwargs["body"])
+    assert len(batch_body) == 1
+    message: Dict[str, Any] = batch_body[0]
+    assert message["destination"] == "https://example.com"
+    assert message["body"] == '{"test": true}'
+    return message
+
+
+@pytest.mark.parametrize(
+    "redact, expected",
+    [
+        ({"body": True}, "body"),
+        ({"header": True}, "header"),
+        (
+            {"header": ["Authorization", "X-API-Key"]},
+            "header[Authorization],header[X-API-Key]",
+        ),
+        ({"body": True, "header": ["Authorization"]}, "body,header[Authorization]"),
+    ],
+)
+def test_trigger_passes_redact_header(redact: Redact, expected: str) -> None:
+    message = _trigger(redact)
+    assert message["headers"]["Upstash-Redact-Fields"] == expected
+
+
+@pytest.mark.parametrize("redact", [None, {}, {"header": []}])
+def test_trigger_passes_no_redact(redact: Optional[Redact]) -> None:
+    message = _trigger(redact)
+    assert "Upstash-Redact-Fields" not in message["headers"]
+
+
+def test_trigger_sends_workflow_headers_unprefixed() -> None:
+    """Workflow control headers must reach QStash as-is, not as forwarded headers."""
+    message = _trigger(None)
+    headers = message["headers"]
+
+    assert headers["Upstash-Workflow-Init"] == "true"
+    assert headers["Upstash-Workflow-RunId"] == "wfr-test-id"
+    assert headers["Upstash-Workflow-Url"] == "https://example.com"
+    assert headers["Upstash-Feature-Set"] == "LazyFetch,InitialBody,WF_DetectTrigger"
+    assert headers["Upstash-Forward-Upstash-Workflow-Sdk-Version"] == "1"
+    assert headers["Content-Type"] == "application/json"
+    assert "Upstash-Forward-Upstash-Workflow-Init" not in headers
+
+
+def test_trigger_sends_failure_callback_when_failure_url_set() -> None:
+    """A failing first step must also reach the failure function (as in workflow-js)."""
+    message = _trigger(None, failure_url="https://example.com/failure")
+    headers = message["headers"]
+
+    assert headers["Upstash-Failure-Callback"] == "https://example.com/failure"
+    assert headers["Upstash-Failure-Callback-Workflow-Runid"] == "wfr-test-id"
+    assert (
+        headers["Upstash-Failure-Callback-Forward-Upstash-Workflow-Is-Failure"]
+        == "true"
+    )
+
+
+def test_trigger_sends_no_failure_callback_without_failure_url() -> None:
+    message = _trigger(None)
+    assert "Upstash-Failure-Callback" not in message["headers"]

upstash_workflow/asyncio/context/context.py

@@ -9,7 +9,6 @@
     Awaitable,
     TypeVar,
     Any,
-    cast,
     Generic,
 )
 from qstash import AsyncQStash
@@ -165,14 +164,18 @@ async def call(
             )
         )
 
+        body = result["body"]
         try:
-            return CallResponse(
-                status=result["status"],
-                body=json.loads(result["body"]),
-                header=result["header"],
-            )
+            body = json.loads(body)
         except Exception:
-            return cast(CallResponse[Any], result)
+            # not a JSON body, return it as it is
+            pass
+
+        return CallResponse(
+            status=result["status"],
+            body=body,
+            header=result["header"],
+        )
 
     async def _add_step(self, step: _BaseLazyStep[TResult]) -> TResult:
         """

upstash_workflow/asyncio/workflow_requests.py

@@ -18,7 +18,11 @@
 )
 from upstash_workflow.types import StepTypes, Redact
 from upstash_workflow.workflow_types import _AsyncRequest
-from upstash_workflow.workflow_requests import _get_headers, _recreate_user_headers
+from upstash_workflow.workflow_requests import (
+    _get_headers,
+    _recreate_user_headers,
+    _get_first_invocation_batch_body,
+)
 
 if TYPE_CHECKING:
     from upstash_workflow import AsyncWorkflowContext
@@ -33,20 +37,21 @@ async def _trigger_first_invocation(
     retries: int,
     redact: Optional[Redact] = None,
 ) -> None:
-    headers = _get_headers(
-        "true",
+    batch_body = _get_first_invocation_batch_body(
         workflow_context.workflow_run_id,
         workflow_context.url,
         workflow_context.headers,
-        None,
+        workflow_context.request_payload,
         retries,
-    ).headers
+        redact,
+        workflow_context.failure_url,
+    )
 
-    await workflow_context.qstash_client.message.publish_json(
-        url=workflow_context.url,
-        body=workflow_context.request_payload,
-        headers=headers,
-        redact=redact,
+    await workflow_context.qstash_client.http.request(
+        path="/v2/batch",
+        method="POST",
+        headers={"Content-Type": "application/json"},
+        body=json.dumps(batch_body),
     )
 
 
@@ -193,10 +198,14 @@ async def _handle_third_party_call_result(
                 "concurrent": int(concurrent_str),
             }
 
-            await client.message.publish_json(
-                headers=request_headers,
-                body=call_result_step,
-                url=workflow_url,
+            # Not using `client.message.publish_json` because qstash-py (>=3)
+            # prefixes every header with `Upstash-Forward-`, corrupting the
+            # workflow control headers.
+            await client.http.request(
+                path=f"/v2/publish/{workflow_url}",
+                method="POST",
+                headers={"Content-Type": "application/json", **request_headers},
+                body=json.dumps(call_result_step),
             )
 
             return "is-call-return"

upstash_workflow/context/context.py

@@ -8,7 +8,6 @@
     Callable,
     TypeVar,
     Any,
-    cast,
     Generic,
 )
 from qstash import QStash
@@ -164,14 +163,18 @@ def call(
             )
         )
 
+        body = result["body"]
         try:
-            return CallResponse(
-                status=result["status"],
-                body=json.loads(result["body"]),
-                header=result["header"],
-            )
+            body = json.loads(body)
         except Exception:
-            return cast(CallResponse[Any], result)
+            # not a JSON body, return it as it is
+            pass
+
+        return CallResponse(
+            status=result["status"],
+            body=body,
+            header=result["header"],
+        )
 
     def _add_step(self, step: _BaseLazyStep[TResult]) -> TResult:
         """