feat(flow): support worker tag override on AI agent steps (#9513)

* feat(flow): support worker tag override on AI agent steps

* refactor: drop ineffective tag passthrough in nested agent tool path

* chore: regenerate openflow-derived system prompt artifacts

backend/windmill-types/src/flows.rs

@@ -942,6 +942,8 @@ pub enum FlowModuleValue {
     AIAgent {
         input_transforms: HashMap<String, InputTransform>,
         tools: Vec<AgentTool>,
+        #[serde(skip_serializing_if = "is_none_or_empty")]
+        tag: Option<String>,
         #[serde(default, skip_serializing_if = "is_false")]
         omit_output_from_conversation: bool,
     },
@@ -1080,6 +1082,7 @@ impl<'de> Deserialize<'de> for FlowModuleValue {
                 tools: untagged
                     .tools
                     .ok_or_else(|| serde::de::Error::missing_field("tools"))?,
+                tag: untagged.tag,
                 omit_output_from_conversation: untagged
                     .omit_output_from_conversation
                     .unwrap_or(false),
@@ -1233,4 +1236,41 @@ mod tests {
 
         assert!(omit_output_from_conversation);
     }
+
+    #[test]
+    fn ai_agent_tag_round_trips() {
+        let input = json!({
+            "type": "aiagent",
+            "tools": [],
+            "input_transforms": {},
+            "tag": "bedrock"
+        });
+
+        let val: FlowModuleValue = serde_json::from_value(input).unwrap();
+        let FlowModuleValue::AIAgent { ref tag, .. } = val else {
+            panic!("expected aiagent module");
+        };
+        assert_eq!(tag.as_deref(), Some("bedrock"));
+
+        let output = serde_json::to_string(&val).unwrap();
+        assert!(output.contains("\"tag\":\"bedrock\""));
+    }
+
+    #[test]
+    fn ai_agent_tag_defaults_to_none_and_is_omitted_when_serializing() {
+        let input = json!({
+            "type": "aiagent",
+            "tools": [],
+            "input_transforms": {}
+        });
+
+        let val: FlowModuleValue = serde_json::from_value(input).unwrap();
+        let FlowModuleValue::AIAgent { ref tag, .. } = val else {
+            panic!("expected aiagent module");
+        };
+        assert!(tag.is_none());
+
+        let output = serde_json::to_string(&val).unwrap();
+        assert!(!output.contains("tag"));
+    }
 }

backend/windmill-worker/src/ai/tools.rs

@@ -426,6 +426,8 @@ async fn execute_windmill_tool(
                 ));
             }
             let path = format!("{}/tools/{}", ctx.job.runnable_path(), tool_module.id);
+            // tool jobs are pushed with the parent agent job's tag and executed inline on the
+            // same worker, so a tag override on a nested agent tool does not apply here
             JobPayloadWithTag {
                 payload: JobPayload::AIAgent { path },
                 tag: None,

backend/windmill-worker/src/worker_flow.rs

@@ -5058,13 +5058,13 @@ async fn compute_next_flow_transform(
                 NextStatus::NextStep,
             ))
         }
-        FlowModuleValue::AIAgent { .. } => {
+        FlowModuleValue::AIAgent { tag, .. } => {
             let path = get_path(flow_job, status, module);
             let payload = JobPayload::AIAgent { path };
             Ok(NextFlowTransform::Continue(
                 ContinuePayload::SingleJob(JobPayloadWithTag {
                     payload,
-                    tag: None,
+                    tag: tag.filter(|t| !t.trim().is_empty()),
                     delete_after_use,
                     delete_after_secs,
                     timeout: None,

backend/windmill-worker/src/worker_lockfiles.rs

@@ -1248,6 +1248,7 @@ async fn lock_modules(
                 FlowModuleValue::AIAgent {
                     input_transforms,
                     mut tools,
+                    tag,
                     omit_output_from_conversation,
                 } => {
                     // Extract FlowModules from tools and track their original indices
@@ -1299,6 +1300,7 @@ async fn lock_modules(
                     e.value = FlowModuleValue::AIAgent {
                         input_transforms,
                         tools,
+                        tag,
                         omit_output_from_conversation,
                     }
                     .into();

frontend/src/lib/components/flows/content/FlowModuleComponent.svelte

@@ -746,7 +746,7 @@
 						console.log('tagChange', e.detail)
 						if (flowModule.value.type == 'script') {
 							flowModule.value.tag_override = e.detail
-						} else if (flowModule.value.type == 'rawscript') {
+						} else if (flowModule.value.type == 'rawscript' || flowModule.value.type == 'aiagent') {
 							flowModule.value.tag = e.detail
 						}
 					}}

frontend/src/lib/components/flows/content/FlowModuleHeader.svelte

@@ -28,7 +28,8 @@
 	}
 
 	let { module, tag }: Props = $props()
-	const { scriptEditorDrawer, flowEditorDrawer } = getContext<FlowEditorContext>('FlowEditorContext')
+	const { scriptEditorDrawer, flowEditorDrawer } =
+		getContext<FlowEditorContext>('FlowEditorContext')
 
 	const dispatch = createEventDispatcher()
 	let customUi: undefined | FlowBuilderWhitelabelCustomUi = getContext('customUi')
@@ -176,6 +177,17 @@
 		/>
 	{/if}
 
+	{#if module.value.type === 'aiagent' && customUi?.tagEdit != false}
+		<FlowModuleWorkerTagSelect
+			isPreprocessor={false}
+			placeholder={customUi?.tagSelectPlaceholder}
+			noLabel={customUi?.tagSelectNoLabel}
+			nullTag={tag}
+			tag={module.value.tag}
+			on:change={(e) => dispatch('tagChange', e.detail)}
+		/>
+	{/if}
+
 	{#if module.value.type === 'rawscript'}
 		<FlowModuleWorkerTagSelect
 			isPreprocessor={module.id == 'preprocessor'}

openflow.openapi.yaml

@@ -1078,6 +1078,9 @@ components:
           type: string
           enum:
             - aiagent
+        tag:
+          type: string
+          description: Worker group tag for execution routing. If not set, the AI agent step runs on the flow's tag (default `flow`)
         omit_output_from_conversation:
           type: boolean
           default: false