vellum-workflow-server 1.7.10__py3-none-any.whl → 1.9.6.post1__py3-none-any.whl

{vellum_workflow_server-1.7.10.dist-info → vellum_workflow_server-1.9.6.post1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 1.7.10
+Version: 1.9.6.post1
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4
@@ -29,7 +29,7 @@ Requires-Dist: pyjwt (==2.10.0)
 Requires-Dist: python-dotenv (==1.0.1)
 Requires-Dist: retrying (==1.3.4)
 Requires-Dist: sentry-sdk[flask] (==2.20.0)
-Requires-Dist: vellum-ai (==1.7.10)
+Requires-Dist: vellum-ai (==1.9.6)
 Description-Content-Type: text/markdown
 
 # Vellum Workflow Runner Server

{vellum_workflow_server-1.7.10.dist-info → vellum_workflow_server-1.9.6.post1.dist-info}/RECORD

@@ -1,22 +1,24 @@
 workflow_server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-workflow_server/api/auth_middleware.py,sha256=IlZaCiwZ5nwQqk5sYQorvOFj7lt0p1ZSSEqUxfiFaW0,2458
+workflow_server/api/auth_middleware.py,sha256=qHsMZFYtZyaR7dvkJ6xI6ZYdYw2-2y8X4JyfIQqQTSw,2495
 workflow_server/api/healthz_view.py,sha256=itiRvBDBXncrw8Kbbc73UZLwqMAhgHOR3uSre_dAfgY,404
+workflow_server/api/status_view.py,sha256=Jah8dBAVL4uOcRfsjKAOyfVONFyk9HQjXeRfjcIqhmA,514
 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=_9SHdK1t-wwWILZFrnoYe5NhkoLazJnt7K1oH4bnJU0,30355
-workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=GjWbwSmjxb691LmKqaCsAItkcb3QJaHw_OG_H5ETWdc,37432
-workflow_server/api/workflow_view.py,sha256=tSepIRP0LcJUKae_wSHGw6P5Ho8K20_YTj3J-9ZL6bs,20463
+workflow_server/api/tests/test_workflow_view.py,sha256=81kAHpijNp0rvb3ZjvceB5uFEriVWPeWHnK78-xoeTc,32343
+workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=Yrp_DlLbbwZJe5WRLwdlFT17R8CQoCK9-jlQ1jUT_eM,40377
+workflow_server/api/workflow_view.py,sha256=XSVfHYgsy2k_QqTiue9Xx438Z4qprHbD0PeS8JI04sY,24547
 workflow_server/code_exec_runner.py,sha256=DLNNrinCRbnkSvlqVvSZ1wv_etI7r_kKAXNPGMj3jBk,2196
-workflow_server/config.py,sha256=cUdI_lEovV7e7lwCkGJ1eM9R4OZVJw5R5zT1eG1SzQQ,2122
+workflow_server/config.py,sha256=I4hfTsjIbHxoSKylPCjKnrysPV0jO5nfRKwpKvEcfAE,2193
 workflow_server/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/core/cancel_workflow.py,sha256=QcEeYUIrxq4pub-z9BlGi5fLI3gVRml-56rMCW7j5Hc,2212
 workflow_server/core/events.py,sha256=24MA66DVQuaLJJcZrS8IL1Zq4Ohi9CoouKZ5VgoH3Cs,1402
-workflow_server/core/executor.py,sha256=rMpVgP3PqMnjHVYeuTdHvbYGwJW5C9O5TFGW926jFM0,16106
+workflow_server/core/executor.py,sha256=lP69l8ATeSe88DOUPIO5mmwq1iuQ-02smw7Tr471wTY,17754
 workflow_server/core/utils.py,sha256=si0NB4Suurc-mn8NYdn59xM9CkPrfOP1aWEVrZvifDI,1929
-workflow_server/core/workflow_executor_context.py,sha256=Q2R0T2KkYZ1z52v8erDMysJfxSODbzPhDxBxX--k4Zw,3202
-workflow_server/server.py,sha256=lhHPmK1PhRZd6eCkj1C0acK3YwaApZgoPHghMChw0fc,1461
-workflow_server/start.py,sha256=xSIobowtSLoZI86bbMkmEw3pqJHQaFdDyNffk4kGYL8,2544
+workflow_server/core/workflow_executor_context.py,sha256=uUlFF2PIfFzIzhHS25mpvO4wO97UWqQVArg7zC2xVcM,3490
+workflow_server/logging_config.py,sha256=Hvx1t8uhqMMinl-5qcef7ufUvzs6x14VRnCb7YZxEAg,1206
+workflow_server/server.py,sha256=pBl0OQmrLE-PbTDwTgsVmxgz_Ai3TVhFRaMnr6PX6Yk,1849
+workflow_server/start.py,sha256=Ams5ycqVbBorC7s6EI95BYzjpxzlo5mQbBnMNOkJS0w,2753
 workflow_server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/utils/exit_handler.py,sha256=_FacDVi4zc3bfTA3D2mJsISePlJ8jpLrnGVo5-xZQFs,743
 workflow_server/utils/log_proxy.py,sha256=nugi6fOgAYKX2X9DIc39TG366rsmmDUPoEtG3gzma_Y,3088
@@ -28,7 +30,7 @@ workflow_server/utils/tests/test_sentry_integration.py,sha256=14PfuW8AaQNNtqLmBs
 workflow_server/utils/tests/test_system_utils.py,sha256=_4GwXvVvU5BrATxUEWwQIPg0bzQXMWBtiBmjP8MTxJM,4314
 workflow_server/utils/tests/test_utils.py,sha256=0Nq6du8o-iBtTrip9_wgHES53JSiJbVdSXaBnPobw3s,6930
 workflow_server/utils/utils.py,sha256=m7iMJtor5SQLWu7jlJw-X5Q3nmbq69BCxTMv6qnFYrA,4835
-vellum_workflow_server-1.7.10.dist-info/METADATA,sha256=0jfzSrst_oJpWuB5jFawlvAeVZIZWvaBn4t-O2G1M_I,2269
-vellum_workflow_server-1.7.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_workflow_server-1.7.10.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
-vellum_workflow_server-1.7.10.dist-info/RECORD,,
+vellum_workflow_server-1.9.6.post1.dist-info/METADATA,sha256=69SUWrSyCFeBTIuN7NVAhHUTUBaNtUAgz8SrHFlEOag,2273
+vellum_workflow_server-1.9.6.post1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-1.9.6.post1.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-1.9.6.post1.dist-info/RECORD,,

workflow_server/api/auth_middleware.py

@@ -5,7 +5,7 @@ from flask import Flask, Request, Response
 import jwt
 from jwt import ExpiredSignatureError
 
-from workflow_server.config import IS_VPC, NAMESPACE, VEMBDA_PUBLIC_KEY, is_development
+from workflow_server.config import IS_ASYNC_MODE, IS_VPC, NAMESPACE, VEMBDA_PUBLIC_KEY, is_development
 
 
 class AuthMiddleware:
@@ -15,7 +15,7 @@ class AuthMiddleware:
     def __call__(self, environ: Dict[str, Any], start_response: Any) -> Any:
         try:
             request = Request(environ)
-            if not request.path.startswith("/healthz") and not is_development() and not IS_VPC:
+            if not request.path.startswith("/healthz") and not is_development() and not IS_VPC and not IS_ASYNC_MODE:
                 token = request.headers.get("X-Vembda-Signature")
                 if token:
                     decoded = jwt.decode(token, VEMBDA_PUBLIC_KEY, algorithms=["RS256"])

workflow_server/api/status_view.py

@@ -0,0 +1,19 @@
+from typing import Tuple
+
+from flask import Blueprint, Response, jsonify
+
+from workflow_server.config import CONCURRENCY
+from workflow_server.utils.system_utils import get_active_process_count
+
+bp = Blueprint("status", __name__)
+
+
+@bp.route("/is_available", methods=["GET"])
+def is_available() -> Tuple[Response, int]:
+    resp = jsonify(
+        available=get_active_process_count() < CONCURRENCY,
+        process_count=get_active_process_count(),
+        max_concurrency=CONCURRENCY,
+    )
+
+    return resp, 200

workflow_server/api/tests/test_workflow_view.py

@@ -537,6 +537,57 @@ def test_serialize_route__with_invalid_workspace_api_key():
     assert "exec_config" in response.json
 
 
+def test_serialize_route__with_is_new_server_header():
+    """
+    Tests that the serialize route returns the is_new_server header.
+    """
+    # GIVEN a Flask application
+    flask_app = create_app()
+
+    workflow_files = {
+        "__init__.py": "",
+        "workflow.py": (
+            "from vellum.workflows import BaseWorkflow\n\n"
+            "class Workflow(BaseWorkflow):\n"
+            "    class Outputs(BaseWorkflow.Outputs):\n"
+            "        foo = 'hello'\n"
+        ),
+    }
+
+    # WHEN we make a request with is_new_server=True
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/serialize", json={"files": workflow_files, "is_new_server": True})
+
+    # THEN we should get a successful response
+    assert response.status_code == 200
+
+    # AND the response should contain the is_new_server header set to true
+    assert "X-Vellum-Is-New-Server" in response.headers
+    assert response.headers["X-Vellum-Is-New-Server"] == "true"
+
+    # WHEN we make a request with is_new_server=False
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/serialize", json={"files": workflow_files, "is_new_server": False})
+
+    # THEN we should get a successful response
+    assert response.status_code == 200
+
+    # AND the response should contain the is_new_server header set to false
+    assert "X-Vellum-Is-New-Server" in response.headers
+    assert response.headers["X-Vellum-Is-New-Server"] == "false"
+
+    # WHEN we make a request without is_new_server
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/serialize", json={"files": workflow_files})
+
+    # THEN we should get a successful response
+    assert response.status_code == 200
+
+    # AND the response should contain the is_new_server header set to false (default)
+    assert "X-Vellum-Is-New-Server" in response.headers
+    assert response.headers["X-Vellum-Is-New-Server"] == "false"
+
+
 def test_stream_node_route__with_node_id():
     """
     Tests that the stream-node endpoint works with node_id.

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
 
@@ -131,8 +132,11 @@ class Workflow(BaseWorkflow):
         },
     }
 
-    # WHEN we call the stream route
-    status_code, events = both_stream_types(request_body)
+    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
     assert status_code == 200, events
@@ -164,9 +168,37 @@ class Workflow(BaseWorkflow):
     assert isinstance(display_context["workflow_inputs"], dict)
     assert isinstance(display_context["workflow_outputs"], dict)
     assert "foo" in display_context["workflow_outputs"]
+
+    # AND the initiated event should have server_metadata with version info and memory usage
+    assert "server_metadata" in events[1]["body"], events[1]["body"]
+    server_metadata = events[1]["body"]["server_metadata"]
+    assert server_metadata is not None, "server_metadata should not be None"
+    assert "server_version" in server_metadata
+    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"]
 
+    # AND the fulfilled event should have server_metadata with memory usage
+    assert "server_metadata" in events[2]["body"], events[2]["body"]
+    fulfilled_metadata = events[2]["body"]["server_metadata"]
+    assert fulfilled_metadata is not None, "fulfilled server_metadata should not be None"
+    assert "memory_usage_mb" in fulfilled_metadata
+    assert isinstance(fulfilled_metadata["memory_usage_mb"], (int, float))
+
     assert events[3] == {
         "id": mock.ANY,
         "trace_id": events[0]["trace_id"],
@@ -366,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"},
         ],
@@ -405,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",
@@ -419,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"]
@@ -653,7 +701,10 @@ class TimeoutWorkflow(BaseWorkflow):
 
     assert "workflow.execution.rejected" in event_names, "Should emit workflow.execution.rejected on timeout"
     workflow_execution_rejected = next(e for e in events if e["name"] == "workflow.execution.rejected")
-    assert "vellum/workflows/runner/runner.py" in workflow_execution_rejected["body"]["stacktrace"]
+    assert workflow_execution_rejected["body"]["error"]["code"] == "WORKFLOW_TIMEOUT"
+    # TODO: Uncomment once version 1.8.1 is released
+    # assert "stacktrace" in workflow_execution_rejected["body"]
+    # assert "vellum/workflows/runner/runner.py" in workflow_execution_rejected["body"]["stacktrace"]
 
     assert "vembda.execution.fulfilled" in event_names
     vembda_fulfilled = next(e for e in events if e["name"] == "vembda.execution.fulfilled")

workflow_server/api/workflow_view.py

@@ -8,6 +8,7 @@ import os
 import pkgutil
 from queue import Empty
 import sys
+import threading
 import time
 import traceback
 from uuid import uuid4
@@ -71,7 +72,6 @@ WORKFLOW_INITIATION_TIMEOUT_SECONDS = 60
 @bp.route("/stream", methods=["POST"])
 def stream_workflow_route() -> Response:
     data = request.get_json()
-
     try:
         context = WorkflowExecutorContext.model_validate(data)
     except ValidationError as e:
@@ -84,28 +84,7 @@ def stream_workflow_route() -> Response:
             content_type="application/json",
         )
 
-    logger.info(
-        f"Starting Workflow Server Request, trace ID: {context.trace_id}, "
-        f"process count: {get_active_process_count()}, process wrapper: {ENABLE_PROCESS_WRAPPER}"
-    )
-
-    # Create this event up here so timestamps are fully from the start to account for any unknown overhead
-    vembda_initiated_event = VembdaExecutionInitiatedEvent(
-        id=uuid4(),
-        timestamp=datetime.now(),
-        trace_id=context.trace_id,
-        span_id=context.execution_id,
-        body=VembdaExecutionInitiatedBody.model_validate(get_version()),
-        parent=None,
-    )
-
-    process_output_queue: Queue[Union[str, dict]] = Queue()
-
-    headers = {
-        "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
-        "X-Vellum-Server-Version": vembda_initiated_event.body.server_version,
-        "X-Vellum-Events-Emitted": str(is_events_emitting_enabled(context)),
-    }
+    headers = _get_headers(context)
 
     # We can exceed the concurrency count currently with long running workflows due to a knative issue. So here
     # if we detect a memory problem just exit us early
@@ -122,6 +101,183 @@ def stream_workflow_route() -> Response:
             headers=headers,
         )
 
+    start_workflow_state = _start_workflow(context)
+    if isinstance(start_workflow_state, Response):
+        return start_workflow_state
+
+    workflow_events, vembda_initiated_event, process, span_id, headers = start_workflow_state
+
+    def generator() -> Generator[str, None, None]:
+        try:
+            yield "\n"
+            yield vembda_initiated_event.model_dump_json()
+            yield "\n"
+            for row in workflow_events:
+                yield "\n"
+                if isinstance(row, dict):
+                    dump = json.dumps(row)
+                    yield dump
+                else:
+                    yield row
+                yield "\n"
+            # Sometimes the connections get hung after they finish with the vembda fulfilled event
+            # if it happens during a knative scale down event. So we emit an END string so that
+            # we don't have to do string compares on all the events for performance.
+            yield "\n"
+            yield "END"
+            yield "\n"
+
+            logger.info(
+                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
+            # from knative activator gateway timeouts which are caused by idleTimeout or responseStartSeconds
+            # being exceeded.
+            app.logger.error(
+                "Client disconnected in the middle of the Workflow Stream",
+                extra={
+                    "sentry_tags": {
+                        "server_version": vembda_initiated_event.body.server_version,
+                        "sdk_version": vembda_initiated_event.body.sdk_version,
+                    }
+                },
+            )
+            return
+        except Exception as e:
+            logger.exception("Error during workflow response stream generator", extra={"error": e})
+            yield "\n"
+            yield "END"
+            yield "\n"
+            return
+        finally:
+            if ENABLE_PROCESS_WRAPPER:
+                try:
+                    if process and process.is_alive():
+                        process.kill()
+                    if process:
+                        increment_process_count(-1)
+                        remove_active_span_id(span_id)
+                except Exception as e:
+                    logger.error("Failed to kill process", e)
+            else:
+                increment_process_count(-1)
+                remove_active_span_id(span_id)
+
+    resp = Response(
+        stream_with_context(generator()),
+        status=200,
+        content_type="application/x-ndjson",
+        headers=headers,
+    )
+    return resp
+
+
+@bp.route("/async-exec", methods=["POST"])
+def async_exec_workflow() -> Response:
+    data = request.get_json()
+    try:
+        context = WorkflowExecutorContext.model_validate(data)
+    except ValidationError as e:
+        error_message = e.errors()[0]["msg"]
+        error_location = e.errors()[0]["loc"]
+
+        # TODO need to convert this to a vembda event so that trigger'd execs can me notified
+        #  can either do it here in the workflow server or
+        return Response(
+            json.dumps({"detail": f"Invalid context: {error_message} at {error_location}"}),
+            status=400,
+            content_type="application/json",
+        )
+
+    # Reject back to the queue handler if were low on memory here, though maybe we should update the is_available
+    # route to look at memory too. Don't send this response as an event. Though we might want some logic to catch
+    # if they have a workflow server that can never start a workflow because the base image uses so much memory.
+    if not wait_for_available_process():
+        return Response(
+            json.dumps({"detail": f"Server resources low." f"Process count: {get_active_process_count()}"}),
+            status=429,
+            content_type="application/json",
+        )
+
+    def run_workflow_background() -> None:
+        process: Optional[Process] = None
+        span_id: Optional[str] = None
+
+        try:
+            start_workflow_result = _start_workflow(context)
+            if isinstance(start_workflow_result, Response):
+                # TODO same here, should return this response as en event or it will get yeeted to the nether
+                # return start_workflow_result
+                return
+
+            workflow_events, vembda_initiated_event, process, span_id, headers = start_workflow_result
+
+            for _ in workflow_events:
+                # This is way inefficient in process mode since were just having the main proc stream the events
+                # to nowhere wasting memory I/O and cpu.
+                continue
+            logger.info(
+                f"Workflow async exec completed, execution ID: {span_id}, process count: {get_active_process_count()}"
+            )
+        except Exception as e:
+            logger.exception("Error during workflow async background worker", e)
+        finally:
+            if ENABLE_PROCESS_WRAPPER:
+                try:
+                    if process and process.is_alive():
+                        process.kill()
+                    if process:
+                        increment_process_count(-1)
+                        if span_id:
+                            remove_active_span_id(span_id)
+                except Exception as e:
+                    logger.error("Failed to kill process", e)
+            else:
+                increment_process_count(-1)
+                if span_id:
+                    remove_active_span_id(span_id)
+
+    thread = threading.Thread(target=run_workflow_background)
+    thread.start()
+
+    return Response(
+        json.dumps({"success": True}),
+        status=200,
+        content_type="application/json",
+    )
+
+
+def _start_workflow(
+    context: WorkflowExecutorContext,
+) -> Union[
+    Response,
+    tuple[
+        Iterator[Union[str, dict]],
+        VembdaExecutionInitiatedEvent,
+        Optional[Process],
+        str,
+        dict[str, str],
+    ],
+]:
+    headers = _get_headers(context)
+    logger.info(
+        f"Starting Workflow Server Request, trace ID: {context.trace_id}, "
+        f"process count: {get_active_process_count()}, process wrapper: {ENABLE_PROCESS_WRAPPER}"
+    )
+
+    # Create this event up here so timestamps are fully from the start to account for any unknown overhead
+    version_data = get_version()
+    vembda_initiated_event = VembdaExecutionInitiatedEvent(
+        id=uuid4(),
+        timestamp=datetime.now(),
+        trace_id=context.trace_id,
+        span_id=context.execution_id,
+        body=VembdaExecutionInitiatedBody.model_validate(version_data),
+        parent=None,
+    )
+
+    output_queue: Queue[Union[str, dict]] = Queue()
     cancel_signal = MultiprocessingEvent()
     timeout_signal = MultiprocessingEvent()
 
@@ -130,7 +286,7 @@ def stream_workflow_route() -> Response:
         try:
             process = stream_workflow_process_timeout(
                 executor_context=context,
-                queue=process_output_queue,
+                queue=output_queue,
                 cancel_signal=cancel_signal,
                 timeout_signal=timeout_signal,
             )
@@ -138,10 +294,10 @@ def stream_workflow_route() -> Response:
         except Exception as e:
             logger.exception(e)
 
-            process_output_queue.put(create_vembda_rejected_event(context, traceback.format_exc()))
+            output_queue.put(create_vembda_rejected_event(context, traceback.format_exc()))
 
         try:
-            first_item = process_output_queue.get(timeout=WORKFLOW_INITIATION_TIMEOUT_SECONDS)
+            first_item = output_queue.get(timeout=WORKFLOW_INITIATION_TIMEOUT_SECONDS)
         except Empty:
             logger.error("Request timed out trying to initiate the Workflow")
 
@@ -290,72 +446,9 @@ def stream_workflow_route() -> Response:
                 break
             yield event
 
-    workflow_events = process_events(process_output_queue)
-
-    def generator() -> Generator[str, None, None]:
-        try:
-            yield "\n"
-            yield vembda_initiated_event.model_dump_json()
-            yield "\n"
-            for row in workflow_events:
-                yield "\n"
-                if isinstance(row, dict):
-                    dump = json.dumps(row)
-                    yield dump
-                else:
-                    yield row
-                yield "\n"
-            # Sometimes the connections get hung after they finish with the vembda fulfilled event
-            # if it happens during a knative scale down event. So we emit an END string so that
-            # we don't have to do string compares on all the events for performance.
-            yield "\n"
-            yield "END"
-            yield "\n"
-
-            logger.info(
-                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
-            # from knative activator gateway timeouts which are caused by idleTimeout or responseStartSeconds
-            # being exceeded.
-            app.logger.error(
-                "Client disconnected in the middle of the Workflow Stream",
-                extra={
-                    "sentry_tags": {
-                        "server_version": vembda_initiated_event.body.server_version,
-                        "sdk_version": vembda_initiated_event.body.sdk_version,
-                    }
-                },
-            )
-            return
-        except Exception as e:
-            logger.exception("Error during workflow response stream generator", extra={"error": e})
-            yield "\n"
-            yield "END"
-            yield "\n"
-            return
-        finally:
-            if ENABLE_PROCESS_WRAPPER:
-                try:
-                    if process and process.is_alive():
-                        process.kill()
-                    if process:
-                        increment_process_count(-1)
-                        remove_active_span_id(span_id)
-                except Exception as e:
-                    logger.error("Failed to kill process", e)
-            else:
-                increment_process_count(-1)
-                remove_active_span_id(span_id)
+    workflow_events = process_events(output_queue)
 
-    resp = Response(
-        stream_with_context(generator()),
-        status=200,
-        content_type="application/x-ndjson",
-        headers=headers,
-    )
-    return resp
+    return workflow_events, vembda_initiated_event, process, span_id, headers
 
 
 @bp.route("/stream-node", methods=["POST"])
@@ -374,12 +467,13 @@ def stream_node_route() -> Response:
         )
 
     # Create this event up here so timestamps are fully from the start to account for any unknown overhead
+    version_data = get_version()
     vembda_initiated_event = VembdaExecutionInitiatedEvent(
         id=uuid4(),
         timestamp=datetime.now(),
         trace_id=context.trace_id,
         span_id=context.execution_id,
-        body=VembdaExecutionInitiatedBody.model_validate(get_version()),
+        body=VembdaExecutionInitiatedBody.model_validate(version_data),
         parent=None,
     )
 
@@ -433,6 +527,7 @@ def serialize_route() -> Response:
 
     files = data.get("files", {})
     workspace_api_key = data.get("workspace_api_key")
+    is_new_server = data.get("is_new_server", False)
 
     if not files:
         return Response(
@@ -446,6 +541,11 @@ def serialize_route() -> Response:
     # Generate a unique namespace for this serialization request
     namespace = get_random_namespace()
     virtual_finder = VirtualFileFinder(files, namespace)
+
+    headers = {
+        "X-Vellum-Is-New-Server": str(is_new_server).lower(),
+    }
+
     try:
         sys.meta_path.append(virtual_finder)
         result = BaseWorkflowDisplay.serialize_module(namespace, client=client, dry_run=True)
@@ -454,6 +554,7 @@ def serialize_route() -> Response:
             json.dumps(result.model_dump()),
             status=200,
             content_type="application/json",
+            headers=headers,
         )
 
     except WorkflowInitializationException as e:
@@ -463,6 +564,7 @@ def serialize_route() -> Response:
             json.dumps({"detail": error_message}),
             status=400,
             content_type="application/json",
+            headers=headers,
         )
 
     except Exception as e:
@@ -471,6 +573,7 @@ def serialize_route() -> Response:
             json.dumps({"detail": f"Serialization failed: {str(e)}"}),
             status=500,
             content_type="application/json",
+            headers=headers,
         )
 
     finally:
@@ -553,3 +656,12 @@ def startup_error_generator(
             },
         )
         return
+
+
+def _get_headers(context: WorkflowExecutorContext) -> dict[str, Union[str, Any]]:
+    headers = {
+        "X-Vellum-SDK-Version": get_version()["sdk_version"],
+        "X-Vellum-Server-Version": get_version()["server_version"],
+        "X-Vellum-Events-Emitted": str(is_events_emitting_enabled(context)),
+    }
+    return headers

workflow_server/config.py

@@ -42,6 +42,8 @@ LOCAL_WORKFLOW_MODULE = os.getenv("LOCAL_WORKFLOW_MODULE")
 # The deployment name to match against when using local mode so you can still run your normal workflow
 LOCAL_DEPLOYMENT = os.getenv("LOCAL_DEPLOYMENT")
 
+IS_ASYNC_MODE = os.getenv("IS_ASYNC_MODE", "false").lower() == "true"
+
 
 def is_development() -> bool:
     return os.getenv("FLASK_ENV", "local") == "local"

workflow_server/core/executor.py

@@ -1,5 +1,4 @@
-from datetime import datetime
-import importlib
+from datetime import datetime, timezone
 from io import StringIO
 import json
 import logging
@@ -18,10 +17,12 @@ from vellum_ee.workflows.display.utils.events import event_enricher
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
 from vellum.workflows import BaseWorkflow
+from vellum.workflows.context import execution_context
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.emitters.vellum_emitter import VellumEmitter
 from vellum.workflows.events.exception_handling import stream_initialization_exception
 from vellum.workflows.events.types import BaseEvent
+from vellum.workflows.events.workflow import WorkflowEvent
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import BaseNode
@@ -30,8 +31,9 @@ from vellum.workflows.resolvers.base import BaseWorkflowResolver
 from vellum.workflows.resolvers.resolver import VellumResolver
 from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.state.store import EmptyStore
+from vellum.workflows.triggers import BaseTrigger
 from vellum.workflows.types import CancelSignal
-from vellum.workflows.workflows.event_filters import all_workflow_event_filter
+from vellum.workflows.workflows.event_filters import workflow_sandbox_event_filter
 from workflow_server.config import LOCAL_DEPLOYMENT, LOCAL_WORKFLOW_MODULE
 from workflow_server.core.cancel_workflow import CancelWorkflowWatcherThread
 from workflow_server.core.events import (
@@ -51,6 +53,8 @@ from workflow_server.core.workflow_executor_context import (
     WorkflowExecutorContext,
 )
 from workflow_server.utils.log_proxy import redirect_log
+from workflow_server.utils.system_utils import get_memory_in_use_mb
+from workflow_server.utils.utils import get_version
 
 logger = logging.getLogger(__name__)
 
@@ -146,7 +150,21 @@ def stream_workflow(
     cancel_watcher_kill_switch = ThreadingEvent()
     try:
         workflow, namespace = _create_workflow(executor_context)
-        workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
+
+        trigger_id = executor_context.trigger_id
+
+        inputs_or_trigger = workflow.deserialize_trigger(trigger_id=trigger_id, inputs=executor_context.inputs)
+
+        # Determine whether we have inputs or a trigger
+        if isinstance(inputs_or_trigger, BaseInputs):
+            workflow_inputs = inputs_or_trigger
+            trigger = None
+        elif isinstance(inputs_or_trigger, BaseTrigger):
+            workflow_inputs = None
+            trigger = inputs_or_trigger
+        else:
+            workflow_inputs = None
+            trigger = None
 
         workflow_state = (
             workflow.deserialize_state(
@@ -167,14 +185,22 @@ def stream_workflow(
             inputs=workflow_inputs,
             state=workflow_state,
             node_output_mocks=node_output_mocks,
-            event_filter=all_workflow_event_filter,
+            event_filter=workflow_sandbox_event_filter,
             cancel_signal=cancel_signal,
             entrypoint_nodes=[executor_context.node_id] if executor_context.node_id else None,
             previous_execution_id=executor_context.previous_execution_id,
+            timeout=executor_context.timeout,
+            trigger=trigger,
+            execution_id=executor_context.workflow_span_id,
         )
     except WorkflowInitializationException as e:
         cancel_watcher_kill_switch.set()
-        initialization_exception_stream = stream_initialization_exception(e)
+
+        with execution_context(
+            parent_context=executor_context.execution_context.parent_context,
+            trace_id=executor_context.execution_context.trace_id,
+        ):
+            initialization_exception_stream = stream_initialization_exception(e)
 
         def _stream_generator() -> Generator[dict[str, Any], Any, None]:
             for event in initialization_exception_stream:
@@ -401,10 +427,55 @@ def get_random_namespace() -> str:
     return "workflow_tmp_" + "".join(random.choice(string.ascii_letters + string.digits) for i in range(14))
 
 
+def _enrich_event(event: WorkflowEvent, executor_context: Optional[BaseExecutorContext] = None) -> WorkflowEvent:
+    """
+    Enrich an event with metadata based on the event type.
+
+    For initiated events, include server and SDK versions.
+    For fulfilled events with WORKFLOW_DEPLOYMENT parent, include memory usage.
+    """
+    metadata: Optional[dict] = None
+
+    try:
+        is_deployment = event.parent and event.parent.type in ("WORKFLOW_DEPLOYMENT", "EXTERNAL")
+
+        if event.name == "workflow.execution.initiated" and is_deployment:
+            metadata = {
+                **get_version(),
+            }
+
+            memory_mb = get_memory_in_use_mb()
+            if memory_mb is not None:
+                metadata["memory_usage_mb"] = memory_mb
+
+            if executor_context is not None:
+                metadata["is_new_server"] = executor_context.is_new_server
+
+                if executor_context.vembda_service_initiated_timestamp is not None and event.timestamp is not None:
+                    event_ts = event.timestamp
+                    if event_ts.tzinfo is None:
+                        event_ts = event_ts.replace(tzinfo=timezone.utc)
+                    event_ts_ns = int(event_ts.timestamp() * 1_000_000_000)
+                    initiated_latency = event_ts_ns - executor_context.vembda_service_initiated_timestamp
+                    metadata["initiated_latency"] = initiated_latency
+        elif event.name == "workflow.execution.fulfilled" and is_deployment:
+            metadata = {}
+            memory_mb = get_memory_in_use_mb()
+            if memory_mb is not None:
+                metadata["memory_usage_mb"] = memory_mb
+    except Exception:
+        pass
+
+    vellum_client = executor_context.vellum_client if executor_context else None
+    return event_enricher(event, vellum_client, metadata=metadata)
+
+
 def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext) -> dict:
     module_base = executor_context.module.split(".")
+
     dump = event.model_dump(
-        mode="json", context={"event_enricher": lambda event: event_enricher(event, executor_context.vellum_client)}
+        mode="json",
+        context={"event_enricher": lambda event: _enrich_event(event, executor_context)},
     )
     if dump["name"] in {
         "workflow.execution.initiated",
@@ -426,38 +497,3 @@ def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext) -> dict
         dump["body"]["node_definition"]["module"] = module_base + dump["body"]["node_definition"]["module"][1:]
 
     return dump
-
-
-def _get_workflow_inputs(
-    executor_context: BaseExecutorContext, workflow_class: Type[BaseWorkflow]
-) -> Optional[BaseInputs]:
-    if not executor_context.inputs:
-        return None
-
-    if not executor_context.files.get("inputs.py"):
-        return None
-
-    namespace = _get_file_namespace(executor_context)
-    inputs_module_path = f"{namespace}.inputs"
-    try:
-        inputs_module = importlib.import_module(inputs_module_path)
-    except Exception as e:
-        raise WorkflowInitializationException(
-            message=f"Failed to initialize workflow inputs: {e}",
-            workflow_definition=workflow_class,
-        ) from e
-
-    if not hasattr(inputs_module, "Inputs"):
-        raise WorkflowInitializationException(
-            message=f"Inputs module {inputs_module_path} does not have a required Inputs class",
-            workflow_definition=workflow_class,
-        )
-
-    if not issubclass(inputs_module.Inputs, BaseInputs):
-        raise WorkflowInitializationException(
-            message=f"""The class {inputs_module_path}.Inputs was expected to be a subclass of BaseInputs, \
-but found {inputs_module.Inputs.__class__.__name__}""",
-            workflow_definition=workflow_class,
-        )
-
-    return inputs_module.Inputs(**executor_context.inputs)

workflow_server/core/workflow_executor_context.py

@@ -35,6 +35,12 @@ class BaseExecutorContext(UniversalBaseModel):
     environment_variables: Optional[dict[str, str]] = None
     previous_execution_id: Optional[UUID] = None
     feature_flags: Optional[dict[str, bool]] = None
+    is_new_server: bool = False
+    trigger_id: Optional[UUID] = None
+    # The actual 'execution id' of the workflow that we pass into the workflow
+    # when running in async mode.
+    workflow_span_id: Optional[UUID] = None
+    vembda_service_initiated_timestamp: Optional[int] = None
 
     @field_validator("inputs", mode="before")
     @classmethod

workflow_server/logging_config.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+import json
+import logging
+
+
+class GCPJsonFormatter(logging.Formatter):
+    """
+    Custom JSON formatter for Google Cloud Platform logging.
+
+    Outputs logs in JSON format with a 'severity' field that GCP Cloud Logging
+    can properly parse. This ensures INFO logs show up as INFO in GCP instead of ERROR.
+
+    See: https://cloud.google.com/logging/docs/structured-logging
+    """
+
+    SEVERITY_MAP = {
+        "DEBUG": "DEBUG",
+        "INFO": "INFO",
+        "WARNING": "WARNING",
+        "ERROR": "ERROR",
+        "CRITICAL": "CRITICAL",
+    }
+
+    def format(self, record: logging.LogRecord) -> str:
+        log_obj = {
+            "severity": self.SEVERITY_MAP.get(record.levelname, "DEFAULT"),
+            "message": record.getMessage(),
+            "timestamp": datetime.utcfromtimestamp(record.created).isoformat() + "Z",
+            "logging.googleapis.com/sourceLocation": {
+                "file": record.pathname,
+                "line": record.lineno,
+                "function": record.funcName,
+            },
+        }
+
+        if record.exc_info:
+            log_obj["exception"] = self.formatException(record.exc_info)
+
+        return json.dumps(log_obj)

workflow_server/server.py

@@ -5,8 +5,10 @@ from flask import Flask
 
 from workflow_server.api.auth_middleware import AuthMiddleware
 from workflow_server.api.healthz_view import bp as healthz_bp
+from workflow_server.api.status_view import bp as status_bp
 from workflow_server.api.workflow_view import bp as workflow_bp
 from workflow_server.config import is_development
+from workflow_server.logging_config import GCPJsonFormatter
 from workflow_server.utils.sentry import init_sentry
 from workflow_server.utils.utils import get_version
 
@@ -14,11 +16,18 @@ from workflow_server.utils.utils import get_version
 # enable_log_proxy()
 
 logger = logging.getLogger(__name__)
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-)
+
+if is_development():
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+else:
+    handler = logging.StreamHandler()
+    handler.setFormatter(GCPJsonFormatter())
+    logging.root.addHandler(handler)
+    logging.root.setLevel(logging.INFO)
 
 
 def create_app() -> Flask:
@@ -40,6 +49,7 @@ def create_app() -> Flask:
 
     # Register blueprints
     app.register_blueprint(healthz_bp, url_prefix="/healthz")
+    app.register_blueprint(status_bp, url_prefix="/status")
     app.register_blueprint(workflow_bp, url_prefix="/workflow")
 
     logger.info(is_development())

workflow_server/start.py

@@ -33,6 +33,7 @@ class CustomGunicornLogger(glogging.Logger):
         logger = logging.getLogger("gunicorn.access")
         logger.addFilter(HealthCheckFilter())
         logger.addFilter(SignalFilter())
+        logger.addFilter(StatusIsAvailableFilter())
 
 
 class HealthCheckFilter(logging.Filter):
@@ -45,6 +46,11 @@ class SignalFilter(logging.Filter):
         return "SIGTERM" not in record.getMessage()
 
 
+class StatusIsAvailableFilter(logging.Filter):
+    def filter(self, record: Any) -> bool:
+        return "/status/is_available" not in record.getMessage()
+
+
 def start() -> None:
     if not is_development():
         start_oom_killer_worker()