vellum-workflow-server 1.3.4__tar.gz → 1.3.4.post1__tar.gz

{vellum_workflow_server-1.3.4 → vellum_workflow_server-1.3.4.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 1.3.4
+Version: 1.3.4.post1
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4

{vellum_workflow_server-1.3.4 → vellum_workflow_server-1.3.4.post1}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
 
 [tool.poetry]
 name = "vellum-workflow-server"
-version = "1.3.4"
+version = "1.3.4.post1"
 description = ""
 readme = "README.md"
 authors = []

{vellum_workflow_server-1.3.4 → vellum_workflow_server-1.3.4.post1}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -4,6 +4,7 @@ from importlib.metadata import version
 import io
 import json
 from queue import Empty
+import re
 from unittest import mock
 from uuid import uuid4
 
@@ -408,7 +409,16 @@ from vellum.workflows.inputs import BaseInputs
     assert len(events) == 2
 
 
-def test_stream_workflow_route__cancel(both_stream_types):
+@pytest.mark.parametrize(
+    ["execute_workflow_stream", "assert_last_request"],
+    [
+        (flask_stream, False),  # Unfortunately, can't make assertions on requests made in a subprocess.
+        (code_exec_stream, True),
+        (flask_stream_disable_process_wrapper, True),
+    ],
+    ids=["flask_stream", "code_exec_stream", "flask_stream_disable_process_wrapper"],
+)
+def test_stream_workflow_route__cancel(execute_workflow_stream, assert_last_request):
     # GIVEN a valid request body
     span_id = uuid4()
     request_body = {
@@ -447,10 +457,10 @@ class BasicCancellableWorkflow(BaseWorkflow):
 
     # WHEN we call the stream route with a mock cancelled return true
     with requests_mock.Mocker() as mocker:
-        mocker.get(
-            f"http://test.biz/vembda-public/cancel-workflow-execution-status/{span_id}", json={"cancelled": True}
+        response_mock = mocker.get(
+            re.compile("http://test.biz/vembda-public/cancel-workflow-execution-status"), json={"cancelled": True}
         )
-        status_code, events = both_stream_types(request_body)
+        status_code, events = execute_workflow_stream(request_body)
 
     # THEN we get a 200 response
     assert status_code == 200, events
@@ -473,14 +483,20 @@ class BasicCancellableWorkflow(BaseWorkflow):
 
     assert events[1]["name"] == "workflow.execution.initiated", events[1]
     assert "display_context" in events[1]["body"], events[1]["body"]
-    cancelled_event = None
-    for event in events:
-        if event["name"] == "workflow.execution.rejected":
-            cancelled_event = event
-            break
 
+    cancelled_event = events[-2]
+    assert cancelled_event["name"] == "workflow.execution.rejected"
     assert cancelled_event["body"]["error"]["message"] == "Workflow run cancelled"
 
+    # AND we called the cancel endpoint with the correct execution id
+    workflow_span_id = events[1]["span_id"]
+    if assert_last_request:
+        assert response_mock.last_request
+        assert (
+            response_mock.last_request.url
+            == f"http://test.biz/vembda-public/cancel-workflow-execution-status/{workflow_span_id}"
+        )
+
 
 def test_stream_workflow_route__very_large_events(both_stream_types):
     # GIVEN a valid request body

{vellum_workflow_server-1.3.4 → vellum_workflow_server-1.3.4.post1}/src/workflow_server/core/executor.py

@@ -192,19 +192,6 @@ def stream_workflow(
 
     cancel_watcher_kill_switch = ThreadingEvent()
     cancel_signal = cancel_signal or ThreadingEvent()
-    cancel_watcher = CancelWorkflowWatcherThread(
-        kill_switch=cancel_watcher_kill_switch,
-        execution_id=executor_context.execution_id,
-        timeout_seconds=executor_context.timeout,
-        vembda_public_url=executor_context.vembda_public_url,
-        cancel_signal=cancel_signal,
-    )
-
-    try:
-        if executor_context.vembda_public_url:
-            cancel_watcher.start()
-    except Exception:
-        logger.exception("Failed to start cancel watcher")
 
     try:
         stream = workflow.stream(
@@ -216,12 +203,25 @@ def stream_workflow(
             entrypoint_nodes=[run_from_node] if run_from_node else None,
             previous_execution_id=executor_context.previous_execution_id,
         )
-
     except Exception:
         cancel_watcher_kill_switch.set()
         logger.exception("Failed to generate Workflow Stream")
         raise
 
+    cancel_watcher = CancelWorkflowWatcherThread(
+        kill_switch=cancel_watcher_kill_switch,
+        execution_id=stream.span_id,
+        timeout_seconds=executor_context.timeout,
+        vembda_public_url=executor_context.vembda_public_url,
+        cancel_signal=cancel_signal,
+    )
+
+    try:
+        if executor_context.vembda_public_url:
+            cancel_watcher.start()
+    except Exception:
+        logger.exception("Failed to start cancel watcher")
+
     def call_workflow() -> Generator[dict[str, Any], Any, None]:
         try:
             first = True