vellum-workflow-server 1.8.2__py3-none-any.whl → 1.10.7__py3-none-any.whl

workflow_server/api/tests/test_workflow_view_async_exec.py

@@ -0,0 +1,410 @@
+import pytest
+import logging
+import time
+from uuid import uuid4
+
+from workflow_server.server import create_app
+from workflow_server.utils.system_utils import get_active_process_count
+
+
+@pytest.fixture(autouse=True)
+def drain_background_threads():
+    """
+    Ensures background threads from previous tests complete before starting the next test.
+    This prevents cross-test interference in process count assertions.
+    """
+    baseline = get_active_process_count()
+    yield
+
+    deadline = time.time() + 15
+    while time.time() < deadline:
+        current_count = get_active_process_count()
+        if current_count == baseline:
+            break
+        time.sleep(0.1)
+
+
+def test_async_exec_route__happy_path():
+    """
+    Tests that the async-exec route successfully accepts a valid workflow and returns immediately.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    # AND a valid workflow request
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+        },
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 200 response
+    assert response.status_code == 200
+
+    # AND the response should indicate success
+    assert response.json == {"success": True}
+
+
+def test_async_exec_route__with_inputs():
+    """
+    Tests that the async-exec route handles workflows with inputs correctly.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    # AND a valid workflow request with inputs
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [
+            {"name": "foo", "type": "STRING", "value": "hello"},
+        ],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+from .inputs import Inputs
+
+class Workflow(BaseWorkflow[Inputs, BaseState]):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+            "inputs.py": """\
+from vellum.workflows.inputs import BaseInputs
+
+class Inputs(BaseInputs):
+    foo: str
+""",
+        },
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 200 response
+    assert response.status_code == 200
+
+    # AND the response should indicate success
+    assert response.json == {"success": True}
+
+
+def test_async_exec_route__with_state():
+    """
+    Tests that the async-exec route handles workflows with state correctly.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    # AND a valid workflow request with state
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "state": {"foo": "bar"},
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs import BaseInputs
+from .state import State
+
+class Workflow(BaseWorkflow[BaseInputs, State]):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = State.foo
+""",
+            "state.py": """\
+from vellum.workflows.state import BaseState
+
+class State(BaseState):
+    foo: str
+""",
+        },
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 200 response
+    assert response.status_code == 200
+
+    # AND the response should indicate success
+    assert response.json == {"success": True}
+
+
+def test_async_exec_route__invalid_context():
+    """
+    Tests that the async-exec route returns 400 for invalid request context.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    # AND an invalid request missing required fields
+    request_body = {
+        "inputs": [],
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 400 response
+    assert response.status_code == 400
+
+    # AND the response should contain error details
+    assert "detail" in response.json
+    assert "Invalid context" in response.json["detail"]
+
+
+def test_async_exec_route__missing_files():
+    """
+    Tests that the async-exec route returns 400 when files are missing.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 400 response
+    assert response.status_code == 400
+
+    # AND the response should contain error details
+    assert "detail" in response.json
+    assert "Invalid context" in response.json["detail"]
+
+
+def test_async_exec_route__with_syntax_error_in_workflow():
+    """
+    Tests that the async-exec route handles workflows with syntax errors gracefully.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow)
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+        },
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 200 response (async execution is accepted)
+    assert response.status_code == 200
+
+    # AND the response should indicate success
+    assert response.json == {"success": True}
+
+
+def test_async_exec_route__with_invalid_inputs():
+    """
+    Tests that the async-exec route handles workflows with invalid inputs gracefully.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+from .inputs import Inputs
+
+class Workflow(BaseWorkflow[Inputs, BaseState]):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+            "inputs.py": """\
+from vellum.workflows.inputs import BaseInputs
+
+class Inputs(BaseInputs):
+    foo: str
+""",
+        },
+    }
+
+    # WHEN we make a request to the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we should get a 200 response (async execution is accepted)
+    assert response.status_code == 200
+
+    # AND the response should indicate success
+    assert response.json == {"success": True}
+
+
+def test_async_exec_route__background_thread_completes(caplog):
+    """
+    Verifies that the async background worker thread runs to completion.
+    """
+    # GIVEN a Flask application with log capture enabled
+    caplog.set_level(logging.INFO, logger="workflow_server.api.workflow_view")
+    flask_app = create_app()
+
+    baseline = get_active_process_count()
+
+    # AND a valid workflow request
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+        },
+    }
+
+    # WHEN we call the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we get immediate acceptance
+    assert response.status_code == 200
+    assert response.json == {"success": True}
+
+    # AND the background thread should complete
+    completion_deadline = time.time() + 15
+    saw_completion_log = False
+    while time.time() < completion_deadline:
+        if any("Workflow async exec completed" in rec.message for rec in caplog.records):
+            saw_completion_log = True
+            break
+        time.sleep(0.1)
+
+    # THEN we should observe the completion log
+    assert saw_completion_log, "Did not observe background completion log within 15 seconds"
+
+    cleanup_deadline = time.time() + 15
+    process_count_returned = False
+    while time.time() < cleanup_deadline:
+        current_count = get_active_process_count()
+        if current_count == baseline:
+            process_count_returned = True
+            break
+        time.sleep(0.1)
+
+    current_count = get_active_process_count()
+    assert process_count_returned, (
+        f"Process count did not return to baseline within 15 seconds after completion log. "
+        f"Expected: {baseline}, Current: {current_count}"
+    )
+
+
+def test_async_exec_route__background_thread_completes_on_error(caplog):
+    """
+    Verifies that the background worker completes even when the workflow fails early.
+    """
+    # GIVEN a Flask application with log capture enabled
+    caplog.set_level(logging.INFO, logger="workflow_server.api.workflow_view")
+    flask_app = create_app()
+
+    baseline = get_active_process_count()
+
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow)
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+        },
+    }
+
+    # WHEN we call the async-exec route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/async-exec", json=request_body)
+
+    # THEN we get immediate acceptance
+    assert response.status_code == 200
+    assert response.json == {"success": True}
+
+    # AND the background thread should complete and clean up resources
+    deadline = time.time() + 15
+    process_count_returned = False
+    while time.time() < deadline:
+        current_count = get_active_process_count()
+        if current_count == baseline:
+            process_count_returned = True
+            break
+        time.sleep(0.1)
+
+    current_count = get_active_process_count()
+    assert process_count_returned, (
+        f"Process count did not return to baseline on error within 15 seconds. "
+        f"Expected: {baseline}, Current: {current_count}"
+    )

workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -5,6 +5,7 @@ import io
 import json
 from queue import Empty
 import re
+import time
 from unittest import mock
 from uuid import uuid4
 
@@ -133,6 +134,8 @@ class Workflow(BaseWorkflow):
 
     with mock.patch("builtins.open", mock.mock_open(read_data="104857600")):
         # WHEN we call the stream route
+        ts_ns = time.time_ns()
+        request_body["vembda_service_initiated_timestamp"] = ts_ns
         status_code, events = both_stream_types(request_body)
 
     # THEN we get a 200 response
@@ -174,6 +177,17 @@ class Workflow(BaseWorkflow):
     assert "sdk_version" in server_metadata
     assert "memory_usage_mb" in server_metadata
     assert isinstance(server_metadata["memory_usage_mb"], (int, float))
+    assert "is_new_server" in server_metadata
+    assert server_metadata["is_new_server"] is False
+
+    # AND the initiated event should have initiated_latency within a reasonable range
+    assert "initiated_latency" in server_metadata, "initiated_latency should be present in server_metadata"
+    initiated_latency = server_metadata["initiated_latency"]
+    assert isinstance(initiated_latency, int), "initiated_latency should be an integer (nanoseconds)"
+    # Latency should be positive and less than 60 seconds (60_000_000_000 nanoseconds) for CI
+    assert (
+        0 < initiated_latency < 60_000_000_000
+    ), f"initiated_latency should be between 0 and 60 seconds, got {initiated_latency} ns"
 
     assert events[2]["name"] == "workflow.execution.fulfilled", events[2]
     assert events[2]["body"]["workflow_definition"]["module"] == ["test", "workflow"]
@@ -384,9 +398,15 @@ class State(BaseState):
 def test_stream_workflow_route__bad_indent_in_inputs_file(both_stream_types):
     # GIVEN a valid request body
     span_id = uuid4()
+    trace_id = uuid4()
+    parent_span_id = uuid4()
     request_body = {
         "timeout": 360,
         "execution_id": str(span_id),
+        "execution_context": {
+            "trace_id": str(trace_id),
+            "parent_context": {"span_id": str(parent_span_id)},
+        },
         "inputs": [
             {"name": "foo", "type": "STRING", "value": "hello"},
         ],
@@ -423,7 +443,7 @@ from vellum.workflows.inputs import BaseInputs
 
     assert events[0] == {
         "id": mock.ANY,
-        "trace_id": mock.ANY,
+        "trace_id": str(trace_id),
         "span_id": str(span_id),
         "timestamp": mock.ANY,
         "api_version": "2024-10-25",
@@ -437,9 +457,19 @@ from vellum.workflows.inputs import BaseInputs
     }
 
     assert events[1]["name"] == "workflow.execution.initiated"
+    assert events[1]["trace_id"] == str(trace_id), "workflow initiated event should use request trace_id"
+    assert events[1]["parent"] is not None, "workflow initiated event should have parent context"
+    assert events[1]["parent"]["span_id"] == str(
+        parent_span_id
+    ), "workflow initiated event parent should match request parent_context"
 
     assert events[2]["name"] == "workflow.execution.rejected"
+    assert events[2]["trace_id"] == str(trace_id), "workflow rejected event should use request trace_id"
     assert events[2]["span_id"] == events[1]["span_id"]
+    assert events[2]["parent"] is not None, "workflow rejected event should have parent context"
+    assert events[2]["parent"]["span_id"] == str(
+        parent_span_id
+    ), "workflow rejected event parent should match request parent_context"
     assert (
         "Syntax Error raised while loading Workflow: "
         "unexpected indent (inputs.py, line 3)" in events[2]["body"]["error"]["message"]
@@ -1211,3 +1241,72 @@ class InvalidWorkflow(BaseWorkflow):
     assert events[3]["name"] == "vembda.execution.fulfilled"
     assert events[3]["span_id"] == str(span_id)
     assert events[3]["body"]["exit_code"] == 0
+
+
+@mock.patch("workflow_server.api.workflow_view.get_is_oom_killed")
+def test_stream_workflow_route__oom_does_not_set_timed_out_flag(mock_get_is_oom_killed):
+    """
+    Tests that when an OOM error occurs, we don't set the timed_out flag in the vembda fulfilled event.
+    """
+    # GIVEN a workflow that takes some time to execute
+    span_id = uuid4()
+    request_body = {
+        "timeout": 10,
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+import time
+
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.workflows.base import BaseWorkflow
+
+
+class SlowNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        value: str
+
+    def run(self) -> Outputs:
+        time.sleep(2)
+        return self.Outputs(value="hello world")
+
+
+class OOMWorkflow(BaseWorkflow):
+    graph = SlowNode
+    class Outputs(BaseWorkflow.Outputs):
+        final_value = SlowNode.Outputs.value
+
+""",
+        },
+    }
+
+    # WHEN we mock the OOM killer to trigger after a few checks
+    call_count = [0]
+
+    def mock_oom_side_effect():
+        call_count[0] += 1
+        if call_count[0] > 3:
+            return True
+        return False
+
+    mock_get_is_oom_killed.side_effect = mock_oom_side_effect
+
+    # AND we call the stream route
+    status_code, events = flask_stream(request_body)
+
+    # THEN we get a 200 response
+    assert status_code == 200
+
+    # AND we get the expected events
+    event_names = [e["name"] for e in events]
+
+    assert "vembda.execution.initiated" in event_names
+
+    # THEN the key assertion: if there's a vembda.execution.fulfilled event, it should NOT have timed_out=True
+    vembda_fulfilled_event = next(e for e in events if e["name"] == "vembda.execution.fulfilled")
+    assert (
+        vembda_fulfilled_event["body"].get("timed_out") is not True
+    ), "timed_out flag should not be set when OOM occurs"