From cf4bec611f6df8a45b238afe5b8da3e3a5748bf9 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 6 Oct 2026 17:51:07 +0200 Subject: [PATCH] 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) * fix: also strip NUL from tool call ids and citations Co-Authored-By: Claude Opus 5.5 (1M context) --------- Co-authored-by: Claude Opus 5.5 (1M context) --- backend/tests/ai_agent_nul.rs | 125 +++++++++++++++++++ backend/windmill-ai/src/query_builder.rs | 37 +++++- backend/windmill-queue/src/jobs.rs | 14 ++- backend/windmill-worker/src/ai/compaction.rs | 1 + backend/windmill-worker/src/ai_executor.rs | 2 +- 5 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 backend/tests/ai_agent_nul.rs diff --git a/backend/tests/ai_agent_nul.rs b/backend/tests/ai_agent_nul.rs new file mode 100644 index 0000000000..1af46ecc82 --- /dev/null +++ b/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) -> ([(&'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, modules: Value) -> anyhow::Result { + 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) -> 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) -> 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(()) +} diff --git a/backend/windmill-ai/src/query_builder.rs b/backend/windmill-ai/src/query_builder.rs index 986c433f42..a08d01dadd 100644 --- a/backend/windmill-ai/src/query_builder.rs +++ b/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] diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 49aca2cba7..207c8ee176 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/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, + Json(sanitized_args) as Json>, job_kind.clone() as JobKind, trigger_path.flatten(), language as Option, diff --git a/backend/windmill-worker/src/ai/compaction.rs b/backend/windmill-worker/src/ai/compaction.rs index c88ecb8178..b933348a2e 100644 --- a/backend/windmill-worker/src/ai/compaction.rs +++ b/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 { diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index d6089746ec..133a0a0b0d 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/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 {