fix: strip NUL from ai agent output and job args before jsonb writes (#11568)

* fix: strip NUL from ai agent output and job args before jsonb writes

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* fix: also strip NUL from tool call ids and citations

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

backend/tests/ai_agent_nul.rs

@@ -0,0 +1,125 @@
+use axum::{routing::post, Json, Router};
+use serde_json::{json, Value};
+use sqlx::{Pool, Postgres};
+use std::time::Duration;
+use windmill_common::{flows::FlowValue, jobs::JobPayload};
+use windmill_test_utils::{
+    completed_job, in_test_worker, listen_for_completed_jobs, ApiServer, RunJob, StreamFind,
+};
+
+const MEMORY_ID: &str = "0199c0de-0000-7000-8000-000000000001";
+
+/// A model that puts a U+0000 in its tool call arguments, then in its answer.
+async fn model(Json(body): Json<Value>) -> ([(&'static str, &'static str); 1], String) {
+    let finished = body["messages"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .any(|m| m["role"] == "tool");
+    let delta = if finished {
+        json!({"role":"assistant", "content":"are\u{0}tired"})
+    } else {
+        json!({"role":"assistant", "tool_calls":[{
+            "index":0, "id":"call_0", "type":"function",
+            "function":{"name":"echo", "arguments":"{\"x\":\"a\\u0000b\"}"}
+        }]})
+    };
+    let event = json!({"choices":[{"index":0, "delta":delta, "finish_reason":null}]});
+    let end = json!({"choices":[{"index":0, "delta":{}, "finish_reason":if finished {"stop"} else {"tool_calls"}}]});
+    (
+        [("content-type", "text/event-stream")],
+        format!("data: {event}\n\ndata: {end}\n\ndata: [DONE]\n\n"),
+    )
+}
+
+async fn run_flow(db: &Pool<Postgres>, modules: Value) -> anyhow::Result<Value> {
+    let server = ApiServer::start(db.clone()).await?;
+    let flow: FlowValue = serde_json::from_value(json!({ "modules": modules }))?;
+    let id = RunJob::from(JobPayload::RawFlow {
+        value: flow,
+        path: Some("u/test/nul".into()),
+        restarted_from: None,
+    })
+    .push(db)
+    .await;
+    let notifications = listen_for_completed_jobs(db).await;
+    tokio::time::timeout(
+        Duration::from_secs(45),
+        in_test_worker(db, notifications.find(&id), server.addr.port()),
+    )
+    .await?;
+    let job = completed_job(id, db).await;
+    server.close().await?;
+    assert!(job.success, "{:?}", job.result);
+    Ok(job.json_result().expect("flow result"))
+}
+
+#[sqlx::test(fixtures("base"))]
+#[serial_test::serial]
+async fn agent_step_survives_nul_in_model_output(db: Pool<Postgres>) -> anyhow::Result<()> {
+    std::env::set_var("ALLOW_PRIVATE_AI_BASE_URLS", "true");
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let base = format!("http://{}", listener.local_addr()?);
+    let stub = tokio::spawn(async move {
+        let router = Router::new().route("/v1/chat/completions", post(model));
+        axum::serve(listener, router).await.unwrap();
+    });
+    let result = run_flow(
+        &db,
+        json!([{"id":"agent","value":{
+            "type":"aiagent",
+            "tools":[{"id":"t0","summary":"echo","value":{
+                "type":"rawscript", "language":"bun", "input_transforms":{},
+                "content":"export async function main(x: string) { return x }"
+            }}],
+            "input_transforms":{
+                "provider":{"type":"static","value":{"kind":"customai","model":"nul-test","resource":{"base_url":format!("{base}/v1")}}},
+                "user_message":{"type":"static","value":"echo"},
+                "memory":{"type":"static","value":{"kind":"auto","context_length":10,"memory_id":MEMORY_ID}}
+            }
+        }}]),
+    )
+    .await?;
+    stub.abort();
+    assert_eq!(result["output"], "aretired");
+    let tool_message = result["messages"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .find(|m| m["role"] == "tool")
+        .expect("tool message");
+    assert_eq!(tool_message["content"], "\"ab\"");
+    // The model's own words are also written to a jsonb column the job result never passes through.
+    let memory: String = sqlx::query_scalar(
+        "SELECT messages::text FROM ai_agent_memory WHERE conversation_id = $1::uuid",
+    )
+    .bind(MEMORY_ID)
+    .fetch_one(&db)
+    .await?;
+    assert!(memory.contains("aretired"), "{memory}");
+    Ok(())
+}
+
+/// A step's args are evaluated from the previous step's result as the worker holds it, not as
+/// the completed job stored it.
+#[sqlx::test(fixtures("base"))]
+#[serial_test::serial]
+async fn flow_step_args_survive_nul_in_previous_result(db: Pool<Postgres>) -> anyhow::Result<()> {
+    let result = run_flow(
+        &db,
+        json!([
+            {"id":"a","value":{
+                "type":"rawscript", "language":"bun", "input_transforms":{},
+                "content":"export async function main() { return 'a\\u0000b' }"
+            }},
+            {"id":"b","value":{
+                "type":"rawscript", "language":"bun",
+                "input_transforms":{"x":{"type":"javascript","expr":"results.a"}},
+                "content":"export async function main(x: string) { return x }"
+            }}
+        ]),
+    )
+    .await?;
+    assert_eq!(result, "ab");
+    Ok(())
+}

backend/windmill-ai/src/query_builder.rs

@@ -1,5 +1,7 @@
+use std::borrow::Cow;
+
 use async_trait::async_trait;
-use windmill_common::{client::AuthedClient, error::Error};
+use windmill_common::{client::AuthedClient, error::Error, utils::strip_json_nul};
 use windmill_types::s3::S3Object;
 
 use crate::ai_types::OpenAIToolCall;
@@ -92,6 +94,39 @@ pub enum ParsedResponse {
     },
 }
 
+impl ParsedResponse {
+    /// Drops every U+0000 the model emitted. Postgres stores neither a NUL in `text` nor its
+    /// `\u0000` escape in `jsonb`, and this output is written to both: tool job args, the flow
+    /// status, agent memory, conversation rows.
+    pub fn without_nul(mut self) -> Self {
+        fn strip(s: &mut String) {
+            if s.contains('\0') {
+                s.retain(|c| c != '\0');
+            }
+        }
+        if let ParsedResponse::Text { content, reasoning, tool_calls, annotations, .. } = &mut self
+        {
+            for text in [content, reasoning].into_iter().flatten() {
+                strip(text);
+            }
+            for annotation in annotations {
+                strip(&mut annotation.url);
+                annotation.title.as_mut().into_iter().for_each(strip);
+            }
+            for tool_call in tool_calls {
+                strip(&mut tool_call.id);
+                strip(&mut tool_call.function.name);
+                // Arguments are serialized JSON, where a NUL is spelled as its escape.
+                strip(&mut tool_call.function.arguments);
+                if let Cow::Owned(stripped) = strip_json_nul(&tool_call.function.arguments) {
+                    tool_call.function.arguments = stripped;
+                }
+            }
+        }
+        self
+    }
+}
+
 /// Trait for streaming AI events to a sink (e.g., database persistence).
 /// Implemented by the worker's StreamEventProcessor.
 #[async_trait]

backend/windmill-queue/src/jobs.rs

@@ -7523,6 +7523,18 @@ async fn push_inner<'c, 'd>(
                 concurrency_settings.concurrency_time_window_s,
             )
         };
+    // jsonb rejects the `\u0000` escape (22P05), which would fail the whole push, and args are
+    // not always what a caller typed: a flow step's are evaluated from the previous step's
+    // in-memory result, an agent tool's are written by the model.
+    let serialized_args = serde_json::value::to_raw_value(&args).map_err(|e| {
+        Error::internal_err(format!("Could not serialize args of job {job_id}: {e:#}"))
+    })?;
+    let sanitized_args = match strip_json_nul(serialized_args.get()) {
+        Cow::Owned(stripped) => RawValue::from_string(stripped).map_err(|e| {
+            Error::internal_err(format!("Could not sanitize args of job {job_id}: {e:#}"))
+        })?,
+        Cow::Borrowed(_) => serialized_args,
+    };
     sqlx::query!(
         "WITH inserted_job AS (
             INSERT INTO v2_job (
@@ -7589,7 +7601,7 @@ async fn push_inner<'c, 'd>(
         permissioned_as,
         runnable_id,
         runnable_path.clone(),
-        Json(args) as Json<PushArgs>,
+        Json(sanitized_args) as Json<Box<RawValue>>,
         job_kind.clone() as JobKind,
         trigger_path.flatten(),
         language as Option<ScriptLang>,

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

@@ -651,6 +651,7 @@ async fn summarize_prefix(
             .query_builder
             .parse_streaming_response(resp, StreamEventProcessor::new_silent().boxed_sink())
             .await?
+            .without_nul()
     };
 
     match parsed {

backend/windmill-worker/src/ai_executor.rs

@@ -1915,7 +1915,7 @@ pub async fn run_agent(
                         query_builder.parse_image_response(resp).await?
                     }
                 };
-                Ok(parsed)
+                Ok(parsed.without_nul())
             }.await;
 
             match attempt {