refactor: render $flow_expr tag values through the $args path lookup

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

backend/windmill-queue/src/jobs.rs

@@ -4657,6 +4657,23 @@ pub fn tag_reads_flow_expr(tag: &str) -> bool {
     tag.contains("$flow_expr[")
 }
 
+/// Renders the value at the dotted `path` below `root` as a dynamic tag component, shared by
+/// `$args[...]` and `$flow_expr[...]`: its JSON text with surrounding quotes trimmed, and empty
+/// once a segment is missing. Only object keys are followed, never array indexes.
+pub fn render_tag_path(root: Option<&RawValue>, path: &str) -> String {
+    let mut value = root.map(|x| x.get()).unwrap_or_default().to_string();
+    for part in path.split('.').filter(|p| !p.is_empty()) {
+        match serde_json::from_str::<serde_json::Value>(&value) {
+            Ok(obj) => value = obj.get(part).map(|v| v.to_string()).unwrap_or_default(),
+            Err(_) => {
+                value = String::new();
+                break;
+            }
+        }
+    }
+    value.trim_matches('"').to_string()
+}
+
 pub fn interpolate_args(x: String, args: &PushArgs, workspace_id: &str) -> String {
     // Save this value to avoid parsing twice
     let workspaced = x.as_str().replace("$workspace", workspace_id).to_string();
@@ -4664,40 +4681,12 @@ pub fn interpolate_args(x: String, args: &PushArgs, workspace_id: &str) -> Strin
         let mut interpolated = workspaced.clone();
         for cap in RE_ARG_TAG.captures_iter(&workspaced) {
             let arg_name = cap.get(1).unwrap().as_str();
-            let arg_value = if arg_name.contains('.') {
-                let parts: Vec<&str> = arg_name.split('.').collect();
-                let root = parts[0];
-                let mut value = args
-                    .args
-                    .get(root)
-                    .or(args.extra.as_ref().and_then(|x| x.get(root)))
-                    .map(|x| x.get())
-                    .unwrap_or_default()
-                    .to_string();
-
-                for part in parts.iter().skip(1) {
-                    if let Ok(obj) = serde_json::from_str::<serde_json::Value>(&value) {
-                        value = obj
-                            .get(part)
-                            .and_then(|v| Some(v.to_string()))
-                            .unwrap_or_default()
-                            .as_str()
-                            .to_string();
-                    } else {
-                        value = "".to_string(); // Invalid JSON or missing field
-                        break;
-                    }
-                }
-                value.trim_matches('"').to_string()
-            } else {
-                args.args
-                    .get(arg_name)
-                    .or(args.extra.as_ref().and_then(|x| x.get(arg_name)))
-                    .map(|x| x.get())
-                    .unwrap_or_default()
-                    .trim_matches('"')
-                    .to_string()
-            };
+            let (root, rest) = arg_name.split_once('.').unwrap_or((arg_name, ""));
+            let root_value = args
+                .args
+                .get(root)
+                .or(args.extra.as_ref().and_then(|x| x.get(root)));
+            let arg_value = render_tag_path(root_value.map(|x| &**x), rest);
             interpolated =
                 interpolated.replace(format!("$args[{}]", arg_name).as_str(), &arg_value);
         }
@@ -7934,3 +7923,28 @@ mod result_metadata_tests {
         assert_eq!(meta.wm_failure.as_deref(), Some("boom"));
     }
 }
+
+#[cfg(test)]
+mod render_tag_path_tests {
+    use super::render_tag_path;
+    use serde_json::value::RawValue;
+
+    fn render(root: &str, path: &str) -> String {
+        render_tag_path(
+            Some(&RawValue::from_string(root.to_string()).unwrap()),
+            path,
+        )
+    }
+
+    // Existing `$args[...]` tags route on exactly these renderings.
+    #[test]
+    fn renders_like_args_tags() {
+        assert_eq!(render(r#""eu""#, ""), "eu");
+        assert_eq!(render(r#"{"a": {"b": "eu"}}"#, "a.b"), "eu");
+        assert_eq!(render(r#"{"n": 4}"#, "n"), "4");
+        assert_eq!(render("null", ""), "null");
+        assert_eq!(render(r#"{"a": 1}"#, "b.c"), "");
+        assert_eq!(render(r#"{"a": ["eu"]}"#, "a.0"), "");
+        assert_eq!(render_tag_path(None, "a"), "");
+    }
+}

backend/windmill-worker/src/worker_flow.rs

@@ -69,7 +69,7 @@ use windmill_common::{
 use windmill_queue::schedule::get_schedule_opt;
 use windmill_queue::{
     add_completed_job, add_completed_job_error, append_logs, get_mini_pulled_job,
-    insert_concurrency_key_capped, interpolate_args,
+    insert_concurrency_key_capped, interpolate_args, render_tag_path,
     report_error_to_workspace_handler_or_critical_side_channel, tag_reads_args,
     tag_reads_flow_expr, try_schedule_next_job, CanceledBy, FlowRunners, MiniCompletedJob,
     MiniPulledJob, PushArgs, PushIsolationLevel, SameWorkerPayload, WrappedError, RE_FLOW_EXPR_TAG,
@@ -3114,9 +3114,8 @@ fn resolve_flow_step_tag(
     }
 }
 
-/// Resolves each `$flow_expr[root.path]` of a step tag by reading `path` from `results` (keyed by
-/// step id), `flow_input` or `flow_env`. As for `$args[...]`, a path that reaches nothing renders
-/// empty, a string renders bare and any other value as its JSON text.
+/// Resolves each `$flow_expr[root.key.path]` of a step tag by reading `key.path` from `results`
+/// (where `key` is a step id), `flow_input` or `flow_env`, rendered as `$args[key.path]` would be.
 async fn interpolate_flow_expr_tag(
     tag: &str,
     db: &DB,
@@ -3139,42 +3138,34 @@ async fn interpolate_flow_expr_tag(
         if rendered.contains_key(path) {
             continue;
         }
-        let mut segments = path.split('.');
-        let root = segments.next().unwrap_or_default();
-        let segments = segments.collect::<Vec<_>>();
-        let value = match (root, segments.split_first()) {
-            ("flow_input", _) => read_flow_expr_path(Some(flow_input), &segments),
-            ("flow_env", _) => read_flow_expr_path(flow_env, &segments),
-            ("results", Some((step_id, rest))) => {
-                match windmill_queue::get_result_by_id(
-                    db.clone(),
-                    flow_job.workspace_id.clone(),
-                    flow_job.id,
-                    step_id.to_string(),
-                    (!rest.is_empty()).then(|| rest.join(".")),
-                )
-                .await
-                {
-                    Ok(result) => serde_json::from_str(result.get()).unwrap_or_default(),
-                    Err(Error::NotFound(_)) => Value::Null,
-                    Err(e) => {
-                        return Err(Error::ExecutionErr(format!(
-                            "Could not resolve the step tag `{tag}`: {e}"
-                        )))
-                    }
+        let (root, key_path) = path.split_once('.').unwrap_or((path, ""));
+        let (key, rest) = key_path.split_once('.').unwrap_or((key_path, ""));
+        if key.is_empty() || !matches!(root, "results" | "flow_input" | "flow_env") {
+            return Err(Error::ExecutionErr(format!(
+                "Could not resolve the step tag `{tag}`: `{path}` must start with \
+                 `results.<step_id>`, `flow_input.<key>` or `flow_env.<key>`"
+            )));
+        }
+        let value = match root {
+            "flow_input" => render_tag_path(flow_input.get(key).map(|x| &**x), rest),
+            "flow_env" => render_tag_path(flow_env.and_then(|e| e.get(key)).map(|x| &**x), rest),
+            _ => match windmill_queue::get_result_by_id(
+                db.clone(),
+                flow_job.workspace_id.clone(),
+                flow_job.id,
+                key.to_string(),
+                None,
+            )
+            .await
+            {
+                Ok(result) => render_tag_path(Some(&*result), rest),
+                Err(Error::NotFound(_)) => String::new(),
+                Err(e) => {
+                    return Err(Error::ExecutionErr(format!(
+                        "Could not resolve the step tag `{tag}`: {e}"
+                    )))
                 }
-            }
-            _ => {
-                return Err(Error::ExecutionErr(format!(
-                    "Could not resolve the step tag `{tag}`: `{path}` must start with \
-                     `results.<step_id>`, `flow_input` or `flow_env`"
-                )))
-            }
-        };
-        let value = match value {
-            Value::String(s) => s,
-            Value::Null => String::new(),
-            v => v.to_string(),
+            },
         };
         rendered.insert(path, value);
     }
@@ -3183,19 +3174,6 @@ async fn interpolate_flow_expr_tag(
         .into_owned())
 }
 
-fn read_flow_expr_path(map: Option<&HashMap<String, Box<RawValue>>>, segments: &[&str]) -> Value {
-    let Some((key, rest)) = segments.split_first() else {
-        return Value::Null;
-    };
-    map.and_then(|m| m.get(*key))
-        .and_then(|raw| serde_json::from_str::<Value>(raw.get()).ok())
-        .and_then(|v| {
-            v.pointer(&rest.iter().map(|s| format!("/{s}")).collect::<String>())
-                .cloned()
-        })
-        .unwrap_or_default()
-}
-
 #[cfg(test)]
 mod tag_resolution_tests {
     use super::resolve_flow_step_tag;