From a6d4390790d21d535df1e9d525bffd577c50d8dc Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 9 Mar 2026 19:39:24 +0000 Subject: [PATCH] feat: workflow-as-code (WAC) v2 (#8172) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: workflow-as-code v2 with @task decorator API Replace ctx.step("name", "script") API with @task decorators where functions are called directly. Users no longer need to pass WorkflowCtx or use string-based step names/script paths. Python: @task decorator with contextvars-based implicit context TypeScript: task() wrapper with module-level context variable Parsers: detect @task function calls instead of ctx.step() calls Worker: updated wrappers to set implicit context Co-Authored-By: Claude Opus 4.6 * feat: WAC v2 checkpoint/replay with _executing_key child dispatch - Rust-side orchestration: parent dispatches child jobs, suspends, resumes on completion - _executing_key in checkpoint tells child which step to execute directly - task() throws StepSuspend(mode="step_complete") after executing target step - result_processor handles child completion and updates parent checkpoint - WacGraph.svelte for runtime execution visualization - Sequential and parallel workflows tested end-to-end Co-Authored-By: Claude Opus 4.6 * fix: WAC v2 bundle cache, globalThis ctx sharing, description optional - Disable bun bundle caching for WAC v2 scripts (wrapper needs windmill-client from node_modules, not available in bundle mode) - Use Reflect.set/get(globalThis, "__wmill_wf_ctx") to share workflow context across dual module instances (wrapper vs user script) - Never-resolving thenable for non-matching steps in child job mode prevents Promise.all race conditions - Make description field optional in NewScript API (defaults to "") Co-Authored-By: Claude Opus 4.6 * feat: add step() primitive for inline checkpointed steps step() executes a function inline (no child job) and persists the result to the checkpoint. On replay, the cached value is returned — ensuring deterministic behavior for non-deterministic operations like Date.now() or Math.random(). - TypeScript: step(name, fn) — executes inline, throws StepSuspend with mode "inline_checkpoint" to persist before continuing - Rust: InlineCheckpoint variant in WacOutput, saves to checkpoint and resets running=false for immediate re-pickup (no zombie wait) - Shared step counter between task() and step() via _allocKey() Co-Authored-By: Claude Opus 4.6 * feat: add Python WAC v2 support with task(), step(), workflow() - Python SDK: WorkflowCtx with _executing_key child mode, _alloc_key shared counter, _run_inline_step for step(), _execute_directly and _never_resolve for child mode, step() async function - Python executor: WAC v2 detection, checkpoint.json writing, WAC wrapper.py generation calling _run_workflow(), post-execution hook into shared handle_wac_v2_output() - Make handle_wac_v2_output pub so both bun and python executors share the same dispatch/suspend/inline-checkpoint logic - 17 Python tests covering dispatch, replay, parallel, conditional, inline checkpoint, and child mode Co-Authored-By: Claude Opus 4.6 * chore: update sqlx prepared queries Co-Authored-By: Claude Opus 4.6 * fix: WacGraph Tooltip→Popover, simplify wacToFlow parsers - Fix type error: Tooltip doesn't accept text snippet, use Popover - Extract shared helpers for task matching and block collection - Replace linear tasks.find() with Map lookups - Remove mutable module-level counter Co-Authored-By: Claude Opus 4.6 * fix: Box::pin WAC v2 output handler to prevent stack overflow handle_python_job's async state machine was too large when combined with handle_wac_v2_output. Box::pin heap-allocates the future. Co-Authored-By: Claude Opus 4.6 * fix: merge WAC v1 and v2 task decorators to preserve backward compat The v2 @task decorator was shadowing the v1 one, breaking WAC v1 scripts that rely on HTTP-based dispatch via /workflow_as_code/ API. The merged decorator handles three modes: - v2: inside @workflow context → checkpoint/replay dispatch - v1: WM_JOB_ID set, no @workflow → HTTP API dispatch + wait_job - standalone: no Windmill env → execute function body directly Co-Authored-By: Claude Opus 4.6 * fix: skip no_main_func detection for WAC v2 scripts in TS and Python parsers Co-Authored-By: Claude Opus 4.6 * fix: prevent empty/noop dispatch causing infinite requeue loop - Validate steps.len() > 0 in WAC dispatch handler (issue 3) - Replace noop StepSuspend throw with never-resolving promise so it can't reach the backend as an empty dispatch (issue 4) Co-Authored-By: Claude Opus 4.6 * fix: Python task wrapper now converts positional args to kwargs in v2 mode Previously only **kwargs were passed to _next_step(), silently dropping positional arguments. Extract shared _merge_args() helper used by both v1 and v2 paths. Co-Authored-By: Claude Opus 4.6 * fix: replace unwrap() with proper error propagation in WAC arg serialization Co-Authored-By: Claude Opus 4.6 * fix: add workspace_id filter to v2_job queries in WAC dispatch Co-Authored-By: Claude Opus 4.6 * fix: prevent race condition in WAC child dispatch Restructure dispatch to save checkpoint + suspend parent + seed child checkpoints in a single transaction BEFORE pushing child jobs. This ensures a fast child can't complete before the parent is suspended. Also wrap InlineCheckpoint save + running reset in a transaction to prevent corrupted state on crash. Use ULID for pre-generated child job IDs (consistent with rest of API). Co-Authored-By: Claude Opus 4.6 * fix: include step key and child job ID in WAC error propagation Move step_key lookup before the success check so failed child errors include which task failed, the child job ID, and the original error. Co-Authored-By: Claude Opus 4.6 * docs: document WAC determinism contract and step dispatch semantics - Document that workflow functions must be deterministic across replays - Document that WacStepDispatch.script/args are metadata, not dispatch targets - Add comments on counter-based key allocation Co-Authored-By: Claude Opus 4.6 * fix: tighten WAC v2 detection to reduce false positives Replace naive substring matching with line-aware checks that skip comments and look for specific patterns: - TS: import from "windmill-client" containing workflow/task - Python: @workflow and @task decorators with wmill import Extracted shared helpers in wac_executor.rs used by both executors. Co-Authored-By: Claude Opus 4.6 * fix: show failed steps in WacGraph when workflow completes with errors When flowDone is true and a pending step isn't in completedSteps, mark it as 'failed' instead of 'running'. The failed state CSS and XCircle icon were already defined but never triggered. Co-Authored-By: Claude Opus 4.6 * fix: unsuspend and fail parent when WAC child push fails Previously if a child push failed mid-batch, the parent remained suspended with suspend = num_steps but fewer children, hanging until the 14-day timeout. Now the push loop catches errors and unsuspends the parent before returning the error. Also adds source hash validation: if the script content changes between replays, the job fails with a clear error instead of silently feeding stale checkpoint data into wrong steps. Co-Authored-By: Claude Opus 4.6 * fix: clear suspend_until when unsuspending WAC parent Set suspend_until = NULL alongside suspend = 0 in both the child failure and all-children-complete paths, so the parent doesn't rely on subtle pull query invariants to be re-picked-up. Co-Authored-By: Claude Opus 4.6 * test: add exhaustive edge case tests for WAC v2 SDK fix: make TS task wrapper non-async to fix unawaited task flush The async wrapper caused microtask-based thenable auto-resolution that fired .then() and threw StepSuspend before _flushPending() could capture unawaited steps — making the flush mechanism completely broken. Now the thenable is returned directly without async wrapping. Backward compatible with v1 (all code paths still return awaitables). Tests added (59 TS + 66 Python) covering: full sequential lifecycle, step after parallel, parallel after parallel, conditional on step result, empty/single-task workflows, 10+ steps, falsy value preservation, inline steps, mixed step/task, unawaited flush, child mode with parallel, key determinism, large parallel groups, and complex mixed patterns. Co-Authored-By: Claude Opus 4.6 * fix: atomic checkpoint updates to prevent parallel child race condition Replace read-modify-write pattern in handle_wac_child_completion with atomic SQL operations: - completed_steps merged via jsonb_set(... || jsonb_build_object(...)) so concurrent children on different workers don't overwrite each other - suspend counter decremented atomically with RETURNING to determine "all done" condition (instead of checking completed_steps in memory) - suspend_until cleared in the same atomic decrement statement Before this fix, two parallel children completing simultaneously could both load the same checkpoint, each add their step, and save — the second write would overwrite the first, silently losing a child result and leaving the parent suspended forever. Co-Authored-By: Claude Opus 4.6 * fix: cancel already-pushed children on partial WAC dispatch failure When pushing child jobs sequentially, if pushing child N fails, children 1..N-1 are already running. Previously the error handler only unsuspended the parent, leaving orphaned children that would complete and corrupt the checkpoint state (decrementing suspend on an already-unsuspended parent, potentially causing duplicate step execution on re-run). Now on partial failure: 1. Cancel all already-pushed children (prevents them from completing and corrupting checkpoint state) 2. Clear pending_steps from checkpoint (so parent doesn't think children are outstanding on re-run) 3. Then unsuspend parent (so the error propagates) Co-Authored-By: Claude Opus 4.6 * fix: skip WAC duration write and child check for non-WAC parents The duration write to workflow_as_code_status was running for every non-flow child with a parent (error handlers, success handlers, run_script children), even though it was only intended for WAC jobs. Add WHERE workflow_as_code_status IS NOT NULL to skip non-WAC parents entirely. Piggyback RETURNING pending_steps.job_ids on the same query so WAC v2 child completion needs zero extra DB round-trips on the success path. Co-Authored-By: Claude Opus 4.6 * fix: seed child checkpoint in same transaction as push The child checkpoint insert was happening before the child job was pushed, violating the FK constraint on v2_job_status. Move it into the push transaction so the job row exists and the child can't be picked up before its checkpoint is ready. Co-Authored-By: Claude Opus 4.6 * fix: set running=false when WAC parent suspends for child dispatch The parent job kept running=true after suspending, so workers wouldn't pick it up when children completed and suspend reached 0. The parent only advanced when the zombie job detector reset it (~90s). Now the dispatch suspend sets running=false so the parent is immediately eligible for pickup. Co-Authored-By: Claude Opus 4.6 * fix: WAC parent suspend/unsuspend lifecycle Keep running=true when suspending the parent so the normal pull query (WHERE running=false) never picks it up. Keep suspend_until non-null when decrementing suspend to 0 so the suspended pull query (WHERE suspend_until IS NOT NULL AND suspend<=0) picks it up. Previously: setting running=false caused infinite restart loops because the normal pull query has no suspend check and would immediately re-pick the parent. Clearing suspend_until on the last child prevented the suspended pull from ever seeing it, requiring the 90s zombie detector. Co-Authored-By: Claude Opus 4.6 * feat: add approval primitive, flow child completion, timeline fixes for WAC v2 Co-Authored-By: Claude Opus 4.6 * feat: add error propagation, task options, sleep, and parallel for WAC v2 Co-Authored-By: Claude Opus 4.6 * test: fix python SDK tests to use name-based keys and add new test coverage Co-Authored-By: Claude Opus 4.6 * fix: address WAC v2 review findings (sleep timing, error marker, atomicity) - Fix sleep using suspend=1 instead of 0 to enforce actual delay - Add approval/sleep resume injection to Python executor - Fix TS SDK concurrency_limit mapping (was reading wrong property) - Namespace error marker as __wmill_error to avoid user data collision - Wrap child completion SQL in transaction for atomicity - Decrement suspend even when step key is missing (prevents hang) - Expand TASK_RE to handle export const, let, var, generics - Validate step key uniqueness before dispatch - Log warning on checkpoint deserialization failure - Remove unimplemented delete_after_use from SDKs - Add TaskError exception class to Python SDK with diagnostic context - Fix extra positional args handling and add functools.wraps - Improve getParamNames to handle typed/destructured params Co-Authored-By: Claude Opus 4.6 * sqlx * sqlx * test: add WAC v1 e2e integration tests for TS and Python Co-Authored-By: Claude Opus 4.6 * fix: revert fake test versions in typescript-client Co-Authored-By: Claude Opus 4.6 * refactor: remove unused WacGraph component and strip wacToFlow to isWorkflowAsCode Co-Authored-By: Claude Opus 4.6 * refactor: extract shared approval/sleep resume logic into wac_executor Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- ...b6efdae570e3416ec6c2493cc04f75c32a699.json | 16 + ...b7ab1f2fc29a2ba79a39576551bdf66b592b6.json | 15 + ...dacbabed4c5ae28101e3ae2694f96fd055a91.json | 2 +- ...50d9d258c288883b2b5b0ab286f5cb50850b5.json | 16 - ...153c43903f929ae5d62fbba12610f89c36d55.json | 2 +- ...e82c45be45409f881292f0d4a3316362ba1f4.json | 14 + ...3e6e08f808f423c8f2d58b9c849aba7d176f5.json | 14 + ...1e91093233cf16af0dae666b4743f3878b22e.json | 14 + ...5be9d9f51071978c0e48df4284d8b90000a4a.json | 22 + ...933aa54523bf44d63e304440e53b9eadd5340.json | 24 + ...e0c9c8c61e1bf609be60ac5dc5d438189353.json} | 10 +- ...706d78a6f24cb0e614d7d81ba1b643805bf06.json | 2 +- ...2808e90cb068d1048717f82992f476377cc20.json | 15 + backend/Cargo.lock | 16 + backend/Cargo.toml | 2 + backend/parsers/windmill-parser-py/src/lib.rs | 5 +- backend/parsers/windmill-parser-ts/src/lib.rs | 14 +- .../parsers/windmill-parser-wac/Cargo.toml | 21 + .../parsers/windmill-parser-wac/src/dag.rs | 44 + .../parsers/windmill-parser-wac/src/lib.rs | 32 + .../parsers/windmill-parser-wac/src/python.rs | 717 ++++++++ .../windmill-parser-wac/src/typescript.rs | 739 ++++++++ .../windmill-parser-wac/src/validation.rs | 64 + .../windmill-parser-wac/tests/python_tests.rs | 266 +++ .../windmill-parser-wac/tests/ts_tests.rs | 245 +++ .../parsers/windmill-parser-wasm/Cargo.toml | 2 + backend/parsers/windmill-parser-wasm/build.nu | 6 + .../parsers/windmill-parser-wasm/src/lib.rs | 7 + backend/test_wac_e2e.sh | 157 ++ backend/windmill-api/openapi.yaml | 1 - backend/windmill-api/src/jobs.rs | 39 +- backend/windmill-common/src/error.rs | 3 + backend/windmill-queue/src/jobs.rs | 35 +- .../nsjail/run.bun.config.proto | 7 + .../windmill-worker/src/ai/image_handler.rs | 2 +- .../windmill-worker/src/ai/query_builder.rs | 4 +- backend/windmill-worker/src/ai/types.rs | 17 +- .../windmill-worker/src/bigquery_executor.rs | 2 +- backend/windmill-worker/src/bun_executor.rs | 963 +++++++++- backend/windmill-worker/src/lib.rs | 1 + .../windmill-worker/src/python_executor.rs | 96 +- .../windmill-worker/src/result_processor.rs | 243 ++- .../windmill-worker/src/snowflake_executor.rs | 2 +- backend/windmill-worker/src/wac_executor.rs | 339 ++++ backend/windmill-worker/src/worker.rs | 7 + backend/windmill-worker/src/worker_flow.rs | 29 +- .../src/lib/components/ScriptBuilder.svelte | 3 +- .../src/lib/components/TimelineBar.svelte | 7 +- .../lib/components/WorkflowTimeline.svelte | 11 +- .../src/lib/components/graph/wacToFlow.ts | 16 + .../lib/components/runs/JobRunsPreview.svelte | 7 +- .../components/scriptEditor/LogPanel.svelte | 8 +- .../(root)/(logged)/run/[...run]/+page.svelte | 21 +- python-client/wmill/pyproject.toml | 6 + python-client/wmill/tests/test_workflow.py | 1114 ++++++++++++ python-client/wmill/wmill/client.py | 553 +++++- typescript-client/build.sh | 24 +- typescript-client/client.ts | 492 ++++- typescript-client/package-lock.json | 4 +- typescript-client/tests/e2e_wac.py | 182 ++ typescript-client/tests/e2e_wac_v1.py | 190 ++ typescript-client/tests/workflow.test.ts | 1596 +++++++++++++++++ 62 files changed, 8318 insertions(+), 209 deletions(-) create mode 100644 backend/.sqlx/query-0cd9cad7109340edc81a5a40620b6efdae570e3416ec6c2493cc04f75c32a699.json create mode 100644 backend/.sqlx/query-10af387fce25f6ea7af275e8e93b7ab1f2fc29a2ba79a39576551bdf66b592b6.json delete mode 100644 backend/.sqlx/query-56f7325e3b0316866714e76d94b50d9d258c288883b2b5b0ab286f5cb50850b5.json create mode 100644 backend/.sqlx/query-a35164456ade8e79cb8f5418c8fe82c45be45409f881292f0d4a3316362ba1f4.json create mode 100644 backend/.sqlx/query-a684f160d1a366c1928fef27c613e6e08f808f423c8f2d58b9c849aba7d176f5.json create mode 100644 backend/.sqlx/query-a72081cb042f09034338dcb49381e91093233cf16af0dae666b4743f3878b22e.json create mode 100644 backend/.sqlx/query-b663e6baf2f8da00c6d94e5b8e35be9d9f51071978c0e48df4284d8b90000a4a.json create mode 100644 backend/.sqlx/query-c331609952e0b98d36f605bd5d2933aa54523bf44d63e304440e53b9eadd5340.json rename backend/.sqlx/{query-1a0ab65bbf2751f702fc696c1e32a7dd9524cdd806be1ad8e9ab88d4c88d3f82.json => query-dbc7e74e259b502e700491ee0248e0c9c8c61e1bf609be60ac5dc5d438189353.json} (65%) create mode 100644 backend/.sqlx/query-f56c58fea9f27d2e55d33720e032808e90cb068d1048717f82992f476377cc20.json create mode 100644 backend/parsers/windmill-parser-wac/Cargo.toml create mode 100644 backend/parsers/windmill-parser-wac/src/dag.rs create mode 100644 backend/parsers/windmill-parser-wac/src/lib.rs create mode 100644 backend/parsers/windmill-parser-wac/src/python.rs create mode 100644 backend/parsers/windmill-parser-wac/src/typescript.rs create mode 100644 backend/parsers/windmill-parser-wac/src/validation.rs create mode 100644 backend/parsers/windmill-parser-wac/tests/python_tests.rs create mode 100644 backend/parsers/windmill-parser-wac/tests/ts_tests.rs create mode 100755 backend/test_wac_e2e.sh create mode 100644 backend/windmill-worker/src/wac_executor.rs create mode 100644 frontend/src/lib/components/graph/wacToFlow.ts create mode 100644 python-client/wmill/tests/test_workflow.py create mode 100644 typescript-client/tests/e2e_wac.py create mode 100644 typescript-client/tests/e2e_wac_v1.py create mode 100644 typescript-client/tests/workflow.test.ts diff --git a/backend/.sqlx/query-0cd9cad7109340edc81a5a40620b6efdae570e3416ec6c2493cc04f75c32a699.json b/backend/.sqlx/query-0cd9cad7109340edc81a5a40620b6efdae570e3416ec6c2493cc04f75c32a699.json new file mode 100644 index 0000000000..b8e52cdbe7 --- /dev/null +++ b/backend/.sqlx/query-0cd9cad7109340edc81a5a40620b6efdae570e3416ec6c2493cc04f75c32a699.json @@ -0,0 +1,16 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET canceled_by = $2, canceled_reason = $3 WHERE id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Varchar", + "Text" + ] + }, + "nullable": [] + }, + "hash": "0cd9cad7109340edc81a5a40620b6efdae570e3416ec6c2493cc04f75c32a699" +} diff --git a/backend/.sqlx/query-10af387fce25f6ea7af275e8e93b7ab1f2fc29a2ba79a39576551bdf66b592b6.json b/backend/.sqlx/query-10af387fce25f6ea7af275e8e93b7ab1f2fc29a2ba79a39576551bdf66b592b6.json new file mode 100644 index 0000000000..3f39982319 --- /dev/null +++ b/backend/.sqlx/query-10af387fce25f6ea7af275e8e93b7ab1f2fc29a2ba79a39576551bdf66b592b6.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET suspend = $2, suspend_until = now() + interval '14 day' WHERE id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Int4" + ] + }, + "nullable": [] + }, + "hash": "10af387fce25f6ea7af275e8e93b7ab1f2fc29a2ba79a39576551bdf66b592b6" +} diff --git a/backend/.sqlx/query-2d6607b3c38fe72b5663c32de58dacbabed4c5ae28101e3ae2694f96fd055a91.json b/backend/.sqlx/query-2d6607b3c38fe72b5663c32de58dacbabed4c5ae28101e3ae2694f96fd055a91.json index da0ce60709..90f38c05a7 100644 --- a/backend/.sqlx/query-2d6607b3c38fe72b5663c32de58dacbabed4c5ae28101e3ae2694f96fd055a91.json +++ b/backend/.sqlx/query-2d6607b3c38fe72b5663c32de58dacbabed4c5ae28101e3ae2694f96fd055a91.json @@ -15,7 +15,7 @@ ] }, "nullable": [ - false + true ] }, "hash": "2d6607b3c38fe72b5663c32de58dacbabed4c5ae28101e3ae2694f96fd055a91" diff --git a/backend/.sqlx/query-56f7325e3b0316866714e76d94b50d9d258c288883b2b5b0ab286f5cb50850b5.json b/backend/.sqlx/query-56f7325e3b0316866714e76d94b50d9d258c288883b2b5b0ab286f5cb50850b5.json deleted file mode 100644 index 6b47103c3a..0000000000 --- a/backend/.sqlx/query-56f7325e3b0316866714e76d94b50d9d258c288883b2b5b0ab286f5cb50850b5.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "UPDATE v2_job_status SET\n workflow_as_code_status = jsonb_set(\n jsonb_set(\n COALESCE(workflow_as_code_status, '{}'::jsonb),\n array[$1],\n COALESCE(workflow_as_code_status->$1, '{}'::jsonb)\n ),\n array[$1, 'duration_ms'],\n to_jsonb($2::bigint)\n )\n WHERE id = $3", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Text", - "Int8", - "Uuid" - ] - }, - "nullable": [] - }, - "hash": "56f7325e3b0316866714e76d94b50d9d258c288883b2b5b0ab286f5cb50850b5" -} diff --git a/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json b/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json index 36ddb8ab9f..713ccb9dd3 100644 --- a/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json +++ b/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json @@ -15,7 +15,7 @@ ] }, "nullable": [ - true + null ] }, "hash": "5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55" diff --git a/backend/.sqlx/query-a35164456ade8e79cb8f5418c8fe82c45be45409f881292f0d4a3316362ba1f4.json b/backend/.sqlx/query-a35164456ade8e79cb8f5418c8fe82c45be45409f881292f0d4a3316362ba1f4.json new file mode 100644 index 0000000000..aedbbf424e --- /dev/null +++ b/backend/.sqlx/query-a35164456ade8e79cb8f5418c8fe82c45be45409f881292f0d4a3316362ba1f4.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET suspend = GREATEST(suspend - 1, 0) WHERE id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [] + }, + "hash": "a35164456ade8e79cb8f5418c8fe82c45be45409f881292f0d4a3316362ba1f4" +} diff --git a/backend/.sqlx/query-a684f160d1a366c1928fef27c613e6e08f808f423c8f2d58b9c849aba7d176f5.json b/backend/.sqlx/query-a684f160d1a366c1928fef27c613e6e08f808f423c8f2d58b9c849aba7d176f5.json new file mode 100644 index 0000000000..7a45e6c402 --- /dev/null +++ b/backend/.sqlx/query-a684f160d1a366c1928fef27c613e6e08f808f423c8f2d58b9c849aba7d176f5.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET running = false, started_at = null WHERE id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [] + }, + "hash": "a684f160d1a366c1928fef27c613e6e08f808f423c8f2d58b9c849aba7d176f5" +} diff --git a/backend/.sqlx/query-a72081cb042f09034338dcb49381e91093233cf16af0dae666b4743f3878b22e.json b/backend/.sqlx/query-a72081cb042f09034338dcb49381e91093233cf16af0dae666b4743f3878b22e.json new file mode 100644 index 0000000000..4fa871c594 --- /dev/null +++ b/backend/.sqlx/query-a72081cb042f09034338dcb49381e91093233cf16af0dae666b4743f3878b22e.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET suspend = 0, suspend_until = NULL WHERE id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [] + }, + "hash": "a72081cb042f09034338dcb49381e91093233cf16af0dae666b4743f3878b22e" +} diff --git a/backend/.sqlx/query-b663e6baf2f8da00c6d94e5b8e35be9d9f51071978c0e48df4284d8b90000a4a.json b/backend/.sqlx/query-b663e6baf2f8da00c6d94e5b8e35be9d9f51071978c0e48df4284d8b90000a4a.json new file mode 100644 index 0000000000..8d09036772 --- /dev/null +++ b/backend/.sqlx/query-b663e6baf2f8da00c6d94e5b8e35be9d9f51071978c0e48df4284d8b90000a4a.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET suspend = GREATEST(suspend - 1, 0) WHERE id = $1 RETURNING suspend", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "suspend", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [ + false + ] + }, + "hash": "b663e6baf2f8da00c6d94e5b8e35be9d9f51071978c0e48df4284d8b90000a4a" +} diff --git a/backend/.sqlx/query-c331609952e0b98d36f605bd5d2933aa54523bf44d63e304440e53b9eadd5340.json b/backend/.sqlx/query-c331609952e0b98d36f605bd5d2933aa54523bf44d63e304440e53b9eadd5340.json new file mode 100644 index 0000000000..efd03ae26e --- /dev/null +++ b/backend/.sqlx/query-c331609952e0b98d36f605bd5d2933aa54523bf44d63e304440e53b9eadd5340.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_status SET\n workflow_as_code_status = jsonb_set(\n jsonb_set(\n workflow_as_code_status,\n array[$1],\n COALESCE(workflow_as_code_status->$1, '{}'::jsonb)\n ),\n array[$1, 'duration_ms'],\n to_jsonb($2::bigint)\n )\n WHERE id = $3 AND workflow_as_code_status IS NOT NULL\n RETURNING workflow_as_code_status->'_checkpoint'->'pending_steps'->'job_ids' AS \"job_ids: serde_json::Value\"", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_ids: serde_json::Value", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [ + "Text", + "Int8", + "Uuid" + ] + }, + "nullable": [ + null + ] + }, + "hash": "c331609952e0b98d36f605bd5d2933aa54523bf44d63e304440e53b9eadd5340" +} diff --git a/backend/.sqlx/query-1a0ab65bbf2751f702fc696c1e32a7dd9524cdd806be1ad8e9ab88d4c88d3f82.json b/backend/.sqlx/query-dbc7e74e259b502e700491ee0248e0c9c8c61e1bf609be60ac5dc5d438189353.json similarity index 65% rename from backend/.sqlx/query-1a0ab65bbf2751f702fc696c1e32a7dd9524cdd806be1ad8e9ab88d4c88d3f82.json rename to backend/.sqlx/query-dbc7e74e259b502e700491ee0248e0c9c8c61e1bf609be60ac5dc5d438189353.json index 5b82c5288f..5791090fc1 100644 --- a/backend/.sqlx/query-1a0ab65bbf2751f702fc696c1e32a7dd9524cdd806be1ad8e9ab88d4c88d3f82.json +++ b/backend/.sqlx/query-dbc7e74e259b502e700491ee0248e0c9c8c61e1bf609be60ac5dc5d438189353.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n WITH job_info AS (\n SELECT id, kind::text AS kind, parent_job\n FROM v2_job\n WHERE id = $1\n )\n SELECT\n q.id AS \"id!\",\n s.flow_status,\n q.suspend AS \"suspend!\",\n j.runnable_path AS script_path,\n j.permissioned_as_email AS email,\n (ji.kind IN ('flow', 'flowpreview')) AS \"is_flow_level!\"\n FROM job_info ji\n JOIN v2_job_queue q ON q.id = CASE\n WHEN ji.kind IN ('flow', 'flowpreview') THEN ji.id\n ELSE ji.parent_job\n END\n JOIN v2_job j ON j.id = q.id\n JOIN v2_job_status s ON s.id = q.id\n FOR UPDATE OF q\n ", + "query": "\n WITH job_info AS (\n SELECT id, kind::text AS kind, parent_job\n FROM v2_job\n WHERE id = $1\n )\n SELECT\n q.id AS \"id!\",\n s.flow_status,\n q.suspend AS \"suspend!\",\n j.runnable_path AS script_path,\n j.permissioned_as_email AS email,\n (ji.kind IN ('flow', 'flowpreview')) AS \"is_flow_level!\",\n (ji.kind NOT IN ('flow', 'flowpreview') AND q.id = ji.id) AS \"is_wac!\"\n FROM job_info ji\n JOIN v2_job_queue q ON q.id = CASE\n WHEN ji.kind IN ('flow', 'flowpreview') THEN ji.id\n ELSE COALESCE(ji.parent_job, ji.id)\n END\n JOIN v2_job j ON j.id = q.id\n LEFT JOIN v2_job_status s ON s.id = q.id\n FOR UPDATE OF q\n ", "describe": { "columns": [ { @@ -32,6 +32,11 @@ "ordinal": 5, "name": "is_flow_level!", "type_info": "Bool" + }, + { + "ordinal": 6, + "name": "is_wac!", + "type_info": "Bool" } ], "parameters": { @@ -45,8 +50,9 @@ false, true, false, + null, null ] }, - "hash": "1a0ab65bbf2751f702fc696c1e32a7dd9524cdd806be1ad8e9ab88d4c88d3f82" + "hash": "dbc7e74e259b502e700491ee0248e0c9c8c61e1bf609be60ac5dc5d438189353" } diff --git a/backend/.sqlx/query-eba16eb819e2644284fb073c891706d78a6f24cb0e614d7d81ba1b643805bf06.json b/backend/.sqlx/query-eba16eb819e2644284fb073c891706d78a6f24cb0e614d7d81ba1b643805bf06.json index c96961eac4..0a06188897 100644 --- a/backend/.sqlx/query-eba16eb819e2644284fb073c891706d78a6f24cb0e614d7d81ba1b643805bf06.json +++ b/backend/.sqlx/query-eba16eb819e2644284fb073c891706d78a6f24cb0e614d7d81ba1b643805bf06.json @@ -15,7 +15,7 @@ ] }, "nullable": [ - false + true ] }, "hash": "eba16eb819e2644284fb073c891706d78a6f24cb0e614d7d81ba1b643805bf06" diff --git a/backend/.sqlx/query-f56c58fea9f27d2e55d33720e032808e90cb068d1048717f82992f476377cc20.json b/backend/.sqlx/query-f56c58fea9f27d2e55d33720e032808e90cb068d1048717f82992f476377cc20.json new file mode 100644 index 0000000000..72acab6120 --- /dev/null +++ b/backend/.sqlx/query-f56c58fea9f27d2e55d33720e032808e90cb068d1048717f82992f476377cc20.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "UPDATE v2_job_queue SET suspend = 1, suspend_until = now() + make_interval(secs => $2) WHERE id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Float8" + ] + }, + "nullable": [] + }, + "hash": "f56c58fea9f27d2e55d33720e032808e90cb068d1048717f82992f476377cc20" +} diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 3f0ee056f9..1ab271d577 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -16956,6 +16956,22 @@ dependencies = [ "windmill-parser-sql", ] +[[package]] +name = "windmill-parser-wac" +version = "1.651.1" +dependencies = [ + "anyhow", + "rustpython-ast", + "rustpython-parser", + "serde", + "serde_json", + "sha2 0.10.9", + "swc_common", + "swc_ecma_ast", + "swc_ecma_parser", + "swc_ecma_visit", +] + [[package]] name = "windmill-parser-yaml" version = "1.651.1" diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 7ad583f31e..0bbdc65bd7 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -68,6 +68,7 @@ members = [ "./parsers/windmill-parser-bash", "./parsers/windmill-parser-py", "./parsers/windmill-parser-py-imports", + "./parsers/windmill-parser-wac", "./parsers/windmill-sql-datatype-parser-wasm", "./parsers/windmill-parser-yaml", "windmill-macros", "parsers/windmill-parser-nu", "./windmill-worker-volumes", @@ -332,6 +333,7 @@ windmill-parser-bash = { path = "./parsers/windmill-parser-bash" } windmill-parser-sql = { path = "./parsers/windmill-parser-sql" } windmill-parser-graphql = { path = "./parsers/windmill-parser-graphql" } windmill-parser-php = { path = "./parsers/windmill-parser-php" } +windmill-parser-wac = { path = "./parsers/windmill-parser-wac" } windmill-jseval = { path = "./windmill-jseval" } windmill-runtime-nativets = { path = "./windmill-runtime-nativets" } windmill-api-client = { path = "./windmill-api-client" } diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index c299bff4af..655a1a4ea9 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/backend/parsers/windmill-parser-py/src/lib.rs @@ -296,11 +296,14 @@ pub fn parse_python_signature( // Check if main function was found if params.is_none() { + let is_wac_v2 = (code.contains("@workflow") || code.contains("workflow(")) + && (code.contains("@task") || code.contains("task(")) + && (code.contains("import wmill") || code.contains("from wmill")); return Ok(MainArgSignature { star_args: false, star_kwargs: false, args: vec![], - no_main_func: Some(true), + no_main_func: Some(!is_wac_v2), has_preprocessor: Some(has_preprocessor), }); } diff --git a/backend/parsers/windmill-parser-ts/src/lib.rs b/backend/parsers/windmill-parser-ts/src/lib.rs index 04dd345b2f..f0928fe984 100644 --- a/backend/parsers/windmill-parser-ts/src/lib.rs +++ b/backend/parsers/windmill-parser-ts/src/lib.rs @@ -261,7 +261,9 @@ pub fn parse_deno_signature( for specifier in &named_export.specifiers { if let swc_ecma_ast::ExportSpecifier::Named(spec) = specifier { let export_name = match &spec.exported { - Some(swc_ecma_ast::ModuleExportName::Ident(ident)) => ident.sym.as_ref(), + Some(swc_ecma_ast::ModuleExportName::Ident(ident)) => { + ident.sym.as_ref() + } Some(swc_ecma_ast::ModuleExportName::Str(s)) => s.value.as_ref(), None => match &spec.orig { swc_ecma_ast::ModuleExportName::Ident(ident) => ident.sym.as_ref(), @@ -315,7 +317,11 @@ pub fn parse_deno_signature( let mut c: u16 = 0; - let no_main_func = entrypoint_params.is_none(); + let is_wac_v2 = entrypoint_params.is_none() + && code.contains("workflow(") + && code.contains("task(") + && code.contains("windmill-client"); + let no_main_func = entrypoint_params.is_none() && !is_wac_v2; let mut type_resolver = HashMap::new(); let r = MainArgSignature { star_args: false, @@ -833,7 +839,9 @@ fn tstype_to_typ( false, ), symbol @ _ if symbol.starts_with("DynMultiselect_") => ( - Typ::DynMultiselect(symbol.strip_prefix("DynMultiselect_").unwrap().to_string()), + Typ::DynMultiselect( + symbol.strip_prefix("DynMultiselect_").unwrap().to_string(), + ), false, ), symbol @ _ => { diff --git a/backend/parsers/windmill-parser-wac/Cargo.toml b/backend/parsers/windmill-parser-wac/Cargo.toml new file mode 100644 index 0000000000..d354b1c653 --- /dev/null +++ b/backend/parsers/windmill-parser-wac/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "windmill-parser-wac" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[lib] +name = "windmill_parser_wac" +path = "./src/lib.rs" + +[dependencies] +rustpython-parser.workspace = true +rustpython-ast = { version = "0.4.0", features = ["visitor"] } +swc_common.workspace = true +swc_ecma_parser.workspace = true +swc_ecma_ast.workspace = true +swc_ecma_visit.workspace = true +serde.workspace = true +serde_json.workspace = true +anyhow.workspace = true +sha2.workspace = true diff --git a/backend/parsers/windmill-parser-wac/src/dag.rs b/backend/parsers/windmill-parser-wac/src/dag.rs new file mode 100644 index 0000000000..662f5a06b6 --- /dev/null +++ b/backend/parsers/windmill-parser-wac/src/dag.rs @@ -0,0 +1,44 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct WorkflowDag { + pub nodes: Vec, + pub edges: Vec, + pub params: Vec, + pub source_hash: String, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct Param { + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub typ: Option, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct DagNode { + pub id: String, + pub node_type: DagNodeType, + pub label: String, + pub line: usize, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(tag = "type")] +pub enum DagNodeType { + Step { name: String, script: String }, + Branch { condition_source: String }, + ParallelStart, + ParallelEnd, + LoopStart { iter_source: String }, + LoopEnd, + Return, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct DagEdge { + pub from: String, + pub to: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub label: Option, +} diff --git a/backend/parsers/windmill-parser-wac/src/lib.rs b/backend/parsers/windmill-parser-wac/src/lib.rs new file mode 100644 index 0000000000..1496b9d6cf --- /dev/null +++ b/backend/parsers/windmill-parser-wac/src/lib.rs @@ -0,0 +1,32 @@ +pub mod dag; +pub mod python; +pub mod typescript; +pub mod validation; + +use dag::WorkflowDag; +use validation::CompileError; + +#[derive(Debug, serde::Serialize)] +#[serde(tag = "type")] +pub enum ParseResult { + #[serde(rename = "success")] + Success(WorkflowDag), + #[serde(rename = "error")] + Error { errors: Vec }, +} + +pub fn parse_workflow(code: &str, language: &str) -> ParseResult { + let result = match language { + "python" | "python3" | "py" => python::parse_python_workflow(code), + "typescript" | "ts" | "deno" | "bun" => typescript::parse_ts_workflow(code), + _ => Err(vec![CompileError { + message: format!("Unsupported language: {language}"), + line: 0, + }]), + }; + + match result { + Ok(dag) => ParseResult::Success(dag), + Err(errors) => ParseResult::Error { errors }, + } +} diff --git a/backend/parsers/windmill-parser-wac/src/python.rs b/backend/parsers/windmill-parser-wac/src/python.rs new file mode 100644 index 0000000000..74f0117376 --- /dev/null +++ b/backend/parsers/windmill-parser-wac/src/python.rs @@ -0,0 +1,717 @@ +use std::collections::HashMap; + +use rustpython_parser::{ + ast::{ + Expr, ExprAwait, ExprCall, ExprName, Stmt, StmtExpr, StmtFor, StmtIf, StmtReturn, StmtTry, + StmtTryStar, StmtWhile, + }, + Parse, +}; + +use crate::dag::{DagEdge, DagNode, DagNodeType, Param, WorkflowDag}; +use crate::validation::{self, CompileError}; + +struct LineIndex { + newline_offsets: Vec, +} + +impl LineIndex { + fn new(source: &str) -> Self { + let mut offsets = vec![0]; + for (i, c) in source.char_indices() { + if c == '\n' { + offsets.push(i + 1); + } + } + Self { newline_offsets: offsets } + } + + fn line_of(&self, byte_offset: usize) -> usize { + match self.newline_offsets.binary_search(&byte_offset) { + Ok(line) => line + 1, + Err(line) => line, + } + } +} + +/// Maps task function name → optional external path (from `@task(path="...")`) +type TaskFunctions = HashMap>; + +/// First pass: scan top-level `@task async def foo(...)` declarations. +fn collect_task_functions(stmts: &[Stmt]) -> TaskFunctions { + let mut tasks = HashMap::new(); + for stmt in stmts { + if let Stmt::AsyncFunctionDef(func) = stmt { + for dec in &func.decorator_list { + match dec { + // @task (bare decorator) + Expr::Name(ExprName { id, .. }) if id.as_str() == "task" => { + tasks.insert(func.name.to_string(), None); + } + // @task(path="...") + Expr::Call(call) => { + if let Expr::Name(ExprName { id, .. }) = call.func.as_ref() { + if id.as_str() == "task" { + let path = extract_task_path_kwarg(call); + tasks.insert(func.name.to_string(), path); + } + } + } + _ => {} + } + } + } + } + tasks +} + +/// Extract the `path=` keyword argument from a `@task(path="...")` call. +fn extract_task_path_kwarg(call: &ExprCall) -> Option { + for kw in &call.keywords { + if let Some(ref arg) = kw.arg { + if arg.as_str() == "path" { + if let Expr::Constant(c) = &kw.value { + if let rustpython_parser::ast::Constant::Str(s) = &c.value { + return Some(s.to_string()); + } + } + } + } + } + None +} + +struct WacWalker { + nodes: Vec, + edges: Vec, + errors: Vec, + node_counter: usize, + line_index: LineIndex, + task_functions: TaskFunctions, + in_try: bool, + in_while: bool, + in_nested_func: bool, + in_comprehension: bool, +} + +impl WacWalker { + fn new(source: &str, task_functions: TaskFunctions) -> Self { + Self { + nodes: Vec::new(), + edges: Vec::new(), + errors: Vec::new(), + node_counter: 0, + line_index: LineIndex::new(source), + task_functions, + in_try: false, + in_while: false, + in_nested_func: false, + in_comprehension: false, + } + } + + fn next_id(&mut self) -> String { + let id = format!("step_{}", self.node_counter); + self.node_counter += 1; + id + } + + fn add_node(&mut self, node: DagNode) -> String { + let id = node.id.clone(); + self.nodes.push(node); + id + } + + fn add_edge(&mut self, from: &str, to: &str, label: Option) { + self.edges + .push(DagEdge { from: from.to_string(), to: to.to_string(), label }); + } + + fn line_of_expr(&self, expr: &Expr) -> usize { + let offset = match expr { + Expr::Call(c) => c.range.start().to_usize(), + Expr::Await(a) => a.range.start().to_usize(), + Expr::Attribute(a) => a.range.start().to_usize(), + Expr::Name(n) => n.range.start().to_usize(), + _ => 0, + }; + self.line_index.line_of(offset) + } + + fn line_of_stmt(&self, stmt: &Stmt) -> usize { + let offset = match stmt { + Stmt::If(s) => s.range.start().to_usize(), + Stmt::For(s) => s.range.start().to_usize(), + Stmt::While(s) => s.range.start().to_usize(), + Stmt::Return(s) => s.range.start().to_usize(), + Stmt::Expr(s) => s.range.start().to_usize(), + Stmt::Try(s) => s.range.start().to_usize(), + Stmt::TryStar(s) => s.range.start().to_usize(), + Stmt::Assign(s) => s.range.start().to_usize(), + Stmt::AnnAssign(s) => s.range.start().to_usize(), + Stmt::FunctionDef(s) => s.range.start().to_usize(), + Stmt::AsyncFunctionDef(s) => s.range.start().to_usize(), + _ => 0, + }; + self.line_index.line_of(offset) + } + + /// Check if an expression is a call to a known @task function + fn is_task_fn_call(&self, expr: &Expr) -> bool { + if let Expr::Call(call) = expr { + if let Expr::Name(ExprName { id, .. }) = call.func.as_ref() { + return self.task_functions.contains_key(id.as_str()); + } + } + false + } + + /// Check if an expression is `asyncio.gather(...)` call + fn is_asyncio_gather_call(expr: &Expr) -> bool { + if let Expr::Call(call) = expr { + if let Expr::Attribute(rustpython_parser::ast::ExprAttribute { value, attr, .. }) = + call.func.as_ref() + { + if attr.as_str() == "gather" { + if let Expr::Name(ExprName { id, .. }) = value.as_ref() { + return id.as_str() == "asyncio"; + } + } + } + } + false + } + + /// Extract step name and script from a task function call. + /// Name = function name, script = task_path or function name. + fn extract_step_info_from_task_call(&self, call: &ExprCall) -> Option<(String, String)> { + if let Expr::Name(ExprName { id, .. }) = call.func.as_ref() { + let name = id.to_string(); + let script = self + .task_functions + .get(id.as_str()) + .and_then(|p| p.clone()) + .unwrap_or_else(|| name.clone()); + Some((name, script)) + } else { + None + } + } + + fn expr_to_source(expr: &Expr) -> String { + match expr { + Expr::Compare(c) => { + let left = Self::expr_to_source(&c.left); + if let Some(comparator) = c.comparators.first() { + let right = Self::expr_to_source(comparator); + let op = match c.ops.first() { + Some(rustpython_parser::ast::CmpOp::Gt) => ">", + Some(rustpython_parser::ast::CmpOp::Lt) => "<", + Some(rustpython_parser::ast::CmpOp::GtE) => ">=", + Some(rustpython_parser::ast::CmpOp::LtE) => "<=", + Some(rustpython_parser::ast::CmpOp::Eq) => "==", + Some(rustpython_parser::ast::CmpOp::NotEq) => "!=", + Some(rustpython_parser::ast::CmpOp::In) => "in", + Some(rustpython_parser::ast::CmpOp::NotIn) => "not in", + Some(rustpython_parser::ast::CmpOp::Is) => "is", + Some(rustpython_parser::ast::CmpOp::IsNot) => "is not", + None => "?", + }; + format!("{left} {op} {right}") + } else { + left + } + } + Expr::Subscript(s) => { + let value = Self::expr_to_source(&s.value); + let slice = Self::expr_to_source(&s.slice); + format!("{value}[{slice}]") + } + Expr::Attribute(a) => { + let value = Self::expr_to_source(&a.value); + format!("{value}.{}", a.attr) + } + Expr::Name(n) => n.id.to_string(), + Expr::Constant(c) => match &c.value { + rustpython_parser::ast::Constant::Str(s) => format!("\"{s}\""), + rustpython_parser::ast::Constant::Int(i) => i.to_string(), + rustpython_parser::ast::Constant::Float(f) => f.to_string(), + rustpython_parser::ast::Constant::Bool(b) => b.to_string(), + rustpython_parser::ast::Constant::None => "None".to_string(), + _ => "...".to_string(), + }, + _ => "...".to_string(), + } + } + + /// Check if a statement body contains any task function calls (recursively) + fn body_contains_step(&self, body: &[Stmt]) -> bool { + for stmt in body { + if self.stmt_contains_step(stmt) { + return true; + } + } + false + } + + fn stmt_contains_step(&self, stmt: &Stmt) -> bool { + match stmt { + Stmt::Expr(StmtExpr { value, .. }) => self.expr_contains_step(value), + Stmt::Assign(a) => self.expr_contains_step(&a.value), + Stmt::If(s) => self.body_contains_step(&s.body) || self.body_contains_step(&s.orelse), + Stmt::For(s) => self.body_contains_step(&s.body) || self.body_contains_step(&s.orelse), + Stmt::While(s) => { + self.body_contains_step(&s.body) || self.body_contains_step(&s.orelse) + } + Stmt::Try(s) => { + self.body_contains_step(&s.body) + || self.body_contains_step(&s.orelse) + || self.body_contains_step(&s.finalbody) + || s.handlers.iter().any(|h| match h { + rustpython_parser::ast::ExceptHandler::ExceptHandler(eh) => { + self.body_contains_step(&eh.body) + } + }) + } + Stmt::TryStar(s) => { + self.body_contains_step(&s.body) + || self.body_contains_step(&s.orelse) + || self.body_contains_step(&s.finalbody) + || s.handlers.iter().any(|h| match h { + rustpython_parser::ast::ExceptHandler::ExceptHandler(eh) => { + self.body_contains_step(&eh.body) + } + }) + } + Stmt::Return(_) => false, + _ => false, + } + } + + fn expr_contains_step(&self, expr: &Expr) -> bool { + if self.is_task_fn_call(expr) { + return true; + } + match expr { + Expr::Await(ExprAwait { value, .. }) => self.expr_contains_step(value), + Expr::Call(call) => { + if self.is_task_fn_call(&Expr::Call(call.clone())) { + return true; + } + if Self::is_asyncio_gather_call(&Expr::Call(call.clone())) { + return call.args.iter().any(|a| self.expr_contains_step(a)); + } + false + } + _ => false, + } + } + + /// Walk a list of statements, returning (first_node_id, last_node_id) + fn walk_body(&mut self, body: &[Stmt]) -> Option<(String, String)> { + let mut first_id: Option = None; + let mut prev_id: Option = None; + + for stmt in body { + if let Some((stmt_first, stmt_last)) = self.walk_stmt(stmt) { + if let Some(ref prev) = prev_id { + self.add_edge(prev, &stmt_first, None); + } + if first_id.is_none() { + first_id = Some(stmt_first); + } + prev_id = Some(stmt_last); + } + } + + match (first_id, prev_id) { + (Some(f), Some(l)) => Some((f, l)), + _ => None, + } + } + + fn walk_stmt(&mut self, stmt: &Stmt) -> Option<(String, String)> { + match stmt { + Stmt::Expr(StmtExpr { value, .. }) => self.walk_expr_stmt(value), + Stmt::Assign(a) => self.walk_expr_stmt(&a.value), + Stmt::If(if_stmt) => self.walk_if(if_stmt), + Stmt::For(for_stmt) => self.walk_for(for_stmt), + Stmt::While(while_stmt) => self.walk_while(while_stmt), + Stmt::Try(try_stmt) => self.walk_try(try_stmt), + Stmt::TryStar(try_stmt) => self.walk_try_star(try_stmt), + Stmt::Return(ret) => self.walk_return(ret), + Stmt::FunctionDef(_) | Stmt::AsyncFunctionDef(_) => { + if self.stmt_contains_step(stmt) { + self.errors.push(validation::error_step_in_nested_function( + self.line_of_stmt(stmt), + )); + } + None + } + _ => None, + } + } + + fn walk_expr_stmt(&mut self, expr: &Expr) -> Option<(String, String)> { + // await task_fn(...) + if let Expr::Await(ExprAwait { value, .. }) = expr { + // await task_fn(...) + if let Expr::Call(call) = value.as_ref() { + if self.is_task_fn_call(&Expr::Call(call.clone())) { + return self.emit_step(call, expr); + } + } + // await asyncio.gather(task_fn(...), task_fn(...), ...) + if Self::is_asyncio_gather_call(value) { + if let Expr::Call(gather_call) = value.as_ref() { + return self.emit_parallel(gather_call, expr); + } + } + } + + // Bare task_fn() without await — validation error + if self.is_task_fn_call(expr) { + self.errors + .push(validation::error_missing_await(self.line_of_expr(expr))); + } + + None + } + + fn emit_step(&mut self, call: &ExprCall, expr: &Expr) -> Option<(String, String)> { + if self.in_try { + self.errors + .push(validation::error_step_in_try(self.line_of_expr(expr))); + return None; + } + if self.in_while { + self.errors + .push(validation::error_step_in_while(self.line_of_expr(expr))); + return None; + } + if self.in_nested_func { + self.errors.push(validation::error_step_in_nested_function( + self.line_of_expr(expr), + )); + return None; + } + if self.in_comprehension { + self.errors.push(validation::error_step_in_comprehension( + self.line_of_expr(expr), + )); + return None; + } + + let (name, script) = self + .extract_step_info_from_task_call(call) + .unwrap_or(("unknown".into(), "unknown".into())); + let id = self.next_id(); + let node_id = self.add_node(DagNode { + id: id.clone(), + node_type: DagNodeType::Step { name: name.clone(), script }, + label: name, + line: self.line_of_expr(expr), + }); + Some((node_id.clone(), node_id)) + } + + fn emit_parallel(&mut self, gather_call: &ExprCall, expr: &Expr) -> Option<(String, String)> { + if self.in_try { + self.errors + .push(validation::error_step_in_try(self.line_of_expr(expr))); + return None; + } + if self.in_while { + self.errors + .push(validation::error_step_in_while(self.line_of_expr(expr))); + return None; + } + + let line = self.line_of_expr(expr); + let start_id = self.next_id(); + let start_node_id = self.add_node(DagNode { + id: start_id.clone(), + node_type: DagNodeType::ParallelStart, + label: "parallel".to_string(), + line, + }); + + let mut step_ids = Vec::new(); + for arg in &gather_call.args { + // Each arg should be task_fn(...) + if let Expr::Call(call) = arg { + if self.is_task_fn_call(&Expr::Call(call.clone())) { + let (name, script) = self + .extract_step_info_from_task_call(call) + .unwrap_or(("unknown".into(), "unknown".into())); + let step_id = self.next_id(); + let node_id = self.add_node(DagNode { + id: step_id.clone(), + node_type: DagNodeType::Step { name: name.clone(), script }, + label: name, + line: self.line_of_expr(arg), + }); + self.add_edge(&start_node_id, &node_id, None); + step_ids.push(node_id); + } + } + } + + let end_id = self.next_id(); + let end_node_id = self.add_node(DagNode { + id: end_id.clone(), + node_type: DagNodeType::ParallelEnd, + label: "join".to_string(), + line, + }); + + for step_id in &step_ids { + self.add_edge(step_id, &end_node_id, None); + } + + Some((start_node_id, end_node_id)) + } + + fn walk_if(&mut self, if_stmt: &StmtIf) -> Option<(String, String)> { + let has_steps_in_body = self.body_contains_step(&if_stmt.body); + let has_steps_in_else = self.body_contains_step(&if_stmt.orelse); + + if !has_steps_in_body && !has_steps_in_else { + return None; + } + + let line = self.line_index.line_of(if_stmt.range.start().to_usize()); + let condition_source = Self::expr_to_source(&if_stmt.test); + + let branch_id = self.next_id(); + let branch_node_id = self.add_node(DagNode { + id: branch_id.clone(), + node_type: DagNodeType::Branch { condition_source }, + label: "if".to_string(), + line, + }); + + let merge_id = format!("{branch_id}_merge"); + + let mut last_ids = Vec::new(); + + if let Some((true_first, true_last)) = self.walk_body(&if_stmt.body) { + self.add_edge(&branch_node_id, &true_first, Some("true".to_string())); + last_ids.push(true_last); + } else { + last_ids.push(branch_node_id.clone()); + } + + if !if_stmt.orelse.is_empty() { + if let Some((else_first, else_last)) = self.walk_body(&if_stmt.orelse) { + self.add_edge(&branch_node_id, &else_first, Some("false".to_string())); + last_ids.push(else_last); + } else { + last_ids.push(branch_node_id.clone()); + } + } + + if last_ids.len() == 1 { + Some((branch_node_id, last_ids.into_iter().next().unwrap())) + } else { + Some((branch_node_id, merge_id)) + } + } + + fn walk_for(&mut self, for_stmt: &StmtFor) -> Option<(String, String)> { + if !self.body_contains_step(&for_stmt.body) { + return None; + } + + let line = self.line_index.line_of(for_stmt.range.start().to_usize()); + let iter_source = Self::expr_to_source(&for_stmt.iter); + + let start_id = self.next_id(); + let start_node_id = self.add_node(DagNode { + id: start_id.clone(), + node_type: DagNodeType::LoopStart { iter_source }, + label: "for".to_string(), + line, + }); + + if let Some((body_first, body_last)) = self.walk_body(&for_stmt.body) { + self.add_edge(&start_node_id, &body_first, None); + self.add_edge(&body_last, &start_node_id, Some("next".to_string())); + } + + let end_id = self.next_id(); + let end_node_id = self.add_node(DagNode { + id: end_id.clone(), + node_type: DagNodeType::LoopEnd, + label: "end for".to_string(), + line, + }); + self.add_edge(&start_node_id, &end_node_id, Some("done".to_string())); + + Some((start_node_id, end_node_id)) + } + + fn walk_while(&mut self, while_stmt: &StmtWhile) -> Option<(String, String)> { + if self.body_contains_step(&while_stmt.body) { + let line = self.line_index.line_of(while_stmt.range.start().to_usize()); + self.errors.push(validation::error_step_in_while(line)); + } + None + } + + fn walk_try(&mut self, try_stmt: &StmtTry) -> Option<(String, String)> { + let has_steps = self.body_contains_step(&try_stmt.body) + || self.body_contains_step(&try_stmt.orelse) + || self.body_contains_step(&try_stmt.finalbody) + || try_stmt.handlers.iter().any(|h| match h { + rustpython_parser::ast::ExceptHandler::ExceptHandler(eh) => { + self.body_contains_step(&eh.body) + } + }); + + if has_steps { + let line = self.line_index.line_of(try_stmt.range.start().to_usize()); + self.errors.push(validation::error_step_in_try(line)); + } + None + } + + fn walk_try_star(&mut self, try_stmt: &StmtTryStar) -> Option<(String, String)> { + let has_steps = self.body_contains_step(&try_stmt.body) + || self.body_contains_step(&try_stmt.orelse) + || self.body_contains_step(&try_stmt.finalbody) + || try_stmt.handlers.iter().any(|h| match h { + rustpython_parser::ast::ExceptHandler::ExceptHandler(eh) => { + self.body_contains_step(&eh.body) + } + }); + + if has_steps { + let line = self.line_index.line_of(try_stmt.range.start().to_usize()); + self.errors.push(validation::error_step_in_try(line)); + } + None + } + + fn walk_return(&mut self, ret: &StmtReturn) -> Option<(String, String)> { + let line = self.line_index.line_of(ret.range.start().to_usize()); + let id = self.next_id(); + let node_id = self.add_node(DagNode { + id: id.clone(), + node_type: DagNodeType::Return, + label: "return".to_string(), + line, + }); + Some((node_id.clone(), node_id)) + } +} + +/// Extract workflow function parameters (no longer skips ctx) +fn extract_params(args: &rustpython_parser::ast::Arguments) -> Vec { + let mut params = Vec::new(); + for arg_with_default in args.args.iter().chain(args.posonlyargs.iter()) { + let name = arg_with_default.def.arg.to_string(); + let typ = arg_with_default + .def + .annotation + .as_ref() + .map(|ann| WacWalker::expr_to_source(ann)); + params.push(Param { name, typ }); + } + params +} + +pub fn parse_python_workflow(code: &str) -> Result> { + let ast = rustpython_parser::ast::Suite::parse(code, "") + .map_err(|e| vec![CompileError { message: format!("Parse error: {e}"), line: 0 }])?; + + // First pass: collect @task functions + let task_functions = collect_task_functions(&ast); + + // Find the @workflow async def + let workflow_fn = ast.iter().find_map(|stmt| { + if let Stmt::AsyncFunctionDef(func) = stmt { + let has_workflow_decorator = func.decorator_list.iter().any(|dec| { + if let Expr::Name(ExprName { id, .. }) = dec { + id.as_str() == "workflow" + } else { + false + } + }); + if has_workflow_decorator { + return Some(func); + } + } + // Also check non-async for error reporting + if let Stmt::FunctionDef(func) = stmt { + let has_workflow_decorator = func.decorator_list.iter().any(|dec| { + if let Expr::Name(ExprName { id, .. }) = dec { + id.as_str() == "workflow" + } else { + false + } + }); + if has_workflow_decorator { + return None; // Will be reported as not-async below + } + } + None + }); + + // Check for non-async workflow function + let non_async_workflow = ast.iter().find_map(|stmt| { + if let Stmt::FunctionDef(func) = stmt { + let has_workflow_decorator = func.decorator_list.iter().any(|dec| { + if let Expr::Name(ExprName { id, .. }) = dec { + id.as_str() == "workflow" + } else { + false + } + }); + if has_workflow_decorator { + let line_index = LineIndex::new(code); + return Some(line_index.line_of(func.range.start().to_usize())); + } + } + None + }); + + if let Some(line) = non_async_workflow { + if workflow_fn.is_none() { + return Err(vec![validation::error_not_async(line)]); + } + } + + let workflow_fn = workflow_fn.ok_or_else(|| { + vec![CompileError { message: "No @workflow async function found.".to_string(), line: 0 }] + })?; + + let params = extract_params(&workflow_fn.args); + let source_hash = compute_source_hash(code); + + let mut walker = WacWalker::new(code, task_functions); + walker.walk_body(&workflow_fn.body); + + if !walker.errors.is_empty() { + return Err(walker.errors); + } + + Ok(WorkflowDag { nodes: walker.nodes, edges: walker.edges, params, source_hash }) +} + +fn compute_source_hash(code: &str) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(code.as_bytes()); + format!("{:x}", hasher.finalize()) +} + +trait ToUsize { + fn to_usize(self) -> usize; +} + +impl ToUsize for rustpython_parser::text_size::TextSize { + fn to_usize(self) -> usize { + u32::from(self) as usize + } +} diff --git a/backend/parsers/windmill-parser-wac/src/typescript.rs b/backend/parsers/windmill-parser-wac/src/typescript.rs new file mode 100644 index 0000000000..bd777749ea --- /dev/null +++ b/backend/parsers/windmill-parser-wac/src/typescript.rs @@ -0,0 +1,739 @@ +use std::collections::HashMap; + +use swc_common::{sync::Lrc, FileName, SourceMap, SourceMapper, Spanned}; +use swc_ecma_ast::*; +use swc_ecma_parser::{lexer::Lexer, Parser, StringInput, Syntax, TsSyntax}; + +use crate::dag::{DagEdge, DagNode, DagNodeType, Param, WorkflowDag}; +use crate::validation::{self, CompileError}; + +/// Maps task function name → optional external path (from `task("f/path", ...)`) +type TaskFunctions = HashMap>; + +/// First pass: scan top-level `const foo = task(async (...) => {})` or +/// `const foo = task("f/path", async (...) => {})` declarations. +fn collect_task_functions(module: &Module) -> TaskFunctions { + let mut tasks = HashMap::new(); + for item in &module.body { + // const foo = task(async (...) => { ... }) + // const foo = task("f/path", async (...) => { ... }) + if let ModuleItem::Stmt(Stmt::Decl(Decl::Var(var_decl))) = item { + for decl in &var_decl.decls { + if let (Some(name), Some(init)) = (extract_var_name(&decl.name), &decl.init) { + if let Some(path) = extract_task_call_info(init) { + tasks.insert(name, path); + } + } + } + } + // export const foo = task(...) + if let ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(export)) = item { + if let Decl::Var(var_decl) = &export.decl { + for decl in &var_decl.decls { + if let (Some(name), Some(init)) = (extract_var_name(&decl.name), &decl.init) { + if let Some(path) = extract_task_call_info(init) { + tasks.insert(name, path); + } + } + } + } + } + } + tasks +} + +/// Extract variable name from a pattern (simple ident case) +fn extract_var_name(pat: &Pat) -> Option { + if let Pat::Ident(BindingIdent { id, .. }) = pat { + Some(id.sym.to_string()) + } else { + None + } +} + +/// Check if expr is `task(async fn)` or `task("path", async fn)`. +/// Returns Some(optional_path) if it is a task() call. +fn extract_task_call_info(expr: &Expr) -> Option> { + if let Expr::Call(call) = expr { + if let Callee::Expr(callee) = &call.callee { + if let Expr::Ident(ident) = callee.as_ref() { + if ident.sym.as_ref() == "task" { + // task("f/path", async fn) or task(async fn) + if call.args.len() == 2 { + // task("f/path", async fn) + let path = extract_string_lit(&call.args[0].expr); + return Some(path); + } else if call.args.len() == 1 { + // task(async fn) + return Some(None); + } + } + } + } + } + None +} + +struct TsWacWalker { + nodes: Vec, + edges: Vec, + errors: Vec, + node_counter: usize, + cm: Lrc, + task_functions: TaskFunctions, + in_try: bool, + in_while: bool, + in_nested_func: bool, +} + +impl TsWacWalker { + fn new(cm: Lrc, task_functions: TaskFunctions) -> Self { + Self { + nodes: Vec::new(), + edges: Vec::new(), + errors: Vec::new(), + node_counter: 0, + cm, + task_functions, + in_try: false, + in_while: false, + in_nested_func: false, + } + } + + fn next_id(&mut self) -> String { + let id = format!("step_{}", self.node_counter); + self.node_counter += 1; + id + } + + fn add_node(&mut self, node: DagNode) -> String { + let id = node.id.clone(); + self.nodes.push(node); + id + } + + fn add_edge(&mut self, from: &str, to: &str, label: Option) { + self.edges + .push(DagEdge { from: from.to_string(), to: to.to_string(), label }); + } + + fn span_line(&self, span: swc_common::Span) -> usize { + let loc = self.cm.lookup_char_pos(span.lo); + loc.line + } + + /// Check if expr is a call to a known task function + fn is_task_call(&self, expr: &Expr) -> bool { + if let Expr::Call(call) = expr { + if let Callee::Expr(callee) = &call.callee { + if let Expr::Ident(ident) = callee.as_ref() { + return self.task_functions.contains_key(ident.sym.as_ref()); + } + } + } + false + } + + /// Check if expr is `Promise.all([...])` + fn is_promise_all(expr: &Expr) -> bool { + if let Expr::Call(call) = expr { + if let Callee::Expr(callee) = &call.callee { + if let Expr::Member(MemberExpr { obj, prop: MemberProp::Ident(prop), .. }) = + callee.as_ref() + { + if prop.sym.as_ref() == "all" { + if let Expr::Ident(ident) = obj.as_ref() { + return ident.sym.as_ref() == "Promise"; + } + } + } + } + } + false + } + + /// Extract step name and script from a task function call. + /// Name = function name, script = task_path or function name. + fn extract_step_info_from_task_call(&self, call: &CallExpr) -> Option<(String, String)> { + if let Callee::Expr(callee) = &call.callee { + if let Expr::Ident(ident) = callee.as_ref() { + let name = ident.sym.to_string(); + let script = self + .task_functions + .get(ident.sym.as_ref()) + .and_then(|p| p.clone()) + .unwrap_or_else(|| name.clone()); + return Some((name, script)); + } + } + None + } + + fn expr_to_source(&self, expr: &Expr) -> String { + let span = expr.span(); + self.cm + .span_to_snippet(span) + .unwrap_or_else(|_| "...".to_string()) + } + + fn body_contains_step(&self, stmts: &[Stmt]) -> bool { + stmts.iter().any(|s| self.stmt_contains_step(s)) + } + + fn stmt_contains_step(&self, stmt: &Stmt) -> bool { + match stmt { + Stmt::Expr(expr_stmt) => self.expr_contains_step(&expr_stmt.expr), + Stmt::Decl(Decl::Var(var_decl)) => var_decl.decls.iter().any(|d| { + d.init + .as_ref() + .map_or(false, |init| self.expr_contains_step(init)) + }), + Stmt::If(if_stmt) => { + self.stmt_contains_step(&if_stmt.cons) + || if_stmt + .alt + .as_ref() + .map_or(false, |alt| self.stmt_contains_step(alt)) + } + Stmt::Block(block) => self.body_contains_step(&block.stmts), + Stmt::For(for_stmt) => self.stmt_contains_step(&for_stmt.body), + Stmt::ForIn(for_in) => self.stmt_contains_step(&for_in.body), + Stmt::ForOf(for_of) => self.stmt_contains_step(&for_of.body), + Stmt::While(while_stmt) => self.stmt_contains_step(&while_stmt.body), + Stmt::Try(try_stmt) => { + self.body_contains_step(&try_stmt.block.stmts) + || try_stmt + .handler + .as_ref() + .map_or(false, |h| self.body_contains_step(&h.body.stmts)) + || try_stmt + .finalizer + .as_ref() + .map_or(false, |f| self.body_contains_step(&f.stmts)) + } + Stmt::Return(ret) => ret + .arg + .as_ref() + .map_or(false, |arg| self.expr_contains_step(arg)), + _ => false, + } + } + + fn expr_contains_step(&self, expr: &Expr) -> bool { + if self.is_task_call(expr) { + return true; + } + match expr { + Expr::Await(await_expr) => self.expr_contains_step(&await_expr.arg), + Expr::Call(call) => { + if Self::is_promise_all(&Expr::Call(call.clone())) { + return call.args.iter().any(|a| self.expr_contains_step(&a.expr)); + } + false + } + Expr::Paren(p) => self.expr_contains_step(&p.expr), + _ => false, + } + } + + fn walk_body(&mut self, stmts: &[Stmt]) -> Option<(String, String)> { + let mut first_id: Option = None; + let mut prev_id: Option = None; + + for stmt in stmts { + if let Some((stmt_first, stmt_last)) = self.walk_stmt(stmt) { + if let Some(ref prev) = prev_id { + self.add_edge(prev, &stmt_first, None); + } + if first_id.is_none() { + first_id = Some(stmt_first); + } + prev_id = Some(stmt_last); + } + } + + match (first_id, prev_id) { + (Some(f), Some(l)) => Some((f, l)), + _ => None, + } + } + + fn walk_stmt(&mut self, stmt: &Stmt) -> Option<(String, String)> { + match stmt { + Stmt::Expr(expr_stmt) => self.walk_expr_stmt(&expr_stmt.expr), + Stmt::Decl(Decl::Var(var_decl)) => { + // const result = await task_fn(...) + for decl in &var_decl.decls { + if let Some(init) = &decl.init { + if let Some(result) = self.walk_expr_stmt(init) { + return Some(result); + } + } + } + None + } + Stmt::If(if_stmt) => self.walk_if(if_stmt), + Stmt::For(for_stmt) => self.walk_for_stmt(for_stmt), + Stmt::ForIn(for_in) => self.walk_for_in(for_in), + Stmt::ForOf(for_of) => self.walk_for_of(for_of), + Stmt::While(while_stmt) => self.walk_while(while_stmt), + Stmt::Try(try_stmt) => self.walk_try(try_stmt), + Stmt::Block(block) => self.walk_body(&block.stmts), + Stmt::Return(ret) => self.walk_return(ret), + Stmt::Decl(Decl::Fn(_)) => { + if self.stmt_contains_step(stmt) { + self.errors.push(validation::error_step_in_nested_function( + self.span_line(stmt.span()), + )); + } + None + } + _ => None, + } + } + + fn walk_expr_stmt(&mut self, expr: &Expr) -> Option<(String, String)> { + // await task_fn(...) + if let Expr::Await(await_expr) = expr { + if let Expr::Call(call) = await_expr.arg.as_ref() { + if self.is_task_call(&Expr::Call(call.clone())) { + return self.emit_step(call, expr); + } + } + // await Promise.all([task_fn(...), ...]) + if Self::is_promise_all(&await_expr.arg) { + if let Expr::Call(promise_call) = await_expr.arg.as_ref() { + return self.emit_parallel(promise_call, expr); + } + } + } + + // Bare task_fn() without await + if self.is_task_call(expr) { + self.errors + .push(validation::error_missing_await(self.span_line(expr.span()))); + } + + None + } + + fn emit_step(&mut self, call: &CallExpr, expr: &Expr) -> Option<(String, String)> { + if self.in_try { + self.errors + .push(validation::error_step_in_catch(self.span_line(expr.span()))); + return None; + } + if self.in_while { + self.errors + .push(validation::error_step_in_while(self.span_line(expr.span()))); + return None; + } + if self.in_nested_func { + self.errors.push(validation::error_step_in_nested_function( + self.span_line(expr.span()), + )); + return None; + } + + let (name, script) = self + .extract_step_info_from_task_call(call) + .unwrap_or(("unknown".into(), "unknown".into())); + let id = self.next_id(); + let node_id = self.add_node(DagNode { + id: id.clone(), + node_type: DagNodeType::Step { name: name.clone(), script }, + label: name, + line: self.span_line(expr.span()), + }); + Some((node_id.clone(), node_id)) + } + + fn emit_parallel(&mut self, promise_call: &CallExpr, expr: &Expr) -> Option<(String, String)> { + if self.in_try { + self.errors + .push(validation::error_step_in_catch(self.span_line(expr.span()))); + return None; + } + if self.in_while { + self.errors + .push(validation::error_step_in_while(self.span_line(expr.span()))); + return None; + } + + let line = self.span_line(expr.span()); + let start_id = self.next_id(); + let start_node_id = self.add_node(DagNode { + id: start_id.clone(), + node_type: DagNodeType::ParallelStart, + label: "parallel".to_string(), + line, + }); + + let mut step_ids = Vec::new(); + + // Promise.all takes an array as first argument + if let Some(first_arg) = promise_call.args.first() { + if let Expr::Array(ArrayLit { elems, .. }) = first_arg.expr.as_ref() { + for elem in elems.iter().flatten() { + if let Expr::Call(call) = elem.expr.as_ref() { + if self.is_task_call(&Expr::Call(call.clone())) { + let (name, script) = self + .extract_step_info_from_task_call(call) + .unwrap_or(("unknown".into(), "unknown".into())); + let step_id = self.next_id(); + let node_id = self.add_node(DagNode { + id: step_id.clone(), + node_type: DagNodeType::Step { name: name.clone(), script }, + label: name, + line: self.span_line(elem.expr.span()), + }); + self.add_edge(&start_node_id, &node_id, None); + step_ids.push(node_id); + } + } + } + } + } + + let end_id = self.next_id(); + let end_node_id = self.add_node(DagNode { + id: end_id.clone(), + node_type: DagNodeType::ParallelEnd, + label: "join".to_string(), + line, + }); + + for step_id in &step_ids { + self.add_edge(step_id, &end_node_id, None); + } + + Some((start_node_id, end_node_id)) + } + + fn walk_if(&mut self, if_stmt: &IfStmt) -> Option<(String, String)> { + let has_steps_cons = self.stmt_contains_step(&if_stmt.cons); + let has_steps_alt = if_stmt + .alt + .as_ref() + .map_or(false, |a| self.stmt_contains_step(a)); + + if !has_steps_cons && !has_steps_alt { + return None; + } + + let line = self.span_line(if_stmt.span); + let condition_source = self.expr_to_source(&if_stmt.test); + + let branch_id = self.next_id(); + let branch_node_id = self.add_node(DagNode { + id: branch_id.clone(), + node_type: DagNodeType::Branch { condition_source }, + label: "if".to_string(), + line, + }); + + let mut last_ids = Vec::new(); + + // True branch + if let Some((true_first, true_last)) = self.walk_stmt(&if_stmt.cons) { + self.add_edge(&branch_node_id, &true_first, Some("true".to_string())); + last_ids.push(true_last); + } else { + last_ids.push(branch_node_id.clone()); + } + + // False branch + if let Some(alt) = &if_stmt.alt { + if let Some((else_first, else_last)) = self.walk_stmt(alt) { + self.add_edge(&branch_node_id, &else_first, Some("false".to_string())); + last_ids.push(else_last); + } else { + last_ids.push(branch_node_id.clone()); + } + } + + if last_ids.len() == 1 { + Some((branch_node_id, last_ids.into_iter().next().unwrap())) + } else { + let merge_id = format!("{branch_id}_merge"); + Some((branch_node_id, merge_id)) + } + } + + fn walk_for_stmt(&mut self, for_stmt: &ForStmt) -> Option<(String, String)> { + if !self.stmt_contains_step(&for_stmt.body) { + return None; + } + self.walk_loop_body(&for_stmt.body, for_stmt.span, "for") + } + + fn walk_for_in(&mut self, for_in: &ForInStmt) -> Option<(String, String)> { + if !self.stmt_contains_step(&for_in.body) { + return None; + } + let iter_source = self.expr_to_source(&for_in.right); + self.walk_loop_body_with_iter(&for_in.body, for_in.span, &iter_source) + } + + fn walk_for_of(&mut self, for_of: &ForOfStmt) -> Option<(String, String)> { + if !self.stmt_contains_step(&for_of.body) { + return None; + } + let iter_source = self.expr_to_source(&for_of.right); + self.walk_loop_body_with_iter(&for_of.body, for_of.span, &iter_source) + } + + fn walk_loop_body( + &mut self, + body: &Stmt, + span: swc_common::Span, + _label: &str, + ) -> Option<(String, String)> { + self.walk_loop_body_with_iter(body, span, "...") + } + + fn walk_loop_body_with_iter( + &mut self, + body: &Stmt, + span: swc_common::Span, + iter_source: &str, + ) -> Option<(String, String)> { + let line = self.span_line(span); + let start_id = self.next_id(); + let start_node_id = self.add_node(DagNode { + id: start_id.clone(), + node_type: DagNodeType::LoopStart { iter_source: iter_source.to_string() }, + label: "for".to_string(), + line, + }); + + if let Some((body_first, body_last)) = self.walk_stmt(body) { + self.add_edge(&start_node_id, &body_first, None); + self.add_edge(&body_last, &start_node_id, Some("next".to_string())); + } + + let end_id = self.next_id(); + let end_node_id = self.add_node(DagNode { + id: end_id.clone(), + node_type: DagNodeType::LoopEnd, + label: "end for".to_string(), + line, + }); + self.add_edge(&start_node_id, &end_node_id, Some("done".to_string())); + + Some((start_node_id, end_node_id)) + } + + fn walk_while(&mut self, while_stmt: &WhileStmt) -> Option<(String, String)> { + if self.stmt_contains_step(&while_stmt.body) { + self.errors.push(validation::error_step_in_while( + self.span_line(while_stmt.span), + )); + } + None + } + + fn walk_try(&mut self, try_stmt: &TryStmt) -> Option<(String, String)> { + let has_steps = self.body_contains_step(&try_stmt.block.stmts) + || try_stmt + .handler + .as_ref() + .map_or(false, |h| self.body_contains_step(&h.body.stmts)) + || try_stmt + .finalizer + .as_ref() + .map_or(false, |f| self.body_contains_step(&f.stmts)); + + if has_steps { + self.errors.push(validation::error_step_in_catch( + self.span_line(try_stmt.span), + )); + } + None + } + + fn walk_return(&mut self, ret: &ReturnStmt) -> Option<(String, String)> { + let line = self.span_line(ret.span); + let id = self.next_id(); + let node_id = self.add_node(DagNode { + id: id.clone(), + node_type: DagNodeType::Return, + label: "return".to_string(), + line, + }); + Some((node_id.clone(), node_id)) + } +} + +/// Extract workflow function params (no longer skips ctx) +fn extract_ts_params(params: &[swc_ecma_ast::Param], cm: &Lrc) -> Vec { + let mut result = Vec::new(); + for param in params { + let (name, typ) = match ¶m.pat { + Pat::Ident(BindingIdent { id, type_ann, .. }) => { + let name = id.sym.to_string(); + let typ = type_ann.as_ref().map(|ann| { + cm.span_to_snippet(ann.type_ann.span()) + .unwrap_or_else(|_| "unknown".to_string()) + }); + (name, typ) + } + _ => continue, + }; + result.push(Param { name, typ }); + } + result +} + +pub fn parse_ts_workflow(code: &str) -> Result> { + let cm: Lrc = Default::default(); + let fm = cm.new_source_file(FileName::Custom("workflow.ts".into()).into(), code.into()); + let lexer = Lexer::new( + Syntax::Typescript(TsSyntax::default()), + Default::default(), + StringInput::from(&*fm), + None, + ); + + let mut parser = Parser::new_from(lexer); + let module = parser + .parse_module() + .map_err(|e| vec![CompileError { message: format!("Parse error: {e:?}"), line: 0 }])?; + + // First pass: collect task functions + let task_functions = collect_task_functions(&module); + + // Find: export default workflow(async (...) => { ... }) + // or: export default workflow(async function(...) { ... }) + let mut workflow_body: Option<(&[Stmt], Vec)> = None; + + for item in &module.body { + // export default workflow(async (...) => { ... }) + if let ModuleItem::ModuleDecl(ModuleDecl::ExportDefaultExpr(export)) = item { + if let Some(result) = find_workflow_call(&export.expr, &cm) { + workflow_body = Some(result); + break; + } + } + // const wf = workflow(async (...) => { ... }); export default wf; + if let ModuleItem::ModuleDecl(ModuleDecl::ExportDefaultDecl(export)) = item { + if let DefaultDecl::Fn(_) = &export.decl { + // `export default async function(...) { ... }` — not wrapped in workflow(), skip + } + } + if let ModuleItem::Stmt(Stmt::Decl(Decl::Var(var_decl))) = item { + for decl in &var_decl.decls { + if let Some(init) = &decl.init { + if let Some(result) = find_workflow_call(init, &cm) { + workflow_body = Some(result); + break; + } + } + } + } + } + + let (stmts, params) = workflow_body.ok_or_else(|| { + vec![CompileError { + message: "No workflow() wrapped async function found.".to_string(), + line: 0, + }] + })?; + + let source_hash = compute_source_hash(code); + + let mut walker = TsWacWalker::new(cm, task_functions); + walker.walk_body(stmts); + + if !walker.errors.is_empty() { + return Err(walker.errors); + } + + Ok(WorkflowDag { nodes: walker.nodes, edges: walker.edges, params, source_hash }) +} + +/// Find workflow(async (...) => { ... }) or workflow(async function(...) { ... }) +fn find_workflow_call<'a>(expr: &'a Expr, cm: &Lrc) -> Option<(&'a [Stmt], Vec)> { + if let Expr::Call(call) = expr { + // Check if callee is `workflow` + let is_workflow = match &call.callee { + Callee::Expr(callee_expr) => { + if let Expr::Ident(ident) = callee_expr.as_ref() { + ident.sym.as_ref() == "workflow" + } else { + false + } + } + _ => false, + }; + + if is_workflow { + if let Some(first_arg) = call.args.first() { + return extract_async_fn_body(&first_arg.expr, cm); + } + } + } + None +} + +fn extract_async_fn_body<'a>( + expr: &'a Expr, + cm: &Lrc, +) -> Option<(&'a [Stmt], Vec)> { + match expr { + Expr::Arrow(arrow) if arrow.is_async => { + let params = extract_arrow_params(&arrow.params, cm); + match &*arrow.body { + BlockStmtOrExpr::BlockStmt(block) => Some((&block.stmts, params)), + _ => None, + } + } + Expr::Fn(fn_expr) if fn_expr.function.is_async => { + let params = extract_ts_params(&fn_expr.function.params, cm); + fn_expr + .function + .body + .as_ref() + .map(|body| (body.stmts.as_slice(), params)) + } + Expr::Paren(p) => extract_async_fn_body(&p.expr, cm), + _ => None, + } +} + +/// Extract arrow function params (no longer skips ctx) +fn extract_arrow_params(pats: &[Pat], cm: &Lrc) -> Vec { + let mut result = Vec::new(); + for pat in pats { + match pat { + Pat::Ident(BindingIdent { id, type_ann, .. }) => { + let name = id.sym.to_string(); + let typ = type_ann.as_ref().map(|ann| { + cm.span_to_snippet(ann.type_ann.span()) + .unwrap_or_else(|_| "unknown".to_string()) + }); + result.push(Param { name, typ }); + } + _ => {} + } + } + result +} + +fn extract_string_lit(expr: &Expr) -> Option { + match expr { + Expr::Lit(Lit::Str(s)) => Some(s.value.to_string()), + Expr::Tpl(tpl) if tpl.exprs.is_empty() && tpl.quasis.len() == 1 => { + tpl.quasis.first().map(|q| q.raw.to_string()) + } + _ => None, + } +} + +fn compute_source_hash(code: &str) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(code.as_bytes()); + format!("{:x}", hasher.finalize()) +} diff --git a/backend/parsers/windmill-parser-wac/src/validation.rs b/backend/parsers/windmill-parser-wac/src/validation.rs new file mode 100644 index 0000000000..e3a57c6811 --- /dev/null +++ b/backend/parsers/windmill-parser-wac/src/validation.rs @@ -0,0 +1,64 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct CompileError { + pub message: String, + pub line: usize, +} + +impl std::fmt::Display for CompileError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "line {}: {}", self.line, self.message) + } +} + +pub fn error_step_in_try(line: usize) -> CompileError { + CompileError { + message: + "Task calls inside try/except are not allowed. Steps have built-in error handling." + .to_string(), + line, + } +} + +pub fn error_step_in_while(line: usize) -> CompileError { + CompileError { + message: "Task calls inside while loops are not allowed. Use for loops instead." + .to_string(), + line, + } +} + +pub fn error_step_in_nested_function(line: usize) -> CompileError { + CompileError { + message: "Task calls inside nested functions, closures, or lambdas are not allowed." + .to_string(), + line, + } +} + +pub fn error_step_in_comprehension(line: usize) -> CompileError { + CompileError { message: "Task calls inside comprehensions are not allowed.".to_string(), line } +} + +pub fn error_not_async(line: usize) -> CompileError { + CompileError { message: "Workflow function must be async.".to_string(), line } +} + +pub fn error_missing_await(line: usize) -> CompileError { + CompileError { + message: + "Task calls must be awaited directly or used inside asyncio.gather()/Promise.all()." + .to_string(), + line, + } +} + +pub fn error_step_in_catch(line: usize) -> CompileError { + CompileError { + message: + "Task calls inside catch blocks are not allowed. Steps have built-in error handling." + .to_string(), + line, + } +} diff --git a/backend/parsers/windmill-parser-wac/tests/python_tests.rs b/backend/parsers/windmill-parser-wac/tests/python_tests.rs new file mode 100644 index 0000000000..59f0b59f5c --- /dev/null +++ b/backend/parsers/windmill-parser-wac/tests/python_tests.rs @@ -0,0 +1,266 @@ +use windmill_parser_wac::dag::DagNodeType; +use windmill_parser_wac::python::parse_python_workflow; + +#[test] +fn test_simple_sequential_workflow() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def extract_data(url: str): ... +@task +async def load_data(data: list): ... + +@workflow +async def my_etl(url: str): + raw = await extract_data(url=url) + await load_data(data=raw) + return {"status": "done"} +"#; + + let dag = parse_python_workflow(code).expect("should parse"); + assert_eq!(dag.nodes.len(), 3); // 2 steps + 1 return + assert_eq!(dag.edges.len(), 2); // step0->step1, step1->return + + // Check params (url — no ctx to skip) + assert_eq!(dag.params.len(), 1); + assert_eq!(dag.params[0].name, "url"); + assert_eq!(dag.params[0].typ.as_deref(), Some("str")); + + // Check first step + match &dag.nodes[0].node_type { + DagNodeType::Step { name, script } => { + assert_eq!(name, "extract_data"); + assert_eq!(script, "extract_data"); + } + _ => panic!("expected Step node"), + } + + // Check second step + match &dag.nodes[1].node_type { + DagNodeType::Step { name, script } => { + assert_eq!(name, "load_data"); + assert_eq!(script, "load_data"); + } + _ => panic!("expected Step node"), + } + + // Check return + assert!(matches!(dag.nodes[2].node_type, DagNodeType::Return)); + + // Check source hash is non-empty + assert!(!dag.source_hash.is_empty()); +} + +#[test] +fn test_parallel_workflow() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def extract_data(url: str): ... +@task +async def clean_data(data: list): ... +@task +async def compute_stats(data: list): ... +@task +async def load_to_warehouse(rows: list): ... + +@workflow +async def my_etl(url: str): + raw = await extract_data(url=url) + cleaned, stats = await asyncio.gather( + clean_data(data=raw), + compute_stats(data=raw), + ) + await load_to_warehouse(rows=cleaned) + return {"status": "done"} +"#; + + let dag = parse_python_workflow(code).expect("should parse"); + + // extract, ParallelStart, clean, stats, ParallelEnd, load, return = 7 + assert_eq!(dag.nodes.len(), 7); + + assert!(matches!(dag.nodes[0].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[1].node_type, DagNodeType::ParallelStart)); + assert!(matches!(dag.nodes[2].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[3].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[4].node_type, DagNodeType::ParallelEnd)); + assert!(matches!(dag.nodes[5].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[6].node_type, DagNodeType::Return)); +} + +#[test] +fn test_conditional_workflow() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def send_alert(msg: str): ... +@task +async def load_data(): ... + +@workflow +async def my_etl(count: int): + if count > 100: + await send_alert(msg="large") + await load_data() + return {"done": True} +"#; + + let dag = parse_python_workflow(code).expect("should parse"); + // Branch, notify step, load step, return = 4 + assert_eq!(dag.nodes.len(), 4); + assert!(matches!(dag.nodes[0].node_type, DagNodeType::Branch { .. })); + assert!(matches!(dag.nodes[1].node_type, DagNodeType::Step { .. })); +} + +#[test] +fn test_for_loop_workflow() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def process_item(item: str): ... + +@workflow +async def my_etl(items: list): + for item in items: + await process_item(item=item) + return {"done": True} +"#; + + let dag = parse_python_workflow(code).expect("should parse"); + // LoopStart, step, LoopEnd, return = 4 + assert_eq!(dag.nodes.len(), 4); + assert!(matches!( + dag.nodes[0].node_type, + DagNodeType::LoopStart { .. } + )); + assert!(matches!(dag.nodes[1].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[2].node_type, DagNodeType::LoopEnd)); +} + +#[test] +fn test_reject_step_in_try() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def extract_data(): ... + +@workflow +async def my_etl(): + try: + await extract_data() + except Exception: + pass +"#; + + let result = parse_python_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("try/except")); +} + +#[test] +fn test_reject_step_in_while() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def extract_data(): ... + +@workflow +async def my_etl(): + while True: + await extract_data() +"#; + + let result = parse_python_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("while")); +} + +#[test] +fn test_reject_non_async() { + let code = r#" +from wmill import workflow + +@workflow +def my_etl(): + pass +"#; + + let result = parse_python_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("async")); +} + +#[test] +fn test_reject_missing_await() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task +async def extract_data(): ... + +@workflow +async def my_etl(): + extract_data() +"#; + + let result = parse_python_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("awaited")); +} + +#[test] +fn test_no_workflow_function() { + let code = r#" +async def my_func(): + pass +"#; + + let result = parse_python_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("No @workflow")); +} + +#[test] +fn test_task_with_external_path() { + let code = r#" +import asyncio +from wmill import workflow, task + +@task(path="f/external_script") +async def run_external(x: int): ... + +@workflow +async def my_wf(x: int): + result = await run_external(x=x) + return result +"#; + + let dag = parse_python_workflow(code).expect("should parse"); + assert_eq!(dag.nodes.len(), 2); // 1 step + 1 return (bare `return` is not a step node but walk_return creates one) + match &dag.nodes[0].node_type { + DagNodeType::Step { name, script } => { + assert_eq!(name, "run_external"); + assert_eq!(script, "f/external_script"); + } + _ => panic!("expected Step node"), + } +} diff --git a/backend/parsers/windmill-parser-wac/tests/ts_tests.rs b/backend/parsers/windmill-parser-wac/tests/ts_tests.rs new file mode 100644 index 0000000000..949f326b90 --- /dev/null +++ b/backend/parsers/windmill-parser-wac/tests/ts_tests.rs @@ -0,0 +1,245 @@ +use windmill_parser_wac::dag::DagNodeType; +use windmill_parser_wac::typescript::parse_ts_workflow; + +#[test] +fn test_simple_sequential_ts_workflow() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const extract_data = task(async (url: string) => {}); +const load_data = task(async (data: any) => {}); + +export default workflow(async (url: string) => { + const raw = await extract_data(url); + await load_data(raw); + return { status: "done" }; +}); +"#; + + let dag = parse_ts_workflow(code).expect("should parse"); + assert_eq!(dag.nodes.len(), 3); // 2 steps + 1 return + assert_eq!(dag.edges.len(), 2); + + // Check params (url — no ctx to skip) + assert_eq!(dag.params.len(), 1); + assert_eq!(dag.params[0].name, "url"); + assert_eq!(dag.params[0].typ.as_deref(), Some("string")); + + match &dag.nodes[0].node_type { + DagNodeType::Step { name, script } => { + assert_eq!(name, "extract_data"); + assert_eq!(script, "extract_data"); + } + _ => panic!("expected Step node"), + } + + match &dag.nodes[1].node_type { + DagNodeType::Step { name, script } => { + assert_eq!(name, "load_data"); + assert_eq!(script, "load_data"); + } + _ => panic!("expected Step node"), + } + + assert!(matches!(dag.nodes[2].node_type, DagNodeType::Return)); + assert!(!dag.source_hash.is_empty()); +} + +#[test] +fn test_parallel_ts_workflow() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const extract_data = task(async (url: string) => {}); +const clean_data = task(async (data: any) => {}); +const compute_stats = task(async (data: any) => {}); +const load_to_warehouse = task(async (rows: any) => {}); + +export default workflow(async (url: string) => { + const raw = await extract_data(url); + + const [cleaned, stats] = await Promise.all([ + clean_data(raw), + compute_stats(raw), + ]); + + await load_to_warehouse(cleaned); + return { status: "done", rows: stats.rowCount }; +}); +"#; + + let dag = parse_ts_workflow(code).expect("should parse"); + // extract, ParallelStart, clean, stats, ParallelEnd, load, return = 7 + assert_eq!(dag.nodes.len(), 7); + + assert!(matches!(dag.nodes[0].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[1].node_type, DagNodeType::ParallelStart)); + assert!(matches!(dag.nodes[2].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[3].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[4].node_type, DagNodeType::ParallelEnd)); + assert!(matches!(dag.nodes[5].node_type, DagNodeType::Step { .. })); + assert!(matches!(dag.nodes[6].node_type, DagNodeType::Return)); +} + +#[test] +fn test_conditional_ts_workflow() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const send_alert = task(async (msg: string) => {}); +const load_data = task(async () => {}); + +export default workflow(async (count: number) => { + if (count > 100) { + await send_alert("large"); + } + await load_data(); + return { done: true }; +}); +"#; + + let dag = parse_ts_workflow(code).expect("should parse"); + // Branch, notify, load, return = 4 + assert_eq!(dag.nodes.len(), 4); + assert!(matches!(dag.nodes[0].node_type, DagNodeType::Branch { .. })); +} + +#[test] +fn test_for_of_ts_workflow() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const process_item = task(async (item: string) => {}); + +export default workflow(async (items: string[]) => { + for (const item of items) { + await process_item(item); + } + return { done: true }; +}); +"#; + + let dag = parse_ts_workflow(code).expect("should parse"); + // LoopStart, step, LoopEnd, return = 4 + assert_eq!(dag.nodes.len(), 4); + assert!(matches!( + dag.nodes[0].node_type, + DagNodeType::LoopStart { .. } + )); +} + +#[test] +fn test_reject_step_in_try_catch() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const extract_data = task(async () => {}); + +export default workflow(async () => { + try { + await extract_data(); + } catch (e) { + console.log(e); + } +}); +"#; + + let result = parse_ts_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("catch")); +} + +#[test] +fn test_reject_step_in_while_ts() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const extract_data = task(async () => {}); + +export default workflow(async () => { + while (true) { + await extract_data(); + } +}); +"#; + + let result = parse_ts_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("while")); +} + +#[test] +fn test_reject_missing_await_ts() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const extract_data = task(async () => {}); + +export default workflow(async () => { + extract_data(); +}); +"#; + + let result = parse_ts_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("awaited")); +} + +#[test] +fn test_no_workflow_wrapper() { + let code = r#" +export default async function main(ctx: any) { + return {}; +} +"#; + + let result = parse_ts_workflow(code); + assert!(result.is_err()); + let errors = result.unwrap_err(); + assert!(errors[0].message.contains("No workflow()")); +} + +#[test] +fn test_variable_declaration_with_step() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const compute = task(async () => {}); + +export default workflow(async () => { + const result = await compute(); + return result; +}); +"#; + + let dag = parse_ts_workflow(code).expect("should parse"); + assert_eq!(dag.nodes.len(), 2); // step + return + assert!(matches!(dag.nodes[0].node_type, DagNodeType::Step { .. })); +} + +#[test] +fn test_task_with_external_path() { + let code = r#" +import { workflow, task } from "windmill-client"; + +const run_external = task("f/external_script", async (x: number) => {}); + +export default workflow(async (x: number) => { + const result = await run_external(x); + return result; +}); +"#; + + let dag = parse_ts_workflow(code).expect("should parse"); + assert_eq!(dag.nodes.len(), 2); // step + return + match &dag.nodes[0].node_type { + DagNodeType::Step { name, script } => { + assert_eq!(name, "run_external"); + assert_eq!(script, "f/external_script"); + } + _ => panic!("expected Step node"), + } +} diff --git a/backend/parsers/windmill-parser-wasm/Cargo.toml b/backend/parsers/windmill-parser-wasm/Cargo.toml index 1a0d425704..c4891c091b 100644 --- a/backend/parsers/windmill-parser-wasm/Cargo.toml +++ b/backend/parsers/windmill-parser-wasm/Cargo.toml @@ -38,6 +38,7 @@ csharp-parser = [ "dep:windmill-parser-csharp"] nu-parser = [ "dep:windmill-parser-nu"] java-parser = [ "dep:windmill-parser-java"] ruby-parser = [ "dep:windmill-parser-ruby"] +wac-parser = [ "dep:windmill-parser-wac"] [dependencies] anyhow.workspace = true @@ -55,6 +56,7 @@ windmill-parser-csharp = { workspace = true, optional = true } windmill-parser-nu = { workspace = true, optional = true } windmill-parser-java = { workspace = true, optional = true } windmill-parser-ruby = { workspace = true, optional = true } +windmill-parser-wac = { workspace = true, optional = true } wasm-bindgen.workspace = true serde_json.workspace = true diff --git a/backend/parsers/windmill-parser-wasm/build.nu b/backend/parsers/windmill-parser-wasm/build.nu index d954366113..cbed629c58 100755 --- a/backend/parsers/windmill-parser-wasm/build.nu +++ b/backend/parsers/windmill-parser-wasm/build.nu @@ -56,6 +56,12 @@ const targets = [ features: "ruby-parser", env: "tree-sitter", }, + { + ident: "wac", + desc: "Workflow-as-Code", + features: "wac-parser", + env: "default", + }, # ^^^ Add new entry here ^^^ ]; # NOTE: This is legacy command for building all, but it is not more used diff --git a/backend/parsers/windmill-parser-wasm/src/lib.rs b/backend/parsers/windmill-parser-wasm/src/lib.rs index 61cc11d81d..634f2348e1 100644 --- a/backend/parsers/windmill-parser-wasm/src/lib.rs +++ b/backend/parsers/windmill-parser-wasm/src/lib.rs @@ -223,4 +223,11 @@ pub fn parse_assets_ansible(code: &str) -> String { } } +#[cfg(feature = "wac-parser")] +#[wasm_bindgen] +pub fn parse_workflow_as_code(code: &str, language: &str) -> String { + let result = windmill_parser_wac::parse_workflow(code, language); + serde_json::to_string(&result).unwrap_or_else(|_| "{\"type\": \"error\"}".to_string()) +} + // for related places search: ADD_NEW_LANG diff --git a/backend/test_wac_e2e.sh b/backend/test_wac_e2e.sh new file mode 100755 index 0000000000..4379ffde42 --- /dev/null +++ b/backend/test_wac_e2e.sh @@ -0,0 +1,157 @@ +#!/usr/bin/env bash +# E2E test for WAC v2 workflow-as-code suspend/resume lifecycle +set -euo pipefail + +BASE_URL="${BASE_URL:-http://localhost:8070}" +TOKEN="${WM_TOKEN:-}" +WORKSPACE="dev" +TIMEOUT=60 # seconds + +# Get auth token if not set +if [ -z "$TOKEN" ]; then + TOKEN=$(curl -s "${BASE_URL}/api/auth/login" \ + -H "Content-Type: application/json" \ + -d '{"email":"admin@windmill.dev","password":"changeme"}' | tr -d '"') +fi + +echo "=== WAC v2 E2E Test ===" +echo "Base URL: $BASE_URL" +echo "" + +WAC_CODE='import { task, workflow } from "windmill-client@1.999.19"; + +const double = task(async (x: number): Promise => { + console.log("[double] START at " + new Date().toISOString()); + await new Promise(r => setTimeout(r, 2000)); + console.log("[double] END at " + new Date().toISOString()); + return x * 2; +}); + +const increment = task(async (x: number): Promise => { + console.log("[increment] START at " + new Date().toISOString()); + await new Promise(r => setTimeout(r, 2000)); + console.log("[increment] END at " + new Date().toISOString()); + return x + 1; +}); + +export const main = workflow(async (x: number = 10) => { + const [doubled, incremented] = await Promise.all([ + double(x), + increment(x), + ]); + const final_result = await double(incremented); + return { doubled, incremented, final_result }; +});' + +echo "Step 1: Submitting preview job..." +JOB_ID=$(curl -s "${BASE_URL}/api/w/${WORKSPACE}/jobs/run/preview" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $TOKEN" \ + -d "$(jq -n --arg code "$WAC_CODE" '{ + content: $code, + language: "bun", + args: {"x": 10} + }')" | tr -d '"') + +echo "Job ID: $JOB_ID" + +if [ -z "$JOB_ID" ] || [ "$JOB_ID" = "null" ]; then + echo "FAIL: Could not create job" + exit 1 +fi + +echo "" +echo "Step 2: Polling for completion (timeout: ${TIMEOUT}s)..." + +START=$SECONDS +LAST_STATUS="" +while true; do + ELAPSED=$((SECONDS - START)) + if [ $ELAPSED -gt $TIMEOUT ]; then + echo "FAIL: Timed out after ${TIMEOUT}s" + # Dump job state for debugging + echo "" + echo "=== Debug info ===" + echo "Parent job queue state:" + source /home/rfiszel/windmill__worktrees/workflows-as-code-v2/.env.local + psql "$DATABASE_URL" -c "SELECT id, running, suspend, suspend_until, canceled_by FROM v2_job_queue WHERE id = '$JOB_ID'::uuid" 2>/dev/null + echo "Child jobs:" + psql "$DATABASE_URL" -c "SELECT id, running, suspend, created_at FROM v2_job_queue WHERE parent_job = '$JOB_ID'::uuid ORDER BY created_at" 2>/dev/null + echo "Completed children:" + psql "$DATABASE_URL" -c "SELECT id FROM completed_job WHERE parent_job = '$JOB_ID'::uuid" 2>/dev/null + echo "Checkpoint:" + psql "$DATABASE_URL" -c "SELECT workflow_as_code_status->'_checkpoint' FROM v2_job_status WHERE id = '$JOB_ID'::uuid" 2>/dev/null + echo "Total child count:" + psql "$DATABASE_URL" -c "SELECT count(*) FROM v2_job WHERE parent_job = '$JOB_ID'::uuid" 2>/dev/null + exit 1 + fi + + # Check completed job + RESULT=$(curl -s "${BASE_URL}/api/w/${WORKSPACE}/jobs_u/completed/get_result/${JOB_ID}" \ + -H "Authorization: Bearer $TOKEN" 2>/dev/null) + HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" "${BASE_URL}/api/w/${WORKSPACE}/jobs_u/completed/get_result/${JOB_ID}" \ + -H "Authorization: Bearer $TOKEN" 2>/dev/null) + + if [ "$HTTP_CODE" = "200" ]; then + echo "Job completed in ${ELAPSED}s!" + echo "" + echo "Step 3: Checking result..." + echo "Result: $RESULT" + + # Validate + DOUBLED=$(echo "$RESULT" | jq -r '.doubled // empty') + INCREMENTED=$(echo "$RESULT" | jq -r '.incremented // empty') + FINAL=$(echo "$RESULT" | jq -r '.final_result // empty') + + PASS=true + if [ "$DOUBLED" != "20" ]; then + echo "FAIL: doubled = $DOUBLED, expected 20" + PASS=false + fi + if [ "$INCREMENTED" != "11" ]; then + echo "FAIL: incremented = $INCREMENTED, expected 11" + PASS=false + fi + if [ "$FINAL" != "22" ]; then + echo "FAIL: final_result = $FINAL, expected 22" + PASS=false + fi + + if $PASS; then + echo "PASS: All values correct!" + # Check no excessive child jobs + source /home/rfiszel/windmill__worktrees/workflows-as-code-v2/.env.local 2>/dev/null + CHILD_COUNT=$(psql -tA "$DATABASE_URL" -c "SELECT count(*) FROM v2_job WHERE parent_job = '$JOB_ID'::uuid" 2>/dev/null) + echo "Total child jobs created: $CHILD_COUNT (expected: 3)" + if [ "$CHILD_COUNT" -gt "3" ]; then + echo "WARN: More children than expected ($CHILD_COUNT > 3)" + fi + exit 0 + else + exit 1 + fi + fi + + # Show progress + STATUS=$(curl -s "${BASE_URL}/api/w/${WORKSPACE}/jobs_u/get/${JOB_ID}" \ + -H "Authorization: Bearer $TOKEN" 2>/dev/null | jq -r '.type // empty') + if [ "$STATUS" != "$LAST_STATUS" ]; then + echo " [${ELAPSED}s] Status: $STATUS" + LAST_STATUS="$STATUS" + fi + + # Check for runaway child creation + source /home/rfiszel/windmill__worktrees/workflows-as-code-v2/.env.local 2>/dev/null + CHILD_COUNT=$(psql -tA "$DATABASE_URL" -c "SELECT count(*) FROM v2_job WHERE parent_job = '$JOB_ID'::uuid" 2>/dev/null) + if [ "$CHILD_COUNT" -gt "10" ]; then + echo "FAIL: Runaway child creation detected! $CHILD_COUNT children (expected 3)" + echo "" + echo "=== Debug info ===" + psql "$DATABASE_URL" -c "SELECT id, running, suspend, suspend_until FROM v2_job_queue WHERE id = '$JOB_ID'::uuid" 2>/dev/null + psql "$DATABASE_URL" -c "SELECT id, running, suspend, created_at FROM v2_job_queue WHERE parent_job = '$JOB_ID'::uuid ORDER BY created_at LIMIT 20" 2>/dev/null + psql "$DATABASE_URL" -c "SELECT workflow_as_code_status->'_checkpoint' FROM v2_job_status WHERE id = '$JOB_ID'::uuid" 2>/dev/null + exit 1 + fi + + sleep 1 +done diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index 1750276778..8960dfb2a1 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -18803,7 +18803,6 @@ components: required: - path - summary - - description - content - language diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index ac5a9a306b..184602dc87 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -2255,12 +2255,13 @@ async fn resume_suspended_job_internal( let value = value.unwrap_or(serde_json::Value::Null); verify_suspended_secret(&w_id, &db, job_id, resume_id, &approver, secret).await?; - // Get flow info - works for both step-level (job_id is a step) and flow-level (job_id is the flow) - let (flow_info, is_flow_level) = get_flow_info_for_resume(job_id, &db).await?; + // Get flow info - works for step-level, flow-level, and WAC approval + let (flow_info, is_flow_level, is_wac) = get_flow_info_for_resume(job_id, &db).await?; // For step-level resumes, verify user auth and flow status // For flow-level resumes (pre-approvals), the flow might not be at a suspended step yet - if !is_flow_level { + // For WAC approvals, skip flow status checks (there is no flow) + if !is_flow_level && !is_wac { let parent_flow = GetQuery::new() .without_logs() .without_code() @@ -2322,6 +2323,16 @@ async fn resume_suspended_job_internal( ) .execute(&mut *tx) .await?; + } else if is_wac { + // WAC approval: decrement suspend counter directly on the WAC parent job + if flow_info.suspend > 0 { + sqlx::query!( + "UPDATE v2_job_queue SET suspend = GREATEST(suspend - 1, 0) WHERE id = $1", + flow_info.id, + ) + .execute(&mut *tx) + .await?; + } } else if is_flow_level { // For flow-level resumes, decrement the suspend counter if the flow is currently suspended // The approval will be matched when the worker checks for resumes (both step-level and flow-level) @@ -2479,10 +2490,15 @@ struct FlowInfo { email: Option, } -/// Get flow info from either a step job (by looking up its parent) or a flow job directly. -/// Returns (FlowInfo, is_flow_level) where is_flow_level indicates if job_id was a flow job. -async fn get_flow_info_for_resume(job_id: Uuid, db: &DB) -> error::Result<(FlowInfo, bool)> { - // Single query that determines if job_id is a flow or step, and fetches the appropriate flow info +/// Get flow info from either a step job (by looking up its parent), a flow job directly, +/// or a WAC workflow job (self-suspended for approval). +/// Returns (FlowInfo, is_flow_level, is_wac) where: +/// - is_flow_level: job_id was a flow job (pre-approval) +/// - is_wac: job_id is a WAC workflow suspended for approval (target is itself) +async fn get_flow_info_for_resume(job_id: Uuid, db: &DB) -> error::Result<(FlowInfo, bool, bool)> { + // Single query that determines if job_id is a flow, step, or WAC job, + // and fetches the appropriate suspended job info. + // For WAC jobs (no parent, not a flow), the job itself is the suspended target. let result = sqlx::query!( r#" WITH job_info AS ( @@ -2496,14 +2512,15 @@ async fn get_flow_info_for_resume(job_id: Uuid, db: &DB) -> error::Result<(FlowI q.suspend AS "suspend!", j.runnable_path AS script_path, j.permissioned_as_email AS email, - (ji.kind IN ('flow', 'flowpreview')) AS "is_flow_level!" + (ji.kind IN ('flow', 'flowpreview')) AS "is_flow_level!", + (ji.kind NOT IN ('flow', 'flowpreview') AND q.id = ji.id) AS "is_wac!" FROM job_info ji JOIN v2_job_queue q ON q.id = CASE WHEN ji.kind IN ('flow', 'flowpreview') THEN ji.id - ELSE ji.parent_job + ELSE COALESCE(ji.parent_job, ji.id) END JOIN v2_job j ON j.id = q.id - JOIN v2_job_status s ON s.id = q.id + LEFT JOIN v2_job_status s ON s.id = q.id FOR UPDATE OF q "#, job_id, @@ -2520,7 +2537,7 @@ async fn get_flow_info_for_resume(job_id: Uuid, db: &DB) -> error::Result<(FlowI email: Some(result.email), }; - Ok((flow_info, result.is_flow_level)) + Ok((flow_info, result.is_flow_level, result.is_wac)) } async fn get_suspended_flow_info<'c>( diff --git a/backend/windmill-common/src/error.rs b/backend/windmill-common/src/error.rs index eaf7d857c1..49dcf25450 100644 --- a/backend/windmill-common/src/error.rs +++ b/backend/windmill-common/src/error.rs @@ -78,6 +78,8 @@ pub enum Error { AIError(String), #[error("{0}")] AlreadyCompleted(String), + #[error("WAC job suspended: {0}")] + WacSuspended(String), #[error("Find python error: {0}")] FindPythonError(String), #[error("Problem with arguments: {0}")] @@ -108,6 +110,7 @@ impl Error { Self::JsonErr(_) => "JsonErr", Self::AIError(_) => "AIError", Self::AlreadyCompleted(_) => "AlreadyCompleted", + Self::WacSuspended(_) => "WacSuspended", Self::FindPythonError(_) => "FindPythonError", Self::ArgumentErr(_) => "ArgumentErr", Self::Generic(_, _) => "Generic", diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 806e3e99a2..d8cb3f4de4 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -818,7 +818,7 @@ pub async fn add_completed_job( flow_is_done: bool, duration: Option, from_cache: bool, -) -> Result<(Uuid, i64), Error> { +) -> Result<(Uuid, i64, Option), Error> { // tracing::error!("Start"); // let start = tokio::time::Instant::now(); @@ -830,7 +830,7 @@ pub async fn add_completed_job( } let result_columns = result_columns.as_ref(); - let (opt_uuid, duration, _skip_downstream_error_handlers) = (|| { + let (opt_uuid, duration, _skip_downstream_error_handlers, wac_job_ids) = (|| { commit_completed_job( db, completed_job, @@ -866,7 +866,7 @@ pub async fn add_completed_job( // if scheduling next job failed, return the job_id early to ensure the job get retried after a timeout if let Some(job_id) = opt_uuid { - return Ok((job_id, duration)); + return Ok((job_id, duration, None)); } #[cfg(feature = "cloud")] @@ -887,7 +887,7 @@ pub async fn add_completed_job( // tracing::error!("4 {:?}", start.elapsed()); - Ok((completed_job.id, duration)) + Ok((completed_job.id, duration, wac_job_ids)) } async fn commit_completed_job( @@ -902,7 +902,7 @@ async fn commit_completed_job( flow_is_done: bool, duration: Option, from_cache: bool, -) -> windmill_common::error::Result<(Option, i64, bool)> { +) -> windmill_common::error::Result<(Option, i64, bool, Option)> { // let start = std::time::Instant::now(); let mut tx = db.begin().warn_after_seconds(10).await?; @@ -1003,25 +1003,31 @@ async fn commit_completed_job( .map_err(|e| Error::InternalErr(format!("Could not update job labels: {e:#}")))?; } + let mut wac_job_ids: Option = None; if !completed_job.is_flow_step() { if let Some(parent_job) = completed_job.parent_job { - let _ = sqlx::query_scalar!( - "UPDATE v2_job_status SET + // Only update WAC parents (v1 or v2). The WHERE condition skips + // non-WAC parents entirely (error handlers, run_script children, etc.). + // Also returns pending_steps.job_ids so WAC v2 child completion + // doesn't need a separate read. + let row = sqlx::query_scalar!( + r#"UPDATE v2_job_status SET workflow_as_code_status = jsonb_set( jsonb_set( - COALESCE(workflow_as_code_status, '{}'::jsonb), + workflow_as_code_status, array[$1], COALESCE(workflow_as_code_status->$1, '{}'::jsonb) ), array[$1, 'duration_ms'], to_jsonb($2::bigint) ) - WHERE id = $3", + WHERE id = $3 AND workflow_as_code_status IS NOT NULL + RETURNING workflow_as_code_status->'_checkpoint'->'pending_steps'->'job_ids' AS "job_ids: serde_json::Value""#, &completed_job.id.to_string(), duration, parent_job ) - .execute(&mut *tx) + .fetch_optional(&mut *tx) .warn_after_seconds(10) .await .inspect_err(|e| { @@ -1029,7 +1035,10 @@ async fn commit_completed_job( "Could not update parent job `duration_ms` in workflow as code status: {}", e, ) - }); + }) + .ok() + .flatten(); + wac_job_ids = row.flatten(); } } // tracing::error!("Added completed job {:#?}", queued_job); @@ -1250,14 +1259,14 @@ async fn commit_completed_job( completed_job.id ); // tracing::info!("completed job: {:?}", start.elapsed().as_micros()); - Ok((None, duration, _skip_downstream_error_handlers)) + Ok((None, duration, _skip_downstream_error_handlers, wac_job_ids)) } async fn check_result_size( db: &Pool, queued_job: &MiniCompletedJob, result: Json<&T>, -) -> Option, i64, bool), Error>> { +) -> Option, i64, bool, Option), Error>> { let result_size = result.size() / 1024 / 1024; if result_size > 2 { if result_size > *MAX_RESULT_SIZE_MB { diff --git a/backend/windmill-worker/nsjail/run.bun.config.proto b/backend/windmill-worker/nsjail/run.bun.config.proto index 3ba8c73257..afd5c42ba9 100644 --- a/backend/windmill-worker/nsjail/run.bun.config.proto +++ b/backend/windmill-worker/nsjail/run.bun.config.proto @@ -142,6 +142,13 @@ mount { rw: true } +mount { + src: "{JOB_DIR}/checkpoint.json" + dst: "/tmp/{LANG}/checkpoint.json" + is_bind: true + mandatory: false +} + mount { src: "{JOB_DIR}/result.json" dst: "/tmp/{LANG}/result.json" diff --git a/backend/windmill-worker/src/ai/image_handler.rs b/backend/windmill-worker/src/ai/image_handler.rs index 7eb6bb6597..63d8aeaec3 100644 --- a/backend/windmill-worker/src/ai/image_handler.rs +++ b/backend/windmill-worker/src/ai/image_handler.rs @@ -2,8 +2,8 @@ use base64::Engine; use futures; use ulid; use windmill_common::{client::AuthedClient, error::Error}; -use windmill_types::s3::S3Object; use windmill_queue::MiniPulledJob; +use windmill_types::s3::S3Object; use crate::ai::types::*; diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs index 04f1b4b548..73010b5ba1 100644 --- a/backend/windmill-worker/src/ai/query_builder.rs +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -1,7 +1,5 @@ use async_trait::async_trait; -use windmill_common::{ - client::AuthedClient, error::Error, worker::Connection, -}; +use windmill_common::{client::AuthedClient, error::Error, worker::Connection}; use windmill_queue::MiniPulledJob; use windmill_types::s3::S3Object; diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs index eebe09bc77..631a8ab113 100644 --- a/backend/windmill-worker/src/ai/types.rs +++ b/backend/windmill-worker/src/ai/types.rs @@ -20,8 +20,8 @@ use windmill_common::{ flow_status::AgentAction, flows::FlowModule, }; -use windmill_types::s3::S3Object; use windmill_parser::Typ; +use windmill_types::s3::S3Object; // Re-export shared types from windmill_common::ai_types pub use windmill_common::ai_types::{ @@ -1603,10 +1603,7 @@ mod tests { schema.sanitize_for_google(); - assert!( - schema.multiple_of.is_none(), - "multipleOf should be removed" - ); + assert!(schema.multiple_of.is_none(), "multipleOf should be removed"); } #[test] @@ -1639,7 +1636,10 @@ mod tests { assert!(schema.default.is_none()); let value_prop = schema.properties.as_ref().unwrap().get("value").unwrap(); - assert!(value_prop.default.is_none(), "nested default should be removed"); + assert!( + value_prop.default.is_none(), + "nested default should be removed" + ); assert!( value_prop.exclusive_minimum.is_none(), "nested exclusiveMinimum should be removed" @@ -1652,7 +1652,10 @@ mod tests { value_prop.multiple_of.is_none(), "nested multipleOf should be removed" ); - assert!(value_prop.r#const.is_none(), "nested const should be removed"); + assert!( + value_prop.r#const.is_none(), + "nested const should be removed" + ); assert!(schema.properties.is_some()); assert!(matches!(&schema.r#type, Some(SchemaType::Single(t)) if t == "object")); diff --git a/backend/windmill-worker/src/bigquery_executor.rs b/backend/windmill-worker/src/bigquery_executor.rs index 2f748fe22c..b092e135b6 100644 --- a/backend/windmill-worker/src/bigquery_executor.rs +++ b/backend/windmill-worker/src/bigquery_executor.rs @@ -6,9 +6,9 @@ use reqwest::Client; use serde_json::{json, value::RawValue, Value}; use windmill_common::client::AuthedClient; use windmill_common::error::to_anyhow; -use windmill_object_store::convert_json_line_stream; use windmill_common::worker::{Connection, SqlResultCollectionStrategy}; use windmill_common::{error::Error, worker::to_raw_value}; +use windmill_object_store::convert_json_line_stream; use windmill_parser_sql::{ parse_bigquery_sig, parse_db_resource, parse_s3_mode, parse_sql_blocks, parse_sql_statement_named_params, diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 6ada66c4bb..16900dad8b 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -1051,6 +1051,37 @@ pub async fn handle_bun_job( let apply_preprocessor = job.flow_step_id.as_deref() != Some("preprocessor") && job.preprocessed == Some(false); + let is_wac_v2 = main_override.is_none() && crate::wac_executor::is_wac_v2_ts(inner_content); + + // For WAC v2, inject variable names into unnamed task() calls so the + // runtime can use them for step naming (timeline, graph). + // `const double = task(async ...` → `const double = task("double", async ...` + // Also handles: export const, let, var, and optional generic type parameters. + // Skips calls that already have a string argument: `task("path", async ...` + let inner_content = if is_wac_v2 { + use regex::Regex; + use std::borrow::Cow; + lazy_static::lazy_static! { + static ref TASK_RE: Regex = + Regex::new(r#"(?m)((?:export\s+)?(?:const|let|var)\s+)(\w+)(\s*=\s*task\s*(?:<[^>]*>)?\s*\(\s*)(async\b)"#).unwrap(); + } + let replaced = TASK_RE.replace_all(inner_content, r#"${1}${2}${3}"${2}", ${4}"#); + match replaced { + Cow::Borrowed(_) => inner_content.to_string(), + Cow::Owned(s) => s, + } + } else { + inner_content.to_string() + }; + let inner_content = inner_content.as_str(); + + // WAC v2 scripts can't use bundle caching because the wrapper imports + // windmill-client from node_modules, which isn't available in bundle mode + if is_wac_v2 && has_bundle_cache { + has_bundle_cache = false; + let _ = write_file(job_dir, "main.ts", inner_content)?; + } + let mut format = BundleFormat::Cjs; if has_bundle_cache { let target; @@ -1184,13 +1215,20 @@ pub async fn handle_bun_job( return Ok(()) as error::Result<()>; } // let mut start = Instant::now(); - let args = windmill_parser_ts::parse_deno_signature( - inner_content, - true, - false, - main_override.map(ToString::to_string), - )? - .args; + let args = if is_wac_v2 { + // For WAC v2, try to parse "main" args; if that fails, try the default export + windmill_parser_ts::parse_deno_signature(inner_content, true, false, None) + .unwrap_or_default() + .args + } else { + windmill_parser_ts::parse_deno_signature( + inner_content, + true, + false, + main_override.map(ToString::to_string), + )? + .args + }; let pre_args = if apply_preprocessor { Some( @@ -1227,6 +1265,14 @@ pub async fn handle_bun_job( // we cannot use Bun.read and Bun.write because it results in an EBADF error on cloud let main_name = main_override.unwrap_or("main"); + // For WAC child jobs where the parser can't find params (task-wrapped consts), + // fall back to passing arg values directly (filtering out internal fields) + let child_spread = if spread.is_empty() && main_override.is_some() { + "Object.values(Object.fromEntries(Object.entries(args).filter(([k]) => !k.startsWith('_'))))".to_string() + } else { + "argsObjToArr(args)".to_string() + }; + let main_import = if codebase.is_some() || has_bundle_cache { "./main.js" } else { @@ -1250,8 +1296,116 @@ pub async fn handle_bun_job( "".to_string() }; - let wrapper_content = format!( - r#" + let wac_spread = if spread.is_empty() { + "Object.values(args)".to_string() + } else { + format!("argsObjToArr(args)") + }; + + let wrapper_content = if is_wac_v2 { + format!( + r#" +import * as Main from "{main_import}"; +import {{ WorkflowCtx, StepSuspend, setWorkflowCtx }} from "windmill-client"; + +import * as fs from "fs/promises"; + +let args = await fs.readFile('args.json', {{ encoding: 'utf8' }}).then(JSON.parse); +const checkpoint = JSON.parse(await fs.readFile('checkpoint.json', {{ encoding: 'utf8' }})); + +function argsObjToArr({{ {spread} }}) {{ + return [ {spread} ]; +}} + +BigInt.prototype.toJSON = function () {{ + return this.toString(); +}}; + +// Find the workflow entrypoint (export default) +let workflowFn = Main.default; +if (!workflowFn || !workflowFn._is_workflow) {{ + for (const key of Object.keys(Main)) {{ + if (Main[key]?._is_workflow) {{ + workflowFn = Main[key]; + break; + }} + }} +}} +if (!workflowFn) {{ + throw new Error("No workflow() entrypoint found. Wrap your main function with workflow()."); +}} + +async function run() {{ + {dates} + {preprocessor} + const argsArr = {wac_spread}; + + const ctx = new WorkflowCtx(checkpoint); + setWorkflowCtx(ctx); + + try {{ + const result = await workflowFn(...argsArr); + setWorkflowCtx(null); + // Flush any unawaited tasks (e.g. forgotten await on last statement) + const trailing = ctx._flushPending(); + if (trailing.length > 0) {{ + return {{ type: "dispatch", mode: trailing.length > 1 ? "parallel" : "sequential", steps: trailing }}; + }} + return {{ type: "complete", result: result ?? null }}; + }} catch (e) {{ + setWorkflowCtx(null); + if (e?.name === "StepSuspend" || e instanceof StepSuspend) {{ + const dispatch = e.dispatchInfo ?? e.dispatch_info ?? {{}}; + if (dispatch.mode === "step_complete") {{ + return {{ type: "complete", result: dispatch.result ?? null }}; + }} + if (dispatch.mode === "inline_checkpoint") {{ + return {{ type: "inline_checkpoint", key: dispatch.key, result: dispatch.result ?? null }}; + }} + if (dispatch.mode === "approval") {{ + return {{ type: "approval", key: dispatch.key, timeout: dispatch.timeout, form: dispatch.form }}; + }} + if (dispatch.mode === "sleep") {{ + return {{ type: "sleep", key: dispatch.key, seconds: dispatch.seconds }}; + }} + return {{ type: "dispatch", mode: dispatch.mode ?? "sequential", steps: dispatch.steps ?? [] }}; + }} + throw e; + }} +}} + +try {{ + const output = await run(); + const output_json = JSON.stringify(output, (key, value) => + typeof value === 'undefined' ? null : value + ); + await fs.writeFile("result.json", output_json); + process.exit(0); +}} catch(e) {{ + console.error(e); + let err = {{ message: e.message, name: e.name, stack: e.stack }}; + let step_id = process.env.WM_FLOW_STEP_ID; + if (step_id) {{ + err["step_id"] = step_id; + }} + const extra = {{}}; + Object.getOwnPropertyNames(e).forEach((key) => {{ + if (['line', 'name', 'stack', 'column', 'message', 'sourceURL', 'originalLine', 'originalColumn'].includes(key)) {{ + return; + }} + extra[key] = e[key]; + }}); + if (Object.keys(extra).length > 0) {{ + err["extra"] = extra; + }} + await fs.writeFile("result.json", JSON.stringify(err)); + process.exit(1); +}} + "#, + ) + } else { + format!( + r#" import * as Main from "{main_import}"; import * as fs from "fs/promises"; @@ -1273,11 +1427,14 @@ BigInt.prototype.toJSON = function () {{ async function run() {{ {dates} {preprocessor} - const argsArr = argsObjToArr(args); + // If the entrypoint has no parsed params (spread is empty), pass values directly + // This handles WAC child jobs where tasks are const-wrapped functions + const argsArr = {child_spread}; if (Main.{main_name} === undefined || typeof Main.{main_name} !== 'function') {{ throw new Error("{main_name} function is missing"); }} - let res = await Main.{main_name}(...argsArr); + let entrypoint = Main.{main_name}; + let res = await entrypoint(...argsArr); if (isAsyncIterable(res)) {{ for await (const chunk of res) {{ console.log("WM_STREAM: " + chunk.replace(/\n/g, '\\n')); @@ -1311,7 +1468,8 @@ try {{ process.exit(1); }} "#, - ); + ) + }; write_file(job_dir, "wrapper.mjs", &wrapper_content)?; Ok(()) as error::Result<()> }; @@ -1336,6 +1494,7 @@ try {{ && !annotation.nobundling && !*DISABLE_BUNDLING && !codebase.is_some() + && !is_wac_v2 && (maybe_lock.get_lock().is_some() || annotation.native); let write_loader_f = async { @@ -1381,6 +1540,23 @@ try {{ write_wrapper_f, write_loader_f )?; + + // For WAC v2, write checkpoint.json before bun runs + if is_wac_v2 { + if let Connection::Sql(db) = conn { + let checkpoint = crate::wac_executor::load_checkpoint(db, &job.id).await?; + let checkpoint = + crate::wac_executor::prepare_checkpoint_for_resume(db, &job.id, checkpoint).await?; + + let checkpoint_json = serde_json::to_string(&checkpoint).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize checkpoint: {e}")) + })?; + write_file(job_dir, "checkpoint.json", &checkpoint_json)?; + } else { + write_file(job_dir, "checkpoint.json", r#"{"completed_steps":{}}"#)?; + } + } + if !codebase.is_some() && !has_bundle_cache { if build_cache { generate_bun_bundle( @@ -1476,7 +1652,7 @@ try {{ let result = crate::js_eval::eval_fetch_timeout( env_code, - inner_content.clone(), + inner_content.to_string(), js_code, job_args, job.script_entrypoint_override.clone(), @@ -1679,7 +1855,766 @@ try {{ })?; *new_args = Some(args.clone()); } - read_result(job_dir, handle_result.result_stream).await + + let result = read_result(job_dir, handle_result.result_stream).await?; + + // WAC v2 post-execution: parse output and handle dispatch/suspend + if is_wac_v2 { + return handle_wac_v2_output(result, job, conn).await; + } + + Ok(result) +} + +/// Handle WAC v2 output after bun/python exits. Parse result as WacOutput, +/// dispatch child jobs on suspend, or return the final result. +pub async fn handle_wac_v2_output( + result: Box, + job: &MiniPulledJob, + conn: &Connection, +) -> error::Result> { + use crate::wac_executor::{ + add_completed_step, load_checkpoint, parse_wac_output, update_checkpoint_for_dispatch, + WacOutput, + }; + use serde_json::Value; + use windmill_common::get_latest_flow_version_info_for_path; + use windmill_common::jobs::{script_path_to_payload, JobKind, JobPayload, RawCode}; + use windmill_common::runnable_settings::{ + ConcurrencySettings, ConcurrencySettingsWithCustom, DebouncingSettings, + }; + use windmill_queue::{push, PushArgs, PushIsolationLevel}; + + let output = parse_wac_output(&result)?; + + match output { + WacOutput::Complete { result: value } => { + // Workflow completed — return the inner result value + let raw = serde_json::value::to_raw_value(&value).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize WAC result: {e}")) + })?; + Ok(raw) + } + WacOutput::Dispatch { mode, steps } => { + if steps.is_empty() { + return Err(error::Error::internal_err( + "WAC v2 dispatch with no steps — this is a bug in the workflow SDK".to_string(), + )); + } + let db = match conn { + Connection::Sql(db) => db, + _ => { + return Err(error::Error::internal_err( + "WAC v2 dispatch requires SQL connection".to_string(), + )) + } + }; + + let mut checkpoint = load_checkpoint(db, &job.id).await?; + + // Source hash validation: detect if code changed between replays + let current_hash = job.runnable_id.map(|h| h.0.to_string()).unwrap_or_default(); + if !current_hash.is_empty() { + if checkpoint.source_hash.is_empty() { + checkpoint.source_hash = current_hash.clone(); + } else if checkpoint.source_hash != current_hash { + return Err(error::Error::ExecutionErr( + "Workflow source code changed between replays. \ + Cannot safely resume from checkpoint — step keys may have shifted. \ + Please restart this workflow." + .to_string(), + )); + } + } + let num_steps = steps.len(); + + tracing::info!( + job_id = %job.id, + mode = %mode, + num_steps = num_steps, + steps = ?steps.iter().map(|s| &s.name).collect::>(), + "WAC v2 dispatching child jobs" + ); + + // Create child jobs for each step. + // Each child re-runs the full workflow with a checkpoint containing + // _executing_key = step_key, so only that step runs its inner function. + // + // IMPORTANT: To prevent a race condition where a fast child completes + // before the parent is suspended, we: + // 1. Pre-generate child UUIDs + // 2. Save checkpoint + suspend parent + seed child checkpoints + // 3. THEN push the child jobs (making them visible to workers) + + // Read the parent's original args for the child jobs + let parent_args: HashMap> = { + let stored: serde_json::Map = checkpoint.input_args.clone(); + if stored.is_empty() { + // First dispatch — read from the parent job's args + let row: Option = sqlx::query_scalar( + "SELECT args FROM v2_job WHERE id = $1 AND workspace_id = $2", + ) + .bind(&job.id) + .bind(&job.workspace_id) + .fetch_optional(db) + .await?; + let args_val = row.unwrap_or(Value::Object(Default::default())); + if let Value::Object(map) = args_val { + // Store for future re-runs + checkpoint.input_args = map.clone(); + map.into_iter() + .map(|(k, v)| { + let raw = serde_json::value::to_raw_value(&v).map_err(|e| { + error::Error::internal_err(format!( + "Failed to serialize arg '{k}': {e}" + )) + })?; + Ok((k, raw)) + }) + .collect::>>()? + } else { + HashMap::new() + } + } else { + stored + .into_iter() + .map(|(k, v)| { + let raw = serde_json::value::to_raw_value(&v).map_err(|e| { + error::Error::internal_err(format!( + "Failed to serialize arg '{k}': {e}" + )) + })?; + Ok((k, raw)) + }) + .collect::>>()? + } + }; + + // Pre-generate child UUIDs so we can save them in the checkpoint + // before the children become visible to workers. + // Validate key uniqueness — duplicate keys would cause one child's + // UUID to be overwritten in the job_ids map, making it unmappable + // on completion (the parent would hang). + { + let mut seen_keys = std::collections::HashSet::new(); + for s in &steps { + if !seen_keys.insert(&s.key) { + return Err(error::Error::internal_err(format!( + "WAC v2 duplicate step key '{}' — each task call must produce a unique key", + s.key + ))); + } + } + } + let job_ids: Vec<(String, Uuid)> = steps + .iter() + .map(|s| (s.key.clone(), ulid::Ulid::new().into())) + .collect(); + + // Resolve job_payload once (same for all children since they re-run + // the parent script) + let job_payload_template = match job.kind { + JobKind::Script => { + if let Some(hash) = job.runnable_id { + Ok(JobPayload::ScriptHash { + hash, + path: job.runnable_path.clone().unwrap_or_default(), + cache_ttl: job.cache_ttl, + cache_ignore_s3_path: job.cache_ignore_s3_path, + dedicated_worker: None, + language: job.script_lang.unwrap_or(ScriptLang::Bun), + priority: job.priority, + apply_preprocessor: false, + concurrency_settings: ConcurrencySettings::default(), + debouncing_settings: DebouncingSettings::default(), + }) + } else { + Err(error::Error::internal_err( + "WAC v2 Script job missing runnable_id".to_string(), + )) + } + } + JobKind::Preview => { + let row: Option<(Option, Option)> = sqlx::query_as( + "SELECT raw_code, raw_lock FROM v2_job WHERE id = $1 AND workspace_id = $2", + ) + .bind(&job.id) + .bind(&job.workspace_id) + .fetch_optional(db) + .await?; + let (code, lock) = row.unwrap_or_default(); + Ok(JobPayload::Code(RawCode { + content: code.unwrap_or_default(), + path: job.runnable_path.clone(), + hash: None, + language: job.script_lang.unwrap_or(ScriptLang::Bun), + lock: lock, + cache_ttl: job.cache_ttl, + cache_ignore_s3_path: job.cache_ignore_s3_path, + dedicated_worker: None, + concurrency_settings: ConcurrencySettingsWithCustom::default(), + debouncing_settings: DebouncingSettings::default(), + })) + } + _ => Err(error::Error::internal_err(format!( + "WAC v2 unsupported job kind: {:?}", + job.kind + ))), + }?; + + // Step 1: Save checkpoint, suspend parent, and seed child checkpoints + // in a single transaction — all BEFORE children become visible. + { + let mut tx = db.begin().await?; + + // Update checkpoint with pending steps + update_checkpoint_for_dispatch(&mut checkpoint, &steps, &mode, &job_ids); + let status_json = serde_json::to_value(&checkpoint).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize checkpoint: {e}")) + })?; + sqlx::query( + "INSERT INTO v2_job_status (id, workflow_as_code_status) + VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb)) + ON CONFLICT (id) DO UPDATE SET + workflow_as_code_status = jsonb_set( + COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb), + '{_checkpoint}', + $2::jsonb + )", + ) + .bind(&job.id) + .bind(&status_json) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!("Failed to save WAC checkpoint: {e}")) + })?; + + // Store per-child-job info for the WorkflowTimeline UI + for (step, (_, child_id)) in steps.iter().zip(job_ids.iter()) { + let child_id_str = child_id.to_string(); + let timeline_val = serde_json::json!({ + "scheduled_for": chrono::Utc::now().to_rfc3339(), + "name": step.key, + }); + sqlx::query( + "UPDATE v2_job_status SET workflow_as_code_status = jsonb_set( + COALESCE(workflow_as_code_status, '{}'::jsonb), + ARRAY[$2], + $3 + ) WHERE id = $1", + ) + .bind(&job.id) + .bind(&child_id_str) + .bind(&timeline_val) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!( + "Failed to update WAC timeline status: {e}" + )) + })?; + } + + // Suspend parent before children become visible. + // Keep running = true so the normal pull query ignores it. + // The suspended pull query picks it up when suspend reaches 0 + // (it checks: suspend_until IS NOT NULL AND suspend <= 0). + let suspend_count = num_steps as i32; + sqlx::query!( + "UPDATE v2_job_queue SET suspend = $2, suspend_until = now() + interval '14 day' WHERE id = $1", + job.id, + suspend_count, + ) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!( + "Failed to suspend WAC parent job {}: {e}", + job.id + )) + })?; + + tx.commit().await?; + } + + // Step 2: Push child jobs (now visible to workers). + // Parent is already suspended, so child completions are safe. + // Track successfully pushed children so we can cancel them on + // partial failure (e.g. pushing child 3 of 5 fails). + let mut pushed_ids: Vec = Vec::with_capacity(num_steps); + let push_result: error::Result<()> = async { + for (step, (_, child_uuid)) in steps.iter().zip(job_ids.iter()) { + // Resolve job payload based on dispatch_type + let (job_payload, child_args, is_external) = match step.dispatch_type.as_str() { + "script" => { + // Resolve script path to job payload (handles hash, lang, etc.) + let (payload, _, _, _, _) = script_path_to_payload( + &step.script, + None, // no authed db for background workers + db.clone(), + &job.workspace_id, + Some(true), // skip preprocessor + ) + .await?; + let step_args: HashMap> = step + .args + .iter() + .map(|(k, v)| { + let raw = serde_json::value::to_raw_value(v).unwrap(); + (k.clone(), raw) + }) + .collect(); + (payload, step_args, true) + } + "flow" => { + let flow_info = get_latest_flow_version_info_for_path( + None, + db, + &job.workspace_id, + &step.script, + true, + ) + .await?; + let payload = JobPayload::Flow { + path: step.script.clone(), + dedicated_worker: flow_info.dedicated_worker, + apply_preprocessor: false, + version: flow_info.version, + }; + let step_args: HashMap> = step + .args + .iter() + .map(|(k, v)| { + let raw = serde_json::value::to_raw_value(v).unwrap(); + (k.clone(), raw) + }) + .collect(); + (payload, step_args, true) + } + _ => { + // "inline" — re-run parent with _executing_key + (job_payload_template.clone(), parent_args.clone(), false) + } + }; + + let push_args = PushArgs { args: &child_args, extra: None }; + + // Apply step-level overrides to payload (cache, concurrency) + let mut job_payload = job_payload; + if let Some(cache_ttl) = step.cache_ttl { + match &mut job_payload { + JobPayload::ScriptHash { cache_ttl: ref mut ct, .. } => { + *ct = Some(cache_ttl) + } + JobPayload::Code(ref mut code) => code.cache_ttl = Some(cache_ttl), + _ => {} + } + } + if step.concurrent_limit.is_some() + || step.concurrency_key.is_some() + || step.concurrency_time_window_s.is_some() + { + match &mut job_payload { + JobPayload::ScriptHash { concurrency_settings: ref mut cs, .. } => { + if let Some(limit) = step.concurrent_limit { + cs.concurrent_limit = Some(limit); + } + if let Some(ref key) = step.concurrency_key { + cs.concurrency_key = Some(key.clone()); + } + if let Some(window) = step.concurrency_time_window_s { + cs.concurrency_time_window_s = Some(window); + } + } + JobPayload::Code(ref mut code) => { + if let Some(limit) = step.concurrent_limit { + code.concurrency_settings.concurrent_limit = Some(limit); + } + if let Some(ref key) = step.concurrency_key { + code.concurrency_settings.custom_concurrency_key = + Some(key.clone()); + } + if let Some(window) = step.concurrency_time_window_s { + code.concurrency_settings.concurrency_time_window_s = + Some(window); + } + } + _ => {} + } + } + + let (_, mut tx) = push( + db, + PushIsolationLevel::IsolatedRoot(db.clone()), + &job.workspace_id, + job_payload, + push_args, + &job.created_by, + &job.permissioned_as_email, + job.permissioned_as.clone(), + None, + None, + None, + Some(job.id), // parent_job + job.root_job.or(Some(job.id)), // root_job + job.flow_innermost_root_job, + Some(*child_uuid), // pre-generated job_id + false, // is_flow_step + false, // same_worker + None, // pre_run_error + job.visible_to_owner, + step.tag.clone().or_else(|| Some(job.tag.clone())), + step.timeout.or(job.timeout), + None, // flow_step_id + step.priority, // priority_override + None, // authed + false, // running + None, // end_user_email + None, // trigger + None, // suspended_mode + ) + .await?; + + // Seed child checkpoint only for inline tasks (they need + // _executing_key to know which step to run). External + // scripts/flows don't need a WAC checkpoint. + if !is_external { + let child_checkpoint_json = serde_json::json!({ + "completed_steps": &checkpoint.completed_steps, + "_executing_key": &step.key, + }); + sqlx::query( + "INSERT INTO v2_job_status (id, workflow_as_code_status) + VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb)) + ON CONFLICT (id) DO UPDATE SET + workflow_as_code_status = jsonb_set( + COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb), + '{_checkpoint}', + $2::jsonb + )", + ) + .bind(child_uuid) + .bind(&child_checkpoint_json) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!( + "Failed to seed child checkpoint: {e}" + )) + })?; + } + + tx.commit().await.map_err(|e| { + error::Error::internal_err(format!("Failed to commit child push: {e}")) + })?; + + pushed_ids.push(*child_uuid); + + tracing::info!( + parent_job = %job.id, + child_job = %child_uuid, + step_name = %step.name, + step_key = %step.key, + "WAC v2 dispatched child job" + ); + } + Ok(()) + } + .await; + + if let Err(e) = push_result { + tracing::error!( + job_id = %job.id, + error = %e, + pushed_count = pushed_ids.len(), + total_count = num_steps, + "WAC v2 failed to push child jobs, cleaning up" + ); + + // Cancel already-pushed children so they don't complete and + // corrupt the checkpoint (they'd decrement suspend on a parent + // that's about to be unsuspended and re-run). + for child_id in &pushed_ids { + let _ = sqlx::query!( + "UPDATE v2_job_queue SET canceled_by = $2, canceled_reason = $3 WHERE id = $1", + child_id, + "system", + "WAC dispatch failed: not all children could be pushed", + ) + .execute(db) + .await; + } + + // Clear pending_steps from checkpoint so the parent doesn't + // think children are outstanding when it re-runs. + let _ = sqlx::query( + "UPDATE v2_job_status SET workflow_as_code_status = \ + workflow_as_code_status #- '{_checkpoint,pending_steps}' \ + WHERE id = $1", + ) + .bind(&job.id) + .execute(db) + .await; + + // Unsuspend parent so the error propagates instead of a 14-day hang + let _ = sqlx::query!( + "UPDATE v2_job_queue SET suspend = 0, suspend_until = NULL WHERE id = $1", + job.id, + ) + .execute(db) + .await; + return Err(e); + } + + tracing::info!( + job_id = %job.id, + num_steps = num_steps, + "WAC v2 parent job suspended" + ); + + Err(error::Error::WacSuspended(format!( + "WAC v2 job {} suspended waiting for {} child job(s)", + job.id, num_steps + ))) + } + WacOutput::Approval { key, timeout, form } => { + let db = match conn { + Connection::Sql(db) => db, + _ => { + return Err(error::Error::internal_err( + "WAC v2 approval requires SQL connection".to_string(), + )) + } + }; + + let mut checkpoint = load_checkpoint(db, &job.id).await?; + let timeout_secs = timeout.unwrap_or(1800) as f64; + + // Mark this step as pending approval + checkpoint.pending_steps = Some(crate::wac_executor::WacPendingSteps { + mode: "approval".to_string(), + keys: vec![key.clone()], + job_ids: serde_json::Map::new(), + }); + + let mut tx = db.begin().await?; + + // Save checkpoint + let status_json = serde_json::to_value(&checkpoint).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize checkpoint: {e}")) + })?; + sqlx::query( + "INSERT INTO v2_job_status (id, workflow_as_code_status) + VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb)) + ON CONFLICT (id) DO UPDATE SET + workflow_as_code_status = jsonb_set( + COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb), + '{_checkpoint}', + $2::jsonb + )", + ) + .bind(&job.id) + .bind(&status_json) + .execute(&mut *tx) + .await + .map_err(|e| error::Error::internal_err(format!("Failed to save checkpoint: {e}")))?; + + // Store approval form metadata for the approval page endpoint + let approval_meta = serde_json::json!({ + "key": key, + "form": form, + "timeout": timeout_secs as u32, + }); + sqlx::query( + "UPDATE v2_job_status SET workflow_as_code_status = jsonb_set( + COALESCE(workflow_as_code_status, '{}'::jsonb), + '{_approval}', + $2::jsonb + ) WHERE id = $1", + ) + .bind(&job.id) + .bind(&approval_meta) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!("Failed to save approval meta: {e}")) + })?; + + // Suspend parent with suspend=1 (waiting for 1 approval event) + sqlx::query!( + "UPDATE v2_job_queue SET suspend = 1, suspend_until = now() + make_interval(secs => $2) WHERE id = $1", + job.id, + timeout_secs, + ) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + tracing::info!( + job_id = %job.id, + approval_key = %key, + timeout_secs = timeout_secs, + "WAC v2 parent job suspended waiting for approval" + ); + + Err(error::Error::WacSuspended(format!( + "WAC v2 job {} suspended waiting for approval (key: {})", + job.id, key + ))) + } + WacOutput::Sleep { key, seconds } => { + let db = match conn { + Connection::Sql(db) => db, + _ => { + return Err(error::Error::internal_err( + "WAC v2 sleep requires SQL connection".to_string(), + )) + } + }; + + let mut checkpoint = load_checkpoint(db, &job.id).await?; + let sleep_secs = seconds.max(1) as f64; + + // Mark this step as pending sleep + checkpoint.pending_steps = Some(crate::wac_executor::WacPendingSteps { + mode: "sleep".to_string(), + keys: vec![key.clone()], + job_ids: serde_json::Map::new(), + }); + + let mut tx = db.begin().await?; + + // Save checkpoint + let status_json = serde_json::to_value(&checkpoint).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize checkpoint: {e}")) + })?; + sqlx::query( + "INSERT INTO v2_job_status (id, workflow_as_code_status) + VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb)) + ON CONFLICT (id) DO UPDATE SET + workflow_as_code_status = jsonb_set( + COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb), + '{_checkpoint}', + $2::jsonb + )", + ) + .bind(&job.id) + .bind(&status_json) + .execute(&mut *tx) + .await + .map_err(|e| error::Error::internal_err(format!("Failed to save checkpoint: {e}")))?; + + // Suspend parent — it will auto-resume when suspend_until passes. + // Use suspend=1 (not 0) so the suspended pull query only picks it up + // when `suspend_until <= now()`, not via `suspend <= 0`. + sqlx::query!( + "UPDATE v2_job_queue SET suspend = 1, suspend_until = now() + make_interval(secs => $2) WHERE id = $1", + job.id, + sleep_secs, + ) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + tracing::info!( + job_id = %job.id, + sleep_key = %key, + sleep_secs = sleep_secs, + "WAC v2 parent job sleeping for {}s", + sleep_secs + ); + + Err(error::Error::WacSuspended(format!( + "WAC v2 job {} sleeping for {}s (key: {})", + job.id, seconds, key + ))) + } + WacOutput::InlineCheckpoint { key, result: value } => { + let db = match conn { + Connection::Sql(db) => db, + _ => { + return Err(error::Error::internal_err( + "WAC v2 inline checkpoint requires SQL connection".to_string(), + )) + } + }; + + let mut checkpoint = load_checkpoint(db, &job.id).await?; + + // Source hash validation (same as Dispatch path) + let current_hash = job.runnable_id.map(|h| h.0.to_string()).unwrap_or_default(); + if !current_hash.is_empty() { + if checkpoint.source_hash.is_empty() { + checkpoint.source_hash = current_hash.clone(); + } else if checkpoint.source_hash != current_hash { + return Err(error::Error::ExecutionErr( + "Workflow source code changed between replays. \ + Cannot safely resume from checkpoint — step keys may have shifted. \ + Please restart this workflow." + .to_string(), + )); + } + } + + tracing::info!( + job_id = %job.id, + step_key = %key, + "WAC v2 inline checkpoint — persisting step result" + ); + + add_completed_step(&mut checkpoint, &key, value); + + // Save checkpoint + reset running in a single transaction + { + let mut tx = db.begin().await?; + let status_json = serde_json::to_value(&checkpoint).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize checkpoint: {e}")) + })?; + sqlx::query( + "INSERT INTO v2_job_status (id, workflow_as_code_status) + VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb)) + ON CONFLICT (id) DO UPDATE SET + workflow_as_code_status = jsonb_set( + COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb), + '{_checkpoint}', + $2::jsonb + )", + ) + .bind(&job.id) + .bind(&status_json) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!("Failed to save WAC checkpoint: {e}")) + })?; + + // Reset running=false so the job is immediately eligible for pickup. + // Unlike dispatch (which sets suspend>0), inline checkpoints don't suspend — + // the job should be re-run right away to continue past the cached step. + sqlx::query!( + "UPDATE v2_job_queue SET running = false, started_at = null WHERE id = $1", + job.id, + ) + .execute(&mut *tx) + .await + .map_err(|e| { + error::Error::internal_err(format!( + "Failed to reset running state for inline checkpoint: {e}" + )) + })?; + + tx.commit().await?; + } + + Err(error::Error::WacSuspended(format!( + "WAC v2 job {} inline checkpoint for step {}", + job.id, key + ))) + } + } } pub async fn get_common_bun_proc_envs(base_internal_url: Option<&str>) -> HashMap { diff --git a/backend/windmill-worker/src/lib.rs b/backend/windmill-worker/src/lib.rs index f6d752558d..2fe56f50b9 100644 --- a/backend/windmill-worker/src/lib.rs +++ b/backend/windmill-worker/src/lib.rs @@ -73,6 +73,7 @@ mod universal_pkg_installer; #[cfg(feature = "private")] mod volume_ee; mod volume_oss; +pub mod wac_executor; mod worker; mod worker_flow; mod worker_lockfiles; diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index 320bcb3e40..c74feb0a0a 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -567,6 +567,9 @@ pub async fn handle_python_job( let annotations = PythonAnnotations::parse(inner_content); + let is_wac_v2 = job.script_entrypoint_override.is_none() + && crate::wac_executor::is_wac_v2_py(inner_content); + if annotations.sandbox && NSJAIL_AVAILABLE.is_none() { return Err(Error::ExecutionErr( "Script has #sandbox annotation but nsjail is not available on this worker. \ @@ -677,8 +680,68 @@ pub async fn handle_python_job( String::new() }; let main_override = main_name.unwrap_or_else(|| "main".to_string()); - let wrapper_content: String = format!( - r#" + let wrapper_content: String = if is_wac_v2 { + format!( + r#" +import os +import json +{import_loader} +{import_base64} +{import_datetime} +import traceback +import sys +from {module_dir_dot} import {last} as inner_script +from wmill.client import _run_workflow + +with open("args.json") as f: + kwargs = json.load(f, strict=False) +args = {{}} +{transforms} + +with open("checkpoint.json") as f: + checkpoint = json.load(f, strict=False) + +result_json = os.path.join(os.path.abspath(os.path.dirname(__file__)), "result.json") + +# Find the @workflow-decorated function +workflow_fn = None +for name in dir(inner_script): + obj = getattr(inner_script, name) + if callable(obj) and getattr(obj, '_is_workflow', False): + workflow_fn = obj + break + +if workflow_fn is None: + raise ValueError("No @workflow function found in script") + +for k, v in list(args.items()): + if v == '': + del args[k] + +try: + output = _run_workflow(workflow_fn, checkpoint, args) + output_json = json.dumps(output, separators=(',', ':'), default=str) + with open(result_json, 'w') as f: + f.write(output_json) +except BaseException as e: + exc_type, exc_value, exc_traceback = sys.exc_info() + tb = traceback.format_tb(exc_traceback) + with open(result_json, 'w') as f: + err = {{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }} + extra = e.__dict__ + if extra and len(extra) > 0: + err['extra'] = extra + flow_node_id = os.environ.get('WM_FLOW_STEP_ID') + if flow_node_id: + err['step_id'] = flow_node_id + err_json = json.dumps(err, separators=(',', ':'), default=str).replace('\n', '') + f.write(err_json) + sys.exit(1) +"#, + ) + } else { + format!( + r#" import os import json {import_loader} @@ -751,9 +814,26 @@ except BaseException as e: f.write(err_json) sys.exit(1) "#, - ); + ) + }; write_file(job_dir, "wrapper.py", &wrapper_content)?; + // For WAC v2, write checkpoint.json before python runs. + if is_wac_v2 { + if let Connection::Sql(db) = conn { + let checkpoint = crate::wac_executor::load_checkpoint(db, &job.id).await?; + let checkpoint = + crate::wac_executor::prepare_checkpoint_for_resume(db, &job.id, checkpoint).await?; + + let checkpoint_json = serde_json::to_string(&checkpoint).map_err(|e| { + error::Error::internal_err(format!("Failed to serialize checkpoint: {e}")) + })?; + write_file(job_dir, "checkpoint.json", &checkpoint_json)?; + } else { + write_file(job_dir, "checkpoint.json", r#"{"completed_steps":{}}"#)?; + } + } + tracing::debug!("Finished writing wrapper"); let mut reserved_variables = @@ -936,7 +1016,15 @@ mount {{ *new_args = Some(args.clone()); } - read_result(job_dir, handle_result.result_stream).await + let result = read_result(job_dir, handle_result.result_stream).await?; + + // WAC v2 post-execution: parse output and handle dispatch/suspend. + // Box::pin to avoid bloating handle_python_job's async state machine (stack overflow). + if is_wac_v2 { + return Box::pin(crate::bun_executor::handle_wac_v2_output(result, job, conn)).await; + } + + Ok(result) } async fn prepare_wrapper( diff --git a/backend/windmill-worker/src/result_processor.rs b/backend/windmill-worker/src/result_processor.rs index c237d8cc02..e26633fc9a 100644 --- a/backend/windmill-worker/src/result_processor.rs +++ b/backend/windmill-worker/src/result_processor.rs @@ -665,7 +665,7 @@ pub async fn process_completed_job( add_time!(bench, "pre add_completed_job"); - let (_, duration) = add_completed_job( + let (_, duration, wac_job_ids) = add_completed_job( db, &job, true, @@ -717,6 +717,29 @@ pub async fn process_completed_job( } return Ok(r); } + } else if let Some(parent_job) = parent_job { + // wac_job_ids is piggybacked from the duration write in + // add_completed_job — no extra query needed. + if let Some(job_ids) = wac_job_ids { + if let Ok(Some(_)) = handle_wac_child_completion( + db, + &job_id, + parent_job, + &workspace_id, + result, + true, + job_ids, + ) + .await + { + if let Some(done_tx) = done_tx { + done_tx + .send(()) + .expect("done receiver should still be alive"); + } + return Ok(None); + } + } } } else { let result = add_completed_job_error( @@ -770,11 +793,229 @@ pub async fn process_completed_job( } return Ok(r); } + } else if let Some(parent_job) = job.parent_job { + // WAC child failed — query job_ids from parent (errors are rare, + // so the extra read is acceptable here). + let job_ids_json: Option> = sqlx::query_scalar( + "SELECT workflow_as_code_status->'_checkpoint'->'pending_steps'->'job_ids' \ + FROM v2_job_status WHERE id = $1", + ) + .bind(&parent_job) + .fetch_optional(db) + .await?; + if let Some(Some(job_ids)) = job_ids_json { + let err_result = Arc::new(serde_json::value::to_raw_value(&result).unwrap()); + if let Ok(Some(_)) = handle_wac_child_completion( + db, + &job.id, + parent_job, + &job.workspace_id, + err_result, + false, + job_ids, + ) + .await + { + if let Some(done_tx) = done_tx { + done_tx + .send(()) + .expect("done receiver should still be alive"); + } + return Ok(None); + } + } } } return Ok(None); } +/// Handle a WAC v2 child job completion. +/// Returns Ok(Some(())) if the parent was a WAC job and was handled, +/// Ok(None) if the parent is not a WAC job (caller should fall through). +/// +/// CONCURRENCY: Multiple parallel children may complete simultaneously on +/// different workers. We use atomic SQL operations throughout: +/// - `completed_steps` is merged via `jsonb_set(... || jsonb_build_object(...))` +/// — PostgreSQL serialises concurrent UPDATEs on the same row, so each +/// worker sees the previous worker's writes. +/// - The suspend counter (set to N at dispatch time) is decremented atomically +/// with `RETURNING` to determine the "all done" condition. +pub(crate) async fn handle_wac_child_completion( + db: &DB, + child_job_id: &Uuid, + parent_job_id: Uuid, + workspace_id: &str, + result: Arc>, + success: bool, + job_ids_value: Value, +) -> error::Result> { + let job_ids = match job_ids_value { + Value::Object(m) => m, + _ => return Ok(None), // Not a WAC parent or no pending steps + }; + + let child_id_str = child_job_id.to_string(); + let step_key = job_ids.iter().find_map(|(key, val)| { + if val.as_str() == Some(&child_id_str) { + Some(key.clone()) + } else { + None + } + }); + + let step_key = match step_key { + Some(k) => k, + None => { + if !success { + // No step key and failed — can't store error, fail parent immediately + tracing::error!( + parent_job = %parent_job_id, + child_job = %child_job_id, + "WAC v2 child job failed but no step key found, failing parent" + ); + sqlx::query!( + "UPDATE v2_job_queue SET suspend = 0, suspend_until = NULL WHERE id = $1", + parent_job_id, + ) + .execute(db) + .await?; + let parent_mini = get_mini_completed_job(&parent_job_id, workspace_id, db).await?; + if let Some(parent_mini) = parent_mini { + let child_err: Value = + serde_json::from_str(result.get()).unwrap_or(Value::Null); + let err_value = json!({ + "message": format!("WAC child job {} failed (no step key)", child_job_id), + "error": child_err, + }); + let _ = windmill_queue::add_completed_job_error( + db, + &parent_mini, + 0, + None, + err_value, + "wac_child_handler", + false, + None, + ) + .await; + } + return Ok(Some(())); + } + tracing::warn!( + parent_job = %parent_job_id, + child_job = %child_job_id, + "WAC v2 child completed but no matching step key found in checkpoint, decrementing suspend to avoid parent hang" + ); + // Still decrement suspend so the parent doesn't hang indefinitely + let _ = sqlx::query_scalar!( + "UPDATE v2_job_queue \ + SET suspend = GREATEST(suspend - 1, 0) \ + WHERE id = $1 \ + RETURNING suspend", + parent_job_id, + ) + .fetch_optional(db) + .await?; + return Ok(Some(())); + } + }; + + // Build result — wrap errors with _error marker so workflow try/catch can handle them + let result_value: Value = if success { + serde_json::from_str(result.get()).unwrap_or(Value::Null) + } else { + let child_err: Value = serde_json::from_str(result.get()).unwrap_or(Value::Null); + tracing::info!( + parent_job = %parent_job_id, + child_job = %child_job_id, + step_key = %step_key, + "WAC v2 child job failed, storing error for workflow try/catch" + ); + json!({ + "__wmill_error": true, + "message": format!("WAC task '{}' failed (child job {})", step_key, child_job_id), + "child_job_id": child_job_id.to_string(), + "step_key": step_key, + "result": child_err, + }) + }; + + tracing::info!( + parent_job = %parent_job_id, + child_job = %child_job_id, + step_key = %step_key, + success = success, + "WAC v2 child job completed" + ); + + // Use a transaction to ensure completed_steps merge + suspend decrement + // are atomic. Without this, a crash between the two could strand the parent. + let result_json = serde_json::to_value(&result_value) + .map_err(|e| error::Error::InternalErr(format!("Failed to serialize step result: {e}")))?; + + let mut tx = db.begin().await?; + + // Merge the completed step into the checkpoint. + // Uses `|| jsonb_build_object(key, value)` so concurrent children on + // different workers don't overwrite each other — PostgreSQL serialises + // concurrent UPDATEs on the same row and each sees the previous write. + sqlx::query( + "UPDATE v2_job_status SET workflow_as_code_status = jsonb_set( + workflow_as_code_status, + '{_checkpoint,completed_steps}', + COALESCE(workflow_as_code_status->'_checkpoint'->'completed_steps', '{}'::jsonb) + || jsonb_build_object($2::text, $3::jsonb) + ) WHERE id = $1", + ) + .bind(&parent_job_id) + .bind(&step_key) + .bind(&result_json) + .execute(&mut *tx) + .await + .map_err(|e| error::Error::InternalErr(format!("Failed to add WAC completed step: {e}")))?; + + // Decrement the suspend counter. The counter was set to N (number of + // children) at dispatch time. When it reaches 0 all children are done. + // Keep suspend_until non-null so the suspended pull query + // (`WHERE suspend_until IS NOT NULL AND suspend <= 0`) picks up the parent. + let new_suspend: Option = sqlx::query_scalar!( + "UPDATE v2_job_queue \ + SET suspend = GREATEST(suspend - 1, 0) \ + WHERE id = $1 \ + RETURNING suspend", + parent_job_id, + ) + .fetch_optional(&mut *tx) + .await?; + + let all_done = new_suspend == Some(0); + + if all_done { + // Clear pending_steps from checkpoint since all children are complete. + // This is cosmetic — the next replay will overwrite it anyway — but + // keeps the checkpoint clean for frontend display. + let _ = sqlx::query( + "UPDATE v2_job_status SET workflow_as_code_status = \ + workflow_as_code_status #- '{_checkpoint,pending_steps}' \ + WHERE id = $1", + ) + .bind(&parent_job_id) + .execute(&mut *tx) + .await; + } + + tx.commit().await?; + + if all_done { + tracing::info!( + parent_job = %parent_job_id, + "WAC v2 all child jobs completed, unsuspending parent" + ); + } + + Ok(Some(())) +} + pub async fn handle_non_flow_job_error( db: &DB, job: &MiniCompletedJob, diff --git a/backend/windmill-worker/src/snowflake_executor.rs b/backend/windmill-worker/src/snowflake_executor.rs index 9f89f2fc44..90d06287bc 100644 --- a/backend/windmill-worker/src/snowflake_executor.rs +++ b/backend/windmill-worker/src/snowflake_executor.rs @@ -9,8 +9,8 @@ use serde_json::{json, value::RawValue, Value}; use sha2::{Digest, Sha256}; use std::collections::HashMap; use windmill_common::error::to_anyhow; -use windmill_object_store::convert_json_line_stream; use windmill_common::worker::{Connection, SqlResultCollectionStrategy}; +use windmill_object_store::convert_json_line_stream; use windmill_common::{error::Error, worker::to_raw_value}; use windmill_parser_sql::{ diff --git a/backend/windmill-worker/src/wac_executor.rs b/backend/windmill-worker/src/wac_executor.rs new file mode 100644 index 0000000000..7c350ac92b --- /dev/null +++ b/backend/windmill-worker/src/wac_executor.rs @@ -0,0 +1,339 @@ +use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; +use serde_json::Value; +use uuid::Uuid; + +use windmill_common::error::{self, Error}; +use windmill_common::DB; + +/// Checkpoint state persisted across workflow invocations. +#[derive(Debug, Serialize, Deserialize, Default, Clone)] +pub struct WacCheckpoint { + #[serde(default)] + pub source_hash: String, + #[serde(default)] + pub completed_steps: serde_json::Map, + #[serde(skip_serializing_if = "Option::is_none")] + pub pending_steps: Option, + #[serde(default)] + pub input_args: serde_json::Map, + /// Accumulated map of step_key → child job UUID across all dispatch rounds. + /// Unlike `pending_steps.job_ids` (cleared after completion), this persists + /// so the frontend can always resolve step keys to child job names. + #[serde(default)] + pub job_ids: serde_json::Map, + /// When set on a child job's checkpoint, indicates which step this child + /// should execute directly (instead of dispatching). + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + pub _executing_key: Option, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct WacPendingSteps { + pub mode: String, + pub keys: Vec, + pub job_ids: serde_json::Map, +} + +/// Output from a single WAC invocation (parsed from result.json). +#[derive(Debug, Deserialize)] +#[serde(tag = "type")] +pub enum WacOutput { + #[serde(rename = "dispatch")] + Dispatch { mode: String, steps: Vec }, + #[serde(rename = "complete")] + Complete { result: Value }, + /// An inline step executed in the parent process — persist result to + /// checkpoint and re-run immediately (no child job, no suspend). + #[serde(rename = "inline_checkpoint")] + InlineCheckpoint { key: String, result: Value }, + /// Suspend the workflow waiting for an external approval event. + /// No child job is dispatched — the parent suspends directly and resumes + /// when a user hits the resume/cancel endpoint. + #[serde(rename = "approval")] + Approval { key: String, timeout: Option, form: Option }, + /// Server-side sleep — suspend the workflow for a duration without holding a worker. + #[serde(rename = "sleep")] + Sleep { key: String, seconds: u32 }, +} + +/// A step dispatched by the WAC SDK. +/// +/// `dispatch_type` determines how the child job is created: +/// - `"inline"` (default): re-runs the parent workflow with `_executing_key` set +/// - `"script"`: runs a separate Windmill script resolved from `script` path +/// - `"flow"`: runs a separate Windmill flow resolved from `script` path +#[derive(Debug, Deserialize, Clone)] +pub struct WacStepDispatch { + pub name: String, + pub script: String, + pub args: serde_json::Map, + pub key: String, + #[serde(default = "default_dispatch_type")] + pub dispatch_type: String, + // Per-task options forwarded to push() + #[serde(default)] + pub timeout: Option, + #[serde(default)] + pub tag: Option, + #[serde(default)] + pub cache_ttl: Option, + #[serde(default)] + pub priority: Option, + #[serde(default)] + pub concurrent_limit: Option, + #[serde(default)] + pub concurrency_key: Option, + #[serde(default)] + pub concurrency_time_window_s: Option, +} + +fn default_dispatch_type() -> String { + "inline".to_string() +} + +/// Load the WAC checkpoint from `v2_job_status.workflow_as_code_status._checkpoint`. +pub async fn load_checkpoint(db: &DB, job_id: &Uuid) -> error::Result { + let row: Option> = sqlx::query_scalar( + "SELECT workflow_as_code_status->'_checkpoint' FROM v2_job_status WHERE id = $1", + ) + .bind(job_id) + .fetch_optional(db) + .await?; + + match row { + Some(Some(status)) => { + let checkpoint: WacCheckpoint = match serde_json::from_value(status) { + Ok(c) => c, + Err(e) => { + tracing::warn!( + job_id = %job_id, + error = %e, + "Failed to deserialize WAC checkpoint, resetting to empty" + ); + WacCheckpoint::default() + } + }; + Ok(checkpoint) + } + _ => Ok(WacCheckpoint::default()), + } +} + +/// Save the WAC checkpoint to `v2_job_status.workflow_as_code_status._checkpoint`. +/// The top level of workflow_as_code_status is reserved for per-child-job timeline data. +pub async fn save_checkpoint( + db: &DB, + job_id: &Uuid, + checkpoint: &WacCheckpoint, +) -> error::Result<()> { + let status_json = serde_json::to_value(checkpoint) + .map_err(|e| Error::InternalErr(format!("Failed to serialize checkpoint: {e}")))?; + + sqlx::query( + "INSERT INTO v2_job_status (id, workflow_as_code_status) + VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb)) + ON CONFLICT (id) DO UPDATE SET + workflow_as_code_status = jsonb_set( + COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb), + '{_checkpoint}', + $2::jsonb + )", + ) + .bind(job_id) + .bind(&status_json) + .execute(db) + .await + .map_err(|e| Error::InternalErr(format!("Failed to save WAC checkpoint: {e}")))?; + + Ok(()) +} + +/// Parse the WAC result from result.json content. +pub fn parse_wac_output(result: &RawValue) -> error::Result { + serde_json::from_str(result.get()) + .map_err(|e| Error::InternalErr(format!("Failed to parse WAC output: {e}"))) +} + +/// Process a "dispatch" result: update checkpoint with pending steps info. +pub fn update_checkpoint_for_dispatch( + checkpoint: &mut WacCheckpoint, + steps: &[WacStepDispatch], + mode: &str, + job_ids: &[(String, Uuid)], +) { + let ids_map: serde_json::Map = job_ids + .iter() + .map(|(key, id)| (key.clone(), Value::String(id.to_string()))) + .collect(); + // Accumulate into persistent job_ids (survives pending_steps clearing) + for (k, v) in ids_map.iter() { + checkpoint.job_ids.insert(k.clone(), v.clone()); + } + let pending = WacPendingSteps { + mode: mode.to_string(), + keys: steps.iter().map(|s| s.key.clone()).collect(), + job_ids: ids_map, + }; + checkpoint.pending_steps = Some(pending); +} + +/// Process a completed child job result: add to checkpoint's completed_steps. +pub fn add_completed_step(checkpoint: &mut WacCheckpoint, step_key: &str, result: Value) { + checkpoint + .completed_steps + .insert(step_key.to_string(), result); + // If all pending steps are complete, clear pending + if let Some(ref pending) = checkpoint.pending_steps { + let all_done = pending + .keys + .iter() + .all(|k| checkpoint.completed_steps.contains_key(k)); + if all_done { + checkpoint.pending_steps = None; + } + } +} + +/// Check if all pending parallel steps are complete. +pub fn all_pending_complete(checkpoint: &WacCheckpoint) -> bool { + match &checkpoint.pending_steps { + None => true, + Some(pending) => pending + .keys + .iter() + .all(|k| checkpoint.completed_steps.contains_key(k)), + } +} + +/// If the checkpoint has a pending approval or sleep, inject the resume result +/// into `completed_steps` and save back to DB. Returns the (possibly modified) checkpoint. +/// +/// Called by both bun and python executors before writing checkpoint.json to disk. +pub async fn prepare_checkpoint_for_resume( + db: &DB, + job_id: &Uuid, + mut checkpoint: WacCheckpoint, +) -> error::Result { + let pending_mode = checkpoint.pending_steps.as_ref().map(|p| p.mode.as_str()); + + match pending_mode { + Some("approval") => { + let approval_key = checkpoint + .pending_steps + .as_ref() + .and_then(|p| p.keys.first().cloned()) + .unwrap_or_default(); + + let resume_row = sqlx::query_as::<_, (sqlx::types::Json>, Option, bool)>( + "SELECT value, approver, approved FROM resume_job WHERE job = $1 ORDER BY created_at ASC LIMIT 1", + ) + .bind(job_id) + .fetch_optional(db) + .await?; + + let approval_result = if let Some((value, approver, approved)) = resume_row { + serde_json::json!({ + "value": serde_json::from_str::(value.get()).unwrap_or(Value::Null), + "approver": approver.unwrap_or_else(|| "anonymous".to_string()), + "approved": approved, + }) + } else { + serde_json::json!({ + "value": null, + "approver": null, + "approved": false, + }) + }; + checkpoint + .completed_steps + .insert(approval_key.clone(), approval_result); + checkpoint.pending_steps = None; + save_checkpoint(db, job_id, &checkpoint).await?; + + tracing::info!( + job_id = %job_id, + approval_key = %approval_key, + "WAC v2 injected approval result into checkpoint" + ); + } + Some("sleep") => { + let sleep_key = checkpoint + .pending_steps + .as_ref() + .and_then(|p| p.keys.first().cloned()) + .unwrap_or_default(); + + checkpoint + .completed_steps + .insert(sleep_key.clone(), Value::Bool(true)); + checkpoint.pending_steps = None; + save_checkpoint(db, job_id, &checkpoint).await?; + + tracing::info!( + job_id = %job_id, + sleep_key = %sleep_key, + "WAC v2 resumed from sleep" + ); + } + _ => {} + } + + Ok(checkpoint) +} + +/// Detect WAC v2 patterns in TypeScript/Bun code. +/// Checks for `import ... from "windmill-client"` containing workflow/task, +/// skipping comment lines. +pub fn is_wac_v2_ts(code: &str) -> bool { + let mut has_wac_import = false; + let mut has_workflow = false; + let mut has_task = false; + for line in code.lines() { + let trimmed = line.trim(); + if trimmed.starts_with("//") { + continue; + } + if trimmed.contains("windmill-client") + && (trimmed.starts_with("import") || trimmed.starts_with("from")) + { + has_wac_import = true; + if trimmed.contains("workflow") { + has_workflow = true; + } + if trimmed.contains("task") { + has_task = true; + } + } + if trimmed.contains("export") && trimmed.contains("workflow(") { + has_workflow = true; + } + } + has_wac_import && has_workflow && has_task +} + +/// Detect WAC v2 patterns in Python code. +/// Checks for `@workflow` decorator and `@task` decorator with wmill import, +/// skipping comment lines. +pub fn is_wac_v2_py(code: &str) -> bool { + let mut has_wmill_import = false; + let mut has_workflow_decorator = false; + let mut has_task_decorator = false; + for line in code.lines() { + let trimmed = line.trim(); + if trimmed.starts_with('#') { + continue; + } + if trimmed.starts_with("import wmill") || trimmed.starts_with("from wmill") { + has_wmill_import = true; + } + if trimmed == "@workflow" || trimmed.starts_with("@workflow(") { + has_workflow_decorator = true; + } + if trimmed == "@task" || trimmed.starts_with("@task(") { + has_task_decorator = true; + } + } + has_wmill_import && has_workflow_decorator && has_task_decorator +} diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 8d7f6051ec..6d86ae2723 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -3528,6 +3528,13 @@ pub async fn handle_queued_job( { return Ok(false); } + if result + .as_ref() + .is_err_and(|err| matches!(err, &Error::WacSuspended(_))) + { + // WAC v2 job suspended while waiting for child jobs — don't complete it + return Ok(true); + } process_result( cjob, result.map(|x| Arc::new(x)), diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 8a45f85a90..24dcfb9b9a 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -1722,8 +1722,8 @@ pub async fn update_flow_status_after_job_completion_internal( chat_ai_info.conversation_id, ) .await?; - let duration = if success { - let (_, duration) = add_completed_job( + let (duration, wac_job_ids) = if success { + let (_, duration, wac_job_ids) = add_completed_job( db, &cflow_job, true, @@ -1737,9 +1737,9 @@ pub async fn update_flow_status_after_job_completion_internal( false, ) .await?; - duration + (duration, wac_job_ids) } else { - let (_, duration) = add_completed_job( + let (_, duration, wac_job_ids) = add_completed_job( db, &cflow_job, false, @@ -1757,11 +1757,30 @@ pub async fn update_flow_status_after_job_completion_internal( false, ) .await?; - duration + (duration, wac_job_ids) }; flow_job_duration = flow_job .started_at .map(|x| FlowJobDuration { started_at: x, duration_ms: duration }); + + // If this flow is a WAC child (not a flow step, has parent), + // notify the WAC parent of completion. + if !flow_job.is_flow_step() { + if let Some(parent_job) = flow_job.parent_job { + if let Some(job_ids) = wac_job_ids { + let _ = crate::result_processor::handle_wac_child_completion( + db, + &flow_job.id, + parent_job, + &flow_job.workspace_id, + nresult.clone(), + success, + job_ids, + ) + .await; + } + } + } } true } else { diff --git a/frontend/src/lib/components/ScriptBuilder.svelte b/frontend/src/lib/components/ScriptBuilder.svelte index e4e166791b..7de0926f0a 100644 --- a/frontend/src/lib/components/ScriptBuilder.svelte +++ b/frontend/src/lib/components/ScriptBuilder.svelte @@ -21,6 +21,7 @@ getPreprocessorFullCode, getMainFunctionPattern } from '$lib/script_helpers' + import { isWorkflowAsCode } from './graph/wacToFlow' import AIFormSettings from './copilot/AIFormSettings.svelte' import { defaultScripts, @@ -591,7 +592,7 @@ if (!disableHistoryChange) { history.replaceState(history.state, '', `/scripts/edit/${script.path}`) } - if (stay || (script.no_main_func && script.kind !== 'preprocessor')) { + if (stay || (script.no_main_func && script.kind !== 'preprocessor' && !isWorkflowAsCode(script.content, script.language))) { script.parent_hash = newHash sendUserToast('Deployed') } else { diff --git a/frontend/src/lib/components/TimelineBar.svelte b/frontend/src/lib/components/TimelineBar.svelte index 7c4e170d93..e4fb511450 100644 --- a/frontend/src/lib/components/TimelineBar.svelte +++ b/frontend/src/lib/components/TimelineBar.svelte @@ -35,7 +35,7 @@ {/if} {/snippet} {#if len > 0} - {#if len}{msToSec(len, 1)}s{/if} {/if} diff --git a/frontend/src/lib/components/WorkflowTimeline.svelte b/frontend/src/lib/components/WorkflowTimeline.svelte index 35cf4445ff..66fbb5b6c5 100644 --- a/frontend/src/lib/components/WorkflowTimeline.svelte +++ b/frontend/src/lib/components/WorkflowTimeline.svelte @@ -3,7 +3,7 @@ import { displayDate, msToSec } from '$lib/utils' import { onDestroy } from 'svelte' import { getDbClockNow } from '$lib/forLater' - import { ExternalLink, Loader2 } from 'lucide-svelte' + import { Loader2 } from 'lucide-svelte' import TimelineBar from './TimelineBar.svelte' import type { WorkflowStatus } from '$lib/gen' @@ -17,13 +17,10 @@ let now = $state(getDbClockNow().getTime()) - let interval = setInterval((x) => { + let interval = setInterval(() => { if (!max) { now = getDbClockNow().getTime() } - if (min && (!max || total == undefined)) { - total = max ? max - min : Math.max(now - min, 2000) - } }, 30) onDestroy(() => { @@ -40,7 +37,7 @@ 0 ) : undefined) - let total = $derived(flowDone && max ? max - min : now - min) + let total = $derived(flowDone && max ? max - min : Math.max(now - min, 2000)) {#if flow_status} @@ -75,7 +72,7 @@
{v.name ?? k} {v.name ?? k}
{#if min && total} diff --git a/frontend/src/lib/components/graph/wacToFlow.ts b/frontend/src/lib/components/graph/wacToFlow.ts new file mode 100644 index 0000000000..91d335ff91 --- /dev/null +++ b/frontend/src/lib/components/graph/wacToFlow.ts @@ -0,0 +1,16 @@ +/** + * Detect whether a script is a workflow-as-code entry point. + */ +export function isWorkflowAsCode(code: string, language: string): boolean { + if (language === 'python3') { + return /^\s*@workflow\s*$/m.test(code) || /from\s+wmill\s+import.*workflow/.test(code) + } + if (language === 'bun' || language === 'deno') { + return ( + /workflow\s*\(/.test(code) && + /task\s*\(/.test(code) && + /import.*(?:workflow|task).*from\s+['"]windmill-client(?:@[^'"]*)?['"]/.test(code) + ) + } + return false +} diff --git a/frontend/src/lib/components/runs/JobRunsPreview.svelte b/frontend/src/lib/components/runs/JobRunsPreview.svelte index 637848c83c..274679d3bc 100644 --- a/frontend/src/lib/components/runs/JobRunsPreview.svelte +++ b/frontend/src/lib/components/runs/JobRunsPreview.svelte @@ -47,7 +47,12 @@ ) function asWorkflowStatus(x: any): Record { - return x as Record + if (!x || typeof x !== 'object') return {} + const result: Record = {} + for (const [k, v] of Object.entries(x)) { + if (!k.startsWith('_')) result[k] = v as WorkflowStatus + } + return result } function handleFilterByConcurrencyKey(key: string) { diff --git a/frontend/src/lib/components/scriptEditor/LogPanel.svelte b/frontend/src/lib/components/scriptEditor/LogPanel.svelte index c47c05c281..63eb9d5a4c 100644 --- a/frontend/src/lib/components/scriptEditor/LogPanel.svelte +++ b/frontend/src/lib/components/scriptEditor/LogPanel.svelte @@ -68,6 +68,7 @@ showCustomResultPanel = false }: Props = $props() + type DContent = { mode: 'json' | Preview['language'] | 'plain' title: string @@ -92,7 +93,12 @@ } function asWorkflowStatus(x: any): Record { - return x as Record + if (!x || typeof x !== 'object') return {} + const result: Record = {} + for (const [k, v] of Object.entries(x)) { + if (!k.startsWith('_')) result[k] = v as WorkflowStatus + } + return result } let forceJson = $state(false) diff --git a/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte b/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte index 9ecedc87ea..ac73b934fd 100644 --- a/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte @@ -281,7 +281,12 @@ let redactSensitive = $state(false) function asWorkflowStatus(x: any): Record { - return x as Record + if (!x || typeof x !== 'object') return {} + const result: Record = {} + for (const [k, v] of Object.entries(x)) { + if (!k.startsWith('_')) result[k] = v as WorkflowStatus + } + return result } function forkPreview() { @@ -779,11 +784,15 @@ {/if}
{#if job?.workflow_as_code_status && job.job_kind !== 'aiagent'} -
- +
+

Workflow Timeline

+
+ +
+
{/if} {#if scriptProgress} diff --git a/python-client/wmill/pyproject.toml b/python-client/wmill/pyproject.toml index 5fcfad1ddd..cf3c7460a5 100644 --- a/python-client/wmill/pyproject.toml +++ b/python-client/wmill/pyproject.toml @@ -22,6 +22,12 @@ httpx = ">=0.24" requires = ["poetry>=1.0.2", "poetry-dynamic-versioning"] build-backend = "poetry.masonry.api" +[dependency-groups] +dev = [ + "httpx>=0.28.1", + "pytest>=9.0.2", +] + [tool.poetry-dynamic-versioning] enable = true vcs = "git" diff --git a/python-client/wmill/tests/test_workflow.py b/python-client/wmill/tests/test_workflow.py new file mode 100644 index 0000000000..43a5f22693 --- /dev/null +++ b/python-client/wmill/tests/test_workflow.py @@ -0,0 +1,1114 @@ +"""Tests for the Workflow-as-Code SDK.""" + +import asyncio +import pytest + +from wmill.client import WorkflowCtx, _StepSuspend, TaskError, workflow, task, step, sleep, parallel, _run_workflow + + +@task +async def extract_data(url: str): + pass # body unused in workflow context + + +@task +async def load_data(data=None): + pass + + +@task +async def clean_data(data=None): + pass + + +@task +async def compute_stats(data=None): + pass + + +@task +async def send_alert(msg: str = ""): + pass + + +@task +async def double(x: int): + return x * 2 + + +@task +async def add_one(x: int): + return x + 1 + + +@task +async def noop_task(): + pass + + +# --- Module-level workflow definitions --- + +@workflow +async def simple_workflow(url: str): + raw = await extract_data(url=url) + result = await load_data(data=raw) + return {"status": "done", "result": result} + + +@workflow +async def parallel_workflow(url: str): + raw = await extract_data(url=url) + cleaned, stats = await asyncio.gather( + clean_data(data=raw), + compute_stats(data=raw), + ) + return {"cleaned": cleaned, "stats": stats} + + +@workflow +async def conditional_workflow(count: int): + if count > 100: + await send_alert(msg="large") + await load_data() + return {"done": True} + + +@workflow +async def step_workflow(x: int): + ts = await step("timestamp", lambda: 1234567890) + doubled = await double(x=x) + rid = await step("random_id", lambda: "abc-123") + return {"ts": ts, "doubled": doubled, "id": rid} + + +# Edge case workflows + +@workflow +async def three_step_wf(n: int): + doubled = await double(x=n) + incremented = await add_one(x=doubled) + final = await double(x=incremented) + return {"doubled": doubled, "incremented": incremented, "final": final} + + +@workflow +async def seq_par_seq_wf(url: str): + raw = await extract_data(url=url) + cleaned, stats = await asyncio.gather( + clean_data(data=raw), + compute_stats(data=raw), + ) + loaded = await load_data(data={"cleaned": cleaned, "stats": stats}) + return loaded + + +@workflow +async def double_parallel_wf(): + a, b = await asyncio.gather(double(x=1), double(x=2)) + c, d = await asyncio.gather(add_one(x=a), add_one(x=b)) + return {"a": a, "b": b, "c": c, "d": d} + + +@workflow +async def cond_on_result_wf(): + val = await double(x=5) + if val > 8: + await send_alert(msg="big") + await load_data(data=val) + return {"val": val} + + +@workflow +async def empty_wf(): + return {"status": "empty"} + + +@workflow +async def single_wf(x: int): + result = await double(x=x) + return result + + +@workflow +async def no_arg_wf(): + result = await noop_task() + return result + + +@workflow +async def many_steps_wf(n: int): + val = n + for _ in range(10): + val = await add_one(x=val) + return val + + +@workflow +async def falsy_wf(): + a = await double(x=0) + b = await load_data(data=a) + c = await extract_data(url="") + return {"a": a, "b": b, "c": c} + + +@task(path="f/external_script") +async def run_external(x: int): + return x * 3 + + +@workflow +async def path_wf(x: int): + result = await run_external(x=x) + return result + + +@workflow +async def mixed_step_task_wf(x: int): + ts = await step("get_time", lambda: 999) + doubled = await double(x=x) + config = await step("get_config", lambda: {"retry": 3}) + added = await add_one(x=doubled) + return {"ts": ts, "doubled": doubled, "config": config, "added": added} + + +@workflow +async def par_child_wf(): + a, b = await asyncio.gather(double(x=3), add_one(x=7)) + return {"a": a, "b": b} + + +@workflow +async def det_wf(n: int): + a = await double(x=n) + b = await add_one(x=a) + c = await double(x=b) + return c + + +@workflow +async def par_args_wf(x: int): + base = await double(x=x) + a, b = await asyncio.gather(add_one(x=base), double(x=base)) + return {"a": a, "b": b} + + +@workflow +async def none_return_wf(): + await double(x=1) + + +@workflow +async def large_par_wf(): + results = await asyncio.gather( + double(x=1), double(x=2), double(x=3), double(x=4), double(x=5) + ) + return list(results) + + +@workflow +async def complex_mixed_wf(): + init = await extract_data(url="start") + a, b = await asyncio.gather(double(x=1), double(x=2)) + mid = await load_data(data={"a": a, "b": b}) + c, d = await asyncio.gather(add_one(x=3), add_one(x=4)) + fin = await clean_data(data={"mid": mid, "c": c, "d": d}) + return fin + + +@workflow +async def pre_par_child_wf(x: int): + base = await double(x=x) + a, b = await asyncio.gather(add_one(x=base), double(x=base)) + return {"a": a, "b": b} + + +# --- Tests --- +# NOTE: Python SDK uses name-based keys (e.g. "double", "double_2") +# not index-based keys (e.g. "step_0", "step_1"). + +class TestWorkflowDecorator: + def test_marks_function(self): + assert hasattr(simple_workflow, "_is_workflow") + assert simple_workflow._is_workflow is True + + +class TestTaskDecorator: + def test_marks_function(self): + assert hasattr(extract_data, "_is_task") + assert extract_data._is_task is True + + def test_standalone_execution(self): + """Outside a workflow, @task runs the function body directly.""" + result = asyncio.run(extract_data(url="https://example.com")) + assert result is None # body returns None + + def test_preserves_function_name(self): + assert extract_data.__name__ == "extract_data" + assert double.__name__ == "double" + + +class TestFirstInvocation: + def test_dispatches_first_step(self): + result = _run_workflow(simple_workflow, {}, {"url": "https://example.com"}) + assert result["type"] == "dispatch" + assert result["mode"] == "sequential" + assert len(result["steps"]) == 1 + assert result["steps"][0]["name"] == "extract_data" + assert result["steps"][0]["script"] == "extract_data" + assert result["steps"][0]["key"] == "extract_data" + assert result["steps"][0]["args"] == {"url": "https://example.com"} + + def test_positional_args_converted_to_kwargs(self): + """Positional args should be mapped to parameter names in dispatch.""" + @workflow + async def pos_workflow(): + await extract_data("https://pos.example.com") + + result = _run_workflow(pos_workflow, {}, {}) + assert result["type"] == "dispatch" + assert result["steps"][0]["args"] == {"url": "https://pos.example.com"} + + +class TestReplayWithCheckpoint: + def test_second_invocation_dispatches_second_step(self): + checkpoint = { + "completed_steps": { + "extract_data": {"data": [1, 2, 3]}, + } + } + result = _run_workflow(simple_workflow, checkpoint, {"url": "https://example.com"}) + assert result["type"] == "dispatch" + assert result["mode"] == "sequential" + assert result["steps"][0]["name"] == "load_data" + assert result["steps"][0]["key"] == "load_data" + + def test_all_steps_complete(self): + checkpoint = { + "completed_steps": { + "extract_data": {"data": [1, 2, 3]}, + "load_data": {"loaded": True}, + } + } + result = _run_workflow(simple_workflow, checkpoint, {"url": "https://example.com"}) + assert result["type"] == "complete" + assert result["result"]["status"] == "done" + assert result["result"]["result"] == {"loaded": True} + + +class TestParallelDispatch: + def test_first_invocation(self): + result = _run_workflow(parallel_workflow, {}, {"url": "https://example.com"}) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "extract_data" + + def test_parallel_dispatch(self): + checkpoint = { + "completed_steps": { + "extract_data": {"raw": "data"}, + } + } + result = _run_workflow(parallel_workflow, checkpoint, {"url": "https://example.com"}) + assert result["type"] == "dispatch" + assert result["mode"] == "parallel" + assert len(result["steps"]) == 2 + assert result["steps"][0]["name"] == "clean_data" + assert result["steps"][1]["name"] == "compute_stats" + + def test_parallel_complete(self): + checkpoint = { + "completed_steps": { + "extract_data": {"raw": "data"}, + "clean_data": {"cleaned": True}, + "compute_stats": {"count": 42}, + } + } + result = _run_workflow(parallel_workflow, checkpoint, {"url": "https://example.com"}) + assert result["type"] == "complete" + assert result["result"]["cleaned"] == {"cleaned": True} + assert result["result"]["stats"] == {"count": 42} + + +class TestConditionalWorkflow: + def test_condition_true(self): + result = _run_workflow(conditional_workflow, {}, {"count": 200}) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "send_alert" + + def test_condition_false(self): + result = _run_workflow(conditional_workflow, {}, {"count": 50}) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "load_data" + + +class TestStepInlineCheckpoint: + def test_first_invocation_returns_inline_checkpoint(self): + result = _run_workflow(step_workflow, {}, {"x": 7}) + assert result["type"] == "inline_checkpoint" + assert result["key"] == "timestamp" + assert result["result"] == 1234567890 + + def test_step_cached_then_task_dispatches(self): + checkpoint = {"completed_steps": {"timestamp": 1234567890}} + result = _run_workflow(step_workflow, checkpoint, {"x": 7}) + assert result["type"] == "dispatch" + assert result["mode"] == "sequential" + assert result["steps"][0]["name"] == "double" + assert result["steps"][0]["key"] == "double" + + def test_step_and_task_cached_then_second_step(self): + checkpoint = {"completed_steps": {"timestamp": 1234567890, "double": 14}} + result = _run_workflow(step_workflow, checkpoint, {"x": 7}) + assert result["type"] == "inline_checkpoint" + assert result["key"] == "random_id" + assert result["result"] == "abc-123" + + def test_all_complete(self): + checkpoint = {"completed_steps": {"timestamp": 1234567890, "double": 14, "random_id": "abc-123"}} + result = _run_workflow(step_workflow, checkpoint, {"x": 7}) + assert result["type"] == "complete" + assert result["result"] == {"ts": 1234567890, "doubled": 14, "id": "abc-123"} + + +class TestUnawaitedTask: + def test_unawaited_last_task_is_flushed(self): + @workflow + async def unawaited_workflow(): + await extract_data(url="x") + load_data(data="y") + + checkpoint = {"completed_steps": {"extract_data": "raw"}} + result = _run_workflow(unawaited_workflow, checkpoint, {}) + assert result["type"] == "dispatch" + assert result["mode"] == "sequential" + assert len(result["steps"]) == 1 + assert result["steps"][0]["name"] == "load_data" + + def test_unawaited_multiple_tasks_flushed_as_parallel(self): + @workflow + async def multi_unawaited_workflow(): + await extract_data(url="x") + clean_data(data="y") + compute_stats(data="y") + + checkpoint = {"completed_steps": {"extract_data": "raw"}} + result = _run_workflow(multi_unawaited_workflow, checkpoint, {}) + assert result["type"] == "dispatch" + assert result["mode"] == "parallel" + assert len(result["steps"]) == 2 + assert result["steps"][0]["name"] == "clean_data" + assert result["steps"][1]["name"] == "compute_stats" + + +class TestChildMode: + def test_child_executes_matching_task(self): + checkpoint = {"completed_steps": {"timestamp": 1234567890}, "_executing_key": "double"} + result = _run_workflow(step_workflow, checkpoint, {"x": 7}) + assert result["type"] == "complete" + assert result["result"] == 14 + + def test_child_replays_cached_steps(self): + checkpoint = { + "completed_steps": {"extract_data": {"data": [1, 2, 3]}}, + "_executing_key": "load_data", + } + result = _run_workflow(simple_workflow, checkpoint, {"url": "https://example.com"}) + assert result["type"] == "complete" + assert result["result"] is None + + +# ===================================================================== +# EDGE CASE TESTS +# ===================================================================== + +class TestFullSequentialLifecycle: + def test_replay_0_dispatches_step_0(self): + result = _run_workflow(three_step_wf, {}, {"n": 5}) + assert result["type"] == "dispatch" + assert result["steps"][0]["key"] == "double" + assert result["steps"][0]["name"] == "double" + assert result["steps"][0]["args"] == {"x": 5} + + def test_replay_1_dispatches_step_1_with_step_0_result(self): + result = _run_workflow(three_step_wf, {"completed_steps": {"double": 10}}, {"n": 5}) + assert result["type"] == "dispatch" + assert result["steps"][0]["key"] == "add_one" + assert result["steps"][0]["name"] == "add_one" + assert result["steps"][0]["args"] == {"x": 10} + + def test_replay_2_dispatches_step_2_with_step_1_result(self): + result = _run_workflow( + three_step_wf, {"completed_steps": {"double": 10, "add_one": 11}}, {"n": 5} + ) + assert result["type"] == "dispatch" + assert result["steps"][0]["key"] == "double_2" + assert result["steps"][0]["name"] == "double" + assert result["steps"][0]["args"] == {"x": 11} + + def test_replay_3_all_complete(self): + result = _run_workflow( + three_step_wf, + {"completed_steps": {"double": 10, "add_one": 11, "double_2": 22}}, + {"n": 5}, + ) + assert result["type"] == "complete" + assert result["result"] == {"doubled": 10, "incremented": 11, "final": 22} + + +class TestStepAfterParallelGroup: + def test_dispatches_first_sequential(self): + result = _run_workflow(seq_par_seq_wf, {}, {"url": "http://x"}) + assert result["steps"][0]["name"] == "extract_data" + + def test_dispatches_parallel_group(self): + result = _run_workflow( + seq_par_seq_wf, {"completed_steps": {"extract_data": "raw"}}, {"url": "http://x"} + ) + assert result["mode"] == "parallel" + assert len(result["steps"]) == 2 + + def test_dispatches_final_step_after_parallel(self): + result = _run_workflow( + seq_par_seq_wf, + {"completed_steps": {"extract_data": "raw", "clean_data": "cleaned", "compute_stats": {"count": 5}}}, + {"url": "http://x"}, + ) + assert result["mode"] == "sequential" + assert result["steps"][0]["name"] == "load_data" + assert result["steps"][0]["key"] == "load_data" + + def test_completes_when_final_step_done(self): + result = _run_workflow( + seq_par_seq_wf, + {"completed_steps": {"extract_data": "raw", "clean_data": "cleaned", "compute_stats": {"count": 5}, "load_data": "final"}}, + {"url": "http://x"}, + ) + assert result["type"] == "complete" + assert result["result"] == "final" + + +class TestParallelAfterParallel: + def test_dispatches_first_parallel(self): + result = _run_workflow(double_parallel_wf, {}, {}) + assert result["mode"] == "parallel" + assert len(result["steps"]) == 2 + assert result["steps"][0]["key"] == "double" + assert result["steps"][1]["key"] == "double_2" + + def test_dispatches_second_parallel(self): + result = _run_workflow( + double_parallel_wf, {"completed_steps": {"double": 2, "double_2": 4}}, {} + ) + assert result["mode"] == "parallel" + assert len(result["steps"]) == 2 + assert result["steps"][0]["name"] == "add_one" + assert result["steps"][0]["args"] == {"x": 2} + assert result["steps"][1]["args"] == {"x": 4} + + def test_completes_all_done(self): + result = _run_workflow( + double_parallel_wf, + {"completed_steps": {"double": 2, "double_2": 4, "add_one": 3, "add_one_2": 5}}, + {}, + ) + assert result["type"] == "complete" + assert result["result"] == {"a": 2, "b": 4, "c": 3, "d": 5} + + +class TestConditionalBasedOnStepResult: + def test_condition_true_path(self): + result = _run_workflow(cond_on_result_wf, {"completed_steps": {"double": 10}}, {}) + assert result["steps"][0]["name"] == "send_alert" + assert result["steps"][0]["key"] == "send_alert" + + def test_condition_false_path(self): + result = _run_workflow(cond_on_result_wf, {"completed_steps": {"double": 4}}, {}) + assert result["steps"][0]["name"] == "load_data" + assert result["steps"][0]["key"] == "load_data" + + def test_condition_true_step_after_alert(self): + result = _run_workflow( + cond_on_result_wf, {"completed_steps": {"double": 10, "send_alert": "alerted"}}, {} + ) + assert result["steps"][0]["name"] == "load_data" + assert result["steps"][0]["key"] == "load_data" + + +class TestEmptyWorkflow: + def test_completes_immediately(self): + result = _run_workflow(empty_wf, {}, {}) + assert result["type"] == "complete" + assert result["result"] == {"status": "empty"} + + +class TestSingleTaskWorkflow: + def test_dispatches_single_step(self): + result = _run_workflow(single_wf, {}, {"x": 7}) + assert result["type"] == "dispatch" + assert len(result["steps"]) == 1 + assert result["steps"][0]["name"] == "double" + + def test_completes_with_result(self): + result = _run_workflow(single_wf, {"completed_steps": {"double": 14}}, {"x": 7}) + assert result["type"] == "complete" + assert result["result"] == 14 + + +class TestTaskWithNoArgs: + def test_dispatches_with_empty_args(self): + result = _run_workflow(no_arg_wf, {}, {}) + assert result["type"] == "dispatch" + assert result["steps"][0]["args"] == {} + + +class TestManySteps: + def test_first_dispatches_step_0(self): + result = _run_workflow(many_steps_wf, {}, {"n": 0}) + assert result["steps"][0]["key"] == "add_one" + + def test_with_5_complete_dispatches_step_5(self): + # add_one, add_one_2, add_one_3, add_one_4, add_one_5 + completed = {} + for i in range(5): + key = "add_one" if i == 0 else f"add_one_{i + 1}" + completed[key] = i + 1 + result = _run_workflow(many_steps_wf, {"completed_steps": completed}, {"n": 0}) + assert result["steps"][0]["key"] == "add_one_6" + assert result["steps"][0]["args"] == {"x": 5} + + def test_all_10_complete(self): + completed = {} + for i in range(10): + key = "add_one" if i == 0 else f"add_one_{i + 1}" + completed[key] = i + 1 + result = _run_workflow(many_steps_wf, {"completed_steps": completed}, {"n": 0}) + assert result["type"] == "complete" + assert result["result"] == 10 + + +class TestFalsyValues: + def test_zero_preserved(self): + result = _run_workflow(falsy_wf, {"completed_steps": {"double": 0}}, {}) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "load_data" + assert result["steps"][0]["args"] == {"data": 0} + + def test_none_preserved(self): + result = _run_workflow(falsy_wf, {"completed_steps": {"double": 0, "load_data": None}}, {}) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "extract_data" + + def test_all_falsy_complete(self): + result = _run_workflow( + falsy_wf, {"completed_steps": {"double": 0, "load_data": None, "extract_data": ""}}, {} + ) + assert result["type"] == "complete" + assert result["result"] == {"a": 0, "b": None, "c": ""} + + def test_false_preserved(self): + @workflow + async def flag_wf(): + val = await load_data(data="check") + if val: + await send_alert(msg="truthy") + return {"val": val} + + result = _run_workflow(flag_wf, {"completed_steps": {"load_data": False}}, {}) + assert result["type"] == "complete" + assert result["result"] == {"val": False} + + +class TestTaskWithExplicitPath: + def test_uses_path_as_script(self): + result = _run_workflow(path_wf, {}, {"x": 42}) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "run_external" + assert result["steps"][0]["script"] == "f/external_script" + assert result["steps"][0]["args"] == {"x": 42} + + +class TestMixedStepAndTask: + def test_step_0_inline(self): + result = _run_workflow(mixed_step_task_wf, {}, {"x": 5}) + assert result["type"] == "inline_checkpoint" + assert result["key"] == "get_time" + assert result["result"] == 999 + + def test_step_1_task_dispatch(self): + result = _run_workflow( + mixed_step_task_wf, {"completed_steps": {"get_time": 999}}, {"x": 5} + ) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "double" + assert result["steps"][0]["key"] == "double" + + def test_step_2_inline(self): + result = _run_workflow( + mixed_step_task_wf, + {"completed_steps": {"get_time": 999, "double": 10}}, + {"x": 5}, + ) + assert result["type"] == "inline_checkpoint" + assert result["key"] == "get_config" + assert result["result"] == {"retry": 3} + + def test_step_3_task_dispatch(self): + result = _run_workflow( + mixed_step_task_wf, + {"completed_steps": {"get_time": 999, "double": 10, "get_config": {"retry": 3}}}, + {"x": 5}, + ) + assert result["type"] == "dispatch" + assert result["steps"][0]["name"] == "add_one" + assert result["steps"][0]["key"] == "add_one" + + def test_all_complete(self): + result = _run_workflow( + mixed_step_task_wf, + {"completed_steps": {"get_time": 999, "double": 10, "get_config": {"retry": 3}, "add_one": 11}}, + {"x": 5}, + ) + assert result["type"] == "complete" + assert result["result"] == {"ts": 999, "doubled": 10, "config": {"retry": 3}, "added": 11} + + +class TestChildModeParallel: + def test_child_executes_first_parallel_step(self): + result = _run_workflow( + par_child_wf, {"completed_steps": {}, "_executing_key": "double"}, {} + ) + assert result["type"] == "complete" + assert result["result"] == 6 + + def test_child_executes_second_parallel_step(self): + result = _run_workflow( + par_child_wf, {"completed_steps": {}, "_executing_key": "add_one"}, {} + ) + assert result["type"] == "complete" + assert result["result"] == 8 + + +class TestKeyDeterminism: + def test_keys_consistent_across_replays(self): + r1 = _run_workflow(det_wf, {}, {"n": 3}) + assert r1["steps"][0]["key"] == "double" + assert r1["steps"][0]["name"] == "double" + + r2 = _run_workflow(det_wf, {"completed_steps": {"double": 6}}, {"n": 3}) + assert r2["steps"][0]["key"] == "add_one" + assert r2["steps"][0]["name"] == "add_one" + + r3 = _run_workflow(det_wf, {"completed_steps": {"double": 6, "add_one": 7}}, {"n": 3}) + assert r3["steps"][0]["key"] == "double_2" + assert r3["steps"][0]["name"] == "double" + + +class TestParallelArgsFromCachedResult: + def test_parallel_steps_receive_cached_args(self): + result = _run_workflow(par_args_wf, {"completed_steps": {"double": 20}}, {"x": 10}) + assert result["mode"] == "parallel" + assert result["steps"][0]["args"] == {"x": 20} + assert result["steps"][1]["args"] == {"x": 20} + + +class TestWorkflowReturningNone: + def test_none_return_captured(self): + result = _run_workflow(none_return_wf, {"completed_steps": {"double": 2}}, {}) + assert result["type"] == "complete" + assert result["result"] is None + + +class TestLargeParallelGroup: + def test_dispatches_5_parallel(self): + result = _run_workflow(large_par_wf, {}, {}) + assert result["mode"] == "parallel" + assert len(result["steps"]) == 5 + keys = [result["steps"][i]["key"] for i in range(5)] + assert keys == ["double", "double_2", "double_3", "double_4", "double_5"] + for i in range(5): + assert result["steps"][i]["args"] == {"x": i + 1} + + +class TestComplexMixedWorkflow: + def test_replay_0_extract(self): + r = _run_workflow(complex_mixed_wf, {}, {}) + assert r["steps"][0]["name"] == "extract_data" + + def test_replay_1_parallel(self): + r = _run_workflow(complex_mixed_wf, {"completed_steps": {"extract_data": "init"}}, {}) + assert r["mode"] == "parallel" + assert len(r["steps"]) == 2 + + def test_replay_2_load(self): + r = _run_workflow( + complex_mixed_wf, + {"completed_steps": {"extract_data": "init", "double": 2, "double_2": 4}}, + {}, + ) + assert r["mode"] == "sequential" + assert r["steps"][0]["name"] == "load_data" + assert r["steps"][0]["key"] == "load_data" + + def test_replay_3_second_parallel(self): + r = _run_workflow( + complex_mixed_wf, + {"completed_steps": {"extract_data": "init", "double": 2, "double_2": 4, "load_data": "mid"}}, + {}, + ) + assert r["mode"] == "parallel" + assert len(r["steps"]) == 2 + assert r["steps"][0]["name"] == "add_one" + + def test_replay_4_clean(self): + r = _run_workflow( + complex_mixed_wf, + {"completed_steps": { + "extract_data": "init", "double": 2, "double_2": 4, + "load_data": "mid", "add_one": 4, "add_one_2": 5, + }}, + {}, + ) + assert r["mode"] == "sequential" + assert r["steps"][0]["name"] == "clean_data" + assert r["steps"][0]["key"] == "clean_data" + + def test_replay_5_all_complete(self): + r = _run_workflow( + complex_mixed_wf, + {"completed_steps": { + "extract_data": "init", "double": 2, "double_2": 4, + "load_data": "mid", "add_one": 4, "add_one_2": 5, "clean_data": "final", + }}, + {}, + ) + assert r["type"] == "complete" + assert r["result"] == "final" + + +class TestChildModeWithCachedStepsBeforeParallel: + def test_child_executes_second_parallel_with_cached_base(self): + result = _run_workflow( + pre_par_child_wf, + {"completed_steps": {"double": 10}, "_executing_key": "double_2"}, + {"x": 5}, + ) + assert result["type"] == "complete" + assert result["result"] == 20 + + def test_child_executes_first_parallel_with_cached_base(self): + result = _run_workflow( + pre_par_child_wf, + {"completed_steps": {"double": 10}, "_executing_key": "add_one"}, + {"x": 5}, + ) + assert result["type"] == "complete" + assert result["result"] == 11 + + +# ===================================================================== +# ERROR PROPAGATION TESTS +# ===================================================================== + + +class TestErrorPropagation: + def test_task_error_is_raised_on_replay(self): + @workflow + async def wf(x: int): + return await double(x=x) + + with pytest.raises(TaskError, match="double"): + _run_workflow( + wf, + { + "completed_steps": { + "double": { + "__wmill_error": True, + "message": "Task 'double' failed", + "result": {"message": "boom"}, + } + } + }, + {"x": 5}, + ) + + def test_error_catchable_with_try_except(self): + @workflow + async def wf(x: int): + try: + result = await double(x=x) + return {"success": True, "result": result} + except Exception as e: + return {"success": False, "error": str(e)} + + r = _run_workflow( + wf, + { + "completed_steps": { + "double": { + "__wmill_error": True, + "message": "Task 'double' failed", + "result": {}, + } + } + }, + {"x": 5}, + ) + assert r["type"] == "complete" + assert r["result"]["success"] is False + assert "double" in r["result"]["error"] + + def test_retry_pattern_with_try_except_loop(self): + @workflow + async def wf(x: int): + for i in range(3): + try: + result = await double(x=x) + return {"result": result, "attempts": i + 1} + except Exception: + if i == 2: + raise + + # First double fails, second succeeds + r = _run_workflow( + wf, + { + "completed_steps": { + "double": {"__wmill_error": True, "message": "temporary", "result": {}}, + "double_2": 10, + } + }, + {"x": 5}, + ) + assert r["type"] == "complete" + assert r["result"]["result"] == 10 + assert r["result"]["attempts"] == 2 + + def test_non_error_object_with_error_false(self): + @workflow + async def wf(): + val = await double(x=5) + return val + + r = _run_workflow( + wf, + {"completed_steps": {"double": {"__wmill_error": False, "data": "ok"}}}, + {}, + ) + assert r["type"] == "complete" + assert r["result"] == {"__wmill_error": False, "data": "ok"} + + def test_inline_step_error(self): + @workflow + async def wf(): + try: + val = await step("risky", lambda: 42) + return {"val": val} + except Exception as e: + return {"caught": str(e)} + + r = _run_workflow( + wf, + {"completed_steps": {"risky": {"__wmill_error": True, "message": "step failed", "result": {}}}}, + {}, + ) + assert r["type"] == "complete" + assert "step failed" in r["result"]["caught"] + + +# ===================================================================== +# TASK OPTIONS TESTS +# ===================================================================== + + +class TestTaskOptions: + def test_options_forwarded_in_dispatch(self): + @task(timeout=600, tag="gpu", cache_ttl=3600, priority=10) + async def heavy(x: int): + return x + + @workflow + async def wf(x: int): + return await heavy(x=x) + + r = _run_workflow(wf, {}, {"x": 42}) + assert r["type"] == "dispatch" + step_info = r["steps"][0] + assert step_info["timeout"] == 600 + assert step_info["tag"] == "gpu" + assert step_info["cache_ttl"] == 3600 + assert step_info["priority"] == 10 + + def test_task_without_options_has_no_extra_fields(self): + @task + async def simple(x: int): + return x + + @workflow + async def wf(x: int): + return await simple(x=x) + + r = _run_workflow(wf, {}, {"x": 1}) + step_info = r["steps"][0] + assert "timeout" not in step_info + assert "tag" not in step_info + + def test_concurrency_options_forwarded(self): + @task(concurrency_limit=5, concurrency_key="my-key", concurrency_time_window_s=60) + async def limited(x: int): + return x + + @workflow + async def wf(x: int): + return await limited(x=x) + + r = _run_workflow(wf, {}, {"x": 1}) + step_info = r["steps"][0] + assert step_info["concurrent_limit"] == 5 + assert step_info["concurrency_key"] == "my-key" + assert step_info["concurrency_time_window_s"] == 60 + + +# ===================================================================== +# SLEEP TESTS +# ===================================================================== + + +class TestSleep: + def test_sleep_returns_sleep_output(self): + @workflow + async def wf(): + await double(x=1) + await sleep(60) + await add_one(x=2) + return "done" + + r = _run_workflow( + wf, + {"completed_steps": {"double": 2}}, + {}, + ) + assert r["type"] == "sleep" + assert r["key"] == "sleep" + assert r["seconds"] == 60 + + def test_sleep_completes_on_replay(self): + @workflow + async def wf(): + await double(x=1) + await sleep(60) + await add_one(x=2) + return "done" + + r = _run_workflow( + wf, + {"completed_steps": {"double": 2, "sleep": True}}, + {}, + ) + assert r["type"] == "dispatch" + assert r["steps"][0]["name"] == "add_one" + assert r["steps"][0]["key"] == "add_one" + + def test_all_steps_with_sleep_complete(self): + @workflow + async def wf(): + await double(x=1) + await sleep(60) + await add_one(x=2) + return "done" + + r = _run_workflow( + wf, + {"completed_steps": {"double": 2, "sleep": True, "add_one": 3}}, + {}, + ) + assert r["type"] == "complete" + assert r["result"] == "done" + + def test_sleep_enforces_minimum(self): + @workflow + async def wf(): + await sleep(0) + return "done" + + r = _run_workflow(wf, {}, {}) + assert r["seconds"] == 1 + + +# ===================================================================== +# PARALLEL UTILITY TESTS +# ===================================================================== + + +class TestParallel: + def test_dispatches_all_items(self): + @workflow + async def wf(): + results = await parallel([1, 2, 3], double) + return results + + r = _run_workflow(wf, {}, {}) + assert r["type"] == "dispatch" + assert r["mode"] == "parallel" + assert len(r["steps"]) == 3 + + def test_completes_with_all_results(self): + @workflow + async def wf(): + results = await parallel([1, 2, 3], double) + return results + + r = _run_workflow( + wf, + {"completed_steps": {"double": 2, "double_2": 4, "double_3": 6}}, + {}, + ) + assert r["type"] == "complete" + assert r["result"] == [2, 4, 6] + + def test_batched_dispatches_first_batch(self): + @workflow + async def wf(): + results = await parallel([1, 2, 3, 4, 5], double, concurrency=2) + return results + + r = _run_workflow(wf, {}, {}) + assert r["type"] == "dispatch" + assert r["mode"] == "parallel" + assert len(r["steps"]) == 2 + + def test_batched_dispatches_second_batch(self): + @workflow + async def wf(): + results = await parallel([1, 2, 3, 4, 5], double, concurrency=2) + return results + + r = _run_workflow( + wf, + {"completed_steps": {"double": 2, "double_2": 4}}, + {}, + ) + assert r["type"] == "dispatch" + assert len(r["steps"]) == 2 + + def test_batched_completes_with_all_results(self): + @workflow + async def wf(): + results = await parallel([1, 2, 3, 4, 5], double, concurrency=2) + return results + + r = _run_workflow( + wf, + {"completed_steps": {"double": 2, "double_2": 4, "double_3": 6, "double_4": 8, "double_5": 10}}, + {}, + ) + assert r["type"] == "complete" + assert r["result"] == [2, 4, 6, 8, 10] + + def test_empty_items_returns_empty(self): + @workflow + async def wf(): + results = await parallel([], double) + return results + + r = _run_workflow(wf, {}, {}) + assert r["type"] == "complete" + assert r["result"] == [] diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index 3b9a9f3f45..f5acd7f58c 100644 --- a/python-client/wmill/wmill/client.py +++ b/python-client/wmill/wmill/client.py @@ -2151,69 +2151,6 @@ def ducklake(name: str = "main") -> DucklakeClient: """ return _client.ducklake(name) -def task(*args, **kwargs): - """Decorator to mark a function as a workflow task. - - When executed inside a Windmill job, the decorated function runs as a - separate workflow step. Outside Windmill, it executes normally. - - Args: - tag: Optional worker tag for execution - - Returns: - Decorated function - """ - from inspect import signature - - def f(func, tag: str | None = None): - if ( - os.environ.get("WM_JOB_ID") is None - or os.environ.get("MAIN_OVERRIDE") == func.__name__ - ): - - def inner(*args, **kwargs): - return func(*args, **kwargs) - - return inner - else: - - def inner(*args, **kwargs): - global _client - if _client is None: - _client = Windmill() - w_id = os.environ.get("WM_WORKSPACE") - job_id = os.environ.get("WM_JOB_ID") - f_name = func.__name__ - json = kwargs - params = list(signature(func).parameters) - for i, arg in enumerate(args): - if i < len(params): - p = params[i] - key = p - if key not in kwargs: - json[key] = arg - - params = {} - if tag is not None: - params["tag"] = tag - w_as_code_response = _client.post( - f"/w/{w_id}/jobs/run/workflow_as_code/{job_id}/{f_name}", - json={"args": json}, - params=params, - ) - job_id = w_as_code_response.text - print(f"Executing task {func.__name__} on job {job_id}") - job_result = _client.wait_job(job_id) - print(f"Task {func.__name__} ({job_id}) completed") - return job_result - - return inner - - if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): - return f(args[0], None) - else: - return lambda x: f(x, kwargs.get("tag")) - def parse_resource_syntax(s: str) -> Optional[str]: """Parse resource syntax from string.""" if s is None: @@ -2413,7 +2350,495 @@ def parse_sql_client_name(name: str) -> tuple[str, Optional[str]]: name = name schema = None if ":" in name: - name, schema = name.split(":", 1) + name, schema = name.split(":", 1) if not name: name = "main" return name, schema + + +# ── Workflow-as-Code SDK ────────────────────────────────────────────── + +import asyncio as _asyncio +import contextvars as _contextvars + + +class _StepSuspend(BaseException): + """Raised to suspend workflow execution. Inherits from BaseException + so it is not caught by bare `except Exception:` blocks.""" + + def __init__(self, dispatch_info: dict): + self.dispatch_info = dispatch_info + + +class TaskError(Exception): + """Raised when a WAC task step failed. + + Attributes: + step_key: The checkpoint key of the failed step. + child_job_id: The UUID of the failed child job. + result: The error result from the child job. + """ + + def __init__(self, message: str, *, step_key: str = "", child_job_id: str = "", result=None): + super().__init__(message) + self.step_key = step_key + self.child_job_id = child_job_id + self.result = result + + +_workflow_ctx: _contextvars.ContextVar["WorkflowCtx"] = _contextvars.ContextVar( + "_workflow_ctx" +) + + +class WorkflowCtx: + """Internal context for workflow replay/suspension. + + Not user-facing — set implicitly by ``@workflow`` via contextvars. + """ + + def __init__(self, checkpoint: dict | None = None): + checkpoint = checkpoint or {} + self._completed: dict = checkpoint.get("completed_steps", {}) + self._counters: dict[str, int] = {} + self._pending: list = [] + self._executing_key: str | None = checkpoint.get("_executing_key") + + def _alloc_key(self, name: str = "step") -> str: + """Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent.""" + n = self._counters.get(name, 0) + 1 + self._counters[name] = n + return name if n == 1 else f"{name}_{n}" + + def _next_step(self, name: str, script: str, func=None, dispatch_type: str = "inline", _task_options: Optional[dict] = None, **kwargs): + """Return an awaitable that either resolves from cache or suspends.""" + key = self._alloc_key(name or script or "step") + + if key in self._completed: + val = self._completed[key] + if isinstance(val, dict) and val.get("__wmill_error"): + raise TaskError( + val.get("message", f"Task '{name}' failed"), + step_key=val.get("step_key", ""), + child_job_id=val.get("child_job_id", ""), + result=val.get("result"), + ) + return self._resolved(val) + + if self._executing_key is not None: + if key == self._executing_key: + return self._execute_directly(func, **kwargs) + else: + return self._never_resolve() + + info = {"name": name or key, "script": script or key, "args": kwargs, "key": key, "dispatch_type": dispatch_type} + if _task_options: + for opt_key in ("timeout", "tag", "cache_ttl", "priority", "concurrent_limit", "concurrency_key", "concurrency_time_window_s"): + if opt_key in _task_options and _task_options[opt_key] is not None: + info[opt_key] = _task_options[opt_key] + self._pending.append(info) + return self._suspend() + + async def _resolved(self, value): + return value + + async def _execute_directly(self, func, **kwargs): + result = func(**kwargs) + if _asyncio.iscoroutine(result): + result = await result + raise _StepSuspend({"mode": "step_complete", "steps": [], "result": result}) + + async def _never_resolve(self): + await _asyncio.Future() + + async def _suspend(self): + steps = list(self._pending) + self._pending.clear() + raise _StepSuspend( + { + "mode": "parallel" if len(steps) > 1 else "sequential", + "steps": steps, + } + ) + + async def _wait_for_approval( + self, timeout: int = 1800, form: dict | None = None + ): + key = self._alloc_key("approval") + + if key in self._completed: + return self._completed[key] + + if self._executing_key is not None: + await _asyncio.Future() + + raise _StepSuspend({ + "mode": "approval", + "key": key, + "timeout": timeout, + "form": form, + "steps": [], + }) + + async def _sleep(self, seconds: int): + key = self._alloc_key("sleep") + + if key in self._completed: + return + + if self._executing_key is not None: + await _asyncio.Future() + + raise _StepSuspend({ + "mode": "sleep", + "key": key, + "seconds": max(1, int(seconds)), + "steps": [], + }) + + async def _run_inline_step(self, name: str, fn): + key = self._alloc_key(name or "step") + + if key in self._completed: + val = self._completed[key] + if isinstance(val, dict) and val.get("__wmill_error"): + raise TaskError( + val.get("message", f"Step '{name}' failed"), + step_key=val.get("step_key", ""), + child_job_id=val.get("child_job_id", ""), + result=val.get("result"), + ) + return val + + if self._executing_key is not None: + await _asyncio.Future() + + result = fn() + if _asyncio.iscoroutine(result): + result = await result + + raise _StepSuspend({ + "mode": "inline_checkpoint", + "steps": [], + "key": key, + "result": result, + }) + + +def task( + _func=None, + *, + path: Optional[str] = None, + tag: Optional[str] = None, + timeout: Optional[int] = None, + cache_ttl: Optional[int] = None, + priority: Optional[int] = None, + concurrency_limit: Optional[int] = None, + concurrency_key: Optional[str] = None, + concurrency_time_window_s: Optional[int] = None, +): + """Decorator that marks a function as a workflow task. + + Works in both WAC v1 (sync, HTTP-based dispatch) and WAC v2 + (async, checkpoint/replay) modes: + + - **v2 (inside @workflow)**: dispatches as a checkpoint step. + - **v1 (WM_JOB_ID set, no @workflow)**: dispatches via HTTP API. + - **Standalone**: executes the function body directly. + + Usage:: + + @task + async def extract_data(url: str): ... + + @task(path="f/external_script", timeout=600, tag="gpu") + async def run_external(x: int): ... + """ + from inspect import signature as _sig + + _task_opts = { + "timeout": timeout, + "tag": tag, + "cache_ttl": cache_ttl, + "priority": priority, + "concurrent_limit": concurrency_limit, + "concurrency_key": concurrency_key, + "concurrency_time_window_s": concurrency_time_window_s, + } + # Remove None values + _task_opts = {k: v for k, v in _task_opts.items() if v is not None} or None + + def decorator(func): + task_path = path + task_name = func.__name__ + + _params_list = list(_sig(func).parameters) + + def _merge_args(args, kwargs): + merged = dict(kwargs) + for i, arg in enumerate(args): + if i < len(_params_list): + key = _params_list[i] + if key not in merged: + merged[key] = arg + else: + merged[f"arg{i}"] = arg + return merged + + @functools.wraps(func) + def wrapper(*args, **kwargs): + # WAC v2: inside a @workflow context + ctx = _workflow_ctx.get(None) + if ctx is not None: + script = task_path if task_path else task_name + merged = _merge_args(args, kwargs) + return ctx._next_step(task_name, script, func, _task_options=_task_opts, **merged) + + # WAC v1: running inside a Windmill job but not in a @workflow + if ( + os.environ.get("WM_JOB_ID") is not None + and os.environ.get("MAIN_OVERRIDE") != func.__name__ + ): + global _client + if _client is None: + _client = Windmill() + w_id = os.environ.get("WM_WORKSPACE") + job_id = os.environ.get("WM_JOB_ID") + json_args = _merge_args(args, kwargs) + api_params = {} + if tag is not None: + api_params["tag"] = tag + resp = _client.post( + f"/w/{w_id}/jobs/run/workflow_as_code/{job_id}/{func.__name__}", + json={"args": json_args}, + params=api_params, + ) + child_job_id = resp.text + print(f"Executing task {func.__name__} on job {child_job_id}") + job_result = _client.wait_job(child_job_id) + print(f"Task {func.__name__} ({child_job_id}) completed") + return job_result + + # Standalone — execute directly + return func(*args, **kwargs) + + wrapper._is_task = True + wrapper._task_path = task_path + return wrapper + + if _func is not None: + # @task without parentheses + return decorator(_func) + # @task() or @task(path="...", tag="...") + return decorator + + +def task_script( + path: str, + *, + timeout: Optional[int] = None, + tag: Optional[str] = None, + cache_ttl: Optional[int] = None, + priority: Optional[int] = None, + concurrency_limit: Optional[int] = None, + concurrency_key: Optional[str] = None, + concurrency_time_window_s: Optional[int] = None, +): + """Create a task that dispatches to a separate Windmill script. + + Usage:: + + extract = task_script("f/data/extract", timeout=600) + + @workflow + async def main(): + data = await extract(url="https://...") + """ + name = path.rsplit("/", 1)[-1] + _opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s}.items() if v is not None} or None + + def wrapper(**kwargs): + ctx = _workflow_ctx.get(None) + if ctx is not None: + return ctx._next_step(name, path, dispatch_type="script", _task_options=_opts, **kwargs) + raise RuntimeError(f'task_script("{path}") can only be called inside a @workflow') + + wrapper.__name__ = name + wrapper._is_task = True + wrapper._task_path = path + return wrapper + + +def task_flow( + path: str, + *, + timeout: Optional[int] = None, + tag: Optional[str] = None, + cache_ttl: Optional[int] = None, + priority: Optional[int] = None, + concurrency_limit: Optional[int] = None, + concurrency_key: Optional[str] = None, + concurrency_time_window_s: Optional[int] = None, +): + """Create a task that dispatches to a separate Windmill flow. + + Usage:: + + pipeline = task_flow("f/etl/pipeline", priority=10) + + @workflow + async def main(): + result = await pipeline(input=data) + """ + name = path.rsplit("/", 1)[-1] + _opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s}.items() if v is not None} or None + + def wrapper(**kwargs): + ctx = _workflow_ctx.get(None) + if ctx is not None: + return ctx._next_step(name, path, dispatch_type="flow", _task_options=_opts, **kwargs) + raise RuntimeError(f'task_flow("{path}") can only be called inside a @workflow') + + wrapper.__name__ = name + wrapper._is_task = True + wrapper._task_path = path + return wrapper + + +def workflow(func): + """Decorator marking an async function as a workflow-as-code entry point. + + The function must be **deterministic**: given the same inputs it must call + tasks in the same order on every replay. Branching on task results is fine + (results are replayed from checkpoint), but branching on external state + (current time, random values, external API calls) must use ``step()`` to + checkpoint the value so replays see the same result. + """ + func._is_workflow = True + return func + + +async def step(name: str, fn): + """Execute ``fn`` inline and checkpoint the result. + + On replay the cached value is returned without re-executing ``fn``. + Use for lightweight deterministic operations (timestamps, random IDs, + config reads) that should not incur the overhead of a child job. + """ + ctx: WorkflowCtx | None = _workflow_ctx.get(None) + if ctx is not None: + return await ctx._run_inline_step(name, fn) + result = fn() + if _asyncio.iscoroutine(result): + result = await result + return result + + +async def sleep(seconds: int): + """Server-side sleep — suspend the workflow for the given duration without holding a worker. + + Inside a @workflow, the parent job suspends and auto-resumes after ``seconds``. + Outside a workflow, falls back to ``asyncio.sleep``. + """ + ctx: WorkflowCtx | None = _workflow_ctx.get(None) + if ctx is not None: + return await ctx._sleep(seconds) + await _asyncio.sleep(seconds) + + +async def wait_for_approval( + timeout: int = 1800, + form: dict | None = None, +) -> dict: + """Suspend the workflow and wait for an external approval. + + Use ``get_resume_urls()`` (wrapped in ``step()``) to obtain + resume/cancel/approval URLs before calling this function. + + Returns a dict with ``value`` (form data), ``approver``, and ``approved``. + + Example:: + + urls = await step("urls", lambda: get_resume_urls()) + await step("notify", lambda: send_email(urls["approvalPage"])) + result = await wait_for_approval(timeout=3600) + """ + ctx: WorkflowCtx | None = _workflow_ctx.get(None) + if ctx is not None: + return await ctx._wait_for_approval(timeout=timeout, form=form) + raise RuntimeError("wait_for_approval can only be called inside a @workflow") + + +async def parallel(items, fn, *, concurrency: Optional[int] = None): + """Process items in parallel with optional concurrency control. + + Each item is processed by calling ``fn(item)``, which should be a @task. + Items are dispatched in batches of ``concurrency`` (default: all at once). + + Example:: + + @task + async def process(item: str): + ... + + results = await parallel(items, process, concurrency=5) + """ + if not items: + return [] + batch_size = concurrency if concurrency and concurrency > 0 else len(items) + results = [] + for i in range(0, len(items), batch_size): + batch = items[i : i + batch_size] + batch_results = await _asyncio.gather(*(fn(item) for item in batch)) + results.extend(batch_results) + return results + + +async def _run_workflow_async(func, checkpoint: dict, input_args: dict): + ctx = WorkflowCtx(checkpoint) + token = _workflow_ctx.set(ctx) + try: + result = await func(**input_args) + # Flush any unawaited tasks (e.g. forgotten await on last statement) + if ctx._pending: + steps = list(ctx._pending) + ctx._pending.clear() + return { + "type": "dispatch", + "mode": "parallel" if len(steps) > 1 else "sequential", + "steps": steps, + } + return {"type": "complete", "result": result} + except _StepSuspend as e: + info = e.dispatch_info + mode = info.get("mode") + if mode == "step_complete": + return {"type": "complete", "result": info.get("result")} + if mode == "inline_checkpoint": + return { + "type": "inline_checkpoint", + "key": info["key"], + "result": info.get("result"), + } + if mode == "approval": + return { + "type": "approval", + "key": info["key"], + "timeout": info.get("timeout"), + "form": info.get("form"), + } + if mode == "sleep": + return { + "type": "sleep", + "key": info["key"], + "seconds": info.get("seconds"), + } + return {"type": "dispatch", **info} + finally: + _workflow_ctx.reset(token) + + +def _run_workflow(func, checkpoint: dict, input_args: dict): + """Synchronous wrapper that runs the workflow coroutine to completion + or until it suspends.""" + return _asyncio.run(_run_workflow_async(func, checkpoint, input_args)) diff --git a/typescript-client/build.sh b/typescript-client/build.sh index bea84d0ebc..f765340bad 100755 --- a/typescript-client/build.sh +++ b/typescript-client/build.sh @@ -40,7 +40,7 @@ cp "${script_dirpath}/sqlUtils.ts" "${script_dirpath}/src/" echo "" >> "${script_dirpath}/src/index.ts" echo 'export type { DenoS3LightClientSettings } from "./s3Types";' >> "${script_dirpath}/src/index.ts" echo "" >> "${script_dirpath}/src/index.ts" -echo 'export { type Base64, setClient, getVariable, setVariable, getResource, setResource, getResumeUrls, setState, setProgress, getProgress, getState, getIdToken, denoS3LightClientSettings, loadS3FileStream, loadS3File, writeS3File, signS3Objects, signS3Object, getPresignedS3PublicUrls, getPresignedS3PublicUrl, task, runScript, runScriptAsync, runScriptByPath, runScriptByHash, runScriptByPathAsync, runScriptByHashAsync, runFlow, runFlowAsync, waitJob, getRootJobId, setFlowUserState, getFlowUserState, usernameToEmail, requestInteractiveSlackApproval, type Sql, requestInteractiveTeamsApproval, appendToResultStream, streamResult, datatable, ducklake, type DatatableSqlTemplateFunction, type SqlTemplateFunction, type S3Object, type S3ObjectRecord, type S3ObjectURI } from "./client";' >> "${script_dirpath}/src/index.ts" +echo 'export { type Base64, setClient, getVariable, setVariable, getResource, setResource, getResumeUrls, setState, setProgress, getProgress, getState, getIdToken, denoS3LightClientSettings, loadS3FileStream, loadS3File, writeS3File, signS3Objects, signS3Object, getPresignedS3PublicUrls, getPresignedS3PublicUrl, task, taskScript, taskFlow, workflow, step, sleep, parallel, waitForApproval, type TaskOptions, WorkflowCtx, _workflowCtx, setWorkflowCtx, StepSuspend, runScript, runScriptAsync, runScriptByPath, runScriptByHash, runScriptByPathAsync, runScriptByHashAsync, runFlow, runFlowAsync, waitJob, getRootJobId, setFlowUserState, getFlowUserState, usernameToEmail, requestInteractiveSlackApproval, type Sql, requestInteractiveTeamsApproval, appendToResultStream, streamResult, datatable, ducklake, type DatatableSqlTemplateFunction, type SqlTemplateFunction, type S3Object, type S3ObjectRecord, type S3ObjectURI } from "./client";' >> "${script_dirpath}/src/index.ts" # Build default export by combining client utilities + services # This preserves backward compatibility for `import wmill from "windmill-client"` @@ -68,6 +68,17 @@ import { getPresignedS3PublicUrls, getPresignedS3PublicUrl, task, + taskScript, + taskFlow, + workflow, + step, + sleep, + parallel, + waitForApproval, + WorkflowCtx, + _workflowCtx, + setWorkflowCtx, + StepSuspend, runScript, runScriptAsync, runScriptByPath, @@ -141,6 +152,17 @@ const wmill = { getPresignedS3PublicUrls, getPresignedS3PublicUrl, task, + taskScript, + taskFlow, + workflow, + step, + sleep, + parallel, + waitForApproval, + WorkflowCtx, + _workflowCtx, + setWorkflowCtx, + StepSuspend, runScript, runScriptAsync, runScriptByPath, diff --git a/typescript-client/client.ts b/typescript-client/client.ts index 1b32a4119e..41054f5334 100644 --- a/typescript-client/client.ts +++ b/typescript-client/client.ts @@ -311,45 +311,45 @@ export async function getResultMaybe(jobId: string): Promise { } const STRIP_COMMENTS = /(\/\/.*$)|(\/\*[\s\S]*?\*\/)|(\s*=[^,\)]*(('(?:\\'|[^'\r\n])*')|("(?:\\"|[^"\r\n])*"))|(\s*=[^,\)]*))/gm; -const ARGUMENT_NAMES = /([^\s,]+)/g; function getParamNames(func: Function): string[] { const fnStr = func.toString().replace(STRIP_COMMENTS, ""); - let result: string[] | null = fnStr - .slice(fnStr.indexOf("(") + 1, fnStr.indexOf(")")) - .match(ARGUMENT_NAMES); - if (result === null) result = []; - return result; -} - -/** - * Wrap a function to execute as a Windmill task within a flow context - * @param f - Function to wrap as a task - * @returns Async wrapper function that executes as a Windmill job - */ -export function task(f: (_: P) => T): (_: P) => Promise { - return async (...y) => { - const args: Record = {}; - const paramNames = getParamNames(f); - y.forEach((x, i) => (args[paramNames[i]] = x)); - let req = await fetch( - `${OpenAPI.BASE}/w/${getWorkspace()}/jobs/run/workflow_as_code/${getEnv( - "WM_JOB_ID" - )}/${f.name}`, - { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${getEnv("WM_TOKEN")}`, - }, - body: JSON.stringify({ args }), - } - ); - let jobId = await req.text(); - console.log(`Started task ${f.name} as job ${jobId}`); - let r = await waitJob(jobId); - console.log(`Task ${f.name} (${jobId}) completed`); - return r; - }; + // Find the matching closing paren for the parameter list, handling nesting + const openIdx = fnStr.indexOf("("); + if (openIdx === -1) return []; + let depth = 1; + let closeIdx = openIdx + 1; + for (; closeIdx < fnStr.length && depth > 0; closeIdx++) { + if (fnStr[closeIdx] === "(") depth++; + else if (fnStr[closeIdx] === ")") depth--; + } + const paramStr = fnStr.slice(openIdx + 1, closeIdx - 1).trim(); + if (!paramStr) return []; + // Split on commas at depth 0 (skip nested parens, angle brackets, braces) + const params: string[] = []; + let current = ""; + let d = 0; + for (const ch of paramStr) { + if ("(<{".includes(ch)) d++; + else if (")>}".includes(ch)) d--; + if (ch === "," && d === 0) { + params.push(current.trim()); + current = ""; + } else { + current += ch; + } + } + if (current.trim()) params.push(current.trim()); + // Extract the parameter name from each param (strip type annotations, destructuring, rest) + return params.map((p) => { + // Remove rest operator + p = p.replace(/^\.\.\./, ""); + // For destructured params like { url, depth }: Config, use a positional fallback + if (p.startsWith("{") || p.startsWith("[")) return ""; + // Strip type annotation (e.g. "x: number" -> "x", "x?: string" -> "x") + const colonIdx = p.indexOf(":"); + if (colonIdx !== -1) p = p.slice(0, colonIdx); + return p.replace(/\?$/, "").trim(); + }).filter(Boolean); } /** @@ -1448,3 +1448,421 @@ export function parseS3Object(s3Object: S3Object): S3ObjectRecord { function parseVariableSyntax(s: string) { if (s.startsWith("var://")) return s.substring(6); } + +// ── Workflow-as-Code SDK ────────────────────────────────────────────── + +export class StepSuspend extends Error { + constructor(public dispatchInfo: Record) { + super("__step_suspend__"); + this.name = "StepSuspend"; + } +} + +export interface TaskOptions { + timeout?: number; + tag?: string; + cache_ttl?: number; + priority?: number; + concurrency_limit?: number; + concurrency_key?: string; + concurrency_time_window_s?: number; +} + +export let _workflowCtx: WorkflowCtx | null = null; +export function setWorkflowCtx(ctx: WorkflowCtx | null) { + _workflowCtx = ctx; + Reflect.set(globalThis, "__wmill_wf_ctx", ctx); +} + + +export class WorkflowCtx { + private completed: Record; + private counters: Record = {}; + private pending: Array<{ + name: string; + script: string; + args: Record; + key: string; + dispatch_type: string; + [k: string]: any; + }> = []; + private _suspended = false; + /** When set, the task matching this key executes its inner function directly */ + _executingKey: string | null; + + constructor(checkpoint: Record = {}) { + this.completed = checkpoint?.completed_steps ?? {}; + this._executingKey = checkpoint?._executing_key ?? null; + } + + /** Name-based key: `double` for first call, `double_2`, `double_3` for subsequent. */ + _allocKey(name: string): string { + const n = (this.counters[name] ?? 0) + 1; + this.counters[name] = n; + return n === 1 ? name : `${name}_${n}`; + } + + _nextStep( + name: string, + script: string, + args: Record = {}, + dispatch_type: string = "inline", + options?: TaskOptions, + ): PromiseLike { + const key = this._allocKey(name || script || "step"); + + if (key in this.completed) { + const value = this.completed[key]; + if (value && typeof value === "object" && (value as any).__wmill_error) { + const err = new Error((value as any).message || `Task '${name}' failed`); + (err as any).result = (value as any).result; + (err as any).step_key = (value as any).step_key; + (err as any).child_job_id = (value as any).child_job_id; + return { then: (_resolve: any, reject?: any) => { if (reject) reject(err); else throw err; } } as PromiseLike; + } + return { then: (resolve: any) => resolve(value) }; + } + + // If this is a child job executing a specific step, return null to signal + // that the task wrapper should run the inner function directly + if (this._executingKey === key) { + return { then: (resolve: any) => resolve(null), _execute_directly: true } as any; + } + + // In child job mode (_executingKey is set), non-matching uncompleted steps + // should never resolve or throw — the matching step will throw step_complete + // which terminates the workflow. Returning a never-resolving thenable prevents + // race conditions where a non-matching step's StepSuspend fires before step_complete. + if (this._executingKey !== null) { + return { then: () => new Promise(() => {}) }; + } + + const stepInfo: any = { name: name || key, script: script || key, args, key, dispatch_type }; + if (options) { + if (options.timeout !== undefined) stepInfo.timeout = options.timeout; + if (options.tag !== undefined) stepInfo.tag = options.tag; + if (options.cache_ttl !== undefined) stepInfo.cache_ttl = options.cache_ttl; + if (options.priority !== undefined) stepInfo.priority = options.priority; + if (options.concurrency_limit !== undefined) stepInfo.concurrent_limit = options.concurrency_limit; + if (options.concurrency_key !== undefined) stepInfo.concurrency_key = options.concurrency_key; + if (options.concurrency_time_window_s !== undefined) stepInfo.concurrency_time_window_s = options.concurrency_time_window_s; + } + this.pending.push(stepInfo); + return { + then: (): never => { + // Only the first .then() call throws with all accumulated steps. + // Subsequent calls (e.g. from Promise.all resolving other thenables) + // also throw (they'll be caught by the same handler). + if (this._suspended) return new Promise(() => {}) as never; + this._suspended = true; + const steps = [...this.pending]; + this.pending = []; + throw new StepSuspend({ + mode: steps.length > 1 ? "parallel" : "sequential", + steps, + }); + }, + }; + } + /** Return and clear any pending (unawaited) steps. */ + _flushPending(): Array<{ name: string; script: string; args: Record; key: string; dispatch_type: string }> { + const steps = [...this.pending]; + this.pending = []; + return steps; + } + + _waitForApproval(options?: { + timeout?: number; + form?: object; + }): PromiseLike<{ value: any; approver: string; approved: boolean }> { + const key = this._allocKey("approval"); + + if (key in this.completed) { + const value = this.completed[key]; + return { then: (resolve: any) => resolve(value) }; + } + + // In child job mode, return never-resolving thenable (same as _nextStep) + if (this._executingKey !== null) { + return { then: () => new Promise(() => {}) }; + } + + // Throw immediately — approval is always a blocking step + throw new StepSuspend({ + mode: "approval", + key, + timeout: options?.timeout ?? 1800, + form: options?.form, + steps: [], + }); + } + + _sleep(seconds: number): PromiseLike { + const key = this._allocKey("sleep"); + + if (key in this.completed) { + return { then: (resolve: any) => resolve(undefined) }; + } + + if (this._executingKey !== null) { + return { then: () => new Promise(() => {}) }; + } + + throw new StepSuspend({ + mode: "sleep", + key, + seconds: Math.max(1, Math.round(seconds)), + steps: [], + }); + } + + async _runInlineStep(name: string, fn: () => T | Promise): Promise { + const key = this._allocKey(name || "step"); + + if (key in this.completed) { + const value = this.completed[key]; + if (value && typeof value === "object" && (value as any).__wmill_error) { + const err = new Error((value as any).message || `Step '${name}' failed`); + (err as any).result = (value as any).result; + (err as any).step_key = (value as any).step_key; + (err as any).child_job_id = (value as any).child_job_id; + throw err; + } + return value as T; + } + + if (this._executingKey !== null) { + return new Promise(() => {}); + } + + const result = await fn(); + throw new StepSuspend({ mode: "inline_checkpoint", steps: [], key, result }); + } +} + +export async function sleep(seconds: number): Promise { + const ctx: WorkflowCtx | null = _workflowCtx ?? Reflect.get(globalThis, "__wmill_wf_ctx"); + if (ctx) { + return ctx._sleep(seconds) as Promise; + } + // Outside workflow context, just wait locally + await new Promise((r) => setTimeout(r, seconds * 1000)); +} + +export async function step(name: string, fn: () => T | Promise): Promise { + const ctx: WorkflowCtx | null = _workflowCtx ?? Reflect.get(globalThis, "__wmill_wf_ctx"); + if (ctx) { + return ctx._runInlineStep(name, fn); + } + return fn(); +} + +/** + * Wrap an async function as a workflow task. + * + * @example + * const extract_data = task(async (url: string) => { ... }); + * const run_external = task("f/external_script", async (x: number) => { ... }); + * + * Inside a `workflow()`, calling a task dispatches it as a step. + * Outside a workflow, the function body executes directly. + */ +export function task Promise>( + fnOrPath: T | string, + maybeFnOrOptions?: T | TaskOptions, + maybeOptions?: TaskOptions, +): T { + let fn: T; + let taskPath: string | undefined; + let taskOptions: TaskOptions | undefined; + + if (typeof fnOrPath === "string") { + taskPath = fnOrPath; + fn = maybeFnOrOptions as T; + taskOptions = maybeOptions; + } else { + fn = fnOrPath; + taskOptions = maybeFnOrOptions as TaskOptions | undefined; + } + + const taskName = fn.name || taskPath || ""; + + // NOT async — in workflow context we return the thenable directly so that + // unawaited task calls leave the step in ctx.pending (for _flushPending). + // An async wrapper would auto-resolve the thenable in a microtask, calling + // .then() which throws StepSuspend and empties pending before the caller + // can flush. + const wrapper = function (...args: any[]) { + const ctx: WorkflowCtx | null = _workflowCtx ?? Reflect.get(globalThis, "__wmill_wf_ctx"); + if (ctx) { + // Inside a workflow with checkpoint/replay context — dispatch as step + const script = taskPath ?? taskName; + const paramNames = getParamNames(fn); + const kwargs: Record = {}; + for (let i = 0; i < args.length; i++) { + if (paramNames[i]) { + kwargs[paramNames[i]] = args[i]; + } else { + kwargs[`arg${i}`] = args[i]; + } + } + const stepResult = ctx._nextStep(taskName, script, kwargs, "inline", taskOptions); + // If this step should execute directly (child job mode), run the inner function + // and throw StepSuspend with mode "step_complete" to signal that we're done + if ((stepResult as any)?._execute_directly) { + return (async () => { + const result = await fn(...args); + throw new StepSuspend({ mode: "step_complete", steps: [], result }); + })(); + } + return stepResult; + } else if (getEnv("WM_JOB_ID") && !getEnv("WM_FLOW_JOB_ID")) { + // Inside a Windmill root job without checkpoint context — v1 HTTP dispatch + // WM_FLOW_JOB_ID is set on child jobs, so we skip dispatch for those + return (async () => { + const paramNames = getParamNames(fn); + const kwargs: Record = {}; + args.forEach((x, i) => (kwargs[paramNames[i]] = x)); + let req = await fetch( + `${OpenAPI.BASE}/w/${getWorkspace()}/jobs/run/workflow_as_code/${getEnv( + "WM_JOB_ID" + )}/${taskName}`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${getEnv("WM_TOKEN")}`, + }, + body: JSON.stringify({ args: kwargs }), + } + ); + let jobId = await req.text(); + console.log(`Started task ${taskName} as job ${jobId}`); + let r = await waitJob(jobId); + console.log(`Task ${taskName} (${jobId}) completed`); + return r; + })(); + } else { + // Standalone — execute directly + return fn(...args); + } + } as unknown as T; + + Object.defineProperty(wrapper, "name", { value: taskName }); + (wrapper as any)._is_task = true; + (wrapper as any)._task_path = taskPath; + return wrapper; +} + +/** + * Create a task that dispatches to a separate Windmill script. + * + * @example + * const extract = taskScript("f/data/extract"); + * // inside workflow: await extract({ url: "https://..." }) + */ +export function taskScript(path: string, options?: TaskOptions): (...args: any[]) => PromiseLike { + const name = path.split("/").pop() || path; + const wrapper = function (...args: any[]) { + const ctx: WorkflowCtx | null = _workflowCtx ?? Reflect.get(globalThis, "__wmill_wf_ctx"); + if (ctx) { + const kwargs = args.length === 1 && typeof args[0] === "object" && args[0] !== null + ? args[0] + : args.reduce((acc, v, i) => { acc[`arg${i}`] = v; return acc; }, {} as Record); + return ctx._nextStep(name, path, kwargs, "script", options); + } + throw new Error(`taskScript("${path}") can only be called inside a workflow()`); + }; + Object.defineProperty(wrapper, "name", { value: name }); + (wrapper as any)._is_task = true; + (wrapper as any)._task_path = path; + return wrapper; +} + +/** + * Create a task that dispatches to a separate Windmill flow. + * + * @example + * const pipeline = taskFlow("f/etl/pipeline"); + * // inside workflow: await pipeline({ input: data }) + */ +export function taskFlow(path: string, options?: TaskOptions): (...args: any[]) => PromiseLike { + const name = path.split("/").pop() || path; + const wrapper = function (...args: any[]) { + const ctx: WorkflowCtx | null = _workflowCtx ?? Reflect.get(globalThis, "__wmill_wf_ctx"); + if (ctx) { + const kwargs = args.length === 1 && typeof args[0] === "object" && args[0] !== null + ? args[0] + : args.reduce((acc, v, i) => { acc[`arg${i}`] = v; return acc; }, {} as Record); + return ctx._nextStep(name, path, kwargs, "flow", options); + } + throw new Error(`taskFlow("${path}") can only be called inside a workflow()`); + }; + Object.defineProperty(wrapper, "name", { value: name }); + (wrapper as any)._is_task = true; + (wrapper as any)._task_path = path; + return wrapper; +} + +/** + * Mark an async function as a workflow-as-code entry point. + * + * The function must be **deterministic**: given the same inputs it must call + * tasks in the same order on every replay. Branching on task results is fine + * (results are replayed from checkpoint), but branching on external state + * (current time, random values, external API calls) must use `step()` to + * checkpoint the value so replays see the same result. + */ +export function workflow(fn: (...args: any[]) => Promise) { + (fn as any)._is_workflow = true; + return fn; +} + +/** + * Suspend the workflow and wait for an external approval. + * + * Use `getResumeUrls()` (wrapped in `step()`) to obtain resume/cancel/approvalPage + * URLs before calling this function. + * + * @example + * const urls = await step("urls", () => getResumeUrls()); + * await step("notify", () => sendEmail(urls.approvalPage)); + * const { value, approver } = await waitForApproval({ timeout: 3600 }); + */ +export function waitForApproval(options?: { + timeout?: number; + form?: object; +}): PromiseLike<{ value: any; approver: string; approved: boolean }> { + const ctx: WorkflowCtx | null = _workflowCtx ?? Reflect.get(globalThis, "__wmill_wf_ctx"); + if (!ctx) { + throw new Error("waitForApproval can only be called inside a workflow()"); + } + return ctx._waitForApproval(options); +} + +/** + * Process items in parallel with optional concurrency control. + * + * Each item is processed by calling `fn(item)`, which should be a task(). + * Items are dispatched in batches of `concurrency` (default: all at once). + * + * @example + * const process = task(async (item: string) => { ... }); + * const results = await parallel(items, process, { concurrency: 5 }); + */ +export async function parallel( + items: T[], + fn: (item: T) => PromiseLike | R, + options?: { concurrency?: number }, +): Promise { + const concurrency = options?.concurrency ?? items.length; + if (concurrency <= 0 || items.length === 0) return []; + const results: R[] = []; + for (let i = 0; i < items.length; i += concurrency) { + const batch = items.slice(i, i + concurrency); + const batchResults = await Promise.all(batch.map((item) => fn(item))); + results.push(...batchResults); + } + return results; +} + diff --git a/typescript-client/package-lock.json b/typescript-client/package-lock.json index e299444082..a25b51cbcb 100644 --- a/typescript-client/package-lock.json +++ b/typescript-client/package-lock.json @@ -1,12 +1,12 @@ { "name": "windmill-client", - "version": "1.618.3", + "version": "1.651.1", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "windmill-client", - "version": "1.618.3", + "version": "1.651.1", "license": "Apache 2.0", "devDependencies": { "@types/node": "^20.17.16", diff --git a/typescript-client/tests/e2e_wac.py b/typescript-client/tests/e2e_wac.py new file mode 100644 index 0000000000..28bb629bbe --- /dev/null +++ b/typescript-client/tests/e2e_wac.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" +E2E test for WAC v2 (Workflow-as-Code) with the _executing_key approach. + +This test: +1. Creates a preview job with a WAC v2 bun script +2. Waits for the parent to suspend (dispatch child jobs) +3. Waits for child jobs to complete +4. Waits for parent to unsuspend and complete +5. Checks the final result + +Usage: + python3 typescript-client/tests/e2e_wac.py +""" +import json +import sys +import time +import urllib.request + +BASE = "http://localhost:8000" +TOKEN = "" # Will be fetched +WORKSPACE = "admins" + +def api(method, path, data=None): + url = f"{BASE}/api{path}" + headers = {"Content-Type": "application/json"} + if TOKEN: + headers["Authorization"] = f"Bearer {TOKEN}" + body = json.dumps(data).encode() if data else None + req = urllib.request.Request(url, data=body, headers=headers, method=method) + try: + with urllib.request.urlopen(req) as resp: + raw = resp.read().decode() + try: + return json.loads(raw, strict=False) + except: + return raw + except urllib.error.HTTPError as e: + body = e.read().decode() + print(f"HTTP {e.code} {method} {path}: {body[:500]}") + raise + +def login(): + global TOKEN + TOKEN = "PdxixPjjfx05H8xJ8kWAll4RtiLGcfXW" + # Verify token works + user = api("GET", "/users/whoami") + print(f"Logged in as: {user.get('email', 'unknown')}") + +def run_preview(code, language="bun"): + """Run a preview job and return the job ID.""" + result = api("POST", f"/w/{WORKSPACE}/jobs/run/preview", { + "content": code, + "language": language, + "args": {"n": 10}, + }) + print(f"Preview job created: {result}") + return result + +def get_job(job_id): + return api("GET", f"/w/{WORKSPACE}/jobs_u/get/{job_id}") + +def get_result(job_id): + return api("GET", f"/w/{WORKSPACE}/jobs_u/completed/get_result/{job_id}") + +def wait_for_job(job_id, timeout=60, check_interval=2): + """Wait for a job to complete. Returns the job object.""" + start = time.time() + while time.time() - start < timeout: + job = get_job(job_id) + job_type = job.get("type", "") + if job_type == "CompletedJob": + return job + # Print status + suspend = job.get("suspend", 0) + status = "suspended" if suspend and suspend > 0 else "running" + print(f" Job {job_id[:8]}... status={status} suspend={suspend} ({time.time()-start:.0f}s)") + time.sleep(check_interval) + raise TimeoutError(f"Job {job_id} did not complete within {timeout}s") + + +WAC_SCRIPT = ''' +import { task, workflow } from "windmill-client"; + +const double = task(async function double(x: number): Promise { + return x * 2; +}); + +const add_one = task(async function add_one(x: number): Promise { + return x + 1; +}); + +export default workflow(async function main(n: number) { + const doubled = await double(n); + const result = await add_one(doubled); + return { doubled, result }; +}); +''' + +def main(): + print("=== WAC v2 E2E Test ===\n") + + # 1. Login + login() + + # 2. Run the WAC preview + print(f"\nRunning WAC v2 preview...") + job_id = run_preview(WAC_SCRIPT) + + # 3. Wait for completion + print(f"\nWaiting for job {job_id} to complete...") + job = wait_for_job(job_id, timeout=120) + + success = job.get("success", False) + result = job.get("result") + + print(f"\nJob completed! success={success}") + print(f"Result: {json.dumps(result, indent=2)}") + + if not success: + print("\nFAILED: Job did not succeed") + # Print logs if available + logs = job.get("logs", "") + if logs: + print(f"\nLogs:\n{logs}") + sys.exit(1) + + # 4. Verify result + expected = {"doubled": 20, "result": 21} + if result == expected: + print(f"\nSUCCESS: Sequential workflow result matches expected {expected}") + else: + print(f"\nFAILED: Expected {expected}, got {result}") + sys.exit(1) + + # 5. Test parallel workflow + print("\n\n=== Parallel Workflow Test ===\n") + parallel_job_id = run_preview(PARALLEL_WAC_SCRIPT) + print(f"\nWaiting for parallel job {parallel_job_id} to complete...") + parallel_job = wait_for_job(parallel_job_id, timeout=120) + + p_success = parallel_job.get("success", False) + p_result = parallel_job.get("result") + + print(f"\nParallel job completed! success={p_success}") + print(f"Result: {json.dumps(p_result, indent=2)}") + + if not p_success: + print("\nFAILED: Parallel job did not succeed") + sys.exit(1) + + p_expected = {"doubled": 20, "incremented": 11, "combined": 31} + if p_result == p_expected: + print(f"\nSUCCESS: Parallel workflow result matches expected {p_expected}") + else: + print(f"\nFAILED: Expected {p_expected}, got {p_result}") + sys.exit(1) + + print("\n\n=== ALL TESTS PASSED ===") + +PARALLEL_WAC_SCRIPT = ''' +import { task, workflow } from "windmill-client"; + +const double = task(async function double(x: number): Promise { + return x * 2; +}); + +const increment = task(async function increment(x: number): Promise { + return x + 1; +}); + +export default workflow(async function main(n: number) { + const [doubled, incremented] = await Promise.all([ + double(n), + increment(n), + ]); + return { doubled, incremented, combined: doubled + incremented }; +}); +''' + +if __name__ == "__main__": + main() diff --git a/typescript-client/tests/e2e_wac_v1.py b/typescript-client/tests/e2e_wac_v1.py new file mode 100644 index 0000000000..64c8f6a0cc --- /dev/null +++ b/typescript-client/tests/e2e_wac_v1.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +""" +E2E test for WAC v1 (Workflow-as-Code) — HTTP-dispatch mode. + +WAC v1 scripts use @task / task() but NOT @workflow / workflow(). +Tasks dispatch via HTTP POST to /jobs/run/workflow_as_code/{job_id}/{task_name}. + +This verifies that v1 still works after the v2 client changes. + +Usage: + python3 typescript-client/tests/e2e_wac_v1.py +""" +import json +import sys +import time +import urllib.request + +BASE = "http://localhost:8000" +TOKEN = "" +WORKSPACE = "dev" + + +def api(method, path, data=None): + url = f"{BASE}/api{path}" + headers = {"Content-Type": "application/json"} + if TOKEN: + headers["Authorization"] = f"Bearer {TOKEN}" + body = json.dumps(data).encode() if data else None + req = urllib.request.Request(url, data=body, headers=headers, method=method) + try: + with urllib.request.urlopen(req) as resp: + raw = resp.read().decode() + try: + return json.loads(raw, strict=False) + except Exception: + return raw + except urllib.error.HTTPError as e: + body = e.read().decode() + print(f"HTTP {e.code} {method} {path}: {body[:500]}") + raise + + +def login(): + global TOKEN + resp = api("POST", "/auth/login", {"email": "admin@windmill.dev", "password": "changeme"}) + TOKEN = resp + user = api("GET", "/users/whoami") + print(f"Logged in as: {user.get('email', 'unknown')}") + + +def run_preview(code, language="bun", args=None): + result = api("POST", f"/w/{WORKSPACE}/jobs/run/preview", { + "content": code, + "language": language, + "args": args or {}, + }) + print(f" Preview job created: {result}") + return result + + +def get_job(job_id): + return api("GET", f"/w/{WORKSPACE}/jobs_u/get/{job_id}") + + +def wait_for_job(job_id, timeout=120, check_interval=2): + start = time.time() + while time.time() - start < timeout: + job = get_job(job_id) + if job.get("type") == "CompletedJob": + return job + elapsed = time.time() - start + print(f" {job_id[:8]}... waiting ({elapsed:.0f}s)") + time.sleep(check_interval) + raise TimeoutError(f"Job {job_id} did not complete within {timeout}s") + + +def check_result(job, expected, label): + success = job.get("success", False) + result = job.get("result") + print(f" success={success} result={json.dumps(result)}") + if not success: + logs = job.get("logs", "") + print(f" FAILED: job did not succeed\n Logs:\n{logs}") + sys.exit(1) + if result != expected: + print(f" FAILED [{label}]: expected {expected}, got {result}") + sys.exit(1) + print(f" PASSED [{label}]") + + +# --------------------------------------------------------------------------- +# WAC v1 TypeScript — no workflow() wrapper, tasks dispatch via HTTP +# --------------------------------------------------------------------------- +TS_V1_SEQUENTIAL = ''' +import { task } from "windmill-client"; + +export const double = task(async function double(x: number): Promise { + return x * 2; +}); + +export const add_one = task(async function add_one(x: number): Promise { + return x + 1; +}); + +export async function main(n: number) { + const doubled = await double(n); + const result = await add_one(doubled); + return { doubled, result }; +} +''' + +TS_V1_MULTI_PARAM = ''' +import { task } from "windmill-client"; + +export const add = task(async function add(a: number, b: number): Promise { + return a + b; +}); + +export async function main(x: number) { + const result = await add(x, 100); + return { result }; +} +''' + +# --------------------------------------------------------------------------- +# WAC v1 Python — no @workflow, tasks dispatch via HTTP +# --------------------------------------------------------------------------- +PY_V1_SEQUENTIAL = ''' +import wmill + +@wmill.task +def double(x: int) -> int: + return x * 2 + +@wmill.task +def add_one(x: int) -> int: + return x + 1 + +def main(n: int): + doubled = double(x=n) + result = add_one(x=doubled) + return {"doubled": doubled, "result": result} +''' + +PY_V1_MULTI_PARAM = ''' +import wmill + +@wmill.task +def add(a: int, b: int) -> int: + return a + b + +def main(x: int): + result = add(a=x, b=100) + return {"result": result} +''' + + +def main(): + print("=== WAC v1 E2E Tests ===\n") + login() + + # --- TypeScript v1: sequential --- + print("\n[1] TypeScript v1 — sequential tasks") + job_id = run_preview(TS_V1_SEQUENTIAL, "bun", {"n": 10}) + job = wait_for_job(job_id) + check_result(job, {"doubled": 20, "result": 21}, "ts_v1_sequential") + + # --- TypeScript v1: multi-param --- + print("\n[2] TypeScript v1 — multi-param task") + job_id = run_preview(TS_V1_MULTI_PARAM, "bun", {"x": 42}) + job = wait_for_job(job_id) + check_result(job, {"result": 142}, "ts_v1_multi_param") + + # --- Python v1: sequential --- + print("\n[3] Python v1 — sequential tasks") + job_id = run_preview(PY_V1_SEQUENTIAL, "python3", {"n": 10}) + job = wait_for_job(job_id) + check_result(job, {"doubled": 20, "result": 21}, "py_v1_sequential") + + # --- Python v1: multi-param --- + print("\n[4] Python v1 — multi-param task") + job_id = run_preview(PY_V1_MULTI_PARAM, "python3", {"x": 42}) + job = wait_for_job(job_id) + check_result(job, {"result": 142}, "py_v1_multi_param") + + print("\n\n=== ALL WAC v1 TESTS PASSED ===") + + +if __name__ == "__main__": + main() diff --git a/typescript-client/tests/workflow.test.ts b/typescript-client/tests/workflow.test.ts new file mode 100644 index 0000000000..a42895ccbb --- /dev/null +++ b/typescript-client/tests/workflow.test.ts @@ -0,0 +1,1596 @@ +/** + * Standalone tests for the Workflow-as-Code TypeScript SDK. + * + * Run with: bun test typescript-client/tests/workflow.test.ts + */ +import { expect, test, describe } from "bun:test"; + +// --- Inline SDK (mirrors client.ts implementation) --- + +class StepSuspend extends Error { + constructor(public dispatchInfo: Record) { + super("__step_suspend__"); + this.name = "StepSuspend"; + } +} + +let _workflowCtx: WorkflowCtx | null = null; + +class WorkflowCtx { + private completed: Record; + private stepIndex = 0; + private pending: Array<{ + name: string; + script: string; + args: Record; + key: string; + }> = []; + private _suspended = false; + _executingKey: string | null; + + constructor(checkpoint: Record = {}) { + this.completed = checkpoint?.completed_steps ?? {}; + this._executingKey = checkpoint?._executing_key ?? null; + } + + _allocKey(): string { + return `step_${this.stepIndex++}`; + } + + _nextStep( + name: string, + script: string, + args: Record = {}, + options?: Record, + ): PromiseLike { + const key = this._allocKey(); + + if (key in this.completed) { + const value = this.completed[key]; + if (value && typeof value === "object" && (value as any).__wmill_error) { + const err = new Error((value as any).message || `Task '${name}' failed`); + (err as any).result = (value as any).result; + (err as any).step_key = (value as any).step_key; + (err as any).child_job_id = (value as any).child_job_id; + return { then: (_resolve: any, reject?: any) => { if (reject) reject(err); else throw err; } }; + } + return { then: (resolve: any) => resolve(value) }; + } + + // Child job mode: execute matching step directly + if (this._executingKey === key) { + return { + then: (resolve: any) => resolve(null), + _execute_directly: true, + } as any; + } + + // Child job mode: non-matching steps never resolve + if (this._executingKey !== null) { + return { then: () => new Promise(() => {}) }; + } + + const stepInfo: any = { name, script, args, key }; + if (options) Object.assign(stepInfo, options); + this.pending.push(stepInfo); + return { + then: (): never => { + if (this._suspended) return new Promise(() => {}) as never; + this._suspended = true; + const steps = [...this.pending]; + this.pending = []; + throw new StepSuspend({ + mode: steps.length > 1 ? "parallel" : "sequential", + steps, + }); + }, + }; + } + + _flushPending(): Array<{ + name: string; + script: string; + args: Record; + key: string; + }> { + const steps = [...this.pending]; + this.pending = []; + return steps; + } + + _sleep(seconds: number): PromiseLike { + const key = this._allocKey(); + if (key in this.completed) { + return { then: (resolve: any) => resolve(undefined) }; + } + if (this._executingKey !== null) { + return { then: () => new Promise(() => {}) }; + } + throw new StepSuspend({ + mode: "sleep", + key, + seconds: Math.max(1, Math.round(seconds)), + steps: [], + }); + } + + async _runInlineStep( + name: string, + fn: () => T | Promise + ): Promise { + const key = this._allocKey(); + + if (key in this.completed) { + const value = this.completed[key]; + if (value && typeof value === "object" && (value as any).__wmill_error) { + const err = new Error((value as any).message || `Step '${name}' failed`); + (err as any).result = (value as any).result; + throw err; + } + return value as T; + } + + if (this._executingKey !== null) { + return new Promise(() => {}); + } + + const result = await fn(); + throw new StepSuspend({ + mode: "inline_checkpoint", + steps: [], + key, + result, + }); + } +} + +function getParamNames(fn: Function): string[] { + const src = fn.toString(); + const match = src.match(/^(?:async\s+)?(?:function\s*\w*)?\s*\(([^)]*)\)/); + if (!match) return []; + return match[1] + .split(",") + .map((p) => p.trim().replace(/[:=].*/s, "").trim()) + .filter(Boolean); +} + +function task Promise>( + fnOrPath: T | string, + maybeFnOrOptions?: T | Record, + maybeOptions?: Record, +): T { + let fn: T; + let taskPath: string | undefined; + let taskOptions: Record | undefined; + + if (typeof fnOrPath === "string") { + taskPath = fnOrPath; + fn = maybeFnOrOptions as T; + taskOptions = maybeOptions; + } else { + fn = fnOrPath; + taskOptions = maybeFnOrOptions as Record | undefined; + } + + const taskName = fn.name || "anonymous"; + + // Non-async wrapper — returns thenable directly in workflow context so + // unawaited calls leave steps in pending for _flushPending. + const wrapper = function (...args: any[]) { + const ctx = _workflowCtx; + if (ctx) { + const script = taskPath ?? taskName; + const paramNames = getParamNames(fn); + const kwargs: Record = {}; + for (let i = 0; i < args.length; i++) { + if (paramNames[i]) { + kwargs[paramNames[i]] = args[i]; + } else { + kwargs[`arg${i}`] = args[i]; + } + } + const stepResult = ctx._nextStep(taskName, script, kwargs, taskOptions); + if ((stepResult as any)?._execute_directly) { + return (async () => { + const result = await fn(...args); + throw new StepSuspend({ + mode: "step_complete", + steps: [], + result, + }); + })(); + } + return stepResult; + } else { + return fn(...args); + } + } as unknown as T; + + Object.defineProperty(wrapper, "name", { value: taskName }); + (wrapper as any)._is_task = true; + (wrapper as any)._task_path = taskPath; + return wrapper; +} + +async function step( + name: string, + fn: () => T | Promise +): Promise { + const ctx = _workflowCtx; + if (ctx) { + return ctx._runInlineStep(name, fn); + } + return fn(); +} + +async function sleep(seconds: number): Promise { + const ctx = _workflowCtx; + if (ctx) { + return ctx._sleep(seconds) as Promise; + } + await new Promise((r) => setTimeout(r, seconds * 1000)); +} + +async function parallel( + items: T[], + fn: (item: T) => PromiseLike | R, + options?: { concurrency?: number }, +): Promise { + const concurrency = options?.concurrency ?? items.length; + if (concurrency <= 0 || items.length === 0) return []; + const results: R[] = []; + for (let i = 0; i < items.length; i += concurrency) { + const batch = items.slice(i, i + concurrency); + const batchResults = await Promise.all(batch.map((item) => fn(item))); + results.push(...batchResults); + } + return results; +} + +function workflow(fn: (...args: any[]) => Promise) { + (fn as any)._is_workflow = true; + return fn; +} + +// --- Helper to run a workflow with a checkpoint --- + +async function runWorkflow( + fn: Function, + checkpoint: Record, + args: any[] +): Promise { + const ctx = new WorkflowCtx(checkpoint); + _workflowCtx = ctx; + try { + const result = await fn(...args); + // Flush unawaited tasks + const pending = ctx._flushPending(); + if (pending.length > 0) { + return { + type: "dispatch", + mode: pending.length > 1 ? "parallel" : "sequential", + steps: pending, + }; + } + return { type: "complete", result }; + } catch (e: any) { + if (e instanceof StepSuspend) { + const info = e.dispatchInfo; + if (info.mode === "step_complete") { + return { type: "complete", result: info.result }; + } + if (info.mode === "inline_checkpoint") { + return { + type: "inline_checkpoint", + key: info.key, + result: info.result, + }; + } + if (info.mode === "approval") { + return { type: "approval", key: info.key, timeout: info.timeout, form: info.form }; + } + if (info.mode === "sleep") { + return { type: "sleep", key: info.key, seconds: info.seconds }; + } + return { type: "dispatch", ...info }; + } + throw e; + } finally { + _workflowCtx = null; + } +} + +// --- Define tasks --- + +const extract_data = task(async function extract_data(url: string) {}); +const load_data = task(async function load_data(data?: any) {}); +const clean_data = task(async function clean_data(data?: any) {}); +const compute_stats = task(async function compute_stats(data?: any) {}); +const send_alert = task(async function send_alert(msg: string) {}); +const double = task(async function double(x: number) { + return x * 2; +}); +const add_one = task(async function add_one(x: number) { + return x + 1; +}); +const noop_task = task(async function noop_task() {}); + +// --- Define workflows --- + +const simple_workflow = workflow(async (url: string) => { + const raw = await extract_data(url); + const result = await load_data(raw); + return { status: "done", result }; +}); + +const parallel_workflow = workflow(async (url: string) => { + const raw = await extract_data(url); + const [cleaned, stats] = await Promise.all([ + clean_data(raw), + compute_stats(raw), + ]); + return { cleaned, stats }; +}); + +const conditional_workflow = workflow(async (count: number) => { + if (count > 100) { + await send_alert("large"); + } + await load_data(); + return { done: true }; +}); + +// ===================================================================== +// TESTS +// ===================================================================== + +describe("task decorator", () => { + test("marks function as task", () => { + expect((extract_data as any)._is_task).toBe(true); + }); + + test("standalone execution runs body directly", async () => { + const result = await extract_data("https://example.com"); + expect(result).toBeUndefined(); + }); + + test("preserves function name", () => { + expect(extract_data.name).toBe("extract_data"); + expect(double.name).toBe("double"); + }); +}); + +describe("workflow decorator", () => { + test("marks function as workflow", () => { + expect((simple_workflow as any)._is_workflow).toBe(true); + }); +}); + +describe("first invocation", () => { + test("dispatches first step", async () => { + const result = await runWorkflow(simple_workflow, {}, [ + "https://example.com", + ]); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("sequential"); + expect(result.steps).toHaveLength(1); + expect(result.steps[0].name).toBe("extract_data"); + expect(result.steps[0].script).toBe("extract_data"); + expect(result.steps[0].key).toBe("step_0"); + expect(result.steps[0].args).toEqual({ url: "https://example.com" }); + }); +}); + +describe("replay with checkpoint", () => { + test("second invocation dispatches second step", async () => { + const checkpoint = { + completed_steps: { step_0: [1, 2, 3] }, + }; + const result = await runWorkflow(simple_workflow, checkpoint, [ + "https://example.com", + ]); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("sequential"); + expect(result.steps[0].name).toBe("load_data"); + expect(result.steps[0].key).toBe("step_1"); + }); + + test("all steps complete returns result", async () => { + const checkpoint = { + completed_steps: { + step_0: [1, 2, 3], + step_1: { loaded: true }, + }, + }; + const result = await runWorkflow(simple_workflow, checkpoint, [ + "https://example.com", + ]); + expect(result.type).toBe("complete"); + expect(result.result.status).toBe("done"); + expect(result.result.result).toEqual({ loaded: true }); + }); +}); + +describe("parallel dispatch", () => { + test("first invocation dispatches extract", async () => { + const result = await runWorkflow(parallel_workflow, {}, [ + "https://example.com", + ]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("extract_data"); + }); + + test("dispatches parallel steps after extract completes", async () => { + const checkpoint = { + completed_steps: { step_0: { raw: "data" } }, + }; + const result = await runWorkflow(parallel_workflow, checkpoint, [ + "https://example.com", + ]); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + expect(result.steps[0].name).toBe("clean_data"); + expect(result.steps[1].name).toBe("compute_stats"); + }); + + test("completes when all parallel steps done", async () => { + const checkpoint = { + completed_steps: { + step_0: { raw: "data" }, + step_1: { cleaned: true }, + step_2: { count: 42 }, + }, + }; + const result = await runWorkflow(parallel_workflow, checkpoint, [ + "https://example.com", + ]); + expect(result.type).toBe("complete"); + expect(result.result.cleaned).toEqual({ cleaned: true }); + expect(result.result.stats).toEqual({ count: 42 }); + }); +}); + +describe("conditional workflow", () => { + test("condition true dispatches send_alert", async () => { + const result = await runWorkflow(conditional_workflow, {}, [200]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("send_alert"); + }); + + test("condition false skips to load_data", async () => { + const result = await runWorkflow(conditional_workflow, {}, [50]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("load_data"); + }); +}); + +describe("task with external path", () => { + const run_external = task( + "f/external_script", + async function run_external(x: number) {} + ); + + test("uses external path as script", async () => { + const wf = workflow(async (x: number) => { + const result = await run_external(x); + return result; + }); + const result = await runWorkflow(wf, {}, [42]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("run_external"); + expect(result.steps[0].script).toBe("f/external_script"); + expect(result.steps[0].args).toEqual({ x: 42 }); + }); +}); + +// ===================================================================== +// EDGE CASE TESTS +// ===================================================================== + +describe("full sequential lifecycle (3 steps)", () => { + const three_step_wf = workflow(async (n: number) => { + const doubled = await double(n); + const incremented = await add_one(doubled); + const final_val = await double(incremented); + return { doubled, incremented, final: final_val }; + }); + + test("replay 0: dispatches step_0", async () => { + const result = await runWorkflow(three_step_wf, {}, [5]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].key).toBe("step_0"); + expect(result.steps[0].name).toBe("double"); + expect(result.steps[0].args).toEqual({ x: 5 }); + }); + + test("replay 1: dispatches step_1 with step_0 result as arg", async () => { + const result = await runWorkflow( + three_step_wf, + { completed_steps: { step_0: 10 } }, + [5] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].key).toBe("step_1"); + expect(result.steps[0].name).toBe("add_one"); + expect(result.steps[0].args).toEqual({ x: 10 }); + }); + + test("replay 2: dispatches step_2 with step_1 result as arg", async () => { + const result = await runWorkflow( + three_step_wf, + { completed_steps: { step_0: 10, step_1: 11 } }, + [5] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].key).toBe("step_2"); + expect(result.steps[0].name).toBe("double"); + expect(result.steps[0].args).toEqual({ x: 11 }); + }); + + test("replay 3: all complete, returns final result", async () => { + const result = await runWorkflow( + three_step_wf, + { completed_steps: { step_0: 10, step_1: 11, step_2: 22 } }, + [5] + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ doubled: 10, incremented: 11, final: 22 }); + }); +}); + +describe("step after parallel group", () => { + const seq_par_seq_wf = workflow(async (url: string) => { + const raw = await extract_data(url); + const [cleaned, stats] = await Promise.all([ + clean_data(raw), + compute_stats(raw), + ]); + const loaded = await load_data({ cleaned, stats }); + return loaded; + }); + + test("dispatches first sequential step", async () => { + const result = await runWorkflow(seq_par_seq_wf, {}, ["http://x"]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("extract_data"); + }); + + test("dispatches parallel group", async () => { + const result = await runWorkflow( + seq_par_seq_wf, + { completed_steps: { step_0: "raw" } }, + ["http://x"] + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + }); + + test("dispatches final step after parallel completes", async () => { + const result = await runWorkflow( + seq_par_seq_wf, + { + completed_steps: { + step_0: "raw", + step_1: "cleaned", + step_2: { count: 5 }, + }, + }, + ["http://x"] + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("sequential"); + expect(result.steps[0].name).toBe("load_data"); + expect(result.steps[0].key).toBe("step_3"); + }); + + test("completes when final step done", async () => { + const result = await runWorkflow( + seq_par_seq_wf, + { + completed_steps: { + step_0: "raw", + step_1: "cleaned", + step_2: { count: 5 }, + step_3: "final", + }, + }, + ["http://x"] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe("final"); + }); +}); + +describe("parallel after parallel (back to back)", () => { + const double_parallel_wf = workflow(async () => { + const [a, b] = await Promise.all([double(1), double(2)]); + const [c, d] = await Promise.all([add_one(a), add_one(b)]); + return { a, b, c, d }; + }); + + test("dispatches first parallel group", async () => { + const result = await runWorkflow(double_parallel_wf, {}, []); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + expect(result.steps[0].name).toBe("double"); + expect(result.steps[1].name).toBe("double"); + expect(result.steps[0].key).toBe("step_0"); + expect(result.steps[1].key).toBe("step_1"); + }); + + test("dispatches second parallel group after first completes", async () => { + const result = await runWorkflow( + double_parallel_wf, + { completed_steps: { step_0: 2, step_1: 4 } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + expect(result.steps[0].name).toBe("add_one"); + expect(result.steps[1].name).toBe("add_one"); + expect(result.steps[0].args).toEqual({ x: 2 }); + expect(result.steps[1].args).toEqual({ x: 4 }); + }); + + test("completes when all done", async () => { + const result = await runWorkflow( + double_parallel_wf, + { completed_steps: { step_0: 2, step_1: 4, step_2: 3, step_3: 5 } }, + [] + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ a: 2, b: 4, c: 3, d: 5 }); + }); +}); + +describe("conditional based on step result", () => { + const cond_on_result = workflow(async () => { + const val = await double(5); + if (val > 8) { + await send_alert("big"); + } + await load_data(val); + return { val }; + }); + + test("condition true path (val=10 > 8)", async () => { + const result = await runWorkflow( + cond_on_result, + { completed_steps: { step_0: 10 } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("send_alert"); + expect(result.steps[0].key).toBe("step_1"); + }); + + test("condition false path (val=4 <= 8)", async () => { + const result = await runWorkflow( + cond_on_result, + { completed_steps: { step_0: 4 } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("load_data"); + // When condition is false, send_alert is skipped so step index for + // load_data is step_1 (not step_2) + expect(result.steps[0].key).toBe("step_1"); + }); + + test("condition true: step after alert has key step_2", async () => { + const result = await runWorkflow( + cond_on_result, + { completed_steps: { step_0: 10, step_1: "alerted" } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("load_data"); + expect(result.steps[0].key).toBe("step_2"); + }); +}); + +describe("empty workflow (no tasks)", () => { + const empty_wf = workflow(async () => { + return { status: "empty" }; + }); + + test("completes immediately with no dispatch", async () => { + const result = await runWorkflow(empty_wf, {}, []); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ status: "empty" }); + }); +}); + +describe("single task workflow", () => { + const single_wf = workflow(async (x: number) => { + const result = await double(x); + return result; + }); + + test("dispatches single step", async () => { + const result = await runWorkflow(single_wf, {}, [7]); + expect(result.type).toBe("dispatch"); + expect(result.steps).toHaveLength(1); + expect(result.steps[0].name).toBe("double"); + }); + + test("completes with single result", async () => { + const result = await runWorkflow( + single_wf, + { completed_steps: { step_0: 14 } }, + [7] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe(14); + }); +}); + +describe("task with no arguments", () => { + const no_arg_wf = workflow(async () => { + const result = await noop_task(); + return result; + }); + + test("dispatches with empty args", async () => { + const result = await runWorkflow(no_arg_wf, {}, []); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].args).toEqual({}); + }); +}); + +describe("many steps (10+)", () => { + const many_steps_wf = workflow(async (n: number) => { + let val = n; + for (let i = 0; i < 10; i++) { + val = await add_one(val); + } + return val; + }); + + test("first invocation dispatches step_0", async () => { + const result = await runWorkflow(many_steps_wf, {}, [0]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].key).toBe("step_0"); + }); + + test("with 5 steps complete, dispatches step_5", async () => { + const completed: Record = {}; + for (let i = 0; i < 5; i++) completed[`step_${i}`] = i + 1; + const result = await runWorkflow( + many_steps_wf, + { completed_steps: completed }, + [0] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].key).toBe("step_5"); + expect(result.steps[0].args).toEqual({ x: 5 }); + }); + + test("all 10 steps complete returns final value", async () => { + const completed: Record = {}; + for (let i = 0; i < 10; i++) completed[`step_${i}`] = i + 1; + const result = await runWorkflow( + many_steps_wf, + { completed_steps: completed }, + [0] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe(10); + }); +}); + +describe("falsy values preserved in checkpoint", () => { + const falsy_wf = workflow(async () => { + const a = await double(0); // result will be 0 + const b = await load_data(a); // result will be null + const c = await extract_data(""); // result will be "" + return { a, b, c }; + }); + + test("zero is preserved", async () => { + const result = await runWorkflow( + falsy_wf, + { completed_steps: { step_0: 0 } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("load_data"); + expect(result.steps[0].args).toEqual({ data: 0 }); + }); + + test("null is preserved", async () => { + const result = await runWorkflow( + falsy_wf, + { completed_steps: { step_0: 0, step_1: null } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("extract_data"); + }); + + test("all falsy values complete correctly", async () => { + const result = await runWorkflow( + falsy_wf, + { completed_steps: { step_0: 0, step_1: null, step_2: "" } }, + [] + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ a: 0, b: null, c: "" }); + }); + + test("false is preserved", async () => { + const flag_wf = workflow(async () => { + const val = await load_data("check"); + if (val) { + await send_alert("truthy"); + } + return { val }; + }); + // false should be treated as completed (key exists), not as missing + const result = await runWorkflow( + flag_wf, + { completed_steps: { step_0: false } }, + [] + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ val: false }); + }); +}); + +describe("inline step (step function)", () => { + const step_wf = workflow(async (x: number) => { + const ts = await step("timestamp", () => 1234567890); + const doubled = await double(x); + const rid = await step("random_id", () => "abc-123"); + return { ts, doubled, id: rid }; + }); + + test("first invocation returns inline_checkpoint", async () => { + const result = await runWorkflow(step_wf, {}, [7]); + expect(result.type).toBe("inline_checkpoint"); + expect(result.key).toBe("step_0"); + expect(result.result).toBe(1234567890); + }); + + test("step cached, dispatches task", async () => { + const result = await runWorkflow( + step_wf, + { completed_steps: { step_0: 1234567890 } }, + [7] + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("double"); + expect(result.steps[0].key).toBe("step_1"); + }); + + test("step + task cached, returns second inline step", async () => { + const result = await runWorkflow( + step_wf, + { completed_steps: { step_0: 1234567890, step_1: 14 } }, + [7] + ); + expect(result.type).toBe("inline_checkpoint"); + expect(result.key).toBe("step_2"); + expect(result.result).toBe("abc-123"); + }); + + test("all complete returns final result", async () => { + const result = await runWorkflow( + step_wf, + { + completed_steps: { + step_0: 1234567890, + step_1: 14, + step_2: "abc-123", + }, + }, + [7] + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ + ts: 1234567890, + doubled: 14, + id: "abc-123", + }); + }); +}); + +describe("unawaited tasks (flush pending)", () => { + test("single unawaited task at end is flushed", async () => { + const wf = workflow(async () => { + await extract_data("x"); + load_data("y"); // forgotten await + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: "raw" } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("sequential"); + expect(result.steps).toHaveLength(1); + expect(result.steps[0].name).toBe("load_data"); + }); + + test("multiple unawaited tasks flushed as parallel", async () => { + const wf = workflow(async () => { + await extract_data("x"); + clean_data("y"); // forgotten await + compute_stats("y"); // forgotten await + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: "raw" } }, + [] + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + expect(result.steps[0].name).toBe("clean_data"); + expect(result.steps[1].name).toBe("compute_stats"); + }); + + test("no unawaited tasks means normal complete", async () => { + const wf = workflow(async () => { + const val = await double(5); + return val; + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 10 } }, + [] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe(10); + }); +}); + +describe("child mode (_executingKey)", () => { + test("executes matching task directly", async () => { + const wf = workflow(async (x: number) => { + const val = await double(x); + return val; + }); + const result = await runWorkflow( + wf, + { completed_steps: {}, _executing_key: "step_0" }, + [7] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe(14); // double(7) = 14 + }); + + test("replays cached steps before executing key", async () => { + const wf = workflow(async (x: number) => { + const doubled = await double(x); + const result = await add_one(doubled); + return result; + }); + const result = await runWorkflow( + wf, + { + completed_steps: { step_0: 10 }, + _executing_key: "step_1", + }, + [5] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe(11); // add_one(10) = 11 + }); + + test("child mode with external path task", async () => { + const ext = task( + "f/external", + async function ext_task(x: number) { + return x * 3; + } + ); + const wf = workflow(async (x: number) => { + const result = await ext(x); + return result; + }); + const result = await runWorkflow( + wf, + { completed_steps: {}, _executing_key: "step_0" }, + [4] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe(12); // 4 * 3 + }); +}); + +describe("key determinism across replays", () => { + const det_wf = workflow(async (n: number) => { + const a = await double(n); + const b = await add_one(a); + const c = await double(b); + return c; + }); + + test("keys are consistent: step_0 always maps to first double", async () => { + // Empty checkpoint + const r1 = await runWorkflow(det_wf, {}, [3]); + expect(r1.steps[0].key).toBe("step_0"); + expect(r1.steps[0].name).toBe("double"); + + // With step_0 completed + const r2 = await runWorkflow( + det_wf, + { completed_steps: { step_0: 6 } }, + [3] + ); + expect(r2.steps[0].key).toBe("step_1"); + expect(r2.steps[0].name).toBe("add_one"); + + // With step_0 and step_1 completed + const r3 = await runWorkflow( + det_wf, + { completed_steps: { step_0: 6, step_1: 7 } }, + [3] + ); + expect(r3.steps[0].key).toBe("step_2"); + expect(r3.steps[0].name).toBe("double"); + }); +}); + +describe("parallel dispatch includes correct args from cached results", () => { + test("parallel steps receive cached parent result as args", async () => { + const wf = workflow(async (x: number) => { + const base = await double(x); + const [a, b] = await Promise.all([add_one(base), double(base)]); + return { a, b }; + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 20 } }, + [10] + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps[0].args).toEqual({ x: 20 }); + expect(result.steps[1].args).toEqual({ x: 20 }); + }); +}); + +describe("inline step with async function", () => { + test("async step function resolves correctly", async () => { + const wf = workflow(async () => { + const val = await step("async_step", async () => { + return 42; + }); + return val; + }); + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("inline_checkpoint"); + expect(result.key).toBe("step_0"); + expect(result.result).toBe(42); + }); +}); + +describe("workflow returning undefined", () => { + test("undefined return value is captured", async () => { + const wf = workflow(async () => { + await double(1); + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2 } }, + [] + ); + expect(result.type).toBe("complete"); + expect(result.result).toBeUndefined(); + }); +}); + +describe("large parallel group", () => { + test("dispatches 5 parallel steps at once", async () => { + const wf = workflow(async () => { + const results = await Promise.all([ + double(1), + double(2), + double(3), + double(4), + double(5), + ]); + return results; + }); + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(5); + for (let i = 0; i < 5; i++) { + expect(result.steps[i].key).toBe(`step_${i}`); + expect(result.steps[i].args).toEqual({ x: i + 1 }); + } + }); +}); + +describe("complex mixed workflow: seq → par → seq → par → seq", () => { + const complex_wf = workflow(async () => { + const init = await extract_data("start"); + const [a, b] = await Promise.all([double(1), double(2)]); + const mid = await load_data({ a, b }); + const [c, d] = await Promise.all([add_one(3), add_one(4)]); + const fin = await clean_data({ mid, c, d }); + return fin; + }); + + test("replay 0: dispatches extract_data", async () => { + const r = await runWorkflow(complex_wf, {}, []); + expect(r.steps[0].name).toBe("extract_data"); + }); + + test("replay 1: dispatches parallel [double, double]", async () => { + const r = await runWorkflow( + complex_wf, + { completed_steps: { step_0: "init" } }, + [] + ); + expect(r.mode).toBe("parallel"); + expect(r.steps).toHaveLength(2); + expect(r.steps[0].name).toBe("double"); + }); + + test("replay 2: dispatches load_data", async () => { + const r = await runWorkflow( + complex_wf, + { completed_steps: { step_0: "init", step_1: 2, step_2: 4 } }, + [] + ); + expect(r.mode).toBe("sequential"); + expect(r.steps[0].name).toBe("load_data"); + expect(r.steps[0].key).toBe("step_3"); + }); + + test("replay 3: dispatches parallel [add_one, add_one]", async () => { + const r = await runWorkflow( + complex_wf, + { + completed_steps: { + step_0: "init", + step_1: 2, + step_2: 4, + step_3: "mid", + }, + }, + [] + ); + expect(r.mode).toBe("parallel"); + expect(r.steps).toHaveLength(2); + expect(r.steps[0].name).toBe("add_one"); + }); + + test("replay 4: dispatches clean_data", async () => { + const r = await runWorkflow( + complex_wf, + { + completed_steps: { + step_0: "init", + step_1: 2, + step_2: 4, + step_3: "mid", + step_4: 4, + step_5: 5, + }, + }, + [] + ); + expect(r.mode).toBe("sequential"); + expect(r.steps[0].name).toBe("clean_data"); + expect(r.steps[0].key).toBe("step_6"); + }); + + test("replay 5: all complete", async () => { + const r = await runWorkflow( + complex_wf, + { + completed_steps: { + step_0: "init", + step_1: 2, + step_2: 4, + step_3: "mid", + step_4: 4, + step_5: 5, + step_6: "final", + }, + }, + [] + ); + expect(r.type).toBe("complete"); + expect(r.result).toBe("final"); + }); +}); + +// ===================================================================== +// ERROR PROPAGATION TESTS +// ===================================================================== + +describe("error propagation via __wmill_error marker", () => { + test("task error is thrown on replay", async () => { + const wf = workflow(async (x: number) => { + const result = await double(x); + return result; + }); + // Simulate child failure stored as __wmill_error marker + const checkpoint = { + completed_steps: { + step_0: { + __wmill_error: true, + message: "WAC task 'double' failed (child job abc-123)", + result: { message: "division by zero" }, + step_key: "double", + child_job_id: "abc-123", + }, + }, + }; + try { + await runWorkflow(wf, checkpoint, [5]); + expect(true).toBe(false); // should not reach here + } catch (e: any) { + expect(e.message).toContain("double"); + expect(e.result).toEqual({ message: "division by zero" }); + expect(e.child_job_id).toBe("abc-123"); + } + }); + + test("error is catchable with try/catch in workflow", async () => { + const wf = workflow(async (x: number) => { + try { + const result = await double(x); + return { success: true, result }; + } catch (e: any) { + return { success: false, error: e.message }; + } + }); + const checkpoint = { + completed_steps: { + step_0: { + __wmill_error: true, + message: "Task 'double' failed", + result: { message: "boom" }, + }, + }, + }; + const result = await runWorkflow(wf, checkpoint, [5]); + expect(result.type).toBe("complete"); + expect(result.result.success).toBe(false); + expect(result.result.error).toContain("double"); + }); + + test("error in parallel — one fails, caught by Promise.all reject", async () => { + const wf = workflow(async () => { + try { + const [a, b] = await Promise.all([double(1), add_one(2)]); + return { a, b }; + } catch (e: any) { + return { caught: true, error: e.message }; + } + }); + const checkpoint = { + completed_steps: { + step_0: { __wmill_error: true, message: "double failed", result: {} }, + step_1: 3, // add_one succeeded + }, + }; + const result = await runWorkflow(wf, checkpoint, []); + expect(result.type).toBe("complete"); + expect(result.result.caught).toBe(true); + }); + + test("retry pattern with try/catch + loop", async () => { + // Simulates: first attempt fails, second succeeds + const wf = workflow(async (x: number) => { + for (let i = 0; i < 3; i++) { + try { + const result = await double(x); + return { result, attempts: i + 1 }; + } catch (e) { + if (i === 2) throw e; + // retry on next iteration + } + } + }); + // First double (step_0) fails, second double (step_1) succeeds + const checkpoint = { + completed_steps: { + step_0: { __wmill_error: true, message: "temporary failure", result: {} }, + step_1: 10, + }, + }; + const result = await runWorkflow(wf, checkpoint, [5]); + expect(result.type).toBe("complete"); + expect(result.result.result).toBe(10); + expect(result.result.attempts).toBe(2); + }); + + test("inline step error is thrown", async () => { + const wf = workflow(async () => { + try { + const val = await step("risky", () => 42); + return { val }; + } catch (e: any) { + return { caught: e.message }; + } + }); + const checkpoint = { + completed_steps: { + step_0: { __wmill_error: true, message: "inline step failed", result: {} }, + }, + }; + const result = await runWorkflow(wf, checkpoint, []); + expect(result.type).toBe("complete"); + expect(result.result.caught).toContain("inline step failed"); + }); + + test("non-error object with __wmill_error field is NOT treated as error", async () => { + // An object with __wmill_error: false should be treated as a normal value + const wf = workflow(async () => { + const val = await double(5); + return val; + }); + const checkpoint = { + completed_steps: { + step_0: { __wmill_error: false, data: "not an error" }, + }, + }; + const result = await runWorkflow(wf, checkpoint, [5]); + expect(result.type).toBe("complete"); + expect(result.result).toEqual({ __wmill_error: false, data: "not an error" }); + }); +}); + +// ===================================================================== +// TASK OPTIONS TESTS +// ===================================================================== + +describe("task options", () => { + test("options are forwarded in dispatch step info", async () => { + const heavy = task( + async function heavy(x: number) { return x; }, + { timeout: 600, tag: "gpu", cache_ttl: 3600, priority: 10 }, + ); + const wf = workflow(async (x: number) => { + return await heavy(x); + }); + const result = await runWorkflow(wf, {}, [42]); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].timeout).toBe(600); + expect(result.steps[0].tag).toBe("gpu"); + expect(result.steps[0].cache_ttl).toBe(3600); + expect(result.steps[0].priority).toBe(10); + }); + + test("task without options has no extra fields", async () => { + const simple = task(async function simple(x: number) { return x; }); + const wf = workflow(async (x: number) => { + return await simple(x); + }); + const result = await runWorkflow(wf, {}, [1]); + expect(result.steps[0].timeout).toBeUndefined(); + expect(result.steps[0].tag).toBeUndefined(); + }); + + test("concurrency options forwarded", async () => { + const limited = task( + async function limited(x: number) { return x; }, + { concurrent_limit: 5, concurrency_key: "my-key", concurrency_time_window_s: 60 }, + ); + const wf = workflow(async (x: number) => { + return await limited(x); + }); + const result = await runWorkflow(wf, {}, [1]); + expect(result.steps[0].concurrent_limit).toBe(5); + expect(result.steps[0].concurrency_key).toBe("my-key"); + expect(result.steps[0].concurrency_time_window_s).toBe(60); + }); + + test("task with path and options", async () => { + const ext = task( + "f/gpu_script", + async function ext(x: number) { return x; }, + { timeout: 300, tag: "gpu" }, + ); + const wf = workflow(async (x: number) => { + return await ext(x); + }); + const result = await runWorkflow(wf, {}, [1]); + expect(result.steps[0].script).toBe("f/gpu_script"); + expect(result.steps[0].timeout).toBe(300); + expect(result.steps[0].tag).toBe("gpu"); + }); +}); + +// ===================================================================== +// SLEEP TESTS +// ===================================================================== + +describe("sleep", () => { + test("first invocation returns sleep output", async () => { + const wf = workflow(async () => { + await double(1); + await sleep(60); + await add_one(2); + return "done"; + }); + // step_0 (double) complete, step_1 is sleep + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2 } }, + [], + ); + expect(result.type).toBe("sleep"); + expect(result.key).toBe("step_1"); + expect(result.seconds).toBe(60); + }); + + test("sleep completes on replay when stored in checkpoint", async () => { + const wf = workflow(async () => { + await double(1); + await sleep(60); + await add_one(2); + return "done"; + }); + // step_0 (double) and step_1 (sleep) complete + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2, step_1: true } }, + [], + ); + expect(result.type).toBe("dispatch"); + expect(result.steps[0].name).toBe("add_one"); + expect(result.steps[0].key).toBe("step_2"); + }); + + test("all steps including sleep complete returns result", async () => { + const wf = workflow(async () => { + await double(1); + await sleep(60); + await add_one(2); + return "done"; + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2, step_1: true, step_2: 3 } }, + [], + ); + expect(result.type).toBe("complete"); + expect(result.result).toBe("done"); + }); + + test("sleep enforces minimum of 1 second", async () => { + const wf = workflow(async () => { + await sleep(0); + return "done"; + }); + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("sleep"); + expect(result.seconds).toBe(1); + }); + + test("sleep rounds to nearest integer", async () => { + const wf = workflow(async () => { + await sleep(3.7); + return "done"; + }); + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("sleep"); + expect(result.seconds).toBe(4); + }); +}); + +// ===================================================================== +// PARALLEL UTILITY TESTS +// ===================================================================== + +describe("parallel utility", () => { + test("processes all items with default concurrency", async () => { + const wf = workflow(async () => { + const items = [1, 2, 3]; + const results = await parallel(items, double); + return results; + }); + // All 3 items dispatched in parallel: step_0, step_1, step_2 + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(3); + expect(result.steps[0].args).toEqual({ x: 1 }); + expect(result.steps[1].args).toEqual({ x: 2 }); + expect(result.steps[2].args).toEqual({ x: 3 }); + }); + + test("completes when all parallel items done", async () => { + const wf = workflow(async () => { + const items = [1, 2, 3]; + const results = await parallel(items, double); + return results; + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2, step_1: 4, step_2: 6 } }, + [], + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual([2, 4, 6]); + }); + + test("batched concurrency dispatches first batch", async () => { + const wf = workflow(async () => { + const items = [1, 2, 3, 4, 5]; + const results = await parallel(items, double, { concurrency: 2 }); + return results; + }); + // First batch: items[0..2] → step_0, step_1 + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + expect(result.steps[0].args).toEqual({ x: 1 }); + expect(result.steps[1].args).toEqual({ x: 2 }); + }); + + test("batched concurrency dispatches second batch after first completes", async () => { + const wf = workflow(async () => { + const items = [1, 2, 3, 4, 5]; + const results = await parallel(items, double, { concurrency: 2 }); + return results; + }); + // First batch done, second batch: items[2..4] → step_2, step_3 + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2, step_1: 4 } }, + [], + ); + expect(result.type).toBe("dispatch"); + expect(result.mode).toBe("parallel"); + expect(result.steps).toHaveLength(2); + expect(result.steps[0].args).toEqual({ x: 3 }); + expect(result.steps[1].args).toEqual({ x: 4 }); + }); + + test("batched concurrency last batch may be smaller", async () => { + const wf = workflow(async () => { + const items = [1, 2, 3, 4, 5]; + const results = await parallel(items, double, { concurrency: 2 }); + return results; + }); + // Two batches done, third batch: items[4..5] → step_4 + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2, step_1: 4, step_2: 6, step_3: 8 } }, + [], + ); + expect(result.type).toBe("dispatch"); + expect(result.steps).toHaveLength(1); + expect(result.steps[0].args).toEqual({ x: 5 }); + }); + + test("batched concurrency completes with all results in order", async () => { + const wf = workflow(async () => { + const items = [1, 2, 3, 4, 5]; + const results = await parallel(items, double, { concurrency: 2 }); + return results; + }); + const result = await runWorkflow( + wf, + { completed_steps: { step_0: 2, step_1: 4, step_2: 6, step_3: 8, step_4: 10 } }, + [], + ); + expect(result.type).toBe("complete"); + expect(result.result).toEqual([2, 4, 6, 8, 10]); + }); + + test("empty items returns empty array", async () => { + const wf = workflow(async () => { + const results = await parallel([], double); + return results; + }); + const result = await runWorkflow(wf, {}, []); + expect(result.type).toBe("complete"); + expect(result.result).toEqual([]); + }); +});