From 120db3f8c305d17f58689e945b1bdb624751d092 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 22 Sep 2026 02:23:53 +0200 Subject: [PATCH] fix: recover a chosen branch's own result on re-entry Co-Authored-By: Claude Opus 5 (1M context) --- backend/tests/ai_decision.rs | 49 +++++++++++++++++++--- backend/windmill-worker/src/worker_flow.rs | 26 +++++++----- 2 files changed, 59 insertions(+), 16 deletions(-) diff --git a/backend/tests/ai_decision.rs b/backend/tests/ai_decision.rs index 2eb7fb5eeb..99c5703497 100644 --- a/backend/tests/ai_decision.rs +++ b/backend/tests/ai_decision.rs @@ -8,6 +8,7 @@ use std::sync::{ use axum::{routing::post, Json, Router}; use serde_json::{json, Value}; use sqlx::{Pool, Postgres}; +use uuid::Uuid; use windmill_common::{ flow_status::{BranchChosen, FlowStatus, FlowStatusModule}, flows::FlowValue, @@ -113,17 +114,55 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks( assert_eq!(job.json_result(), Some(json!("handled refund"))); let status: FlowStatus = serde_json::from_value(job.flow_status.expect("flow status"))?; - let FlowStatusModule::Success { branch_chosen, decision_job, .. } = &status.modules[0] else { + let FlowStatusModule::Success { job: branch_job, branch_chosen, decision_job, .. } = + status.modules[0].clone() + else { panic!("decision step did not succeed: {:?}", status.modules[0]); }; assert!(matches!( branch_chosen, Some(BranchChosen::Branch { branch: 1 }) )); - assert!( - decision_job.is_some(), - "the decision job is kept once the branch has run" - ); + let decision_job = decision_job.expect("the decision job is kept once the branch has run"); + + // Re-entered after a restart, with the step's status written but no result in memory, the + // step returns its branch's result, or the answers when the chosen branch was empty. + for (planted_job, expected) in [ + (branch_job, json!("handled refund")), + (Uuid::nil(), json!("refund")), + ] { + let flow: FlowValue = serde_json::from_value(json!({"modules": [{"id": "d", "value": { + "type": "aidecision", + "input_transforms": decision_inputs, + "branches": [{"expr": "true", "modules": []}] + }}]}))?; + let planted = json!({ + "type": "Success", "id": "d", "job": planted_job, + "branch_chosen": {"type": "branch", "branch": 0}, "decision_job": decision_job + }); + let job = RunJob::from(JobPayload::RawFlow { value: flow, path: None, restarted_from: None }) + .run_until_complete_with(&db, false, port, |id| { + let (db, planted) = (db.clone(), planted.clone()); + async move { + sqlx::query( + "UPDATE v2_job_status SET flow_status = jsonb_set(flow_status, '{modules,0}', $1) WHERE id = $2", + ) + .bind(planted) + .bind(id) + .execute(&db) + .await + .unwrap(); + } + }) + .await; + let result = job.json_result().expect("a result"); + let result = if planted_job.is_nil() { + result["output"]["intent"]["choice"].clone() + } else { + result + }; + assert_eq!(result, expected); + } // A branch condition still reads the step before the decision as `results.`, while // `previous_result` is the answers. Once an empty branch is chosen, `results.d` stays the diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 5f2dd67fe7..67216c5ba1 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -1305,10 +1305,14 @@ pub async fn update_flow_status_after_job_completion_internal( false }; success = true; - // An AI decision whose chosen branch is empty completes with no job of its - // own: its job, and so `results.`, stays the decision's answers. + // An AI decision can complete with no job of its own (an empty chosen branch, + // or re-entry after a restart): it keeps its branch's job, else the decision's, + // so `results.` stays the step's result. let job = match module_status.decision_job() { - Some(decision_job) if job_id_for_status.is_nil() => decision_job, + Some(decision_job) if job_id_for_status.is_nil() => module_status + .job() + .filter(|job| !job.is_nil()) + .unwrap_or(decision_job), _ => *job_id_for_status, }; ( @@ -3308,16 +3312,16 @@ async fn push_next_flow_job( { let result = if flow.modules.is_empty() { to_raw_value(arc_flow_job_args.as_ref()) - } else if let FlowStatusModule::Success { branch_chosen: Some(_), decision_job, .. } = - &status_module + } else if let FlowStatusModule::Success { + branch_chosen: Some(_), job, decision_job, .. + } = &status_module { - match (&last_job_result, decision_job) { + // Re-entered without the in-memory result, after a restart: the chosen branch's job + // holds it, and an AI decision's empty branch (no job) returns its answers. + let stored = Some(*job).filter(|job| !job.is_nil()).or(*decision_job); + match (&last_job_result, stored) { (Some(result), _) => result.as_ref().clone(), - // Re-entered without the in-memory result (the worker restarted after an empty - // branch was chosen): an AI decision's empty branch returns its answers. - (None, Some(decision_job)) => { - completed_job_result(db, &flow_job.workspace_id, decision_job).await? - } + (None, Some(job)) => completed_job_result(db, &flow_job.workspace_id, &job).await?, (None, None) => to_raw_value(&json!("{}")), } } else {