vellum-workflow-server 1.5.1__tar.gz → 1.9.7__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of vellum-workflow-server might be problematic. Click here for more details.

Files changed (37) hide show
  1. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/PKG-INFO +2 -2
  2. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/pyproject.toml +2 -2
  3. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/auth_middleware.py +2 -2
  4. vellum_workflow_server-1.9.7/src/workflow_server/api/status_view.py +19 -0
  5. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/tests/test_workflow_view.py +120 -27
  6. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py +215 -5
  7. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/workflow_view.py +241 -155
  8. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/code_exec_runner.py +2 -2
  9. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/config.py +8 -0
  10. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/core/cancel_workflow.py +7 -5
  11. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/core/executor.py +106 -182
  12. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/core/workflow_executor_context.py +35 -45
  13. vellum_workflow_server-1.9.7/src/workflow_server/logging_config.py +39 -0
  14. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/server.py +22 -5
  15. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/start.py +17 -3
  16. vellum_workflow_server-1.9.7/src/workflow_server/utils/exit_handler.py +56 -0
  17. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/oom_killer.py +6 -2
  18. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/sentry.py +21 -0
  19. vellum_workflow_server-1.9.7/src/workflow_server/utils/tests/test_sentry_integration.py +143 -0
  20. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/utils.py +2 -1
  21. vellum_workflow_server-1.5.1/src/workflow_server/utils/exit_handler.py +0 -27
  22. vellum_workflow_server-1.5.1/src/workflow_server/utils/tests/test_sentry_integration.py +0 -69
  23. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/README.md +0 -0
  24. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/__init__.py +0 -0
  25. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/__init__.py +0 -0
  26. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/healthz_view.py +0 -0
  27. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/tests/__init__.py +0 -0
  28. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/api/tests/test_input_display_mapping.py +0 -0
  29. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/core/__init__.py +0 -0
  30. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/core/events.py +0 -0
  31. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/core/utils.py +0 -0
  32. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/__init__.py +0 -0
  33. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/log_proxy.py +0 -0
  34. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/system_utils.py +0 -0
  35. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/tests/__init__.py +0 -0
  36. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/tests/test_system_utils.py +0 -0
  37. {vellum_workflow_server-1.5.1 → vellum_workflow_server-1.9.7}/src/workflow_server/utils/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-workflow-server
3
- Version: 1.5.1
3
+ Version: 1.9.7
4
4
  Summary:
5
5
  License: AGPL
6
6
  Requires-Python: >=3.9.0,<4
@@ -29,7 +29,7 @@ Requires-Dist: pyjwt (==2.10.0)
29
29
  Requires-Dist: python-dotenv (==1.0.1)
30
30
  Requires-Dist: retrying (==1.3.4)
31
31
  Requires-Dist: sentry-sdk[flask] (==2.20.0)
32
- Requires-Dist: vellum-ai (==1.5.1)
32
+ Requires-Dist: vellum-ai (==1.9.7)
33
33
  Description-Content-Type: text/markdown
34
34
 
35
35
  # Vellum Workflow Runner Server
@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
3
3
 
4
4
  [tool.poetry]
5
5
  name = "vellum-workflow-server"
6
- version = "1.5.1"
6
+ version = "1.9.7"
7
7
  description = ""
8
8
  readme = "README.md"
9
9
  authors = []
@@ -45,7 +45,7 @@ flask = "2.3.3"
45
45
  orderly-set = "5.2.2"
46
46
  pebble = "5.0.7"
47
47
  gunicorn = "23.0.0"
48
- vellum-ai = "1.5.1"
48
+ vellum-ai = "1.9.7"
49
49
  python-dotenv = "1.0.1"
50
50
  retrying = "1.3.4"
51
51
  sentry-sdk = {extras = ["flask"], version = "2.20.0"}
@@ -5,7 +5,7 @@ from flask import Flask, Request, Response
5
5
  import jwt
6
6
  from jwt import ExpiredSignatureError
7
7
 
8
- from workflow_server.config import IS_VPC, NAMESPACE, VEMBDA_PUBLIC_KEY, is_development
8
+ from workflow_server.config import IS_ASYNC_MODE, IS_VPC, NAMESPACE, VEMBDA_PUBLIC_KEY, is_development
9
9
 
10
10
 
11
11
  class AuthMiddleware:
@@ -15,7 +15,7 @@ class AuthMiddleware:
15
15
  def __call__(self, environ: Dict[str, Any], start_response: Any) -> Any:
16
16
  try:
17
17
  request = Request(environ)
18
- if not request.path.startswith("/healthz") and not is_development() and not IS_VPC:
18
+ if not request.path.startswith("/healthz") and not is_development() and not IS_VPC and not IS_ASYNC_MODE:
19
19
  token = request.headers.get("X-Vembda-Signature")
20
20
  if token:
21
21
  decoded = jwt.decode(token, VEMBDA_PUBLIC_KEY, algorithms=["RS256"])
@@ -0,0 +1,19 @@
1
+ from typing import Tuple
2
+
3
+ from flask import Blueprint, Response, jsonify
4
+
5
+ from workflow_server.config import CONCURRENCY
6
+ from workflow_server.utils.system_utils import get_active_process_count
7
+
8
+ bp = Blueprint("status", __name__)
9
+
10
+
11
+ @bp.route("/is_available", methods=["GET"])
12
+ def is_available() -> Tuple[Response, int]:
13
+ resp = jsonify(
14
+ available=get_active_process_count() < CONCURRENCY,
15
+ process_count=get_active_process_count(),
16
+ max_concurrency=CONCURRENCY,
17
+ )
18
+
19
+ return resp, 200
@@ -63,11 +63,11 @@ class TestNode(BaseNode):
63
63
  "comment": {"expanded": True, "value": "A test node for processing data."},
64
64
  "position": {"x": 0.0, "y": 0.0},
65
65
  },
66
- "id": "7a8b251d-f5ca-462a-b293-071d219460fb",
66
+ "id": "6f4c9178-9f46-4723-bcb7-0bd59db54eca",
67
67
  "label": "Test Node",
68
68
  "outputs": [],
69
- "ports": [{"id": "a3a0eefd-45d0-4f13-8c58-a836a9f7f9ed", "name": "default", "type": "DEFAULT"}],
70
- "trigger": {"id": "a022e36c-9852-4772-9be3-3c6c147fd811", "merge_behavior": "AWAIT_ATTRIBUTES"},
69
+ "ports": [{"id": "4394823f-79a8-4dbc-99ae-06a1df6c7408", "name": "default", "type": "DEFAULT"}],
70
+ "trigger": {"id": "07240af1-67c6-4460-b53d-53f0b0f1b90e", "merge_behavior": "AWAIT_ATTRIBUTES"},
71
71
  "type": "GENERIC",
72
72
  }
73
73
 
@@ -127,11 +127,11 @@ class SomeOtherNode(BaseNode):
127
127
  "comment": {"expanded": True, "value": "This is Some Node."},
128
128
  "position": {"x": 0.0, "y": 0.0},
129
129
  },
130
- "id": "1e559c2e-db82-41f0-9ceb-5e89b0c5a0a3",
130
+ "id": "89e84bac-5a5f-4f64-8083-7d3ebec98be1",
131
131
  "label": "Some Node",
132
132
  "outputs": [],
133
- "ports": [{"id": "48e39e97-5fd4-471e-b4f2-51d3baf06456", "name": "default", "type": "DEFAULT"}],
134
- "trigger": {"id": "e3381fb7-61fc-4c46-ae8e-51fc463b6a59", "merge_behavior": "AWAIT_ATTRIBUTES"},
133
+ "ports": [{"id": "2983ea5c-1d29-483a-b896-53098f5de4f1", "name": "default", "type": "DEFAULT"}],
134
+ "trigger": {"id": "6996efb0-5a20-4719-8835-34fe6552764a", "merge_behavior": "AWAIT_ATTRIBUTES"},
135
135
  "type": "GENERIC",
136
136
  }
137
137
 
@@ -150,11 +150,11 @@ class SomeOtherNode(BaseNode):
150
150
  "comment": {"expanded": True, "value": "This is Some Other Node."},
151
151
  "position": {"x": 0.0, "y": 0.0},
152
152
  },
153
- "id": "7aee541b-b245-4c8a-9137-3e4631d5100c",
153
+ "id": "3cdbba02-8a34-4e0f-8b94-770a944dcaa3",
154
154
  "label": "Some Other Node",
155
155
  "outputs": [],
156
- "ports": [{"id": "fb66b46a-d970-4bc9-83ea-70c154c57ddd", "name": "default", "type": "DEFAULT"}],
157
- "trigger": {"id": "13fa2714-20b3-4bc3-ab79-621a188e3bfa", "merge_behavior": "AWAIT_ATTRIBUTES"},
156
+ "ports": [{"id": "1839bde5-2ad4-4723-b21b-2c55fa833a7a", "name": "default", "type": "DEFAULT"}],
157
+ "trigger": {"id": "c36df8a8-5624-45be-99c9-826cf511a951", "merge_behavior": "AWAIT_ATTRIBUTES"},
158
158
  "type": "GENERIC",
159
159
  }
160
160
 
@@ -222,11 +222,11 @@ class HelperClass:
222
222
  "comment": {"expanded": True, "value": "Processes input data."},
223
223
  "position": {"x": 0.0, "y": 0.0},
224
224
  },
225
- "id": "f92c09f0-0434-46cb-829d-a73f801d6343",
225
+ "id": "7121bcb9-98a1-4907-bf9b-9734d773fd15",
226
226
  "label": "Processing Node",
227
227
  "outputs": [],
228
- "ports": [{"id": "abaa2984-b312-4491-b069-e689759f72c8", "name": "default", "type": "DEFAULT"}],
229
- "trigger": {"id": "35378c2b-f089-44af-ac37-efe4ea42c817", "merge_behavior": "AWAIT_ATTRIBUTES"},
228
+ "ports": [{"id": "de27da74-30e9-4e7b-95c2-92bdfc5bf042", "name": "default", "type": "DEFAULT"}],
229
+ "trigger": {"id": "e02bd85e-8b03-4b21-8b3e-f411042334ce", "merge_behavior": "AWAIT_ATTRIBUTES"},
230
230
  "type": "GENERIC",
231
231
  }
232
232
 
@@ -240,11 +240,11 @@ class HelperClass:
240
240
  "comment": {"expanded": True, "value": "Transforms data format."},
241
241
  "position": {"x": 0.0, "y": 0.0},
242
242
  },
243
- "id": "09ca32f7-c8f2-4469-97e5-1f288f85127a",
243
+ "id": "6a785cb0-f631-4f03-94c6-e82331c14c1a",
244
244
  "label": "Transformation Node",
245
245
  "outputs": [],
246
- "ports": [{"id": "88778117-fbfc-4b44-964b-5a4994aa2f24", "name": "default", "type": "DEFAULT"}],
247
- "trigger": {"id": "5d096263-7fbf-490a-83b7-e441852b5fb6", "merge_behavior": "AWAIT_ATTRIBUTES"},
246
+ "ports": [{"id": "67a13ea0-fd6b-44dc-af46-c72da06aa11f", "name": "default", "type": "DEFAULT"}],
247
+ "trigger": {"id": "08d4e317-baa8-478f-b278-99362e50e6b4", "merge_behavior": "AWAIT_ATTRIBUTES"},
248
248
  "type": "GENERIC",
249
249
  }
250
250
 
@@ -306,11 +306,11 @@ class BrokenNode(BaseNode)
306
306
  "comment": {"expanded": True, "value": "This is Some Node."},
307
307
  "position": {"x": 0.0, "y": 0.0},
308
308
  },
309
- "id": "1e559c2e-db82-41f0-9ceb-5e89b0c5a0a3",
309
+ "id": "a2706730-074b-4ea3-968a-25e68af1caed",
310
310
  "label": "Some Node",
311
311
  "outputs": [],
312
- "ports": [{"id": "48e39e97-5fd4-471e-b4f2-51d3baf06456", "name": "default", "type": "DEFAULT"}],
313
- "trigger": {"id": "e3381fb7-61fc-4c46-ae8e-51fc463b6a59", "merge_behavior": "AWAIT_ATTRIBUTES"},
312
+ "ports": [{"id": "e0ee3653-e071-4b91-9dfc-5e1dca9c665b", "name": "default", "type": "DEFAULT"}],
313
+ "trigger": {"id": "8d931b01-30ca-4c0d-b1b7-7c18379c83e6", "merge_behavior": "AWAIT_ATTRIBUTES"},
314
314
  "type": "GENERIC",
315
315
  }
316
316
 
@@ -371,14 +371,14 @@ class MyAdditionNode(BaseNode):
371
371
  "adornments": None,
372
372
  "attributes": [
373
373
  {
374
- "id": "aed3bcbb-d243-4a77-bb5e-409e9a28e868",
374
+ "id": "4223b340-447f-46c2-b35d-30ef16c5ae17",
375
375
  "name": "arg1",
376
- "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
376
+ "value": None,
377
377
  },
378
378
  {
379
- "id": "9225d225-a41b-4642-8964-f28f58dcf4bf",
379
+ "id": "1de0f46a-95f6-4cd0-bb0f-e2414054d507",
380
380
  "name": "arg2",
381
- "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
381
+ "value": None,
382
382
  },
383
383
  ],
384
384
  "base": {"module": ["vellum", "workflows", "nodes", "bases", "base"], "name": "BaseNode"},
@@ -387,11 +387,11 @@ class MyAdditionNode(BaseNode):
387
387
  "comment": {"expanded": True, "value": "Custom node that performs simple addition."},
388
388
  "position": {"x": 0.0, "y": 0.0},
389
389
  },
390
- "id": "195cd69d-3d2d-41e4-a432-16c433cb8d34",
390
+ "id": "2464b610-fb6d-495b-b17c-933ee147f19f",
391
391
  "label": "My Addition Node",
392
- "outputs": [{"id": "3d8e40cb-2aa8-44bd-ae6a-708a9fbc4779", "name": "result", "type": "NUMBER", "value": None}],
393
- "ports": [{"id": "9a9e4ef6-febf-4093-a515-217bbb1373db", "name": "default", "type": "DEFAULT"}],
394
- "trigger": {"id": "a5298668-d808-4a45-a62e-790943948e8a", "merge_behavior": "AWAIT_ATTRIBUTES"},
392
+ "outputs": [{"id": "f39d85c9-e7bf-45e1-bb67-f16225db0118", "name": "result", "type": "NUMBER", "value": None}],
393
+ "ports": [{"id": "bc489295-cd8a-4aa2-88bb-34446374100d", "name": "default", "type": "DEFAULT"}],
394
+ "trigger": {"id": "ff580cad-73d6-44fe-8f2c-4b8dc990ee70", "merge_behavior": "AWAIT_ATTRIBUTES"},
395
395
  "type": "GENERIC",
396
396
  "should_file_merge": True,
397
397
  }
@@ -429,7 +429,7 @@ class BrokenNode(BaseNode) # Missing colon
429
429
  with flask_app.test_client() as test_client:
430
430
  response = test_client.post("/workflow/serialize", json={"files": {"broken_node.py": invalid_content}})
431
431
 
432
- # THEN we should get a 4xx response
432
+ # THEN we should get a 400 response
433
433
  assert response.status_code == 400
434
434
 
435
435
  # AND the response should contain an error message
@@ -537,6 +537,57 @@ def test_serialize_route__with_invalid_workspace_api_key():
537
537
  assert "exec_config" in response.json
538
538
 
539
539
 
540
+ def test_serialize_route__with_is_new_server_header():
541
+ """
542
+ Tests that the serialize route returns the is_new_server header.
543
+ """
544
+ # GIVEN a Flask application
545
+ flask_app = create_app()
546
+
547
+ workflow_files = {
548
+ "__init__.py": "",
549
+ "workflow.py": (
550
+ "from vellum.workflows import BaseWorkflow\n\n"
551
+ "class Workflow(BaseWorkflow):\n"
552
+ " class Outputs(BaseWorkflow.Outputs):\n"
553
+ " foo = 'hello'\n"
554
+ ),
555
+ }
556
+
557
+ # WHEN we make a request with is_new_server=True
558
+ with flask_app.test_client() as test_client:
559
+ response = test_client.post("/workflow/serialize", json={"files": workflow_files, "is_new_server": True})
560
+
561
+ # THEN we should get a successful response
562
+ assert response.status_code == 200
563
+
564
+ # AND the response should contain the is_new_server header set to true
565
+ assert "X-Vellum-Is-New-Server" in response.headers
566
+ assert response.headers["X-Vellum-Is-New-Server"] == "true"
567
+
568
+ # WHEN we make a request with is_new_server=False
569
+ with flask_app.test_client() as test_client:
570
+ response = test_client.post("/workflow/serialize", json={"files": workflow_files, "is_new_server": False})
571
+
572
+ # THEN we should get a successful response
573
+ assert response.status_code == 200
574
+
575
+ # AND the response should contain the is_new_server header set to false
576
+ assert "X-Vellum-Is-New-Server" in response.headers
577
+ assert response.headers["X-Vellum-Is-New-Server"] == "false"
578
+
579
+ # WHEN we make a request without is_new_server
580
+ with flask_app.test_client() as test_client:
581
+ response = test_client.post("/workflow/serialize", json={"files": workflow_files})
582
+
583
+ # THEN we should get a successful response
584
+ assert response.status_code == 200
585
+
586
+ # AND the response should contain the is_new_server header set to false (default)
587
+ assert "X-Vellum-Is-New-Server" in response.headers
588
+ assert response.headers["X-Vellum-Is-New-Server"] == "false"
589
+
590
+
540
591
  def test_stream_node_route__with_node_id():
541
592
  """
542
593
  Tests that the stream-node endpoint works with node_id.
@@ -729,3 +780,45 @@ class Workflow(BaseWorkflow):
729
780
  assert events[0]["name"] == "vembda.execution.initiated"
730
781
  assert events[1]["name"] == "vembda.execution.fulfilled"
731
782
  assert len(events) == 2
783
+
784
+
785
+ def test_serialize_route__with_invalid_nested_set_graph():
786
+ """
787
+ Tests that a workflow with an invalid nested set graph structure raises a clear user-facing exception.
788
+ """
789
+ # GIVEN a Flask application
790
+ flask_app = create_app()
791
+
792
+ invalid_workflow_content = """
793
+ from vellum.workflows import BaseWorkflow
794
+ from vellum.workflows.nodes import BaseNode
795
+
796
+ class TestNode(BaseNode):
797
+ class Outputs(BaseNode.Outputs):
798
+ value = "test"
799
+
800
+ class InvalidWorkflow(BaseWorkflow):
801
+ graph = {TestNode, {TestNode}}
802
+
803
+ class Outputs(BaseWorkflow.Outputs):
804
+ result = TestNode.Outputs.value
805
+ """
806
+
807
+ workflow_files = {
808
+ "__init__.py": "",
809
+ "workflow.py": invalid_workflow_content,
810
+ }
811
+
812
+ # WHEN we make a request to the serialize route
813
+ with flask_app.test_client() as test_client:
814
+ response = test_client.post("/workflow/serialize", json={"files": workflow_files})
815
+
816
+ # THEN we should get a 400 response
817
+ assert response.status_code == 400
818
+
819
+ # AND the response should contain a user-friendly error message
820
+ assert "detail" in response.json
821
+ error_detail = response.json["detail"]
822
+ assert "Serialization failed" in error_detail
823
+ assert "Invalid graph structure detected" in error_detail
824
+ assert "contact Vellum support" in error_detail
@@ -5,6 +5,7 @@ import io
5
5
  import json
6
6
  from queue import Empty
7
7
  import re
8
+ import time
8
9
  from unittest import mock
9
10
  from uuid import uuid4
10
11
 
@@ -117,7 +118,7 @@ def test_stream_workflow_route__happy_path(both_stream_types):
117
118
  "execution_id": str(span_id),
118
119
  "inputs": [],
119
120
  "environment_api_key": "test",
120
- "module": "workflow",
121
+ "module": "test",
121
122
  "timeout": 360,
122
123
  "files": {
123
124
  "__init__.py": "",
@@ -131,8 +132,11 @@ class Workflow(BaseWorkflow):
131
132
  },
132
133
  }
133
134
 
134
- # WHEN we call the stream route
135
- status_code, events = both_stream_types(request_body)
135
+ with mock.patch("builtins.open", mock.mock_open(read_data="104857600")):
136
+ # WHEN we call the stream route
137
+ ts_ns = time.time_ns()
138
+ request_body["vembda_service_initiated_timestamp"] = ts_ns
139
+ status_code, events = both_stream_types(request_body)
136
140
 
137
141
  # THEN we get a 200 response
138
142
  assert status_code == 200, events
@@ -154,8 +158,46 @@ class Workflow(BaseWorkflow):
154
158
  }
155
159
 
156
160
  assert events[1]["name"] == "workflow.execution.initiated", events[1]
161
+ assert events[1]["body"]["workflow_definition"]["module"] == ["test", "workflow"]
157
162
  assert "display_context" in events[1]["body"], events[1]["body"]
163
+ display_context = events[1]["body"]["display_context"]
164
+ assert "node_displays" in display_context
165
+ assert "workflow_inputs" in display_context
166
+ assert "workflow_outputs" in display_context
167
+ assert isinstance(display_context["node_displays"], dict)
168
+ assert isinstance(display_context["workflow_inputs"], dict)
169
+ assert isinstance(display_context["workflow_outputs"], dict)
170
+ assert "foo" in display_context["workflow_outputs"]
171
+
172
+ # AND the initiated event should have server_metadata with version info and memory usage
173
+ assert "server_metadata" in events[1]["body"], events[1]["body"]
174
+ server_metadata = events[1]["body"]["server_metadata"]
175
+ assert server_metadata is not None, "server_metadata should not be None"
176
+ assert "server_version" in server_metadata
177
+ assert "sdk_version" in server_metadata
178
+ assert "memory_usage_mb" in server_metadata
179
+ assert isinstance(server_metadata["memory_usage_mb"], (int, float))
180
+ assert "is_new_server" in server_metadata
181
+ assert server_metadata["is_new_server"] is False
182
+
183
+ # AND the initiated event should have initiated_latency within a reasonable range
184
+ assert "initiated_latency" in server_metadata, "initiated_latency should be present in server_metadata"
185
+ initiated_latency = server_metadata["initiated_latency"]
186
+ assert isinstance(initiated_latency, int), "initiated_latency should be an integer (nanoseconds)"
187
+ # Latency should be positive and less than 60 seconds (60_000_000_000 nanoseconds) for CI
188
+ assert (
189
+ 0 < initiated_latency < 60_000_000_000
190
+ ), f"initiated_latency should be between 0 and 60 seconds, got {initiated_latency} ns"
191
+
158
192
  assert events[2]["name"] == "workflow.execution.fulfilled", events[2]
193
+ assert events[2]["body"]["workflow_definition"]["module"] == ["test", "workflow"]
194
+
195
+ # AND the fulfilled event should have server_metadata with memory usage
196
+ assert "server_metadata" in events[2]["body"], events[2]["body"]
197
+ fulfilled_metadata = events[2]["body"]["server_metadata"]
198
+ assert fulfilled_metadata is not None, "fulfilled server_metadata should not be None"
199
+ assert "memory_usage_mb" in fulfilled_metadata
200
+ assert isinstance(fulfilled_metadata["memory_usage_mb"], (int, float))
159
201
 
160
202
  assert events[3] == {
161
203
  "id": mock.ANY,
@@ -234,6 +276,15 @@ class Inputs(BaseInputs):
234
276
 
235
277
  assert events[1]["name"] == "workflow.execution.initiated", events[1]
236
278
  assert "display_context" in events[1]["body"], events[1]["body"]
279
+ display_context = events[1]["body"]["display_context"]
280
+ assert "node_displays" in display_context
281
+ assert "workflow_inputs" in display_context
282
+ assert "workflow_outputs" in display_context
283
+ assert isinstance(display_context["node_displays"], dict)
284
+ assert isinstance(display_context["workflow_inputs"], dict)
285
+ assert isinstance(display_context["workflow_outputs"], dict)
286
+ assert "foo" in display_context["workflow_inputs"]
287
+ assert "foo" in display_context["workflow_outputs"]
237
288
  assert events[2]["name"] == "workflow.execution.fulfilled", events[2]
238
289
 
239
290
  assert events[3] == {
@@ -311,6 +362,14 @@ class State(BaseState):
311
362
 
312
363
  assert events[1]["name"] == "workflow.execution.initiated", events[1]
313
364
  assert "display_context" in events[1]["body"], events[1]["body"]
365
+ display_context = events[1]["body"]["display_context"]
366
+ assert "node_displays" in display_context
367
+ assert "workflow_inputs" in display_context
368
+ assert "workflow_outputs" in display_context
369
+ assert isinstance(display_context["node_displays"], dict)
370
+ assert isinstance(display_context["workflow_inputs"], dict)
371
+ assert isinstance(display_context["workflow_outputs"], dict)
372
+ assert "foo" in display_context["workflow_outputs"]
314
373
  assert events[2]["name"] == "workflow.execution.fulfilled", events[2]
315
374
  assert events[2]["body"]["outputs"] == {"foo": "bar"}
316
375
 
@@ -339,9 +398,15 @@ class State(BaseState):
339
398
  def test_stream_workflow_route__bad_indent_in_inputs_file(both_stream_types):
340
399
  # GIVEN a valid request body
341
400
  span_id = uuid4()
401
+ trace_id = uuid4()
402
+ parent_span_id = uuid4()
342
403
  request_body = {
343
404
  "timeout": 360,
344
405
  "execution_id": str(span_id),
406
+ "execution_context": {
407
+ "trace_id": str(trace_id),
408
+ "parent_context": {"span_id": str(parent_span_id)},
409
+ },
345
410
  "inputs": [
346
411
  {"name": "foo", "type": "STRING", "value": "hello"},
347
412
  ],
@@ -378,7 +443,7 @@ from vellum.workflows.inputs import BaseInputs
378
443
 
379
444
  assert events[0] == {
380
445
  "id": mock.ANY,
381
- "trace_id": mock.ANY,
446
+ "trace_id": str(trace_id),
382
447
  "span_id": str(span_id),
383
448
  "timestamp": mock.ANY,
384
449
  "api_version": "2024-10-25",
@@ -392,11 +457,22 @@ from vellum.workflows.inputs import BaseInputs
392
457
  }
393
458
 
394
459
  assert events[1]["name"] == "workflow.execution.initiated"
460
+ assert events[1]["trace_id"] == str(trace_id), "workflow initiated event should use request trace_id"
461
+ assert events[1]["parent"] is not None, "workflow initiated event should have parent context"
462
+ assert events[1]["parent"]["span_id"] == str(
463
+ parent_span_id
464
+ ), "workflow initiated event parent should match request parent_context"
395
465
 
396
466
  assert events[2]["name"] == "workflow.execution.rejected"
467
+ assert events[2]["trace_id"] == str(trace_id), "workflow rejected event should use request trace_id"
397
468
  assert events[2]["span_id"] == events[1]["span_id"]
469
+ assert events[2]["parent"] is not None, "workflow rejected event should have parent context"
470
+ assert events[2]["parent"]["span_id"] == str(
471
+ parent_span_id
472
+ ), "workflow rejected event parent should match request parent_context"
398
473
  assert (
399
- "Failed to initialize workflow: unexpected indent (inputs.py, line 3)" in events[2]["body"]["error"]["message"]
474
+ "Syntax Error raised while loading Workflow: "
475
+ "unexpected indent (inputs.py, line 3)" in events[2]["body"]["error"]["message"]
400
476
  )
401
477
 
402
478
  assert events[3] == {
@@ -570,6 +646,71 @@ class BasicCancellableWorkflow(BaseWorkflow):
570
646
  )
571
647
 
572
648
 
649
+ def test_stream_workflow_route__timeout_emits_rejection_events():
650
+ """
651
+ Tests that when a workflow times out, we emit node and workflow rejection events.
652
+ """
653
+
654
+ span_id = uuid4()
655
+ request_body = {
656
+ "timeout": 1,
657
+ "execution_id": str(span_id),
658
+ "inputs": [],
659
+ "environment_api_key": "test",
660
+ "module": "workflow",
661
+ "files": {
662
+ "__init__.py": "",
663
+ "workflow.py": """\
664
+ import time
665
+
666
+ from vellum.workflows.nodes.bases.base import BaseNode
667
+ from vellum.workflows.workflows.base import BaseWorkflow
668
+
669
+
670
+ class LongRunningNode(BaseNode):
671
+ class Outputs(BaseNode.Outputs):
672
+ value: str
673
+
674
+ def run(self) -> Outputs:
675
+ time.sleep(30)
676
+ return self.Outputs(value="hello world")
677
+
678
+
679
+ class TimeoutWorkflow(BaseWorkflow):
680
+ graph = LongRunningNode
681
+ class Outputs(BaseWorkflow.Outputs):
682
+ final_value = LongRunningNode.Outputs.value
683
+
684
+ """,
685
+ },
686
+ }
687
+
688
+ status_code, events = flask_stream(request_body)
689
+
690
+ assert status_code == 200
691
+
692
+ event_names = [e["name"] for e in events]
693
+
694
+ assert "vembda.execution.initiated" in event_names
695
+ assert "workflow.execution.initiated" in event_names
696
+ assert "node.execution.initiated" in event_names
697
+
698
+ assert "node.execution.rejected" in event_names, "Should emit node.execution.rejected on timeout"
699
+ node_execution_rejected = next(e for e in events if e["name"] == "node.execution.rejected")
700
+ assert "vellum/workflows/runner/runner.py" in node_execution_rejected["body"]["stacktrace"]
701
+
702
+ assert "workflow.execution.rejected" in event_names, "Should emit workflow.execution.rejected on timeout"
703
+ workflow_execution_rejected = next(e for e in events if e["name"] == "workflow.execution.rejected")
704
+ assert workflow_execution_rejected["body"]["error"]["code"] == "WORKFLOW_TIMEOUT"
705
+ # TODO: Uncomment once version 1.8.1 is released
706
+ # assert "stacktrace" in workflow_execution_rejected["body"]
707
+ # assert "vellum/workflows/runner/runner.py" in workflow_execution_rejected["body"]["stacktrace"]
708
+
709
+ assert "vembda.execution.fulfilled" in event_names
710
+ vembda_fulfilled = next(e for e in events if e["name"] == "vembda.execution.fulfilled")
711
+ assert vembda_fulfilled["body"]["timed_out"] is True
712
+
713
+
573
714
  def test_stream_workflow_route__very_large_events(both_stream_types):
574
715
  # GIVEN a valid request body
575
716
  span_id = uuid4()
@@ -1031,3 +1172,72 @@ class Workflow(BaseWorkflow):
1031
1172
  assert len(event_names) == 2, "Should include 2 events"
1032
1173
  assert "workflow.execution.initiated" in event_names, "Should include workflow.execution.initiated event"
1033
1174
  assert "workflow.execution.fulfilled" in event_names, "Should include workflow.execution.fulfilled event"
1175
+
1176
+
1177
+ def test_stream_workflow_route__with_invalid_nested_set_graph(both_stream_types):
1178
+ """
1179
+ Tests that a workflow with an invalid nested set graph structure raises a clear error in the stream response.
1180
+ """
1181
+ # GIVEN a Flask application and invalid workflow content with nested set graph
1182
+ span_id = uuid4()
1183
+
1184
+ invalid_workflow_content = """
1185
+ from vellum.workflows import BaseWorkflow
1186
+ from vellum.workflows.nodes import BaseNode
1187
+
1188
+ class TestNode(BaseNode):
1189
+ class Outputs(BaseNode.Outputs):
1190
+ value = "test"
1191
+
1192
+ class InvalidWorkflow(BaseWorkflow):
1193
+ graph = {TestNode, {TestNode}}
1194
+
1195
+ class Outputs(BaseWorkflow.Outputs):
1196
+ result = TestNode.Outputs.value
1197
+ """
1198
+
1199
+ request_body = {
1200
+ "timeout": 360,
1201
+ "execution_id": str(span_id),
1202
+ "inputs": [],
1203
+ "environment_api_key": "test",
1204
+ "module": "workflow",
1205
+ "files": {
1206
+ "__init__.py": "",
1207
+ "workflow.py": invalid_workflow_content,
1208
+ },
1209
+ }
1210
+
1211
+ # WHEN we call the stream route
1212
+ status_code, events = both_stream_types(request_body)
1213
+
1214
+ # THEN we get a 200 response
1215
+ assert status_code == 200, events
1216
+
1217
+ # THEN we get the expected events: vembda initiated, workflow initiated, workflow rejected, vembda fulfilled
1218
+ assert len(events) == 4
1219
+
1220
+ # AND the first event should be vembda execution initiated
1221
+ assert events[0]["name"] == "vembda.execution.initiated"
1222
+ assert events[0]["span_id"] == str(span_id)
1223
+
1224
+ # AND the second event should be workflow execution initiated
1225
+ assert events[1]["name"] == "workflow.execution.initiated"
1226
+
1227
+ # AND the third event should be workflow execution rejected
1228
+ assert events[2]["name"] == "workflow.execution.rejected"
1229
+ assert events[1]["span_id"] == events[2]["span_id"]
1230
+
1231
+ # AND the error message should contain information about the invalid graph structure
1232
+ error_message = events[2]["body"]["error"]["message"]
1233
+ expected_message = (
1234
+ "Invalid graph structure detected. "
1235
+ "Nested sets or unsupported graph types are not allowed. "
1236
+ "Please contact Vellum support for assistance with Workflow configuration."
1237
+ )
1238
+ assert error_message == expected_message
1239
+
1240
+ # AND the fourth event should be vembda execution fulfilled
1241
+ assert events[3]["name"] == "vembda.execution.fulfilled"
1242
+ assert events[3]["span_id"] == str(span_id)
1243
+ assert events[3]["body"]["exit_code"] == 0