feat: add failure functions

Author: CahidArda

tests/asyncio/test_context.py

@@ -31,6 +31,7 @@ async def test_workflow_headers(qstash_client: AsyncQStash) -> None:
         initial_payload="my-payload",
         env=None,
         retries=None,
+        failure_url=None
     )
 
     async def execute() -> None:

tests/test_context.py

@@ -30,6 +30,7 @@ def test_workflow_headers(qstash_client: QStash) -> None:
         initial_payload="my-payload",
         env=None,
         retries=None,
+        failure_url=WORKFLOW_ENDPOINT
     )
 
     def execute() -> None:

upstash_workflow/asyncio/context/auto_executor.py

@@ -84,6 +84,7 @@ async def submit_steps_to_qstash(
                 self.context.retries,
                 lazy_step.retries if isinstance(lazy_step, _LazyCallStep) else None,
                 lazy_step.timeout if isinstance(lazy_step, _LazyCallStep) else None,
+                workflow_failure_url=self.context.failure_url,
             ).headers
 
             will_wait = (

upstash_workflow/asyncio/context/context.py

@@ -47,6 +47,7 @@ def __init__(
         headers: Dict[str, str],
         steps: List[DefaultStep],
         url: str,
+        failure_url: Optional[str],
         initial_payload: TInitialPayload,
         env: Optional[Dict[str, Optional[str]]] = None,
         retries: Optional[int] = None,
@@ -55,6 +56,7 @@ def __init__(
         self.workflow_run_id: str = workflow_run_id
         self._steps: List[DefaultStep] = steps
         self.url: str = url
+        self.failure_url = failure_url
         self.headers: Dict[str, str] = headers
         self.request_payload: TInitialPayload = initial_payload
         self.env: Dict[str, Optional[str]] = env or {}

upstash_workflow/asyncio/serve/authorization.py

@@ -38,6 +38,7 @@ async def try_authentication(
             initial_payload=context.request_payload,
             env=context.env,
             retries=context.retries,
+            failure_url=context.failure_url,
         )
 
         try:

upstash_workflow/asyncio/serve/options.py

@@ -1,14 +1,15 @@
 import os
 import json
 import logging
-from typing import Callable, Dict, Optional, cast, TypeVar
+from typing import Callable, Dict, Optional, cast, TypeVar, Any
 from qstash import AsyncQStash, Receiver
 from upstash_workflow.workflow_types import _Response
 from upstash_workflow.constants import DEFAULT_RETRIES
 from upstash_workflow.types import (
     _FinishCondition,
 )
 from upstash_workflow.asyncio.types import ServeBaseOptions
+from upstash_workflow import AsyncWorkflowContext
 
 _logger = logging.getLogger(__name__)
 
@@ -26,6 +27,10 @@ def _process_options(
     env: Optional[Dict[str, Optional[str]]] = None,
     retries: Optional[int] = DEFAULT_RETRIES,
     url: Optional[str] = None,
+    failure_function: Optional[
+        Callable[[AsyncWorkflowContext, int, str, Dict[str, str]], Any]
+    ] = None,
+    failure_url: Optional[str] = None,
 ) -> ServeBaseOptions[TInitialPayload, TResponse]:
     environment = env if env is not None else dict(os.environ)
 
@@ -94,6 +99,8 @@ def _initial_payload_parser(initial_request: str) -> TInitialPayload:
         env=environment,
         retries=DEFAULT_RETRIES if retries is None else retries,
         url=url,
+        failure_url=failure_url,
+        failure_function=failure_function,
     )
 
 

upstash_workflow/asyncio/serve/serve.py

@@ -5,6 +5,7 @@
 from upstash_workflow.workflow_types import _Response, _AsyncRequest
 from upstash_workflow.asyncio.workflow_parser import (
     _get_payload,
+    _handle_failure,
 )
 from upstash_workflow.workflow_parser import _validate_request, _parse_request
 from upstash_workflow.asyncio.workflow_requests import (
@@ -39,6 +40,10 @@ def _serve_base(
     env: Optional[Dict[str, Optional[str]]] = None,
     retries: Optional[int] = None,
     url: Optional[str] = None,
+    failure_function: Optional[
+        Callable[[AsyncWorkflowContext, int, str, Dict[str, str]], Awaitable[Any]]
+    ] = None,
+    failure_url: Optional[str] = None,
 ) -> Dict[str, Callable[[TRequest], Awaitable[TResponse]]]:
     processed_options = _process_options(
         qstash_client=qstash_client,
@@ -49,6 +54,8 @@ def _serve_base(
         env=env,
         retries=retries,
         url=url,
+        failure_function=failure_function,
+        failure_url=failure_url,
     )
     qstash_client = processed_options.qstash_client
     on_step_finish = processed_options.on_step_finish
@@ -58,9 +65,13 @@ def _serve_base(
     env = processed_options.env
     retries = processed_options.retries
     url = processed_options.url
+    failure_url = processed_options.failure_url
+    failure_function = processed_options.failure_function
 
     async def _handler(request: TRequest) -> TResponse:
-        workflow_url = _determine_urls(cast(_AsyncRequest, request), url, base_url)
+        workflow_url, workflow_failure_url = _determine_urls(
+            cast(_AsyncRequest, request), url, base_url, failure_function, failure_url
+        )
 
         request_payload = await _get_payload(request) or ""
         _verify_request(
@@ -78,6 +89,20 @@ async def _handler(request: TRequest) -> TResponse:
         raw_initial_payload = parse_request_response.raw_initial_payload
         steps = parse_request_response.steps
 
+        failure_check = await _handle_failure(
+            request,
+            request_payload,
+            qstash_client,
+            initial_payload_parser,
+            route_function,
+            failure_function,
+            env,
+            retries,
+        )
+
+        if failure_check == "is-failure-callback":
+            return on_step_finish(workflow_run_id, "failure-callback")
+
         workflow_context = AsyncWorkflowContext(
             qstash_client=qstash_client,
             workflow_run_id=workflow_run_id,
@@ -89,6 +114,7 @@ async def _handler(request: TRequest) -> TResponse:
             url=workflow_url,
             env=env,
             retries=retries,
+            failure_url=workflow_failure_url,
         )
 
         auth_check = await _DisabledWorkflowContext[Any].try_authentication(
@@ -110,6 +136,7 @@ async def _handler(request: TRequest) -> TResponse:
             raw_initial_payload,
             qstash_client,
             workflow_url,
+            workflow_failure_url,
             retries,
         )
 
@@ -153,6 +180,10 @@ def serve(
     env: Optional[Dict[str, Optional[str]]] = None,
     retries: Optional[int] = None,
     url: Optional[str] = None,
+    failure_function: Optional[
+        Callable[[AsyncWorkflowContext, int, str, Dict[str, str]], Awaitable[Any]]
+    ] = None,
+    failure_url: Optional[str] = None,
 ) -> Dict[str, Callable[[TRequest], Awaitable[TResponse]]]:
     """
     Creates a method that handles incoming requests and runs the provided
@@ -178,4 +209,6 @@ def serve(
         env=env,
         retries=retries,
         url=url,
+        failure_function=failure_function,
+        failure_url=failure_url,
     )

upstash_workflow/asyncio/types.py

@@ -4,10 +4,12 @@
     Dict,
     Generic,
     TypeVar,
+    Any,
 )
 from qstash import AsyncQStash, Receiver
 from dataclasses import dataclass
 from upstash_workflow.types import _FinishCondition
+from upstash_workflow import AsyncWorkflowContext
 
 TInitialPayload = TypeVar("TInitialPayload")
 TResponse = TypeVar("TResponse")
@@ -22,6 +24,10 @@ class ServeOptions(Generic[TInitialPayload, TResponse]):
     env: Dict[str, Optional[str]]
     retries: int
     url: Optional[str]
+    failure_function: Optional[
+        Callable[[AsyncWorkflowContext, int, str, Dict[str, str]], Any]
+    ]
+    failure_url: Optional[str]
 
 
 @dataclass

upstash_workflow/asyncio/workflow_parser.py

@@ -1,5 +1,16 @@
-from typing import Optional
+from typing import Optional, cast
 from upstash_workflow.workflow_types import _AsyncRequest
+import json
+from typing import Callable, Dict, Any, Literal, Awaitable, TypeVar
+from upstash_workflow.utils import _decode_base64
+from upstash_workflow.constants import (
+    WORKFLOW_FAILURE_HEADER,
+)
+from upstash_workflow.error import WorkflowError
+from upstash_workflow.workflow_requests import _recreate_user_headers
+from upstash_workflow.asyncio.serve.authorization import _DisabledWorkflowContext
+from qstash import AsyncQStash
+from upstash_workflow import AsyncWorkflowContext
 
 
 async def _get_payload(request: _AsyncRequest) -> Optional[str]:
@@ -13,3 +24,79 @@ async def _get_payload(request: _AsyncRequest) -> Optional[str]:
         return (await request.body()).decode()
     except Exception:
         return None
+
+
+TInitialPayload = TypeVar("TInitialPayload")
+TRequest = TypeVar("TRequest", bound=_AsyncRequest)
+
+
+async def _handle_failure(
+    request: TRequest,
+    request_payload: str,
+    qstash_client: AsyncQStash,
+    initial_payload_parser: Callable[[str], Any],
+    route_function: Callable[[AsyncWorkflowContext[TInitialPayload]], Awaitable[None]],
+    failure_function: Optional[
+        Callable[
+            [AsyncWorkflowContext[TInitialPayload], int, str, Dict[str, str]],
+            Awaitable[None],
+        ]
+    ],
+    env: Dict[str, Any],
+    retries: int,
+) -> Literal["not-failure-callback", "is-failure-callback"]:
+    if request.headers and request.headers.get(WORKFLOW_FAILURE_HEADER) != "true":
+        return "not-failure-callback"
+
+    if not failure_function:
+        raise WorkflowError(
+            "Workflow endpoint is called to handle a failure, "
+            "but a failure_function is not provided in serve options. "
+            "Either provide a failure_url or a failure_function."
+        )
+
+    try:
+        payload = json.loads(request_payload)
+        status = payload["status"]
+        header = payload["header"]
+        body = payload["body"]
+        url = payload["url"]
+        source_body = payload["sourceBody"]
+        workflow_run_id = payload["workflowRunId"]
+
+        decoded_body = _decode_base64(body) if body else "{}"
+        error_payload = json.loads(decoded_body)
+
+        # Create context
+        workflow_context = AsyncWorkflowContext(
+            qstash_client=qstash_client,
+            workflow_run_id=workflow_run_id,
+            initial_payload=initial_payload_parser(_decode_base64(source_body))
+            if source_body
+            else None,
+            headers=_recreate_user_headers(request.headers or {}),
+            steps=[],
+            url=url,
+            failure_url=url,
+            env=env,
+            retries=retries,
+        )
+
+        # Attempt running route_function until the first step
+        auth_check = await _DisabledWorkflowContext[Any].try_authentication(
+            route_function, cast(AsyncWorkflowContext[TInitialPayload], workflow_context)
+        )
+
+        if auth_check == "run-ended":
+            raise WorkflowError("Not authorized to run the failure function.")
+
+        await failure_function(
+            cast(AsyncWorkflowContext[TInitialPayload], workflow_context),
+            status,
+            error_payload.get("message"),
+            header,
+        )
+    except Exception as error:
+        raise error
+
+    return "is-failure-callback"

upstash_workflow/asyncio/workflow_requests.py

@@ -79,6 +79,7 @@ async def _handle_third_party_call_result(
     request_payload: str,
     client: AsyncQStash,
     workflow_url: str,
+    workflow_failure_url: Optional[str],
     retries: int,
 ) -> Literal["call-will-retry", "is-call-return", "continue-workflow"]:
     """
@@ -173,6 +174,7 @@ async def _handle_third_party_call_result(
                 user_headers,
                 None,
                 retries,
+                workflow_failure_url=workflow_failure_url,
             ).headers
 
             call_response = {