fix: poll decision jobs, keep empty-branch answers and deployed defaults

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ruben Fiszel
2026-09-22 01:38:45 +02:00
co-authored by Claude Opus 5
parent a1a08785ec
commit 1859cfca9c
8 changed files with 172 additions and 50 deletions
@@ -1,4 +1,4 @@
-- The credential an AI agent step with decision output runs TypeSafe's Jev with. Seeded under
-- The credential an AI decision step runs TypeSafe's Jev with. Seeded under
-- 'admins' so every workspace sees it, in the shape the hub gives the other AI providers, which
-- the step reads as `api_key` and `base_url`. A hub sync that later publishes the same name
-- updates the schema in place.
@@ -26,7 +26,7 @@ VALUES (
},
"required": ["api_key"]
}',
'API key for TypeSafe''s Jev decision model, which answers typed questions with calibrated probabilities. Used by AI agent steps with decision output.',
'API key for TypeSafe''s Jev decision model, which answers typed questions with calibrated probabilities. Used by AI decision steps and tools.',
'system',
now()
)
+66 -13
View File
@@ -52,24 +52,25 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks(
let server = ApiServer::start(db.clone()).await?;
let port = server.addr.port();
let decision_inputs = json!({
"provider": {"type": "static", "value": {
"kind": "typesafe",
"resource": {"api_key": "key", "base_url": format!("http://127.0.0.1:{stub_port}/v1")},
"model": "jev-latest"
}},
"state": {"type": "static", "value": "I was charged twice"},
"questions": {"type": "static", "value": {"intent": {
"type": "choice",
"instructions": "What does the customer want?",
"criteria": {"refund": "Money back", "bug": "Something is broken"}
}}}
});
let flow: FlowValue = serde_json::from_value(json!({
"modules": [{
"id": "d",
"value": {
"type": "aidecision",
"input_transforms": {
"provider": {"type": "static", "value": {
"kind": "typesafe",
"resource": {"api_key": "key", "base_url": format!("http://127.0.0.1:{stub_port}/v1")},
"model": "jev-latest"
}},
"state": {"type": "static", "value": "I was charged twice"},
"questions": {"type": "static", "value": {"intent": {
"type": "choice",
"instructions": "What does the customer want?",
"criteria": {"refund": "Money back", "bug": "Something is broken"}
}}}
},
"input_transforms": decision_inputs,
"branches": [
{
"expr": "previous_result.output.intent.choice === 'bug'",
@@ -114,5 +115,57 @@ async fn test_ai_decision_runs_the_branch_its_answer_picks(
decision_job.is_some(),
"the decision job is kept once the branch has run"
);
// A branch condition still reads the step before the decision as `results.<id>`, while
// `previous_result` is the answers. Once an empty branch is chosen, `results.d` stays the
// answers for every later step, not only the next one.
let flow: FlowValue = serde_json::from_value(json!({
"modules": [
{"id": "pre", "value": {
"type": "rawscript",
"language": "deno",
"content": "export function main(){ return { allow: true } }",
"input_transforms": {}
}},
{"id": "d", "value": {
"type": "aidecision",
"input_transforms": decision_inputs,
"branches": [{
"expr": "results.pre.allow && previous_result.output.intent.choice === 'refund'",
"modules": []
}]
}},
{"id": "next", "value": {"type": "identity"}},
{"id": "after", "value": {
"type": "rawscript",
"language": "deno",
"content": "export function main(choice){ return choice }",
"input_transforms": {"choice": {
"type": "javascript",
"expr": "results.d.output.intent.choice"
}}
}}
]
}))?;
let job = run_job_in_new_worker_until_complete(
&db,
false,
JobPayload::RawFlow { value: flow, path: None, restarted_from: None },
port,
)
.await;
assert_eq!(job.json_result(), Some(json!("refund")));
let status: FlowStatus = serde_json::from_value(job.flow_status.expect("flow status"))?;
assert!(
matches!(
&status.modules[1],
FlowStatusModule::Success {
branch_chosen: Some(BranchChosen::Branch { branch: 0 }),
..
}
),
"the branch reading the previous step was not chosen: {:?}",
status.modules[1]
);
Ok(())
}
+1 -1
View File
@@ -124,7 +124,7 @@ pub enum AIProvider {
#[serde(rename = "aws_bedrock")]
AWSBedrock,
/// TypeSafe's Jev decision model. It answers typed questions rather than chatting, so only
/// an agent step with decision output runs it (`run_decision`), never the agent loop.
/// an AI decision runs it (`run_systemone`), never the agent loop.
TypeSafe,
}
+40 -11
View File
@@ -1150,13 +1150,29 @@ impl BranchChoiceShell {
FlowModuleValue::BranchOne { branches, default, default_node }
}
BranchChoiceShell::AIDecision { input_transforms, default_node, tag } => {
FlowModuleValue::AIDecision { input_transforms, branches, default, default_node, tag }
FlowModuleValue::AIDecision {
input_transforms,
branches,
default,
default_node,
tag,
}
}
}
}
}
impl FlowModuleValue {
/// Whether this is an AI decision that runs a branch once it has answered. A deployed flow
/// moves a default's steps into `default_node`, leaving `default` empty.
pub fn is_branched_ai_decision(&self) -> bool {
matches!(
self,
FlowModuleValue::AIDecision { branches, default, default_node, .. }
if !branches.is_empty() || !default.is_empty() || default_node.is_some()
)
}
/// Splits a step that runs one of its branches into its branches, its default and the rest,
/// so a pass rewriting branches handles `BranchOne` and AI decisions alike. Any other step is
/// given back unchanged.
@@ -1164,16 +1180,22 @@ impl FlowModuleValue {
self,
) -> Result<(Vec<Branch>, Vec<FlowModule>, BranchChoiceShell), FlowModuleValue> {
match self {
FlowModuleValue::BranchOne { branches, default, default_node } => {
Ok((branches, default, BranchChoiceShell::BranchOne { default_node }))
}
FlowModuleValue::AIDecision { input_transforms, branches, default, default_node, tag } => {
Ok((
branches,
default,
BranchChoiceShell::AIDecision { input_transforms, default_node, tag },
))
}
FlowModuleValue::BranchOne { branches, default, default_node } => Ok((
branches,
default,
BranchChoiceShell::BranchOne { default_node },
)),
FlowModuleValue::AIDecision {
input_transforms,
branches,
default,
default_node,
tag,
} => Ok((
branches,
default,
BranchChoiceShell::AIDecision { input_transforms, default_node, tag },
)),
other => Err(other),
}
}
@@ -1707,6 +1729,13 @@ mod tests {
output["branches"][0]["expr"],
json!("previous_result.output.intent.choice === 'refund'")
);
assert!(val.is_branched_ai_decision());
// A deployed default-only decision keeps its default as a flow node.
let deployed =
json!({"type": "aidecision", "input_transforms": {}, "default": [], "default_node": 5});
let val: FlowModuleValue = serde_json::from_value(deployed).unwrap();
assert!(val.is_branched_ai_decision());
}
#[test]
+1 -3
View File
@@ -491,9 +491,7 @@ async fn execute_windmill_tool(
}
}
// A tool returns its answers to the calling model, so it has no flow to branch into.
FlowModuleValue::AIDecision { branches, default, .. }
if !branches.is_empty() || !default.is_empty() =>
{
ref value @ FlowModuleValue::AIDecision { .. } if value.is_branched_ai_decision() => {
return Err(Error::BadRequest(format!(
"The AI decision tool {} has branches, which only a flow step can run",
tool_call.function.name
+36 -8
View File
@@ -61,7 +61,10 @@ use crate::{
common::{
build_args_map, resolve_job_timeout, transform_json_value, OccupancyMetrics, StreamNotifier,
},
handle_child::{run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome},
handle_child::{
run_future_with_polling_update_job_poller,
run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome,
},
};
lazy_static::lazy_static! {
@@ -631,6 +634,10 @@ pub async fn handle_ai_agent_job(
job,
&local_args,
direct_parent_job_kind == JobKind::AIAgent,
canceled_by,
mem_peak,
occupancy_metrics,
worker_name,
)
.await;
}
@@ -2116,6 +2123,10 @@ async fn handle_ai_decision(
job: &MiniPulledJob,
local_args: &HashMap<String, Box<RawValue>>,
as_tool: bool,
canceled_by: &mut Option<CanceledBy>,
mem_peak: &mut i32,
occupancy_metrics: &mut OccupancyMetrics,
worker_name: &str,
) -> error::Result<Box<RawValue>> {
let args = serde_json::from_str::<AIDecisionArgs>(&serde_json::to_string(local_args)?)?;
if args.provider.kind != AIProvider::TypeSafe {
@@ -2129,12 +2140,24 @@ async fn handle_ai_decision(
let timeout = resolve_job_timeout(conn, &job.workspace_id, job.id, job.timeout)
.await
.0;
let result = run_systemone(
&credentials,
args.provider.get_model(),
state,
questions,
timeout,
// Under the job poller, so the job keeps its heartbeat and a cancel drops the request.
let result = run_future_with_polling_update_job_poller(
job.id,
job.timeout,
conn,
mem_peak,
canceled_by,
run_systemone(
&credentials,
args.provider.get_model(),
state,
questions,
timeout,
),
worker_name,
&job.workspace_id,
&mut Some(occupancy_metrics),
Box::pin(futures::stream::once(async { 0 })),
)
.await?;
windmill_common::feature_usage::log_feature_usage(
@@ -2169,7 +2192,12 @@ fn decision_tool_description(input_transforms: &HashMap<String, InputTransform>)
Some("score") => {
let levels = criteria
.and_then(|c| c.as_array())
.map(|c| c.iter().filter_map(|l| l.as_str()).collect::<Vec<_>>().join(" < "))
.map(|c| {
c.iter()
.filter_map(|l| l.as_str())
.collect::<Vec<_>>()
.join(" < ")
})
.unwrap_or_default();
format!("a score over {levels}")
}
+24 -10
View File
@@ -1304,11 +1304,17 @@ 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.<id>`, stays the decision's answers.
let job = match module_status.decision_job() {
Some(decision_job) if job_id_for_status.is_nil() => decision_job,
_ => *job_id_for_status,
};
(
true,
Some(FlowStatusModule::Success {
id: module_status.id(),
job: job_id_for_status.clone(),
job,
flow_jobs,
flow_jobs_success,
flow_jobs_duration,
@@ -3913,11 +3919,10 @@ async fn push_next_flow_job(
// we get the args from the last failed job, except for a decision with branches:
// it may have failed in its branch, whose sub-flow took the flow's args, so it
// asks again from its own inputs.
if matches!(
module.get_value(),
Ok(FlowModuleValue::AIDecision { ref branches, ref default, .. })
if !branches.is_empty() || !default.is_empty()
) {
if module
.get_value()
.is_ok_and(|value| value.is_branched_ai_decision())
{
None
} else {
status.retry.failed_jobs.last()
@@ -5417,18 +5422,18 @@ async fn compute_next_flow_transform(
}
// The decision job runs as an AI agent job; the arm below takes over once a decision with
// branches has answered.
FlowModuleValue::AIDecision { ref branches, ref default, tag, .. }
ref value @ FlowModuleValue::AIDecision { ref tag, .. }
if !matches!(
status_module,
FlowStatusModule::InProgress { decision_job: Some(_), branch_chosen: None, .. }
) =>
{
let has_branches = !branches.is_empty() || !default.is_empty();
let has_branches = value.is_branched_ai_decision();
let path = get_path(flow_job, status, module);
Ok(NextFlowTransform::Continue(
ContinuePayload::SingleJob(JobPayloadWithTag {
payload: JobPayload::AIAgent { path },
tag: tag.filter(|t| !t.trim().is_empty()),
tag: tag.clone().filter(|t| !t.trim().is_empty()),
delete_after_use,
delete_after_secs,
timeout: None,
@@ -5728,7 +5733,16 @@ async fn compute_next_flow_transform(
..
} => {
let mut branch_chosen = BranchChosen::Default;
let idcontext = get_transform_context(&flow_job, previous_id, &status);
// `results.<previous id>` reads `previous_result`, which for an answered
// decision is its own answers, so the decision is the previous step here.
let predicate_previous_id = match status_module {
FlowStatusModule::InProgress { decision_job: Some(_), .. } => {
module.id.as_str()
}
_ => previous_id,
};
let idcontext =
get_transform_context(&flow_job, predicate_previous_id, &status);
let mut predicate_err: Option<Error> = None;
for (i, b) in branches.iter().enumerate() {
let pred_res = compute_bool_from_expr(
+2 -2
View File
@@ -148,8 +148,8 @@ export const AI_PROVIDERS: Record<ChatAIProvider, AIProviderDetails> = {
}
}
/** TypeSafe's Jev answers typed questions instead of messages, so only an agent step with decision
* output offers it. The pinned version is there for flows tuned against its probabilities. */
/** TypeSafe's Jev answers typed questions instead of messages, so only an AI decision offers it.
* The pinned version is there for flows tuned against its probabilities. */
export const DECISION_AI_PROVIDERS: Record<
Exclude<AIProvider, ChatAIProvider>,
AIProviderDetails