diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 6b759a575f..1abe2a75d9 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -1990,6 +1990,45 @@ fn find_flow_job_index(flow_jobs: &Vec, 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> { + 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()) + ); + } +}