From b3ef4bc26c5696624efee89b5e4e33e77e10cf15 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 13 Apr 2026 12:49:53 -0400 Subject: [PATCH] perf: add inline-persist fast path for WAC v2 step() (#8807) --- ...480f7fec559c88bc384c5475dba5f03526337.json | 22 ++ ...798a860ce977954a856335bb62ae00f615659.json | 15 - ...153c43903f929ae5d62fbba12610f89c36d55.json | 2 +- ...9ef795049e60a823e8ce754a3f9d91ef1b6bb.json | 20 -- ...a5a87c2c94227798f4545fe0931af2818a6fa.json | 66 ---- ...2354577dc5f77169fd4573317454e47361d65.json | 65 ---- ...92adce1be8b897288b655505729c97c6e324c.json | 15 - backend/tests/python_jobs.rs | 150 +++++++++ backend/windmill-api/src/jobs.rs | 87 ++++- backend/windmill-common/src/lib.rs | 1 + backend/windmill-common/src/wac.rs | 300 ++++++++++++++++++ backend/windmill-worker/src/bun_executor.rs | 133 +++----- backend/windmill-worker/src/wac_executor.rs | 112 +------ python-client/wmill/wmill/client.py | 83 ++++- typescript-client/client.ts | 76 +++++ 15 files changed, 762 insertions(+), 385 deletions(-) create mode 100644 backend/.sqlx/query-15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337.json delete mode 100644 backend/.sqlx/query-50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659.json delete mode 100644 backend/.sqlx/query-5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb.json delete mode 100644 backend/.sqlx/query-8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa.json delete mode 100644 backend/.sqlx/query-cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65.json delete mode 100644 backend/.sqlx/query-dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c.json create mode 100644 backend/windmill-common/src/wac.rs diff --git a/backend/.sqlx/query-15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337.json b/backend/.sqlx/query-15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337.json new file mode 100644 index 0000000000..4b1c9dbdcc --- /dev/null +++ b/backend/.sqlx/query-15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT workflow_as_code_status as \"v: serde_json::Value\"\n FROM v2_job_completed WHERE id = $1", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "v: serde_json::Value", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [ + true + ] + }, + "hash": "15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337" +} diff --git a/backend/.sqlx/query-50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659.json b/backend/.sqlx/query-50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659.json deleted file mode 100644 index 91f0ed4c92..0000000000 --- a/backend/.sqlx/query-50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO folder SELECT name, $1, display_name, owners, extra_perms, summary, edited_at, created_by FROM folder WHERE workspace_id = $2", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Varchar", - "Text" - ] - }, - "nullable": [] - }, - "hash": "50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659" -} diff --git a/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json b/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json index 713ccb9dd3..36ddb8ab9f 100644 --- a/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json +++ b/backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json @@ -15,7 +15,7 @@ ] }, "nullable": [ - null + true ] }, "hash": "5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55" diff --git a/backend/.sqlx/query-5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb.json b/backend/.sqlx/query-5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb.json deleted file mode 100644 index d099fc8b2a..0000000000 --- a/backend/.sqlx/query-5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb.json +++ /dev/null @@ -1,20 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO folder (workspace_id, name, display_name, owners, extra_perms, summary, created_by, edited_at) VALUES ($1, $2, $3, $4, $5, $6, $7, now())", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Varchar", - "Varchar", - "Varchar", - "VarcharArray", - "Jsonb", - "Text", - "Varchar" - ] - }, - "nullable": [] - }, - "hash": "5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb" -} diff --git a/backend/.sqlx/query-8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa.json b/backend/.sqlx/query-8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa.json deleted file mode 100644 index cfd79b3cdd..0000000000 --- a/backend/.sqlx/query-8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa.json +++ /dev/null @@ -1,66 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "SELECT workspace_id, name, display_name, owners, extra_perms, summary, created_by, edited_at FROM folder WHERE workspace_id = $1 ORDER BY name asc LIMIT $2 OFFSET $3", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "workspace_id", - "type_info": "Varchar" - }, - { - "ordinal": 1, - "name": "name", - "type_info": "Varchar" - }, - { - "ordinal": 2, - "name": "display_name", - "type_info": "Varchar" - }, - { - "ordinal": 3, - "name": "owners", - "type_info": "VarcharArray" - }, - { - "ordinal": 4, - "name": "extra_perms", - "type_info": "Jsonb" - }, - { - "ordinal": 5, - "name": "summary", - "type_info": "Text" - }, - { - "ordinal": 6, - "name": "created_by", - "type_info": "Varchar" - }, - { - "ordinal": 7, - "name": "edited_at", - "type_info": "Timestamptz" - } - ], - "parameters": { - "Left": [ - "Text", - "Int8", - "Int8" - ] - }, - "nullable": [ - false, - false, - false, - false, - false, - true, - true, - true - ] - }, - "hash": "8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa" -} diff --git a/backend/.sqlx/query-cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65.json b/backend/.sqlx/query-cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65.json deleted file mode 100644 index a33576931c..0000000000 --- a/backend/.sqlx/query-cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65.json +++ /dev/null @@ -1,65 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "SELECT workspace_id, name, display_name, owners, extra_perms, summary, created_by, edited_at FROM folder WHERE name = $1 AND workspace_id = $2", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "workspace_id", - "type_info": "Varchar" - }, - { - "ordinal": 1, - "name": "name", - "type_info": "Varchar" - }, - { - "ordinal": 2, - "name": "display_name", - "type_info": "Varchar" - }, - { - "ordinal": 3, - "name": "owners", - "type_info": "VarcharArray" - }, - { - "ordinal": 4, - "name": "extra_perms", - "type_info": "Jsonb" - }, - { - "ordinal": 5, - "name": "summary", - "type_info": "Text" - }, - { - "ordinal": 6, - "name": "created_by", - "type_info": "Varchar" - }, - { - "ordinal": 7, - "name": "edited_at", - "type_info": "Timestamptz" - } - ], - "parameters": { - "Left": [ - "Text", - "Text" - ] - }, - "nullable": [ - false, - false, - false, - false, - false, - true, - true, - true - ] - }, - "hash": "cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65" -} diff --git a/backend/.sqlx/query-dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c.json b/backend/.sqlx/query-dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c.json deleted file mode 100644 index 8049f4de1a..0000000000 --- a/backend/.sqlx/query-dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO folder (workspace_id, name, display_name, owners, extra_perms, summary, edited_at, created_by)\n SELECT $2, name, display_name, owners, extra_perms, summary, edited_at, created_by\n FROM folder\n WHERE workspace_id = $1", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Text", - "Varchar" - ] - }, - "nullable": [] - }, - "hash": "dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c" -} diff --git a/backend/tests/python_jobs.rs b/backend/tests/python_jobs.rs index 4819a1da60..76bf6cf136 100644 --- a/backend/tests/python_jobs.rs +++ b/backend/tests/python_jobs.rs @@ -1014,3 +1014,153 @@ async def main(item: str, qty: int, email: str): .await; Ok(()) } + +/// End-to-end comparison between the legacy `step()` suspend-and-replay path +/// and the new SDK inline-persist fast path, toggled per-job via the +/// `WM_WAC_INLINE_FAST_PATH` env var which the Python script sets on its own +/// `os.environ` at the top so parallel tests can't race on a global env var. +/// +/// Runs the same 5-step WAC v2 Python workflow twice, asserts both modes +/// produce the same final result and the same `completed_steps` map, and +/// prints a wall-clock benchmark line so CI and manual runs can track the +/// speedup. The fast path is expected to be faster because it avoids N-1 +/// subprocess spawns and N-1 queue round-trips for a workflow with N +/// `step()` calls, but the exact ratio depends on the CI runner so we don't +/// assert a hard threshold — behavioral equivalence is the important check. +#[cfg(feature = "python")] +#[sqlx::test(fixtures("base"))] +async fn test_python_wac_v2_step_inline_fast_path(db: Pool) -> anyhow::Result<()> { + initialize_tracing().await; + let server = ApiServer::start(db.clone()).await?; + let port = server.addr.port(); + + fn workflow_content(fast_path_enabled: bool) -> String { + let flag = if fast_path_enabled { "1" } else { "0" }; + // NB: the `os.environ[...] = ...` line must execute BEFORE the first + // `step()` call inside the workflow subprocess — setting it at module + // scope (before `from wmill import ...` finishes importing is still + // fine because the SDK reads it lazily at call time, not at import). + format!( + r#"import os +os.environ["WM_WAC_INLINE_FAST_PATH"] = "{flag}" +from wmill import workflow, step + +@workflow +async def main(n: int): + a = await step("a", lambda: n) + b = await step("b", lambda: a + 1) + c = await step("c", lambda: b + 1) + d = await step("d", lambda: c + 1) + e = await step("e", lambda: d + 1) + return {{"a": a, "b": b, "c": c, "d": d, "e": e}} +"# + ) + } + + let db_ref = &db; + + async fn run_once( + db: &Pool, + port: u16, + content: String, + ) -> (serde_json::Value, serde_json::Value, std::time::Duration) { + let mut job_id_out: Option = None; + let mut result_out: Option = None; + let t0 = std::time::Instant::now(); + in_test_worker( + db, + async { + let job = Box::pin( + RunJob::from(JobPayload::Code(RawCode { + language: ScriptLang::Python3, + content, + ..RawCode::default() + })) + .arg("n", json!(1)) + .run_until_complete(db, false, port), + ) + .await; + result_out = Some(job.json_result().unwrap_or_else(|| { + panic!("job {} returned no result — raw job = {:?}", job.id, job) + })); + job_id_out = Some(job.id); + }, + port, + ) + .await; + let elapsed = t0.elapsed(); + + let job_id = job_id_out.expect("job id"); + // Fetch the full workflow_as_code_status for diagnostics, then extract + // _checkpoint.completed_steps. A None here means the step() path never + // wrote the checkpoint, which is exactly the signal we want to surface + // clearly (instead of panicking with an opaque Option::unwrap() error). + let full_status: Option = sqlx::query_scalar!( + r#"SELECT workflow_as_code_status as "v: serde_json::Value" + FROM v2_job_completed WHERE id = $1"#, + job_id + ) + .fetch_one(db) + .await + .expect("v2_job_completed row fetch"); + + let ckpt = full_status + .as_ref() + .and_then(|s| s.get("_checkpoint")) + .and_then(|c| c.get("completed_steps")) + .cloned() + .unwrap_or_else(|| { + panic!( + "job {job_id} completed_steps missing — full workflow_as_code_status = {:?}, job_result = {:?}", + full_status, result_out + ) + }); + + (result_out.unwrap(), ckpt, elapsed) + } + + // --- Legacy path: worker-side suspend & replay --- + let (legacy_result, legacy_ckpt, legacy_elapsed) = + run_once(db_ref, port, workflow_content(false)).await; + + // --- Fast path: SDK persists the delta via the new API endpoint --- + let (fast_result, fast_ckpt, fast_elapsed) = + run_once(db_ref, port, workflow_content(true)).await; + + // Behavioral equivalence: same final result and same completed_steps. + assert_eq!( + legacy_result, fast_result, + "legacy and fast path produced different workflow results" + ); + assert_eq!( + legacy_result, + json!({"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}), + "unexpected workflow result" + ); + assert_eq!( + legacy_ckpt, fast_ckpt, + "legacy and fast path stored different completed_steps in the checkpoint" + ); + + // Benchmark output. We log and print but do NOT assert a hard threshold: + // CI runners have variable noise floors and a "fast < legacy" check would + // flake. Behavioral equivalence above is the important invariant. + let legacy_ms = legacy_elapsed.as_millis(); + let fast_ms = fast_elapsed.as_millis(); + let speedup = if fast_ms > 0 { + legacy_ms as f64 / fast_ms as f64 + } else { + f64::INFINITY + }; + tracing::info!( + "WAC v2 step() benchmark: legacy={}ms fast={}ms speedup={:.2}x", + legacy_ms, + fast_ms, + speedup + ); + println!( + "WAC v2 step() benchmark: legacy={}ms fast={}ms speedup={:.2}x", + legacy_ms, fast_ms, speedup + ); + Ok(()) +} diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 4d11b9aa3e..ae8ad533b5 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -63,7 +63,6 @@ use windmill_worker::get_worker_internal_server_inline_utils; use windmill_common::variables::get_workspace_key; -#[cfg(feature = "run_inline")] use crate::db::OptJobAuthed; use crate::triggers::trigger_helpers::{FlowId, ScriptId}; use crate::{ @@ -161,6 +160,10 @@ pub fn workspaced_service() -> Router { .layer(cors.clone()) .layer(ce_headers.clone()), ) + .route( + "/wac/inline_checkpoint/{job_id}", + post(wac_inline_checkpoint).layer(cors.clone()), + ) .route( "/restart/f/{job_id}", post(restart_flow).head(|| async { "" }).layer(cors.clone()), @@ -4445,6 +4448,88 @@ pub async fn run_workflow_as_code( Ok((StatusCode::CREATED, uuid.to_string())) } +#[derive(Deserialize)] +pub struct WacInlineCheckpointPayload { + pub key: String, + pub result: serde_json::Value, + #[serde(default)] + pub started_at: Option, + #[serde(default)] + pub duration_ms: Option, +} + +/// Fast-path endpoint called by the WAC v2 SDKs to persist a single `step()` +/// checkpoint delta without unwinding the parent workflow subprocess. +/// +/// Mirrors the worker-side `WacOutput::InlineCheckpoint` arm in +/// `bun_executor::handle_wac_v2_output` exactly — same `completed_steps` +/// entry, same `_step/` timeline entry, same source-hash validation — +/// but does **not** touch `v2_job_queue`, because the parent subprocess is +/// still live and about to return the next chunk of script output. +/// +/// Auth: requires the job's ephemeral token (the one the worker sets into +/// `WM_TOKEN` when spawning the subprocess), not just any workspace-scoped +/// `ApiAuthed`. Because WAC v2 replays steps from `completed_steps`, a forged +/// entry directly changes the workflow's observed return values — this is +/// execution state, not user-facing metadata. `OptJobAuthed.job_id` is set +/// only when the caller presents a JWT whose `job_id` claim matches the URL +/// path, so rejecting mismatches closes the workspace-wide privilege gap. +/// +/// Any error here causes the SDK to fall back to raising `_StepSuspend`, +/// which then goes through the untouched worker-side path. Old SDKs that +/// never call this endpoint continue to work unchanged. +pub async fn wac_inline_checkpoint( + OptJobAuthed { authed: _, job_id: token_job_id }: OptJobAuthed, + Extension(db): Extension, + Path((w_id, job_id)): Path<(String, Uuid)>, + Json(payload): Json, +) -> error::Result { + // Enforce ephemeral-job-token binding: the presented token must be the + // one issued to *this* specific job. Regular workspace API tokens don't + // have `job_id` populated in their JWT claims, so `token_job_id` is None + // for them — reject unconditionally. + if token_job_id != Some(job_id) { + return Err(error::Error::PermissionDenied( + "wac_inline_checkpoint requires the job's ephemeral token".to_string(), + )); + } + + // Look up the job's script hash for source-hash validation. We deliberately + // use a minimal query here rather than `fetch_queued(...)` — the latter + // pulls in many extra columns we don't need. Restrict to the job kinds + // that actually run user WAC v2 code (`script` and `preview`) so a + // forged/buggy caller can't write a bogus checkpoint onto a flow, + // dependency, or other job kind that shares the workspace. + let row: Option<(Option,)> = sqlx::query_as( + "SELECT runnable_id FROM v2_job + WHERE id = $1 AND workspace_id = $2 + AND kind IN ('script'::job_kind, 'preview'::job_kind)", + ) + .bind(&job_id) + .bind(&w_id) + .fetch_optional(&db) + .await?; + let (runnable_id,) = row.ok_or_else(|| { + error::Error::NotFound(format!("WAC v2 job {job_id} not found in workspace {w_id}")) + })?; + let source_hash = runnable_id.map(|h| h.to_string()); + + let mut tx = db.begin().await?; + windmill_common::wac::persist_inline_checkpoint_delta( + &mut tx, + &job_id, + source_hash.as_deref(), + &payload.key, + payload.result, + payload.started_at.as_deref(), + payload.duration_ms, + ) + .await?; + tx.commit().await?; + + Ok(StatusCode::OK) +} + lazy_static::lazy_static! { static ref JOB_VIEW_AUDIT_LOGS: bool = std::env::var("JOB_VIEW_AUDIT_LOGS") .ok() diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index f2d5dff755..8ac739cebb 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -108,6 +108,7 @@ pub mod usernames; pub mod users; pub mod utils; pub mod variables; +pub mod wac; pub mod webhook; pub mod worker; pub mod worker_group_job_stats; diff --git a/backend/windmill-common/src/wac.rs b/backend/windmill-common/src/wac.rs new file mode 100644 index 0000000000..7f12dfbcdb --- /dev/null +++ b/backend/windmill-common/src/wac.rs @@ -0,0 +1,300 @@ +//! Workflow-as-Code v2 checkpoint model and persistence primitives. +//! +//! Lives in `windmill-common` (not `windmill-worker`) so the API server can +//! write checkpoint deltas directly from the SDK fast path without pulling in +//! the entire worker crate. The worker still re-exports the same symbols from +//! `windmill_worker::wac_executor` for its own historical call sites. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sqlx::{Postgres, Transaction}; +use uuid::Uuid; + +use crate::error::{self, Error}; +use crate::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, +} + +/// 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(()) +} + +/// 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; + } + } +} + +/// Persist a single inline-step checkpoint delta into the given transaction: +/// validate the source hash, add the step to `completed_steps`, save the +/// checkpoint, and write the `_step/` timeline entry. +/// +/// The caller owns the transaction and commits it. This lets the worker-side +/// `WacOutput::InlineCheckpoint` fallback arm add its own `UPDATE v2_job_queue +/// SET running = false` in the same transaction — preserving the original +/// all-or-nothing atomicity — while the API fast path simply commits after +/// the helper returns. +/// +/// ## Concurrency model +/// +/// The helper does a read-modify-write: `SELECT ... FOR UPDATE` → parse +/// `WacCheckpoint` → modify in Rust via `add_completed_step` → write the +/// full serialized `_checkpoint` back via `INSERT ... ON CONFLICT DO UPDATE` +/// plus a separate `UPDATE` for the `_step/` timeline entry. The +/// important property of this pattern: each call **replaces the whole +/// `_checkpoint` object**, not individual `completed_steps[key]` entries. +/// That means distinct step keys do NOT protect concurrent callers from +/// overwriting each other — two writers that start from the same loaded +/// checkpoint will each produce a new serialized object that lacks the +/// other's step. +/// +/// **Steady state (row exists)** — `SELECT ... FOR UPDATE` holds the row +/// lock until commit. The second concurrent caller blocks on the lock, +/// then re-reads the post-commit checkpoint (which already contains the +/// first caller's step), applies its own delta, and writes. No loss. +/// +/// **First write (row does not yet exist)** — `SELECT ... FOR UPDATE` on a +/// WHERE clause that matches zero rows acquires no lock. Two concurrent +/// callers would both see `None`, both build a fresh `WacCheckpoint` from +/// scratch, and then race on the final `INSERT ... ON CONFLICT DO UPDATE`: +/// the second writer's `DO UPDATE SET workflow_as_code_status = jsonb_set( +/// ..., '{_checkpoint}', $2)` replaces the `_checkpoint` the first writer +/// just inserted, so the first writer's step is lost. +/// +/// That race window is closed **on the client side** by the SDKs: +/// `WorkflowCtx._inline_lock` (Python `asyncio.Lock`) and +/// `WorkflowCtx._inlineChain` (TypeScript promise chain) serialize the +/// fast-path POSTs per workflow invocation. The lock wraps only the HTTP +/// call — `fn()` itself still runs in parallel across `asyncio.gather` / +/// `Promise.all` — so the only thing actually ordered is the sequence of +/// API requests, which is exactly what the helper needs to rely on. +/// +/// **Future contributors: do not remove the SDK-side lock without also +/// fixing the server-side first-write guarantee (e.g. via a single-statement +/// merge-UPDATE that's cheap enough — see note below — or a pre-created +/// `v2_job_status` row).** The comment used to claim the SDKs could fire in +/// parallel without client-side serialization; that was wrong, because the +/// helper writes the whole `_checkpoint`. +/// +/// Cross-process concurrency with the worker-side legacy fallback arm is +/// safe by construction: both paths receive the same `_StepSuspend` payload +/// (same `key`, same `result`, same `started_at`, same `duration_ms`), so +/// even if the fast path's commit and the worker arm's commit land out of +/// order for the same step, the worst case is a redundant idempotent write, +/// not a divergence. +/// +/// ## Why not a single-statement merge-UPDATE? +/// +/// A pure-SQL single-statement variant (pushing load-modify-save entirely +/// into `jsonb_set` + `jsonb_build_object` so correctness on the first +/// write comes from Postgres row locking rather than a client-side lock) +/// was prototyped and measured at ~80 ms per call in debug mode — the +/// nested `COALESCE(v2_job_status.workflow_as_code_status->'_checkpoint' +/// ->...)` accesses cause Postgres to evaluate the growing JSONB subtree +/// multiple times per call, and the `||` merges re-serialize the whole +/// object. The two-statement Rust-side load-modify-save below is ~10× +/// faster in practice, so we keep it and rely on the SDK-level lock. +pub async fn persist_inline_checkpoint_delta( + tx: &mut Transaction<'_, Postgres>, + job_id: &Uuid, + source_hash_hint: Option<&str>, + key: &str, + result: Value, + started_at: Option<&str>, + duration_ms: Option, +) -> error::Result<()> { + // Row-lock the existing checkpoint row (if any) for the duration of the + // transaction. NULL if the row doesn't exist yet — see the doc comment + // above for why the first-write race is accepted. + let row: Option> = sqlx::query_scalar( + "SELECT workflow_as_code_status->'_checkpoint' + FROM v2_job_status WHERE id = $1 FOR UPDATE", + ) + .bind(job_id) + .fetch_optional(&mut **tx) + .await?; + + let mut checkpoint: WacCheckpoint = match row.flatten() { + Some(status) => serde_json::from_value(status).unwrap_or_else(|e| { + tracing::warn!( + job_id = %job_id, + error = %e, + "Failed to deserialize WAC checkpoint, resetting to empty" + ); + WacCheckpoint::default() + }), + None => WacCheckpoint::default(), + }; + + // Source hash validation: detect if code changed between replays. + match source_hash_hint { + Some(hint) if !hint.is_empty() => { + if checkpoint.source_hash.is_empty() { + checkpoint.source_hash = hint.to_string(); + } else if checkpoint.source_hash != hint { + return Err(Error::ExecutionErr( + "Workflow source code changed between replays. \ + Cannot safely resume from checkpoint — step keys may have shifted. \ + Please restart this workflow." + .to_string(), + )); + } + } + _ => { + // Preview / inline jobs have no `runnable_id`, so the caller passes + // None (or Some("")). We can't validate drift for these — log once + // so operators can tell which jobs are running unguarded. + tracing::debug!( + job_id = %job_id, + "WAC v2 inline checkpoint without runnable hash — source-hash drift protection is off for this job" + ); + } + } + + tracing::info!( + job_id = %job_id, + step_key = %key, + "WAC v2 inline checkpoint — persisting step result" + ); + + add_completed_step(&mut checkpoint, key, 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(&mut **tx) + .await + .map_err(|e| Error::InternalErr(format!("Failed to save WAC checkpoint: {e}")))?; + + // Write the `_step/` timeline entry. Fall back to now() when the + // client doesn't provide started_at (older SDK versions omit it). + let now_str = chrono::Utc::now().to_rfc3339(); + let sa = started_at.unwrap_or(&now_str); + let mut timeline_val = serde_json::json!({ + "scheduled_for": sa, + "started_at": sa, + "name": key, + }); + if let Some(dur) = duration_ms { + timeline_val["duration_ms"] = serde_json::json!(dur); + } + let step_timeline_key = format!("_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(&step_timeline_key) + .bind(&timeline_val) + .execute(&mut **tx) + .await + .map_err(|e| Error::InternalErr(format!("Failed to write step timeline: {e}")))?; + + Ok(()) +} diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 0e44cc4750..c7893a63e0 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -2305,8 +2305,7 @@ pub async fn handle_wac_v2_output( modules: &Option>, ) -> error::Result> { use crate::wac_executor::{ - add_completed_step, load_checkpoint, parse_wac_output, update_checkpoint_for_dispatch, - WacOutput, + load_checkpoint, parse_wac_output, update_checkpoint_for_dispatch, WacOutput, }; use serde_json::Value; use windmill_common::get_latest_flow_version_info_for_path; @@ -3143,104 +3142,44 @@ pub async fn handle_wac_v2_output( } }; - let mut checkpoint = load_checkpoint(db, &job.id).await?; + // All-or-nothing: the checkpoint save, the `_step/` timeline + // write, and the `running = false` queue reset must commit + // together. If we split them, a crash or failure in the middle + // would leave the job queued with `running = true` but a + // checkpoint that already contains the current step — any retry + // would then skip the step entirely. Passing the caller's `tx` + // into `persist_inline_checkpoint_delta` preserves the original + // atomicity from before the shared-helper refactor. + let source_hash = job.runnable_id.map(|h| h.0.to_string()); + let mut tx = db.begin().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(), - )); - } - } + crate::wac_executor::persist_inline_checkpoint_delta( + &mut tx, + &job.id, + source_hash.as_deref(), + &key, + value, + started_at.as_deref(), + duration_ms, + ) + .await?; - tracing::info!( - job_id = %job.id, - step_key = %key, - "WAC v2 inline checkpoint — persisting step result" - ); + // 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}" + )) + })?; - add_completed_step(&mut checkpoint, &key, value); - - // Save checkpoint + write step timeline entry + 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}")) - })?; - - // Write timeline entry for the inline step (keyed as _step/). - // Fall back to now() when the client doesn't provide started_at - // (older windmill-client versions omit it). - { - let now_str = chrono::Utc::now().to_rfc3339(); - let sa = started_at.as_deref().unwrap_or(&now_str); - let mut timeline_val = serde_json::json!({ - "scheduled_for": sa, - "started_at": sa, - "name": key, - }); - if let Some(dur) = duration_ms { - timeline_val["duration_ms"] = serde_json::json!(dur); - } - let step_timeline_key = format!("_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(&step_timeline_key) - .bind(&timeline_val) - .execute(&mut *tx) - .await - .map_err(|e| { - error::Error::internal_err(format!("Failed to write step timeline: {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?; - } + tx.commit().await?; Err(error::Error::WacSuspended(format!( "WAC v2 job {} inline checkpoint for step {}", diff --git a/backend/windmill-worker/src/wac_executor.rs b/backend/windmill-worker/src/wac_executor.rs index 28de226102..d9479e0669 100644 --- a/backend/windmill-worker/src/wac_executor.rs +++ b/backend/windmill-worker/src/wac_executor.rs @@ -1,4 +1,4 @@ -use serde::{Deserialize, Serialize}; +use serde::Deserialize; use serde_json::value::RawValue; use serde_json::Value; use uuid::Uuid; @@ -6,35 +6,13 @@ 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, -} +// Checkpoint model + persistence primitives live in windmill-common so the +// API server can use them without pulling in the full worker crate. Re-export +// here for historical call sites inside windmill-worker. +pub use windmill_common::wac::{ + load_checkpoint, persist_inline_checkpoint_delta, save_checkpoint, WacCheckpoint, + WacPendingSteps, +}; /// Output from a single WAC invocation (parsed from result.json). #[derive(Debug, Deserialize)] @@ -106,63 +84,6 @@ 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()) @@ -192,23 +113,6 @@ pub fn update_checkpoint_for_dispatch( 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 { diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index d70372f626..18573e297d 100644 --- a/python-client/wmill/wmill/client.py +++ b/python-client/wmill/wmill/client.py @@ -2416,6 +2416,23 @@ class WorkflowCtx: self._counters: dict[str, int] = {} self._pending: list = [] self._executing_key: str | None = checkpoint.get("_executing_key") + # Reuse a single httpx.AsyncClient across all fast-path step() calls + # in this workflow invocation. Instantiating a fresh client per call + # allocates a new connection pool each time — on localhost this adds + # ~15ms per step, dominating the end-to-end cost. Lazily built so no + # client is created for workflows that never hit the fast path. + self._inline_http_client: "httpx.AsyncClient | None" = None + # Serializes fast-path POSTs across concurrent step() calls within + # one workflow invocation. Wraps only the HTTP call, not fn() — so + # `asyncio.gather(step("a", fn_a), step("b", fn_b))` still runs the + # two fn() bodies in parallel, only the API requests are ordered. + # This closes the first-write race window against `SELECT FOR UPDATE` + # on a not-yet-created `v2_job_status` row: concurrent POSTs would + # both see None and both overwrite each other's checkpoint because + # the helper writes the whole serialized `_checkpoint` object, not + # a single `completed_steps[key]`. Lazily built so the ctx can be + # constructed outside an event loop (tests do this). + self._inline_lock: "_asyncio.Lock | None" = None def _alloc_key(self, name: str = "step") -> str: """Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent.""" @@ -2543,6 +2560,54 @@ class WorkflowCtx: result = await result duration_ms = int((_time_mod.monotonic() - t0) * 1000) + # Fast path: POST the delta to the new per-job API endpoint and return + # the result directly, letting the workflow subprocess continue into + # the next step() without unwinding. On any failure — network, auth, + # timeout, source-hash mismatch, old backend without the endpoint — + # fall through to raising _StepSuspend so the worker takes the legacy + # suspend-and-replay path. Gated by WM_WAC_INLINE_FAST_PATH (default + # on) so the old behavior stays reachable for A/B testing and rollback. + _fast_path_flag = os.environ.get("WM_WAC_INLINE_FAST_PATH", "1").strip().lower() + _fast_path_enabled = _fast_path_flag not in ("0", "false", "off", "no") + _job_id = os.environ.get("WM_JOB_ID") + _workspace = os.environ.get("WM_WORKSPACE") + _base = os.environ.get("BASE_INTERNAL_URL") + _token = os.environ.get("WM_TOKEN") + if _fast_path_enabled and _job_id and _workspace and _base and _token: + try: + if self._inline_lock is None: + self._inline_lock = _asyncio.Lock() + # Lock wraps only the POST, not fn() above — concurrent + # step() calls run fn() in parallel, then serialize on + # the API request. + async with self._inline_lock: + if self._inline_http_client is None: + self._inline_http_client = httpx.AsyncClient( + timeout=httpx.Timeout(10.0), + headers={ + "Authorization": f"Bearer {_token}", + "Content-Type": "application/json", + }, + ) + _resp = await self._inline_http_client.post( + f"{_base}/api/w/{_workspace}/jobs/wac/inline_checkpoint/{_job_id}", + json={ + "key": key, + "result": result, + "started_at": started_at, + "duration_ms": duration_ms, + }, + ) + _resp.raise_for_status() + return result + except Exception as _e: + logger.info( + "WAC v2 inline fast path failed for key %s, falling back to suspend: %s", + key, + _e, + ) + # fall through to the legacy suspend path + raise _StepSuspend({ "mode": "inline_checkpoint", "steps": [], @@ -2873,7 +2938,23 @@ async def _run_workflow_async(func, checkpoint: dict, input_args: dict): } return {"type": "dispatch", **info} finally: - _workflow_ctx.reset(token) + # Close the lazily-built fast-path httpx client so we don't emit + # asyncio ResourceWarning('unclosed transport') on shutdown and don't + # leak connection pools when this coroutine is driven from a + # long-lived loop (tests, REPL, embedded callers). + # + # Wrapped in its own try/finally so that asyncio.CancelledError + # (which is a BaseException since Python 3.8) during aclose() does + # not skip the _workflow_ctx.reset(token) below. + try: + if ctx._inline_http_client is not None: + try: + await ctx._inline_http_client.aclose() + except Exception: + pass + ctx._inline_http_client = None + finally: + _workflow_ctx.reset(token) def _run_workflow(func, checkpoint: dict, input_args: dict): diff --git a/typescript-client/client.ts b/typescript-client/client.ts index 41b66e6d96..ce40bd07e8 100644 --- a/typescript-client/client.ts +++ b/typescript-client/client.ts @@ -1494,6 +1494,16 @@ export class WorkflowCtx { private _suspended = false; /** When set, the task matching this key executes its inner function directly */ _executingKey: string | null; + /** Serializes fast-path POSTs across concurrent step() calls within one + * workflow invocation. Wraps only the HTTP call, not fn() — so + * `Promise.all([step("a", fn_a), step("b", fn_b)])` still runs the two + * fn() bodies in parallel, only the API requests are ordered. This + * closes the first-write race window against `SELECT FOR UPDATE` on a + * not-yet-created `v2_job_status` row: two concurrent POSTs would both + * see None and overwrite each other's checkpoint because the helper + * writes the whole serialized `_checkpoint` object, not a single + * `completed_steps[key]`. Initialized to a resolved promise. */ + private _inlineChain: Promise = Promise.resolve(); constructor(checkpoint: Record = {}) { this.completed = checkpoint?.completed_steps ?? {}; @@ -1652,6 +1662,72 @@ export class WorkflowCtx { const t0 = Date.now(); const result = await fn(); const durationMs = Date.now() - t0; + + // Fast path: POST the delta to the new per-job API endpoint and return the + // result directly so the workflow subprocess continues into the next step() + // without unwinding. Concurrent step() calls (e.g. inside Promise.all) run + // their fn() bodies in parallel, then serialize the API POSTs via a + // per-ctx promise chain (`this._inlineChain`). Serializing the POSTs is + // required because the backend helper writes the whole serialized + // `_checkpoint` object per call, and two concurrent writes against a + // not-yet-created `v2_job_status` row would both see None under + // `SELECT FOR UPDATE` and overwrite each other. + // + // On any failure — network, auth, timeout, source-hash mismatch, old + // backend without the endpoint — fall through to throwing StepSuspend so + // the worker takes the legacy suspend-and-replay path. Gated by + // WM_WAC_INLINE_FAST_PATH (default on) so the old behavior stays + // reachable for A/B testing and rollback. + const fastPathFlagRaw = (getEnv("WM_WAC_INLINE_FAST_PATH") ?? "1").trim().toLowerCase(); + const fastPathEnabled = + fastPathFlagRaw !== "0" && + fastPathFlagRaw !== "false" && + fastPathFlagRaw !== "off" && + fastPathFlagRaw !== "no"; + const jobId = getEnv("WM_JOB_ID"); + const workspace = getEnv("WM_WORKSPACE"); + if (fastPathEnabled && jobId && workspace && OpenAPI.BASE && OpenAPI.TOKEN) { + const chainTail = this._inlineChain.then(async () => { + const ctrl = new AbortController(); + const t = setTimeout(() => ctrl.abort(), 10_000); + try { + const resp = await fetch( + `${OpenAPI.BASE}/w/${workspace}/jobs/wac/inline_checkpoint/${jobId}`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${OpenAPI.TOKEN}`, + }, + body: JSON.stringify({ + key, + result, + started_at: startedAt, + duration_ms: durationMs, + }), + signal: ctrl.signal, + }, + ); + if (!resp.ok) { + throw new Error(`inline_checkpoint API ${resp.status}`); + } + } finally { + clearTimeout(t); + } + }); + // Swallow chain errors so a past failure does not poison future awaits. + this._inlineChain = chainTail.catch(() => {}); + try { + await chainTail; + return result as T; + } catch (e) { + console.log( + `WAC v2 inline fast path failed for key ${key}, falling back to suspend: ${e}`, + ); + // fall through to the legacy suspend path below + } + } + throw new StepSuspend({ mode: "inline_checkpoint", steps: [], key, result, started_at: startedAt, duration_ms: durationMs }); } }