fix: recover empty-branch answers, keep decision_job on failure

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ruben Fiszel
2026-09-22 02:10:17 +02:00
co-authored by Claude Opus 5
parent 1859cfca9c
commit 8afe3f056c
9 changed files with 92 additions and 54 deletions
+43 -5
View File
@@ -1,5 +1,10 @@
//! AI decision steps run end to end against a stand-in for TypeSafe's System One endpoint.
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use axum::{routing::post, Json, Router};
use serde_json::{json, Value};
use sqlx::{Pool, Postgres};
@@ -10,11 +15,15 @@ use windmill_common::{
};
use windmill_test_utils::*;
/// Answers every question as a choice of `refund`, the way TypeSafe shapes an answer.
async fn start_typesafe_stub() -> anyhow::Result<u16> {
/// Answers every question as a choice of `refund`, the way TypeSafe shapes an answer, and counts
/// the calls.
async fn start_typesafe_stub() -> anyhow::Result<(u16, Arc<AtomicUsize>)> {
let calls = Arc::new(AtomicUsize::new(0));
let counted = calls.clone();
let app = Router::new().route(
"/v1/systemone",
post(|Json(body): Json<Value>| async move {
post(move |Json(body): Json<Value>| async move {
counted.fetch_add(1, Ordering::SeqCst);
let answers = body["questions"]
.as_object()
.into_iter()
@@ -35,7 +44,7 @@ async fn start_typesafe_stub() -> anyhow::Result<u16> {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let port = listener.local_addr()?.port();
tokio::spawn(async move { axum::serve(listener, app).await });
Ok(port)
Ok((port, calls))
}
/// The decision chooses its branch from its answers, a branch step reads those answers as
@@ -48,7 +57,7 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks(
initialize_tracing().await;
// The stub listens on loopback, which the SSRF check refuses without this.
std::env::set_var("ALLOW_PRIVATE_AI_BASE_URLS", "true");
let stub_port = start_typesafe_stub().await?;
let (stub_port, calls) = start_typesafe_stub().await?;
let server = ApiServer::start(db.clone()).await?;
let port = server.addr.port();
@@ -167,5 +176,34 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks(
"the branch reading the previous step was not chosen: {:?}",
status.modules[1]
);
// A retry after its branch fails asks the decision again from its own inputs, rather than
// from the failed branch's args.
let flow: FlowValue = serde_json::from_value(json!({
"modules": [{
"id": "d",
"retry": {"constant": {"attempts": 1, "seconds": 0}},
"value": {
"type": "aidecision",
"input_transforms": decision_inputs,
"branches": [{"expr": "true", "modules": [{"id": "boom", "value": {
"type": "rawscript",
"language": "deno",
"content": "export function main(){ throw new Error('branch failed') }",
"input_transforms": {}
}}]}]
}
}]
}))?;
let before = calls.load(Ordering::SeqCst);
let job = run_job_in_new_worker_until_complete(
&db,
false,
JobPayload::RawFlow { value: flow, path: None, restarted_from: None },
port,
)
.await;
assert!(!job.success, "the failing branch should fail the flow");
assert_eq!(calls.load(Ordering::SeqCst) - before, 2);
Ok(())
}
+5 -1
View File
@@ -303,6 +303,8 @@ pub enum FlowStatusModule {
agent_actions: Option<Vec<AgentAction>>,
#[serde(skip_serializing_if = "Option::is_none")]
agent_actions_success: Option<Vec<bool>>,
#[serde(skip_serializing_if = "Option::is_none")]
decision_job: Option<Uuid>,
},
}
@@ -391,6 +393,7 @@ impl<'de> Deserialize<'de> for FlowStatusModule {
failed_retries: untagged.failed_retries.unwrap_or_default(),
agent_actions: untagged.agent_actions,
agent_actions_success: untagged.agent_actions_success,
decision_job: untagged.decision_job,
}),
other => Err(serde::de::Error::unknown_variant(
other,
@@ -526,7 +529,8 @@ impl FlowStatusModule {
pub fn decision_job(&self) -> Option<Uuid> {
match self {
FlowStatusModule::InProgress { decision_job, .. }
| FlowStatusModule::Success { decision_job, .. } => *decision_job,
| FlowStatusModule::Success { decision_job, .. }
| FlowStatusModule::Failure { decision_job, .. } => *decision_job,
_ => None,
}
}
+21 -14
View File
@@ -1083,6 +1083,7 @@ pub async fn update_flow_status_after_job_completion_internal(
failed_retries: vec![],
agent_actions: None,
agent_actions_success: None,
decision_job: None,
}
};
let r = sqlx::query_scalar!(
@@ -1370,6 +1371,7 @@ pub async fn update_flow_status_after_job_completion_internal(
failed_retries: old_status.retry.failed_jobs.clone(),
agent_actions: module_status.agent_actions(),
agent_actions_success: module_status.agent_actions_success(),
decision_job: module_status.decision_job(),
}),
)
}
@@ -3304,23 +3306,28 @@ async fn push_next_flow_job(
if (flow.modules.is_empty() && !step.is_preprocessor_step())
|| matches!(status_module, FlowStatusModule::Success { .. })
{
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
{
match (&last_job_result, decision_job) {
(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, None) => to_raw_value(&json!("{}")),
}
} else {
// it has to be an empty for loop event
serde_json::from_str("[]").unwrap()
};
return Ok(PushNextFlowJob::Done(Some(UpdateFlow {
flow: flow_job.id,
success: true,
result: if flow.modules.is_empty() {
to_raw_value(arc_flow_job_args.as_ref())
} else if matches!(
status_module,
FlowStatusModule::Success { branch_chosen: Some(_), .. }
) {
last_job_result
.as_ref()
.map(|x| x.as_ref().clone())
.unwrap_or_else(|| to_raw_value(&json!("{}")))
} else {
// it has to be an empty for loop event
serde_json::from_str("[]").unwrap()
},
result,
stop_early_override: None,
w_id: flow_job.workspace_id.clone(),
worker_dir: worker_dir.to_string(),
@@ -67,13 +67,7 @@ describe('validateAiAgentProviders', () => {
expect(() =>
validateAiAgentProviders(anthropicProvider('claude-something-new'), {
...COMPLETE,
options: [
{
...ANTHROPIC,
modelsAreLive: false,
models: { ids: ['claude-sonnet-5'], complete: false }
}
]
options: [{ ...ANTHROPIC, modelsAreLive: false, models: { ids: ['claude-sonnet-5'], complete: false } }]
})
).not.toThrow()
})
@@ -182,9 +182,7 @@ export function selectAiAgentProviderCandidates(
if (!configuredPaths.has(candidate.resourcePath)) return 2
return candidate.kind === defaultProviderKind ? 0 : 1
}
return candidates.sort(
(a, b) => rank(a) - rank(b) || a.resourcePath.localeCompare(b.resourcePath)
)
return candidates.sort((a, b) => rank(a) - rank(b) || a.resourcePath.localeCompare(b.resourcePath))
}
/** What a set of modules needs from the catalog: `needsCatalog` is false when no AI agent step
@@ -10,7 +10,7 @@
reorderBranches as reorderBranchesOp,
graphBranchIndex
} from '../branchOps'
import { choiceBranches, type BranchChoiceValue } from '../branchChoice'
import { choiceBranches, hasChoiceBranches, type BranchChoiceValue } from '../branchChoice'
import Button from '$lib/components/common/button/Button.svelte'
import TextInput from '$lib/components/text_input/TextInput.svelte'
import { dragHandle, dragHandleZone } from '@windmill-labs/svelte-dnd-action'
@@ -139,7 +139,7 @@
</div>
{/each}
</section>
{#if value.type === 'branchone' || items.length > 0}
{#if hasChoiceBranches(value)}
<div class="flex items-center gap-2 rounded-md bg-surface-tertiary p-3 shadow-sm">
<Badge color="blue" class="text-xs">Default</Badge>
<p class="text-xs italic text-tertiary">Runs if none of the above match</p>
@@ -16,7 +16,12 @@ import {
DUPLICATE_MODULE_PREFIX,
NEW_MODULE_PREFIX
} from './flowDiff'
import { findModuleInFlow, findModuleParent, removeFlowModule, replaceFlowModule } from './flowTree'
import {
findModuleInFlow,
findModuleParent,
removeFlowModule,
replaceFlowModule
} from './flowTree'
import { refreshStateStore } from '$lib/svelte5Utils.svelte'
import type { StateStore } from '$lib/utils'
import type DiffDrawer from '../DiffDrawer.svelte'
@@ -52,7 +52,7 @@ function filterDependentComponents(modules: FlowModule[], id: string): Record<st
modules
.map((mod) => [mod.id, getModuleExprs(mod).filter((expr) => expr.includes(`flow_input`))])
.filter((x) => x[1].length > 0)
)
)
: Object.fromEntries(
modules
.map((mod) => [
@@ -63,7 +63,7 @@ function filterDependentComponents(modules: FlowModule[], id: string): Record<st
})
])
.filter((x) => x[1].length > 0)
)
)
}
function getModuleExprs(x: FlowModule): string[] {
+11 -19
View File
@@ -65,7 +65,9 @@ export function getChildModuleBranches(module: FlowModule): FlowModule[][] {
}
if (module.value.type === 'aiagent' && module.value.tools) {
return [(module.value.tools as FlowModule[]).filter((tool) => isFlowModuleToolCandidate(tool))]
return [
(module.value.tools as FlowModule[]).filter((tool) => isFlowModuleToolCandidate(tool))
]
}
return []
@@ -270,24 +272,17 @@ export function collectFlowNodes(flow: FlowModuleTree): FlowNodeLocation[] {
})
}
visitFlowNodesInModules(
flow.modules ?? [],
(index) => ({ type: 'root', index }),
(match) => {
matches.push({
module: match.module,
location: match.location
})
}
)
visitFlowNodesInModules(flow.modules ?? [], (index) => ({ type: 'root', index }), (match) => {
matches.push({
module: match.module,
location: match.location
})
})
return matches
}
export function findModuleInModules(
modules: FlowModule[],
moduleId: string
): FlowModule | undefined {
export function findModuleInModules(modules: FlowModule[], moduleId: string): FlowModule | undefined {
return (
findFlowNodeInModules(modules, moduleId, (index) => ({ type: 'root', index }))?.module ??
undefined
@@ -320,10 +315,7 @@ export function findModuleInFlow(flow: FlowModuleTree, moduleId: string): FlowMo
return findFlowNode(flow, moduleId)?.module ?? null
}
export function findModuleParent(
flow: FlowModuleTree,
moduleId: string
): ModuleParentLocation | null {
export function findModuleParent(flow: FlowModuleTree, moduleId: string): ModuleParentLocation | null {
return findFlowNode(flow, moduleId)?.location ?? null
}