windmill/docs/windmill-ai-refactor-plan.md

Refactor Plan: windmill-ai Crate

Context

AI provider logic is currently split across three crates with duplicate code:

• windmill-common — base types (ai_types, ai_providers, ai_google, ai_bedrock, ai_cache)
• windmill-api — chat proxy routes (ai.rs), audit logging, caching, and DB-backed credential resolution into ProviderCredentials
• windmill-worker — agent execution (ai/ module) with QueryBuilder trait, SSE parsers, provider implementations

The goal: a single windmill-ai crate with all AI provider logic. Worker agent execution uses QueryBuilder; the API proxy uses QueryBuilder::build_proxy_request for HTTP-forwarding providers and native proxy handlers for providers that need response conversion or SDK execution.

Dependency Direction

windmill-ai  →  windmill-common (for DB, Error, AgentAction, AuthedClient, etc.)
             →  windmill-types (for S3Object)
             →  windmill-parser (for Typ, used in OpenAPISchema)

windmill-api    →  windmill-ai
windmill-worker →  windmill-ai

windmill-common does NOT re-export from windmill-ai (would be circular). All consumers update imports.

Reviewer Note: Keep API Proxy Unification Split

The crate boundary, shared utilities, SSE parsers, image handling, worker provider implementations, provider-specific API proxy transformations, and resolved runtime credential shape are now in windmill-ai. Raw API resources and worker agent provider payloads remain separate input/deserialization shapes and convert into ProviderCredentials at execution boundaries.

Do not jump directly from the current state to full proxy and credential unification in one PR. The API proxy 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. Split the work by risk:

• Introduce shared proxy request and credential types first.
• Move the OpenAI-compatible proxy path into windmill-ai next, while keeping provider-native behavior unchanged.
• Move Anthropic/Vertex, Google AI, and Bedrock in separate follow-up PRs.
• Unify credential resolution only after all proxy request builders use the shared shape.

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.

Completed Phase: Proxy Contract + OpenAI-Compatible Proxy ✅

Goal: introduce the shared API proxy contract in windmill-ai and move the OpenAI-compatible proxy request builder there without changing provider behavior.

Suggested PR title: refactor(ai): move openai-compatible proxy building to windmill-ai.

Scope:

• Add windmill-ai/src/proxy.rs and export it from lib.rs.
• Define ProviderCredentials, ProxyBuildArgs, and ProxyRequest.
• Include all context known to be needed by the current API proxy path: method, path, incoming headers, body, provider, base URL, API key, OAuth access token, organization/user fields, platform, 1M context flag, custom headers, region, and AWS credentials.
• Add a conversion from API-side AIRequestConfig to ProviderCredentials.
• Add QueryBuilder::build_proxy_request with a default unsupported-provider implementation.
• Implement build_proxy_request for OpenAI-compatible providers (OpenAI, AzureOpenAI, Mistral, DeepSeek, Groq, OpenRouter, TogetherAI, CustomAI).
• Route workspace and global API proxy requests for OpenAI-compatible providers through windmill-ai.
• Keep FIM transformation in windmill-api before calling the proxy builder.
• Keep AIRequestConfig::prepare_request for Anthropic/Vertex and remaining fallback paths.

Out of scope:

• Do not move Anthropic/Vertex proxy behavior yet.
• Do not move Google AI or Bedrock proxy behavior yet.
• Do not change credential resolution, audit logging, cache behavior, SSE keepalive behavior, or Bedrock/Google special cases.
• Do not remove windmill-api/src/google.rs, windmill-api/src/bedrock.rs, or AIRequestConfig::prepare_request.

Validation:

• cargo test -p windmill-ai proxy
• cargo test -p windmill-api invalidates_all_cached_providers_for_workspace
• cargo check -p windmill-ai -p windmill-api
• cargo check -p windmill-ai -p windmill-api --features bedrock

Follow-up status: Anthropic/Vertex proxy handling has since moved into windmill-ai, and the dead AIRequestConfig::prepare_request fallback has been removed.

Completed Phase: Proxy Execution Mode + Google AI Proxy Migration ✅

Goal: introduce a shared provider execution classifier before moving Google AI and Bedrock. ProxyRequest is a good contract for HTTP-forwarding providers such as OpenAI-compatible providers and Anthropic, but Google AI also converts responses back to OpenAI shape and Bedrock uses SDK execution. Model that split explicitly before moving those providers, then move the Google AI proxy transformation into windmill-ai as the first native-provider migration.

Suggested PR title: refactor(ai): add provider proxy execution mode.

Scope:

• Add ProxyExecutionMode in windmill-ai::proxy.
• Classify providers as HTTP-forwarding, native Google AI, or native Bedrock.
• Make supports_query_builder_proxy derive from the shared execution mode.
• Use the shared execution mode in windmill-api/src/ai.rs for workspace proxy routing.
• Move Google AI workspace proxy request conversion, streaming/non-streaming response conversion, and model-list normalization into windmill-ai::providers::google_ai.
• Share Google AI GeminiTextRequest and generation-config construction between worker agent requests and API proxy requests.
• Delete the API-local windmill-api/src/google.rs module.
• Keep global proxy behavior, Bedrock native handling, credential resolution, audit logging, caching, and SSE keepalive behavior unchanged.

Out of scope:

• Do not move windmill-api/src/bedrock.rs.
• Do not unify AIRequestConfig and ProviderWithResource.

Validation:

• cargo test -p windmill-ai google_ai
• cargo test -p windmill-ai proxy
• cargo test -p windmill-api invalidates_all_cached_providers_for_workspace
• cargo test -p windmill-ai anthropic

Follow-up status: Bedrock native proxy handling has since moved into windmill-ai, and the API-local windmill-api/src/bedrock.rs module has been removed.

Completed Phase: Bedrock Native Proxy Migration ✅

Goal: move the remaining native-provider API proxy execution out of windmill-api and into windmill-ai, while leaving API-owned routing, credential resolution, auditing, cache behavior, and Axum response conversion in windmill-api.

Suggested PR title: refactor(ai): move bedrock proxy handling to windmill-ai.

Scope:

• Move Bedrock control-plane proxy calls (foundation-models, inference-profiles) into windmill-ai::providers::bedrock.
• Move Bedrock chat proxy OpenAI request parsing, Converse request execution, streaming SSE conversion, non-streaming OpenAI-shaped response conversion, and auth selection into windmill-ai::providers::bedrock.
• Add an Axum-free BedrockProxyResponse shape in windmill-ai; the API route converts it into an Axum body.
• Move the optional aws-sdk-bedrock dependency from windmill-api to windmill-ai.
• Delete the API-local windmill-api/src/bedrock.rs module.

Out of scope:

• Do not unify AIRequestConfig and ProviderWithResource.
• Do not change Bedrock credential resolution, audit logging, request caching, or non-Bedrock proxy behavior.

Validation:

• cargo test -p windmill-ai bedrock --features bedrock
• cargo check -p windmill-ai -p windmill-api
• cargo check -p windmill-ai -p windmill-api --features bedrock

Known Follow-Ups

These are not blockers for the current migration PR because they either preserve existing behavior or need a separate product decision, but they should stay visible for later hardening work.

• Google AI/Gemini native proxy custom headers: the native Google AI proxy path intentionally does not apply AI_HTTP_HEADERS or resource-level custom headers today. Decide whether and how env/resource custom-header injection should apply to Google AI once the proxy behavior is unified further.
• Bedrock SSE tool-call indexing: Bedrock streaming currently increments the OpenAI tool-call index on every Bedrock ContentBlockStop, including text content blocks. This behavior existed before the move from windmill-api to windmill-ai, but a later cleanup should advance the index only when the stopped block was a tool-use block.
• Bedrock SSE keepalives: Bedrock native SSE streams are still returned directly without the API proxy keepalive injection used by other SSE paths. This also preserves the pre-move behavior. A later cleanup can generalize the keepalive wrapper so it works for both reqwest::Error streams and Bedrock's SDK-backed std::io::Error streams.

Completed Phase: Credential Unification Phase 1 ✅

Goal: make ProviderCredentials the shared resolved runtime credential shape without overloading it with raw resource input or model-selection state.

AIRequestConfig and ProviderWithResource are not equivalent concepts: AIRequestConfig is API-side resolved state after DB, variable, OAuth, and resource handling, while ProviderWithResource is worker-side raw agent input that also carries the selected model. Keep raw/deserialization types separate and convert them into ProviderCredentials at execution boundaries.

Suggested PR title: refactor(ai): use provider credentials for worker builders.

Scope:

• Add a worker-side conversion from ProviderWithResource to ProviderCredentials.
• Keep model outside ProviderCredentials; it remains agent request data.
• Keep ProviderWithResource as the backward-compatible deserialization type for existing agent payloads.
• Use ProviderCredentials for worker query-builder creation.
• Collapse create_query_builder and create_proxy_query_builder into one create_query_builder(&ProviderCredentials) factory.

Out of scope:

• Do not remove API-local AIRequestConfig yet.
• Do not change API request-cache behavior.
• Do not change worker agent payload shape or serialized field names.

Validation:

• cargo check -p windmill-ai -p windmill-api -p windmill-worker
• cargo check -p windmill-ai -p windmill-api -p windmill-worker --features bedrock

Completed Phase: Credential Unification Phase 2 ✅

Goal: remove the API-local resolved credential wrapper after worker execution already uses the shared shape.

Suggested PR title: refactor(ai): resolve api proxy credentials directly.

Scope:

• Change API credential resolution to return ProviderCredentials directly.
• Replace ExpiringAIRequestConfig with an expiring ProviderCredentials cache entry.
• Remove AIRequestConfig::into_provider_credentials.
• Delete AIRequestConfig entirely if no API-only behavior remains.

Out of scope:

• Do not merge raw worker resource input into ProviderCredentials.
• Do not put model selection into ProviderCredentials.

Validation:

• cargo check -p windmill-ai -p windmill-api -p windmill-worker
• cargo check -p windmill-ai -p windmill-api -p windmill-worker --features bedrock
• cargo test -p windmill-api invalidates_all_cached_providers_for_workspace

Step-by-Step Plan

Each step produces a compiling, working backend.

---

Step 1: Create windmill-ai crate, move base types from windmill-common ✅

Create backend/windmill-ai/Cargo.toml and backend/windmill-ai/src/lib.rs.

Move from windmill-common/src/ to windmill-ai/src/:

• ai_types.rs — OpenAI-compatible message types
• ai_providers.rs — AIProvider enum, AIPlatform, base URLs, ProviderConfig
• ai_google.rs — Gemini types and OpenAI↔Gemini conversion
• ai_bedrock.rs — Bedrock SDK wrapper (feature-gated on bedrock)
• ai_cache.rs — instance AI config revision tracking

Update all imports (windmill_common::ai_* → windmill_ai::ai_*).

---

Step 2: Move worker AI types to windmill-ai ✅

Move from windmill-worker/src/ai/types.rs to windmill-ai/src/types.rs:

• ProviderWithResource, ProviderResource — credential types
• TokenUsage — token usage tracking
• OutputType, SchemaType, AdditionalProperties — output configuration
• OpenAPISchema — tool parameter schema (depends on windmill-parser::Typ)
• Tool, Message, ResponseFormat, JsonSchemaFormat — agent types
• StreamingEvent — SSE event enum
• AIAgentArgs, AIAgentArgsRaw, AIAgentResult — agent job args
• Memory — agent memory enum
• S3ObjectWithType — S3 image type
• McpToolSource stub (with same #[cfg(feature = "mcp")] pattern)

Worker ai/types.rs becomes a re-export: pub use windmill_ai::types::*.

---

Step 3: Move QueryBuilder trait, ParsedResponse, and StreamEventSink abstraction to windmill-ai ✅

Move from windmill-worker/src/ai/query_builder.rs to windmill-ai/src/query_builder.rs:

• BuildRequestArgs struct
• ParsedResponse enum
• QueryBuilder trait (with all existing methods)

New StreamEventSink trait in windmill-ai:

#[async_trait]
pub trait StreamEventSink: Send + Sync {
    async fn send(&self, event: StreamingEvent, events_str: &mut String) -> Result<(), Error>;
}

StreamEventSink abstracts the worker's StreamEventProcessor so windmill-ai doesn't depend on windmill-queue or the worker's job logger. The worker's StreamEventProcessor implements StreamEventSink. All provider parse_streaming_response methods and SSE parsers accept Box<dyn StreamEventSink>.

---

Step 4: Move SSE parsers to windmill-ai ✅

Move from windmill-worker/src/ai/sse.rs to windmill-ai/src/sse.rs:

• SSEParser trait
• OpenAISSEParser, AnthropicSSEParser, GeminiSSEParser, OpenAIResponsesSSEParser
• All associated types (delta types, usage types, etc.)

---

Step 5: Move provider implementations to windmill-ai ✅

Move from windmill-worker/src/ai/providers/ to windmill-ai/src/providers/:

• anthropic.rs — AnthropicQueryBuilder
• openai.rs — OpenAIQueryBuilder
• google_ai.rs — GoogleAIQueryBuilder
• bedrock.rs — BedrockQueryBuilder (feature-gated)
• other.rs — OtherQueryBuilder (Mistral, DeepSeek, Groq, TogetherAI, CustomAI)
• openrouter.rs — OpenRouterQueryBuilder
• mod.rs with create_query_builder factory

Move utility functions providers depend on:

• should_use_structured_output_tool (from utils.rs)
• extract_text_content (from utils.rs)

---

Step 6: Move image_handler to windmill-ai ✅

Move from windmill-worker/src/ai/image_handler.rs to windmill-ai/src/image_handler.rs:

• download_and_encode_s3_image — no signature change needed
• prepare_messages_for_api — no signature change needed
• upload_image_to_s3 — refactor: (base64_image, workspace_id, job_id, client) instead of (base64_image, &MiniPulledJob, client) to remove windmill-queue dependency

---

Step 7: Move shared utilities to windmill-ai ✅

Move AI_HTTP_HEADERS lazy_static (currently duplicated in windmill-api/src/ai.rs and windmill-worker/src/ai_executor.rs) to windmill_ai::utils. Both consumers import from windmill-ai.

---

Step 8: Add API proxy execution support to windmill-ai ✅

This is the key proxy unification step. HTTP-forwarding providers use QueryBuilder::build_proxy_request:

/// Build a request from a raw OpenAI-format proxy request.
/// Used by the API chat proxy. Handles format conversion for non-OpenAI providers.
fn build_proxy_request(
    &self,
    args: &ProxyBuildArgs<'_>,
) -> Result<ProxyRequest, Error>;

Where ProxyBuildArgs carries the API proxy context that provider implementations need:

pub struct ProxyBuildArgs<'a> {
    pub method: &'a 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:

pub struct ProxyRequest {
    pub method: http::Method,
    pub url: String,
    pub headers: Vec<(String, String)>,
    pub body: Vec<u8>,
}

Provider implementations:

• OpenAI-compatible (OpenAI, Mistral, DeepSeek, Groq, TogetherAI, CustomAI, OpenRouter): Minimal transformation — pass body through, build URL and auth headers.
• Anthropic: Handle standard vs Vertex AI. For Vertex: transform body (extract model, add anthropic_version). For standard: pass through with appropriate headers.
• Google AI: Native execution mode converts OpenAI format → Gemini format and Gemini responses → OpenAI shape. Replaces windmill-api/src/google.rs.
• Bedrock: Native execution mode converts OpenAI format → Bedrock SDK calls and SDK responses → OpenAI shape. Replaces windmill-api/src/bedrock.rs.

Refactor API proxy (windmill-api/src/ai.rs):

1. Parse provider from headers, resolve credentials → ProviderCredentials
2. Create QueryBuilder via create_query_builder
3. Dispatch by ProxyExecutionMode:
   • HTTP-forwarding providers call query_builder.build_proxy_request(&proxy_args) → ProxyRequest
   • Google AI and Bedrock call native handlers in windmill-ai
4. Convert the provider response to the API response body

Remove from windmill-api:

• AIRequestConfig::prepare_request — replaced by QueryBuilder::build_proxy_request
• google.rs — replaced by windmill_ai::providers::google_ai native proxy handlers
• bedrock.rs — replaced by windmill_ai::providers::bedrock native proxy handlers
• transform_anthropic_for_vertex — moved to AnthropicQueryBuilder
• supports_native_fim, transform_fim_to_chat_completions — moved to windmill-ai

Keep in API:

• credential resolution from DB, workspace settings, instance settings, variables, and OAuth into ProviderCredentials
• HTTP routes, audit logging, request caching
• inject_keepalives, is_sse_response helpers
• AIConfig, ExpiringProviderCredentials caching types

---

Step 9: Unify credential resolution

Make ProviderCredentials the single resolved runtime credential shape in windmill-ai, while keeping raw API and worker input/deserialization types at their boundaries.

The API's resolve_provider_credentials resolves credentials from DB, workspace or instance settings, variables, and OAuth. The worker's ProviderWithResource gets raw credentials from the flow module definition and also carries the selected model. Convert both paths into ProviderCredentials; do not make ProviderCredentials carry raw resource state or the model.

Extend windmill_ai::proxy::ProviderCredentials as needed so both can produce it:

pub struct ProviderCredentials {
    pub provider: AIProvider,
    pub base_url: String,
    pub api_key: Option<String>,
    pub access_token: Option<String>,
    pub organization_id: Option<String>,
    pub user: Option<String>,
    pub platform: AIPlatform,
    pub region: Option<String>,
    pub aws_access_key_id: Option<String>,
    pub aws_secret_access_key: Option<String>,
    pub aws_session_token: Option<String>,
    pub enable_1m_context: bool,
    pub custom_headers: HashMap<String, String>,
}

The create_query_builder factory takes &ProviderCredentials instead of &ProviderWithResource.

---

Final Crate Structure

windmill-ai/src/
├── lib.rs              # module exports
├── ai_types.rs         # OpenAI-compatible message types
├── ai_providers.rs     # AIProvider enum, base URLs, config
├── ai_google.rs        # Gemini types and conversions
├── ai_bedrock.rs       # Bedrock SDK wrapper (feature: bedrock)
├── ai_cache.rs         # Instance AI config revision
├── types.rs            # TokenUsage, Tool, OpenAPISchema, etc.
├── proxy.rs            # ProviderCredentials, ProxyBuildArgs, ProxyRequest
├── query_builder.rs    # QueryBuilder trait, BuildRequestArgs, ParsedResponse, StreamEventSink
├── sse.rs              # SSE parsers (OpenAI, Anthropic, Gemini, Responses)
├── image_handler.rs    # S3 image upload/download
├── utils.rs            # extract_text_content, should_use_structured_output_tool
└── providers/
    ├── mod.rs           # create_query_builder factory
    ├── anthropic.rs     # build_request + build_proxy_request
    ├── openai.rs        # build_request + build_proxy_request
    ├── google_ai.rs     # build_request + native proxy handlers
    ├── bedrock.rs       # build_request + native proxy handlers (feature: bedrock)
    ├── other.rs         # build_request + build_proxy_request
    └── openrouter.rs    # build_request + build_proxy_request

windmill-worker keeps: ai_executor.rs, ai/tools.rs, ai/utils.rs (flow/conversation/MCP logic), StreamEventProcessor (impl of StreamEventSink).

windmill-api keeps: HTTP routes (ai.rs proxy endpoints), audit logging, caching, credential resolution from DB. google.rs and bedrock.rs deleted.