vellum-workflow-server 1.3.4__py3-none-any.whl → 1.3.4.post2__py3-none-any.whl

{vellum_workflow_server-1.3.4.dist-info → vellum_workflow_server-1.3.4.post2.dist-info}/METADATA

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

{vellum_workflow_server-1.3.4.dist-info → vellum_workflow_server-1.3.4.post2.dist-info}/RECORD

@@ -5,14 +5,14 @@ workflow_server/api/healthz_view.py,sha256=itiRvBDBXncrw8Kbbc73UZLwqMAhgHOR3uSre
 workflow_server/api/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/api/tests/test_input_display_mapping.py,sha256=drBZqMudFyB5wgiUOcMgRXz7E7ge-Qgxbstw4E4f0zE,2211
 workflow_server/api/tests/test_workflow_view.py,sha256=gwb53E44LBm06XHY1UwfHqKG-UfwwICh_IYPtowN_kE,19085
-workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=1vr9NI894zdJmmnsy_wBSqaJG83RwnYSmqucknYf9bI,26304
-workflow_server/api/workflow_view.py,sha256=osYa14kdP3LI3HZzh2Nw16-GhI0bHXwCFgasUkfRjeI,21104
+workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=Qo8u6mPyRCmE2jamY1yIh8l44hgo4-Nwlq03z61ND5g,27031
+workflow_server/api/workflow_view.py,sha256=gSKqlMR2r2urSBYm_jFlcpXHOk5jQsDMSDwpSahp5A8,21144
 workflow_server/code_exec_runner.py,sha256=lBnMIorPZL8zZBye6TjeCIs06WTJM7P2HR07B1fjJJI,2533
 workflow_server/config.py,sha256=DyTty8NrAwvtx-esM3KthnpsNh-nKdWNlovWQOgiGpg,1417
 workflow_server/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/core/cancel_workflow.py,sha256=Ffkc3mzmrdMEUcD-sHfEhX4IwVrka-E--SxKA1dUfIU,2185
 workflow_server/core/events.py,sha256=24MA66DVQuaLJJcZrS8IL1Zq4Ohi9CoouKZ5VgoH3Cs,1402
-workflow_server/core/executor.py,sha256=M_LyVOcPDzAA3u1WBLMpz9gftjc_RtO3hYida6tkSJo,18383
+workflow_server/core/executor.py,sha256=TTpMifrIZuzXxOUtX3E3CgLOb8eRXul2aGPpRCfWIco,18368
 workflow_server/core/utils.py,sha256=cmwHbKCfXqtUutBD3akGus0Ga7a1xG3zlOw-jEMx6mI,1795
 workflow_server/core/workflow_executor_context.py,sha256=w3OhV_AXpgh7AxpjEsc0vo-IJypgJcgr5DXJCqGptOU,1587
 workflow_server/server.py,sha256=QBU12AaAfAgLqfCDBd24qIJl_mbheiq0-hfcWV7rZM4,1234
@@ -28,7 +28,7 @@ workflow_server/utils/tests/test_sentry_integration.py,sha256=LGmWiaLhFrx-jslrRj
 workflow_server/utils/tests/test_system_utils.py,sha256=_4GwXvVvU5BrATxUEWwQIPg0bzQXMWBtiBmjP8MTxJM,4314
 workflow_server/utils/tests/test_utils.py,sha256=qwK5Rmy3RQyjtlUrYAuGuDlBeRzZKsf1yS-y2IpUizQ,6452
 workflow_server/utils/utils.py,sha256=Wqqn-1l2ugkGgy5paWWdt0AVxAyPMQCYcnRSSOMjXlA,4355
-vellum_workflow_server-1.3.4.dist-info/METADATA,sha256=5S0p5CKKxsab90EaMaRWtY7Wwuh3n_L58SxvMgvokc0,2267
-vellum_workflow_server-1.3.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_workflow_server-1.3.4.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
-vellum_workflow_server-1.3.4.dist-info/RECORD,,
+vellum_workflow_server-1.3.4.post2.dist-info/METADATA,sha256=t3_ELDwFzZUWqnYxfvYTp8uypnAsPAPeOnEpFoeBRvo,2273
+vellum_workflow_server-1.3.4.post2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-1.3.4.post2.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-1.3.4.post2.dist-info/RECORD,,

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

workflow_server/api/workflow_view.py

@@ -81,7 +81,7 @@ def stream_workflow_route() -> Response:
         )
 
     logger.info(
-        f"Starting workflow stream, execution ID: {context.execution_id}, "
+        f"Starting Workflow Server Request, trace ID: {context.trace_id}, "
         f"process count: {get_active_process_count()}, process wrapper: {ENABLE_PROCESS_WRAPPER}"
     )
 
@@ -194,6 +194,8 @@ def stream_workflow_route() -> Response:
             headers=headers,
         )
 
+    logger.info(f"Starting Workflow Stream, execution ID: {span_id}, ")
+
     def process_events(queue: Queue) -> Iterator[Union[str, dict]]:
         event: Union[str, dict]
         loops = 0
@@ -309,8 +311,7 @@ def stream_workflow_route() -> Response:
             yield "\n"
 
             logger.info(
-                f"Workflow stream completed, execution ID: {context.execution_id}, "
-                f"process count: {get_active_process_count()}"
+                f"Workflow stream completed, execution ID: {span_id}, process count: {get_active_process_count()}"
             )
         except GeneratorExit:
             # These can happen either from Vembda disconnects (possibily from predict disconnects) or

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