fix: support windmill chat answer override (#8909)

* fix: support windmill chat answer override

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

* fix: remove output fallback from chat override

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

* fix: handle non-string chat answer overrides

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

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

backend/windmill-worker/src/worker_flow.rs

@@ -1990,6 +1990,45 @@ fn find_flow_job_index(flow_jobs: &Vec<Uuid>, job_id_for_status: &Uuid) -> Optio
     flow_jobs.iter().position(|x| x == job_id_for_status)
 }
 
+fn format_chat_message_value(value: &Value) -> String {
+    match value {
+        Value::Null => "null".to_string(),
+        Value::Bool(_) | Value::Number(_) => value.to_string(),
+        Value::String(text) => text.clone(),
+        Value::Array(_) | Value::Object(_) => serde_json::to_string_pretty(value)
+            .unwrap_or_else(|e| format!("Failed to serialize result: {e}")),
+    }
+}
+
+/// Selects the assistant message persisted for the final non-AI step in a chat-enabled flow.
+/// `windmill_chat_answer` is the explicit override contract:
+/// - `null` suppresses the assistant message
+/// - any other JSON value is rendered as the chat message
+fn extract_chat_message_from_flow_result(result: &RawValue) -> error::Result<Option<String>> {
+    let value: Value = serde_json::from_str(result.get())
+        .map_err(|e| Error::internal_err(format!("Failed to parse flow result: {e}")))?;
+
+    match value {
+        Value::Object(map) => {
+            match map.get("windmill_chat_answer") {
+                Some(Value::Null) => return Ok(None),
+                Some(answer) => return Ok(Some(format_chat_message_value(answer))),
+                _ => {}
+            }
+
+            Ok(Some(
+                serde_json::to_string_pretty(&Value::Object(map))
+                    .unwrap_or_else(|e| format!("Failed to serialize result: {e}")),
+            ))
+        }
+        Value::String(content) => Ok(Some(content)),
+        value => Ok(Some(
+            serde_json::to_string_pretty(&value)
+                .unwrap_or_else(|e| format!("Failed to serialize result: {e}")),
+        )),
+    }
+}
+
 async fn add_tool_message_to_conversation(
     db: &DB,
     job_id: &Uuid,
@@ -2007,28 +2046,8 @@ async fn add_tool_message_to_conversation(
         if let Some(conversation_id) = conversation_id {
             // Only create assistant message if last module is NOT an AI agent, or there was an error
             if !is_ai_agent_step || success == false {
-                let value = serde_json::to_value(result.get())
-                    .map_err(|e| Error::internal_err(format!("Failed to serialize result: {e}")))?;
-
-                let content = match value {
-                    // If it's an Object with "output" key AND the output is a String, return it
-                    serde_json::Value::Object(mut map)
-                        if map.contains_key("output")
-                            && matches!(map.get("output"), Some(serde_json::Value::String(_))) =>
-                    {
-                        if let Some(serde_json::Value::String(s)) = map.remove("output") {
-                            s
-                        } else {
-                            // prettify the whole result
-                            serde_json::to_string_pretty(&map)
-                                .unwrap_or_else(|e| format!("Failed to serialize result: {e}"))
-                        }
-                    }
-                    // Otherwise, if the whole value is a String, return it
-                    serde_json::Value::String(s) => s,
-                    // Otherwise, prettify the whole result
-                    v => serde_json::to_string_pretty(&v)
-                        .unwrap_or_else(|e| format!("Failed to serialize result: {e}")),
+                let Some(content) = extract_chat_message_from_flow_result(result.as_ref())? else {
+                    return Ok(());
                 };
 
                 // Insert new assistant message
@@ -5614,3 +5633,84 @@ pub async fn get_previous_job_result(
         _ => Ok(None),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::extract_chat_message_from_flow_result;
+    use serde_json::{json, value::to_raw_value};
+
+    #[test]
+    fn pretty_prints_full_result_when_no_override_is_present() {
+        let value = json!({
+            "output": "final answer",
+            "metadata": { "foo": "bar" }
+        });
+        let result = to_raw_value(&value).unwrap();
+
+        let message = extract_chat_message_from_flow_result(result.as_ref()).unwrap();
+
+        assert_eq!(message, Some(serde_json::to_string_pretty(&value).unwrap()));
+    }
+
+    #[test]
+    fn uses_windmill_chat_answer_when_it_is_a_string() {
+        let result = to_raw_value(&json!({
+            "windmill_chat_answer": "chat-visible answer",
+            "output": "ignored output",
+            "metadata": { "foo": "bar" }
+        }))
+        .unwrap();
+
+        let message = extract_chat_message_from_flow_result(result.as_ref()).unwrap();
+
+        assert_eq!(message, Some("chat-visible answer".to_string()));
+    }
+
+    #[test]
+    fn skips_persisting_when_windmill_chat_answer_is_null() {
+        let result = to_raw_value(&json!({
+            "windmill_chat_answer": null,
+            "output": "should not be stored"
+        }))
+        .unwrap();
+
+        let message = extract_chat_message_from_flow_result(result.as_ref()).unwrap();
+
+        assert_eq!(message, None);
+    }
+
+    #[test]
+    fn coerces_scalar_windmill_chat_answer_to_a_string() {
+        let result = to_raw_value(&json!({
+            "windmill_chat_answer": 42,
+            "output": "ignored output",
+            "metadata": { "foo": "bar" }
+        }))
+        .unwrap();
+
+        let message = extract_chat_message_from_flow_result(result.as_ref()).unwrap();
+
+        assert_eq!(message, Some("42".to_string()));
+    }
+
+    #[test]
+    fn pretty_prints_structured_windmill_chat_answer_only() {
+        let override_value = json!({
+            "text": "chat-visible answer",
+            "meta": ["a", "b"]
+        });
+        let result = to_raw_value(&json!({
+            "windmill_chat_answer": override_value,
+            "output": "ignored output",
+            "metadata": { "foo": "bar" }
+        }))
+        .unwrap();
+
+        let message = extract_chat_message_from_flow_result(result.as_ref()).unwrap();
+
+        assert_eq!(
+            message,
+            Some(serde_json::to_string_pretty(&override_value).unwrap())
+        );
+    }
+}