diff --git a/tests/asyncio/test_context.py b/tests/asyncio/test_context.py index 96549c8..4b65089 100644 --- a/tests/asyncio/test_context.py +++ b/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", + } + ], ), ) diff --git a/tests/test_context.py b/tests/test_context.py index fe274f0..81504be 100644 --- a/tests/test_context.py +++ b/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", + } + ], ), ) diff --git a/tests/test_redact.py b/tests/test_redact.py index 6b5fd29..78e49f5 100644 --- a/tests/test_redact.py +++ b/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"] diff --git a/upstash_workflow/asyncio/context/context.py b/upstash_workflow/asyncio/context/context.py index 9268133..d50d3f3 100644 --- a/upstash_workflow/asyncio/context/context.py +++ b/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: """ diff --git a/upstash_workflow/asyncio/workflow_requests.py b/upstash_workflow/asyncio/workflow_requests.py index 9a0df16..5b6e5bf 100644 --- a/upstash_workflow/asyncio/workflow_requests.py +++ b/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" diff --git a/upstash_workflow/context/context.py b/upstash_workflow/context/context.py index 8cabaff..c2d530b 100644 --- a/upstash_workflow/context/context.py +++ b/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: """ diff --git a/upstash_workflow/workflow_requests.py b/upstash_workflow/workflow_requests.py index 3aced72..f78b2ea 100644 --- a/upstash_workflow/workflow_requests.py +++ b/upstash_workflow/workflow_requests.py @@ -11,6 +11,8 @@ cast, TypeVar, Dict, + List, + Any, ) from qstash import QStash, Receiver from upstash_workflow.error import WorkflowError, WorkflowAbort @@ -36,25 +38,121 @@ TInitialPayload = TypeVar("TInitialPayload") +def _get_redact_headers(redact: Optional[Redact]) -> Dict[str, str]: + """ + Converts the redact option into the `Upstash-Redact-Fields` header + understood by QStash. + + :param redact: redact configuration + :return: headers to add to the request (empty if nothing to redact) + """ + if redact is None: + return {} + + redact_parts = [] + if redact.get("body"): + redact_parts.append("body") + + header_redact = redact.get("header") + if header_redact is True: + redact_parts.append("header") + elif isinstance(header_redact, list): + redact_parts.extend(f"header[{name}]" for name in header_redact) + + if not redact_parts: + return {} + return {"Upstash-Redact-Fields": ",".join(redact_parts)} + + +def _get_first_invocation_batch_body( + workflow_run_id: str, + workflow_url: str, + user_headers: Dict[str, str], + request_payload: Any, + retries: int, + redact: Optional[Redact] = None, + workflow_failure_url: Optional[str] = None, +) -> List[Dict[str, Any]]: + """ + Builds the body of the `/v2/batch` request which starts a workflow run. + + The workflow control headers (`Upstash-Workflow-Init`, `Upstash-Workflow-RunId`, ...) + are sent as-is inside the batch body. We don't use `qstash.message.publish_json` + because qstash-py (>=3) prefixes every header with `Upstash-Forward-`, which + would turn the control headers into plain forwarded headers and QStash would + never recognize the request as a workflow run start. + + :param workflow_run_id: id of the workflow run + :param workflow_url: url of the workflow endpoint + :param user_headers: headers of the initial request + :param request_payload: initial payload. Sent as-is if it's a string, otherwise + serialized to JSON. + :param retries: number of retries + :param redact: fields to redact in QStash logs + :param workflow_failure_url: failure callback url, so that a failure of the + first step also triggers the failure function + :return: batch body with a single message + """ + headers = _get_headers( + "true", + workflow_run_id, + workflow_url, + user_headers, + None, + retries, + workflow_failure_url=workflow_failure_url, + ).headers + + # QStash doesn't forward content-type when passed in `upstash-forward-content-type` + # so we need to pass it in the headers + content_type = next( + ( + value + for header, value in user_headers.items() + if header.lower() == "content-type" + ), + "application/json", + ) + + body = ( + request_payload + if isinstance(request_payload, str) + else json.dumps(request_payload) + ) + + return [ + { + "destination": workflow_url, + "headers": { + "Content-Type": content_type, + **headers, + **_get_redact_headers(redact), + }, + "body": body, + } + ] + + def _trigger_first_invocation( workflow_context: WorkflowContext[TInitialPayload], 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, + ) - workflow_context.qstash_client.message.publish_json( - url=workflow_context.url, - body=workflow_context.request_payload, - headers=headers, - redact=redact, + workflow_context.qstash_client.http.request( + path="/v2/batch", + method="POST", + headers={"Content-Type": "application/json"}, + body=json.dumps(batch_body), ) @@ -232,10 +330,14 @@ def _handle_third_party_call_result( "concurrent": int(concurrent_str), } - 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. + 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"