refactor: move ai sse plumbing to windmill-ai (#9059)

* docs: refine windmill ai refactor plan

* refactor: move ai sse plumbing to windmill-ai

* refactor: remove ai re-export shims

* fix: update ee ai memory ref

* chore: update ee-repo-ref to d3bc7fa85195b46b7a38d43c2f806520bf8b5454

This commit updates the EE repository reference after PR #560 was merged in windmill-ee-private.

Previous ee-repo-ref: ff35bf7cc198e13884b33654e1d6dbd8a8b314d3

New ee-repo-ref: d3bc7fa85195b46b7a38d43c2f806520bf8b5454

Automated by sync-ee-ref workflow.

---------

Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>

backend/Cargo.lock

@@ -16114,12 +16114,14 @@ dependencies = [
  "aws-sdk-bedrockruntime",
  "aws-smithy-types",
  "base64 0.22.1",
+ "eventsource-stream",
  "lazy_static",
  "reqwest 0.13.1",
  "serde",
  "serde_json",
  "sqlx",
  "tokio",
+ "tokio-stream",
  "tracing",
  "uuid",
  "windmill-common",

backend/ee-repo-ref.txt

@@ -1 +1 @@
-c8d100d74b8de6bd26fc973d5edbd8853d54dd8b
+d3bc7fa85195b46b7a38d43c2f806520bf8b5454

backend/windmill-ai/Cargo.toml

@@ -21,6 +21,7 @@ windmill-mcp = { workspace = true, optional = true }
 
 async-trait.workspace = true
 base64.workspace = true
+eventsource-stream.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
@@ -29,6 +30,7 @@ uuid.workspace = true
 lazy_static.workspace = true
 tracing.workspace = true
 tokio.workspace = true
+tokio-stream.workspace = true
 
 # Bedrock (optional)
 aws-config = { workspace = true, optional = true }

backend/windmill-ai/src/lib.rs

@@ -5,4 +5,6 @@ pub mod ai_google;
 pub mod ai_providers;
 pub mod ai_types;
 pub mod query_builder;
+pub mod sse;
 pub mod types;
+pub mod utils;

backend/windmill-ai/src/sse.rs

@@ -3,17 +3,15 @@ use std::collections::HashMap;
 use eventsource_stream::Eventsource;
 use reqwest::Response;
 use serde::Deserialize;
-use serde_json;
 use tokio_stream::StreamExt;
-use windmill_ai::{
-    ai_google::{parse_gemini_sse_event, GeminiUsageMetadata},
-    ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall},
-};
 use windmill_common::{error::Error, utils::rd_string};
 
-use crate::ai::{
+use crate::{
+    ai_google::{parse_gemini_sse_event, GeminiUsageMetadata},
+    ai_types::UrlCitation,
+    ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall},
     query_builder::StreamEventSink,
-    types::{StreamingEvent, UrlCitation},
+    types::StreamingEvent,
 };
 
 #[derive(Deserialize)]
@@ -64,6 +62,7 @@ lazy_static::lazy_static! {
         .parse::<bool>()
         .unwrap_or(false);
 }
+#[allow(async_fn_in_trait)]
 pub trait SSEParser {
     async fn parse_event_data(&mut self, data: &str) -> Result<(), Error>;
 
@@ -459,11 +458,11 @@ impl SSEParser for AnthropicSSEParser {
 // Gemini SSE Parser
 // ============================================================================
 
-/// Accumulates Gemini streaming events and converts them into the worker's
-/// internal [`OpenAIToolCall`] / [`StreamingEvent`] representation.
+/// Accumulates Gemini streaming events and converts them into the shared
+/// [`OpenAIToolCall`] / [`StreamingEvent`] representation.
 ///
 /// The actual SSE parsing is delegated to [`parse_gemini_sse_event`] from
-/// `windmill_common::ai_google` so the logic can be shared with the API proxy.
+/// `windmill_ai::ai_google` so the logic can be shared with the API proxy.
 pub struct GeminiSSEParser {
     pub accumulated_content: String,
     pub accumulated_tool_calls: HashMap<i64, OpenAIToolCall>,

backend/windmill-ai/src/utils.rs

@@ -0,0 +1,56 @@
+use crate::{
+    ai_providers::AIProvider,
+    ai_types::{ContentPart, OpenAIContent},
+};
+
+lazy_static::lazy_static! {
+    /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples
+    /// Format: "header1: value1, header2: value2"
+    pub static ref AI_HTTP_HEADERS: Vec<(String, String)> = {
+        std::env::var("AI_HTTP_HEADERS")
+            .ok()
+            .map(|headers_str| {
+                headers_str
+                    .split(',')
+                    .filter_map(|header| {
+                        let parts: Vec<&str> = header.splitn(2, ':').collect();
+                        if parts.len() == 2 {
+                            let name = parts[0].trim().to_string();
+                            let value = parts[1].trim().to_string();
+                            if !name.is_empty() && !value.is_empty() {
+                                Some((name, value))
+                            } else {
+                                None
+                            }
+                        } else {
+                            None
+                        }
+                    })
+                    .collect()
+            })
+            .unwrap_or_default()
+    };
+}
+
+/// AWS Bedrock do not handle structured output query param, so we use a tool for structured output. Same for every Claude models.
+pub fn should_use_structured_output_tool(provider: &AIProvider, model: &str) -> bool {
+    model.contains("claude") || provider == &AIProvider::AWSBedrock
+}
+
+/// Extract text content from OpenAIContent, joining parts with space if multiple
+pub fn extract_text_content(content: &OpenAIContent) -> String {
+    match content {
+        OpenAIContent::Text(text) => text.clone(),
+        OpenAIContent::Parts(parts) => parts
+            .iter()
+            .filter_map(|p| {
+                if let ContentPart::Text { text } = p {
+                    Some(text.as_str())
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>()
+            .join(""),
+    }
+}

backend/windmill-api/src/ai.rs

@@ -15,11 +15,12 @@ use serde::{Deserialize, Serialize};
 use serde_json::{json, value::RawValue};
 use std::collections::HashMap;
 use std::time::Duration;
-use windmill_audit::{audit_oss::audit_log, ActionKind};
 use windmill_ai::ai_cache::current_instance_ai_config_revision;
 use windmill_ai::ai_providers::{
     empty_string_as_none, AIPlatform, AIProvider, ProviderConfig, ProviderModel,
 };
+use windmill_ai::utils::AI_HTTP_HEADERS;
+use windmill_audit::{audit_oss::audit_log, ActionKind};
 use windmill_common::db::UserDB;
 use windmill_common::error::{to_anyhow, Error, Result};
 use windmill_common::utils::configure_client;
@@ -101,32 +102,6 @@ lazy_static::lazy_static! {
 
     pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringAIRequestConfig> = Cache::new(500);
 
-    /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples
-    /// Format: "header1: value1, header2: value2"
-    static ref AI_HTTP_HEADERS: Vec<(String, String)> = {
-        std::env::var("AI_HTTP_HEADERS")
-            .ok()
-            .map(|headers_str| {
-                headers_str
-                    .split(',')
-                    .filter_map(|header| {
-                        let parts: Vec<&str> = header.splitn(2, ':').collect();
-                        if parts.len() == 2 {
-                            let name = parts[0].trim().to_string();
-                            let value = parts[1].trim().to_string();
-                            if !name.is_empty() && !value.is_empty() {
-                                Some((name, value))
-                            } else {
-                                None
-                            }
-                        } else {
-                            None
-                        }
-                    })
-                    .collect()
-            })
-            .unwrap_or_default()
-    };
 }
 
 pub(crate) fn invalidate_ai_request_cache_for_workspace(workspace_id: &str) {

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

@@ -1,12 +1,11 @@
 use base64::Engine;
 use futures;
 use ulid;
+use windmill_ai::types::*;
 use windmill_common::{client::AuthedClient, error::Error};
 use windmill_queue::MiniPulledJob;
 use windmill_types::s3::S3Object;
 
-use crate::ai::types::*;
-
 /// Upload image to S3 and return S3Object
 pub async fn upload_image_to_s3(
     base64_image: &str,

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

@@ -4,7 +4,5 @@
 pub mod image_handler;
 pub mod providers;
 pub mod query_builder;
-pub mod sse;
 pub mod tools;
-pub mod types;
 pub mod utils;

backend/windmill-worker/src/ai/providers/anthropic.rs

@@ -1,16 +1,17 @@
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use serde_json::value::RawValue;
-use windmill_ai::{ai_google::parse_data_url, ai_providers::AIProvider};
-use windmill_common::{client::AuthedClient, error::Error};
-
-use crate::ai::{
-    image_handler::prepare_messages_for_api,
+use windmill_ai::{
+    ai_google::parse_data_url,
+    ai_providers::AIProvider,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{AnthropicSSEParser, SSEParser},
     types::*,
     utils::{extract_text_content, should_use_structured_output_tool},
 };
+use windmill_common::{client::AuthedClient, error::Error};
+
+use crate::ai::image_handler::prepare_messages_for_api;
 
 /// Anthropic API version for standard API
 const ANTHROPIC_VERSION_STANDARD: &str = "2023-06-01";

backend/windmill-worker/src/ai/providers/bedrock.rs

@@ -6,25 +6,22 @@
 //! - Stream event parsing
 //! - Helper utilities
 
-use crate::ai::{
-    image_handler::prepare_messages_for_api,
-    query_builder::{ParsedResponse, StreamEventSink},
-    types::StreamingEvent,
-    types::TokenUsage,
-    types::{OpenAIMessage, ToolDef},
-};
+use crate::ai::image_handler::prepare_messages_for_api;
 use std::collections::HashMap;
+use windmill_ai::{
+    query_builder::{ParsedResponse, StreamEventSink},
+    types::{OpenAIMessage, StreamingEvent, TokenUsage, ToolDef},
+};
 use windmill_common::{client::AuthedClient, error::Error};
 
-// Re-export from shared module for use by other parts of the worker
+// Import shared Bedrock helpers for worker-specific orchestration.
 use windmill_ai::ai_bedrock::{
     bedrock_model_supports_prompt_caching, bedrock_stream_event_is_block_stop,
     bedrock_stream_event_to_text, bedrock_stream_event_to_tool_delta,
     bedrock_stream_event_to_tool_start, build_tool_config, create_inference_config,
     format_bedrock_error, openai_messages_to_bedrock, streaming_tool_calls_to_openai,
-    StreamingToolCall,
+    BedrockClient, StreamingToolCall,
 };
-pub use windmill_ai::ai_bedrock::{check_env_credentials, BedrockClient};
 
 // ============================================================================
 // Query Builder (Worker-specific orchestration)

backend/windmill-worker/src/ai/providers/google_ai.rs

@@ -1,17 +1,17 @@
 use async_trait::async_trait;
-use windmill_ai::ai_google::{
-    openai_messages_to_gemini, openai_tools_to_gemini, GeminiGenerationConfig, GeminiImageContent,
-    GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart, GeminiPredictContent,
-    GeminiTextRequest, GeminiTool,
-};
-use windmill_common::{client::AuthedClient, error::Error};
-
-use crate::ai::{
-    image_handler::{download_and_encode_s3_image, prepare_messages_for_api},
+use windmill_ai::{
+    ai_google::{
+        openai_messages_to_gemini, openai_tools_to_gemini, GeminiGenerationConfig,
+        GeminiImageContent, GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart,
+        GeminiPredictContent, GeminiTextRequest, GeminiTool,
+    },
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{GeminiSSEParser, SSEParser},
     types::*,
 };
+use windmill_common::{client::AuthedClient, error::Error};
+
+use crate::ai::image_handler::{download_and_encode_s3_image, prepare_messages_for_api};
 
 // ============================================================================
 // Query Builder Implementation

backend/windmill-worker/src/ai/providers/openai.rs

@@ -1,17 +1,17 @@
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use serde_json::value::RawValue;
-use windmill_ai::ai_providers::AIProvider;
-use windmill_ai::ai_types::OpenAIToolCall;
-use windmill_common::{client::AuthedClient, error::Error};
-
-use crate::ai::{
-    image_handler::{prepare_messages_for_api, s3_object_to_content_part},
+use windmill_ai::{
+    ai_providers::AIProvider,
+    ai_types::OpenAIToolCall,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{OpenAIResponsesSSEParser, SSEParser},
     types::*,
     utils::extract_text_content,
 };
+use windmill_common::{client::AuthedClient, error::Error};
+
+use crate::ai::image_handler::{prepare_messages_for_api, s3_object_to_content_part};
 
 // Responses API structures
 #[derive(Deserialize)]

backend/windmill-worker/src/ai/providers/openrouter.rs

@@ -1,15 +1,14 @@
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use serde_json;
-use windmill_ai::ai_providers::AIProvider;
-use windmill_common::{client::AuthedClient, error::Error};
-
-use crate::ai::{
-    image_handler::prepare_messages_for_api,
-    providers::other::OtherQueryBuilder,
+use windmill_ai::{
+    ai_providers::AIProvider,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     types::*,
 };
+use windmill_common::{client::AuthedClient, error::Error};
+
+use crate::ai::{image_handler::prepare_messages_for_api, providers::other::OtherQueryBuilder};
 
 // OpenRouter-specific types
 #[derive(Serialize)]

backend/windmill-worker/src/ai/providers/other.rs

@@ -1,16 +1,16 @@
 use async_trait::async_trait;
 use serde::Serialize;
 use serde_json;
-use windmill_ai::ai_providers::AIProvider;
-use windmill_common::{client::AuthedClient, error::Error};
-
-use crate::ai::{
-    image_handler::prepare_messages_for_api,
+use windmill_ai::{
+    ai_providers::AIProvider,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{OpenAISSEParser, SSEParser},
     types::*,
     utils::should_use_structured_output_tool,
 };
+use windmill_common::{client::AuthedClient, error::Error};
+
+use crate::ai::image_handler::prepare_messages_for_api;
 
 #[derive(Serialize, Debug, Clone)]
 #[serde(rename_all = "lowercase")]

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

@@ -1,24 +1,19 @@
 use async_trait::async_trait;
+use windmill_ai::{
+    query_builder::{QueryBuilder, StreamEventSink},
+    types::*,
+};
 use windmill_common::{error::Error, worker::Connection};
 use windmill_queue::MiniPulledJob;
 
 use crate::{
-    ai::{
-        providers::{
-            anthropic::AnthropicQueryBuilder, google_ai::GoogleAIQueryBuilder,
-            openai::OpenAIQueryBuilder, openrouter::OpenRouterQueryBuilder,
-            other::OtherQueryBuilder,
-        },
-        types::*,
+    ai::providers::{
+        anthropic::AnthropicQueryBuilder, google_ai::GoogleAIQueryBuilder,
+        openai::OpenAIQueryBuilder, openrouter::OpenRouterQueryBuilder, other::OtherQueryBuilder,
     },
     job_logger::append_result_stream,
 };
 
-// Re-export from windmill_ai
-pub use windmill_ai::query_builder::{
-    BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink,
-};
-
 /// Factory function to create the appropriate query builder for a provider
 pub fn create_query_builder(provider: &ProviderWithResource) -> Box<dyn QueryBuilder> {
     use windmill_ai::ai_providers::AIProvider;

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

@@ -1,6 +1,4 @@
-use crate::ai::query_builder::{StreamEventProcessor, StreamEventSink};
-use crate::ai::types::McpToolSource;
-use crate::ai::types::*;
+use crate::ai::query_builder::StreamEventProcessor;
 use crate::ai::utils::{
     add_message_to_conversation, execute_mcp_tool, get_step_name_from_flow,
     is_completed_input_transform, update_flow_status_module_with_actions,
@@ -20,7 +18,7 @@ use mappable_rc::Marc;
 use serde_json::value::RawValue;
 use std::{collections::HashMap, sync::Arc};
 use uuid::Uuid;
-use windmill_ai::ai_types::OpenAIToolCall;
+use windmill_ai::{ai_types::OpenAIToolCall, query_builder::StreamEventSink, types::*};
 use windmill_common::jobs::JobPayload;
 
 #[cfg(feature = "mcp")]

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

@@ -1,2 +0,0 @@
-// Re-export all types from windmill_ai::types
-pub use windmill_ai::types::*;

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

@@ -1,5 +1,3 @@
-pub use crate::ai::types::McpToolSource;
-use crate::ai::types::ToolDef;
 use anyhow::Context;
 use serde_json::value::RawValue;
 use sqlx::types::Json;
@@ -8,7 +6,7 @@ use std::{
     sync::Arc,
 };
 use uuid::Uuid;
-use windmill_ai::ai_providers::AIProvider;
+use windmill_ai::types::*;
 use windmill_common::flows::FlowModuleValue;
 use windmill_common::{
     db::DB,
@@ -24,7 +22,7 @@ use windmill_common::{
 use windmill_mcp::{McpClient, McpResource, McpTool};
 use windmill_queue::{flow_status::get_step_of_flow_status, MiniPulledJob};
 
-use crate::{ai::types::*, parse_sig_of_lang};
+use crate::parse_sig_of_lang;
 
 pub fn parse_raw_script_schema(
     content: &str,
@@ -323,11 +321,6 @@ pub fn get_step_name_from_flow(
     )
 }
 
-/// AWS Bedrock do not handle structured output query param, so we use a tool for structured output. Same for every Claude models.
-pub fn should_use_structured_output_tool(provider: &AIProvider, model: &str) -> bool {
-    model.contains("claude") || provider == &AIProvider::AWSBedrock
-}
-
 /// Cleanup MCP clients by gracefully shutting down connections
 #[cfg(feature = "mcp")]
 pub async fn cleanup_mcp_clients(mcp_clients: HashMap<String, Arc<McpClient>>) {
@@ -713,21 +706,3 @@ pub fn any_tool_needs_previous_result(tools: &[Tool]) -> bool {
         false
     })
 }
-
-/// Extract text content from OpenAIContent, joining parts with space if multiple
-pub fn extract_text_content(content: &OpenAIContent) -> String {
-    match content {
-        OpenAIContent::Text(text) => text.clone(),
-        OpenAIContent::Parts(parts) => parts
-            .iter()
-            .filter_map(|p| {
-                if let ContentPart::Text { text } = p {
-                    Some(text.as_str())
-                } else {
-                    None
-                }
-            })
-            .collect::<Vec<_>>()
-            .join(""),
-    }
-}

backend/windmill-worker/src/ai_executor.rs

@@ -1,12 +1,10 @@
-#[cfg(feature = "bedrock")]
-use crate::ai::providers::bedrock::check_env_credentials;
 use crate::ai::tools::{execute_tool_calls, ToolAbortHandles, ToolExecutionContext};
 use crate::ai::utils::{
     add_message_to_conversation, any_tool_needs_previous_result, cleanup_mcp_clients,
     filter_schema_by_input_transforms, find_unique_tool_name, get_flow_context,
     get_flow_job_runnable_and_raw_flow, get_step_name_from_flow, load_mcp_tools,
-    parse_raw_script_schema, should_use_structured_output_tool,
-    update_flow_status_module_with_actions, update_flow_status_module_with_actions_success,
+    parse_raw_script_schema, update_flow_status_module_with_actions,
+    update_flow_status_module_with_actions_success,
 };
 use crate::memory_oss::{read_from_memory, write_to_memory};
 use crate::worker_flow::{get_previous_job_result, get_transform_context};
@@ -15,12 +13,19 @@ use regex::Regex;
 use serde_json::value::RawValue;
 use std::{collections::HashMap, sync::Arc};
 use uuid::Uuid;
+#[cfg(feature = "bedrock")]
+use windmill_ai::ai_bedrock::check_env_credentials;
 #[cfg(feature = "mcp")]
 use windmill_mcp::McpClient;
 
 #[cfg(not(feature = "mcp"))]
 use crate::ai::tools::McpClientStub as McpClient;
-use windmill_ai::ai_providers::AIProvider;
+use windmill_ai::{
+    ai_providers::AIProvider,
+    query_builder::{BuildRequestArgs, ParsedResponse},
+    types::*,
+    utils::{should_use_structured_output_tool, AI_HTTP_HEADERS},
+};
 use windmill_common::{
     cache,
     client::AuthedClient,
@@ -40,10 +45,7 @@ use windmill_queue::{cancel_single_job, CanceledBy, MiniPulledJob};
 use crate::{
     ai::{
         image_handler::upload_image_to_s3,
-        query_builder::{
-            create_query_builder, BuildRequestArgs, ParsedResponse, StreamEventProcessor,
-        },
-        types::*,
+        query_builder::{create_query_builder, StreamEventProcessor},
     },
     common::{build_args_map, resolve_job_timeout, OccupancyMetrics, StreamNotifier},
     handle_child::{run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome},
@@ -52,33 +54,6 @@ use crate::{
 lazy_static::lazy_static! {
     static ref TOOL_NAME_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap();
 
-    /// Parse AI_HTTP_HEADERS environment variable into a vector of (header_name, header_value) tuples
-    /// Format: "header1: value1, header2: value2"
-    static ref AI_HTTP_HEADERS: Vec<(String, String)> = {
-        std::env::var("AI_HTTP_HEADERS")
-            .ok()
-            .map(|headers_str| {
-                headers_str
-                    .split(',')
-                    .filter_map(|header| {
-                        let parts: Vec<&str> = header.splitn(2, ':').collect();
-                        if parts.len() == 2 {
-                            let name = parts[0].trim().to_string();
-                            let value = parts[1].trim().to_string();
-                            if !name.is_empty() && !value.is_empty() {
-                                Some((name, value))
-                            } else {
-                                None
-                            }
-                        } else {
-                            None
-                        }
-                    })
-                    .collect()
-            })
-            .unwrap_or_default()
-    };
-
     static ref AI_AGENT_TOOL_SCHEMA: Box<RawValue> = to_raw_value(&serde_json::json!({
         "type": "object",
         "properties": {
@@ -791,7 +766,7 @@ pub async fn run_agent(
 
     let mut actions = vec![];
     let mut content = None;
-    let mut final_usage: Option<crate::ai::types::TokenUsage> = None;
+    let mut final_usage: Option<TokenUsage> = None;
 
     // Check if this provider supports tools with the current output type
     let supports_tools = query_builder.supports_tools_with_output_type(output_type);

backend/windmill-worker/src/memory_common.rs

@@ -1,5 +1,5 @@
-use crate::ai::types::OpenAIMessage;
 use uuid::Uuid;
+use windmill_ai::types::OpenAIMessage;
 use windmill_common::{db::DB, error::Error};
 
 pub const MAX_MEMORY_SIZE_BYTES: usize = 100_000; // 100KB per memory entry in database

backend/windmill-worker/src/memory_oss.rs

@@ -3,7 +3,9 @@
 pub use crate::memory_ee::*;
 
 #[cfg(not(all(feature = "private", feature = "enterprise")))]
-use {crate::ai::types::OpenAIMessage, crate::memory_common, uuid::Uuid, windmill_common::db::DB};
+use {
+    crate::memory_common, uuid::Uuid, windmill_ai::types::OpenAIMessage, windmill_common::db::DB,
+};
 
 /// Read AI agent memory from storage
 /// In OSS: always reads from database

docs/windmill-ai-refactor-plan.md

@@ -23,6 +23,58 @@ windmill-worker →  windmill-ai
 
 windmill-common does **NOT** re-export from windmill-ai (would be circular). All consumers update imports.
 
+## Reviewer Note: Keep the Next PR Small
+
+The first merged PR established the crate boundary; it did not yet remove the duplicated API-vs-worker provider paths. The remaining work should stay split by dependency risk, not by the final desired module layout.
+
+Do not jump directly from the current state to provider moves, proxy unification, and credential unification in one PR. The riskiest part is the API proxy because it combines request transformation, endpoint selection, auth headers, custom headers, OAuth user injection, Azure URL handling, Anthropic Vertex handling, Bedrock SDK calls, and SSE keepalive behavior.
+
+Pull the shared plumbing forward before moving provider implementations:
+- Move tiny shared utilities first, including `AI_HTTP_HEADERS`, `extract_text_content`, and `should_use_structured_output_tool`.
+- Move SSE parsers next, using the existing `StreamEventSink` abstraction, and update callers to import from `windmill_ai` directly.
+- Leave provider implementations, image upload/download handling, API proxy changes, and credential unification out of that PR.
+
+Avoid adding modules whose only purpose is to re-export moved code. Direct imports from `windmill_ai` make ownership and dependency direction clearer at each call site.
+
+Also do not make `build_proxy_request(raw_body, path)` too narrow. The proxy path needs method, incoming headers, resolved credentials, base URL/platform, organization/user fields, custom headers, and Bedrock/Azure/Vertex-specific context. Introduce a structured `ProxyBuildArgs`/`ProviderCredentials` shape before deleting `AIRequestConfig::prepare_request`, `google.rs`, or `bedrock.rs`.
+
+## Next Phase PR: Shared Plumbing Only
+
+Goal: make `windmill-ai` own the provider-independent helper code that later provider moves will need, without changing API proxy behavior or agent request behavior.
+
+Suggested PR title: `refactor(ai): move shared SSE plumbing into windmill-ai`.
+
+Scope:
+- Add `windmill-ai/src/utils.rs`.
+- Move the duplicated `AI_HTTP_HEADERS` parsing into `windmill_ai::utils` with identical parsing behavior.
+- Move `extract_text_content` and `should_use_structured_output_tool` from `windmill-worker/src/ai/utils.rs` to `windmill_ai::utils`.
+- Move `windmill-worker/src/ai/sse.rs` to `windmill-ai/src/sse.rs`.
+- Delete `windmill-worker/src/ai/sse.rs` and update callers to import parser types from `windmill_ai::sse`.
+- Update callers of moved utility functions to import from `windmill_ai::utils` directly.
+- Add the minimal new `windmill-ai` dependencies required by `sse.rs` (`eventsource-stream`, `tokio-stream`) and avoid adding worker/queue dependencies.
+
+Out of scope:
+- Do not move provider implementations.
+- Do not move `image_handler`.
+- Do not change `QueryBuilder` method signatures.
+- Do not add `build_proxy_request`.
+- Do not change API proxy routing, request preparation, credential resolution, audit logging, cache behavior, or Bedrock/Google special cases.
+- Do not remove `windmill-api/src/google.rs`, `windmill-api/src/bedrock.rs`, or `AIRequestConfig::prepare_request`.
+
+Implementation checklist:
+1. Add `utils.rs` to `windmill-ai` and export it from `lib.rs`.
+2. Move `AI_HTTP_HEADERS` exactly once, then update `windmill-api/src/ai.rs` and `windmill-worker/src/ai_executor.rs` to import it.
+3. Move the two provider-independent helper functions into `windmill_ai::utils`; leave worker-specific flow/MCP/conversation utilities in `windmill-worker/src/ai/utils.rs`.
+4. Move `sse.rs` into `windmill-ai`, change imports from `crate::ai::{query_builder, types}` to `crate::{query_builder, types}`, and keep behavior unchanged.
+5. Remove worker `ai/sse.rs` and update provider imports to use `windmill_ai::sse` directly.
+6. Run focused grep checks for duplicate `AI_HTTP_HEADERS`, old local helper definitions, and accidental `windmill_queue`/worker dependencies from `windmill-ai`.
+7. Validate with `cargo check -p windmill-ai`, `cargo check -p windmill-worker`, and `cargo check -p windmill-api`. For `bedrock` builds, also check the existing bedrock feature path.
+
+Review expectations:
+- The diff should be mostly moved code and import updates.
+- The behavior should be byte-for-byte equivalent where practical.
+- Tests are only needed if helper behavior changes. For a pure move, existing backend checks plus manual AI streaming verification are enough.
+
 ## Step-by-Step Plan
 
 Each step produces a compiling, working backend.
@@ -131,18 +183,27 @@ This is the key unification step. Add a new method to the `QueryBuilder` trait:
 /// Used by the API chat proxy. Handles format conversion for non-OpenAI providers.
 fn build_proxy_request(
     &self,
-    raw_body: &[u8],
-    path: &str,
+    args: &ProxyBuildArgs<'_>,
 ) -> Result<ProxyRequest, Error>;
 ```
 
-Where `ProxyRequest` contains the transformed body, endpoint URL, and auth headers:
+Where `ProxyBuildArgs` carries the API proxy context that provider implementations need:
+```rust
+pub struct ProxyBuildArgs<'a> {
+    pub method: http::Method,
+    pub path: &'a str,
+    pub headers: &'a http::HeaderMap,
+    pub body: &'a [u8],
+    pub credentials: &'a ProviderCredentials,
+}
+```
+
+And `ProxyRequest` contains the transformed request:
 ```rust
 pub struct ProxyRequest {
     pub url: String,
     pub body: Vec<u8>,
-    pub auth_headers: Vec<(String, String)>,
-    pub is_sse: bool,
+    pub headers: Vec<(String, String)>,
 }
 ```
 
@@ -153,9 +214,9 @@ pub struct ProxyRequest {
 - **Bedrock**: Convert OpenAI format → Bedrock SDK calls. Replaces `windmill-api/src/bedrock.rs`.
 
 **Refactor API proxy** (`windmill-api/src/ai.rs`):
-1. Parse provider from headers, resolve credentials → `ProviderWithResource`
+1. Parse provider from headers, resolve credentials → `ProviderCredentials`
 2. Create `QueryBuilder` via `create_query_builder`
-3. Call `query_builder.build_proxy_request(body, path)` → `ProxyRequest`
+3. Call `query_builder.build_proxy_request(&proxy_args)` → `ProxyRequest`
 4. Send the request, return response with SSE keepalive injection
 
 **Remove** from windmill-api:
@@ -166,7 +227,7 @@ pub struct ProxyRequest {
 - `supports_native_fim`, `transform_fim_to_chat_completions` — moved to windmill-ai
 
 **Keep** in API:
-- `AIRequestConfig::new` credential resolution (or refactor to produce `ProviderWithResource`)
+- `AIRequestConfig::new` credential resolution until it is refactored to produce `ProviderCredentials`
 - HTTP routes, audit logging, request caching
 - `inject_keepalives`, `is_sse_response` helpers
 - `AIConfig`, `ExpiringAIRequestConfig` caching types