diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 3593ddfd26..eec48074a8 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -13875,6 +13875,7 @@ dependencies = [ "async-trait", "aws-config", "aws-credential-types", + "aws-sdk-bedrock", "aws-sdk-bedrockruntime", "aws-smithy-types", "base64 0.22.1", @@ -13923,13 +13924,8 @@ dependencies = [ "async-stream", "async-trait", "async_zip", - "aws-config", - "aws-credential-types", - "aws-sdk-bedrock", - "aws-sdk-bedrockruntime", "aws-sdk-config", "aws-sigv4", - "aws-smithy-types", "axum 0.8.9", "base32", "base64 0.22.1", diff --git a/backend/windmill-ai/Cargo.toml b/backend/windmill-ai/Cargo.toml index 8f2f679c7e..0276240246 100644 --- a/backend/windmill-ai/Cargo.toml +++ b/backend/windmill-ai/Cargo.toml @@ -6,7 +6,7 @@ edition.workspace = true [features] default = [] -bedrock = ["dep:aws-sdk-bedrockruntime", "dep:aws-credential-types", "dep:aws-smithy-types", "dep:aws-config"] +bedrock = ["dep:aws-sdk-bedrock", "dep:aws-sdk-bedrockruntime", "dep:aws-credential-types", "dep:aws-smithy-types", "dep:aws-config"] mcp = ["dep:windmill-mcp"] [lib] @@ -42,4 +42,5 @@ ulid.workspace = true aws-config = { workspace = true, optional = true } aws-credential-types = { workspace = true, optional = true } aws-smithy-types = { workspace = true, optional = true } +aws-sdk-bedrock = { workspace = true, optional = true } aws-sdk-bedrockruntime = { workspace = true, optional = true } diff --git a/backend/windmill-ai/src/providers/bedrock.rs b/backend/windmill-ai/src/providers/bedrock.rs index 005327c15b..ec6f4fbcd3 100644 --- a/backend/windmill-ai/src/providers/bedrock.rs +++ b/backend/windmill-ai/src/providers/bedrock.rs @@ -7,21 +7,731 @@ //! - Helper utilities use crate::{ + 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, + BearerTokenProvider, BedrockClient, StreamingToolCall, + }, + ai_providers::USE_ENV_REGION, + ai_types::{OpenAIFunction, OpenAIToolCall, ToolDefFunction}, image_handler::prepare_messages_for_api, + proxy::ProxyBuildArgs, query_builder::{ParsedResponse, StreamEventSink}, types::{OpenAIMessage, StreamingEvent, TokenUsage, ToolDef}, }; +use bytes::Bytes; +use futures::{stream::BoxStream, StreamExt}; +use http::{HeaderMap, Method, StatusCode}; +use serde::Deserialize; use std::collections::HashMap; use windmill_common::{client::AuthedClient, error::Error}; -// Import shared Bedrock helpers for provider orchestration. -use crate::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, - BedrockClient, StreamingToolCall, -}; +// ============================================================================ +// Native Proxy Execution +// ============================================================================ + +/// OpenAI-format request body for Bedrock SDK proxy handlers. +#[derive(Deserialize, Debug)] +struct OpenAIRequest { + messages: Vec, + #[serde(default)] + tools: Option>, + #[serde(default)] + tool_choice: Option, + #[serde(default)] + max_tokens: Option, + #[serde(default)] + temperature: Option, +} + +#[derive(Deserialize, Debug)] +struct OpenAIToolDef { + #[serde(default)] + #[allow(dead_code)] + r#type: Option, + function: OpenAIToolFunction, +} + +#[derive(Deserialize, Debug)] +struct OpenAIToolFunction { + name: String, + #[serde(default)] + description: Option, + #[serde(default)] + parameters: Option, +} + +#[derive(Deserialize, Debug)] +struct BedrockProxyChatRequest { + model: String, + #[serde(default)] + stream: bool, +} + +enum BedrockAuthConfig { + BearerToken(String), + IamCredentials { + access_key_id: String, + secret_access_key: String, + session_token: Option, + }, + Environment, +} + +pub enum BedrockProxyResponseBody { + Fixed(Bytes), + Stream(BoxStream<'static, std::result::Result>), +} + +pub struct BedrockProxyResponse { + pub status_code: StatusCode, + pub headers: HeaderMap, + pub body: BedrockProxyResponseBody, +} + +/// Handle a workspace Bedrock proxy request through the AWS SDK. +/// +/// The API still owns credential resolution, route authorization, auditing, and +/// cache behavior. This helper owns Bedrock-specific control-plane and +/// OpenAI-compatible Converse transformations. +pub async fn handle_bedrock_proxy( + args: &ProxyBuildArgs<'_>, +) -> Result { + let region = args.credentials.region.as_deref().unwrap_or(USE_ENV_REGION); + + if *args.method == Method::GET { + return match args.path { + "foundation-models" => list_foundation_models(args, region).await, + "inference-profiles" => list_inference_profiles(args, region).await, + _ => Err(Error::BadRequest(format!( + "Unsupported AWS Bedrock proxy path: {}", + args.path + ))), + }; + } + + if *args.method != Method::POST { + return Err(Error::BadRequest(format!( + "Unsupported AWS Bedrock proxy method: {}", + args.method + ))); + } + + let request: BedrockProxyChatRequest = serde_json::from_slice(args.body) + .map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?; + + if request.stream { + handle_bedrock_sdk_streaming(&request.model, args.body, args, region).await + } else { + handle_bedrock_sdk_non_streaming(&request.model, args.body, args, region).await + } +} + +fn determine_auth_config( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + aws_session_token: Option<&str>, +) -> BedrockAuthConfig { + if let Some(key) = api_key.filter(|k| !k.is_empty()) { + BedrockAuthConfig::BearerToken(key.to_string()) + } else if let (Some(access_key_id), Some(secret_access_key)) = ( + aws_access_key_id.filter(|s| !s.is_empty()), + aws_secret_access_key.filter(|s| !s.is_empty()), + ) { + BedrockAuthConfig::IamCredentials { + access_key_id: access_key_id.to_string(), + secret_access_key: secret_access_key.to_string(), + session_token: aws_session_token + .filter(|token| !token.is_empty()) + .map(str::to_string), + } + } else { + BedrockAuthConfig::Environment + } +} + +async fn create_bedrock_client( + args: &ProxyBuildArgs<'_>, + region: &str, +) -> Result { + match determine_auth_config( + args.credentials.api_key.as_deref(), + args.credentials.aws_access_key_id.as_deref(), + args.credentials.aws_secret_access_key.as_deref(), + args.credentials.aws_session_token.as_deref(), + ) { + BedrockAuthConfig::BearerToken(key) => BedrockClient::from_bearer_token(key, region).await, + BedrockAuthConfig::IamCredentials { access_key_id, secret_access_key, session_token } => { + BedrockClient::from_credentials(access_key_id, secret_access_key, session_token, region) + .await + } + BedrockAuthConfig::Environment => BedrockClient::from_env(region).await, + } +} + +fn build_tool_config_from_request( + tools: Option<&[OpenAIToolDef]>, + tool_choice: Option<&serde_json::Value>, + enable_prompt_caching: bool, +) -> Result, Error> { + if let Some(tools) = tools { + let tool_defs: Vec = tools + .iter() + .map(|t| ToolDef { + r#type: "function".to_string(), + function: ToolDefFunction { + name: t.function.name.clone(), + description: t.function.description.clone(), + parameters: Box::from( + serde_json::value::RawValue::from_string( + serde_json::to_string( + &t.function + .parameters + .clone() + .unwrap_or(serde_json::json!({})), + ) + .unwrap_or_default(), + ) + .unwrap_or_else(|_| { + serde_json::value::RawValue::from_string("{}".to_string()).unwrap() + }), + ), + }, + }) + .collect(); + + let force_tool_use = tool_choice + .map(|tc| tc == "required" || tc.as_str() == Some("required")) + .unwrap_or(false); + + build_tool_config(Some(&tool_defs), force_tool_use, enable_prompt_caching) + } else { + Ok(None) + } +} + +async fn create_bedrock_control_client( + args: &ProxyBuildArgs<'_>, + region: &str, +) -> Result { + use aws_config::BehaviorVersion; + + let region_provider = aws_sdk_bedrock::config::Region::new(region.to_string()); + + match determine_auth_config( + args.credentials.api_key.as_deref(), + args.credentials.aws_access_key_id.as_deref(), + args.credentials.aws_secret_access_key.as_deref(), + args.credentials.aws_session_token.as_deref(), + ) { + BedrockAuthConfig::BearerToken(key) => { + let config = aws_sdk_bedrock::config::Builder::new() + .region(region_provider) + .behavior_version(BehaviorVersion::latest()) + .token_provider(BearerTokenProvider::new(key)) + .build(); + Ok(aws_sdk_bedrock::Client::from_conf(config)) + } + BedrockAuthConfig::IamCredentials { access_key_id, secret_access_key, session_token } => { + let credentials = aws_credential_types::Credentials::new( + access_key_id, + secret_access_key, + session_token, + None, + "windmill", + ); + let config = aws_sdk_bedrock::config::Builder::new() + .region(region_provider) + .behavior_version(BehaviorVersion::latest()) + .credentials_provider(credentials) + .build(); + Ok(aws_sdk_bedrock::Client::from_conf(config)) + } + BedrockAuthConfig::Environment => { + let config = aws_config::defaults(BehaviorVersion::latest()) + .region(region_provider) + .load() + .await; + Ok(aws_sdk_bedrock::Client::new(&config)) + } + } +} + +async fn list_foundation_models( + args: &ProxyBuildArgs<'_>, + region: &str, +) -> Result { + let client = create_bedrock_control_client(args, region).await?; + + let response = client + .list_foundation_models() + .send() + .await + .map_err(|e| Error::internal_err(format!("Failed to list foundation models: {}", e)))?; + + let models: Vec = response + .model_summaries() + .iter() + .map(|m| { + serde_json::json!({ + "modelId": m.model_id(), + "modelName": m.model_name(), + "providerName": m.provider_name(), + "modelArn": m.model_arn(), + "inputModalities": m.input_modalities().iter().map(|i| i.as_str()).collect::>(), + "outputModalities": m.output_modalities().iter().map(|o| o.as_str()).collect::>(), + "responseStreamingSupported": m.response_streaming_supported(), + "inferenceTypesSupported": m.inference_types_supported().iter().map(|i| i.as_str()).collect::>(), + }) + }) + .collect(); + + let body = serde_json::to_vec(&serde_json::json!({ "modelSummaries": models })) + .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; + + Ok(BedrockProxyResponse { + status_code: StatusCode::OK, + headers: json_response_headers(), + body: BedrockProxyResponseBody::Fixed(Bytes::from(body)), + }) +} + +async fn list_inference_profiles( + args: &ProxyBuildArgs<'_>, + region: &str, +) -> Result { + let client = create_bedrock_control_client(args, region).await?; + + let response = + client.list_inference_profiles().send().await.map_err(|e| { + Error::internal_err(format!("Failed to list inference profiles: {}", e)) + })?; + + let profiles: Vec = response + .inference_profile_summaries() + .iter() + .map(|p| { + serde_json::json!({ + "inferenceProfileId": p.inference_profile_id(), + "inferenceProfileName": p.inference_profile_name(), + "inferenceProfileArn": p.inference_profile_arn(), + "description": p.description(), + "status": p.status().as_str(), + "type": p.r#type().as_str(), + }) + }) + .collect(); + + let body = serde_json::to_vec(&serde_json::json!({ "inferenceProfileSummaries": profiles })) + .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; + + Ok(BedrockProxyResponse { + status_code: StatusCode::OK, + headers: json_response_headers(), + body: BedrockProxyResponseBody::Fixed(Bytes::from(body)), + }) +} + +async fn handle_bedrock_sdk_streaming( + model: &str, + body: &[u8], + args: &ProxyBuildArgs<'_>, + region: &str, +) -> Result { + let openai_req: OpenAIRequest = serde_json::from_slice(body) + .map_err(|e| Error::internal_err(format!("Failed to parse OpenAI request: {}", e)))?; + + let bedrock_client = create_bedrock_client(args, region).await?; + let enable_prompt_caching = bedrock_model_supports_prompt_caching(model); + let (bedrock_messages, system_prompts) = + openai_messages_to_bedrock(&openai_req.messages, enable_prompt_caching)?; + let inference_config = create_inference_config(openai_req.temperature, openai_req.max_tokens); + let tool_config = build_tool_config_from_request( + openai_req.tools.as_deref(), + openai_req.tool_choice.as_ref(), + enable_prompt_caching, + )?; + + let mut request_builder = bedrock_client + .client() + .converse_stream() + .model_id(model) + .set_messages(Some(bedrock_messages)); + + if !system_prompts.is_empty() { + request_builder = request_builder.set_system(Some(system_prompts)); + } + + if let Some(config) = inference_config { + request_builder = request_builder.inference_config(config); + } + + if let Some(config) = tool_config { + request_builder = request_builder.set_tool_config(Some(config)); + } + + tracing::debug!("Bedrock SDK streaming: sending converse_stream request"); + let stream_output = request_builder.send().await.map_err(|e| { + let error_msg = format!("Bedrock SDK streaming error: {}", format_bedrock_error(&e)); + tracing::error!("Bedrock SDK streaming failed: {}", error_msg); + Error::internal_err(error_msg) + })?; + tracing::debug!("Bedrock SDK streaming: stream established successfully"); + + Ok(BedrockProxyResponse { + status_code: StatusCode::OK, + headers: event_stream_response_headers(), + body: BedrockProxyResponseBody::Stream( + sdk_stream_to_sse(stream_output.stream, model.to_string()).boxed(), + ), + }) +} + +pub fn sdk_stream_to_sse( + stream: aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver< + aws_sdk_bedrockruntime::types::ConverseStreamOutput, + aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError, + >, + model: String, +) -> impl futures::Stream> + Send { + let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); + let created = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + struct StreamState { + id: String, + model: String, + created: u64, + tool_calls: HashMap, + current_tool_index: usize, + } + + let state = std::sync::Arc::new(tokio::sync::Mutex::new(StreamState { + id, + model, + created, + tool_calls: HashMap::new(), + current_tool_index: 0, + })); + + async_stream::stream! { + let mut stream = stream; + let state = state.clone(); + + loop { + match stream.recv().await { + Ok(Some(event)) => { + let mut state = state.lock().await; + + if let Some(tool_call) = bedrock_stream_event_to_tool_start(&event) { + let index = state.current_tool_index; + state.tool_calls.insert( + index, + (tool_call.id.clone(), tool_call.name.clone(), String::new()), + ); + + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": index, + "id": tool_call.id, + "type": "function", + "function": { + "name": tool_call.name, + "arguments": "" + } + }] + }, + "finish_reason": serde_json::Value::Null + }] + }); + + yield Ok(Bytes::from(format!("data: {}\n\n", chunk))); + } + + if let Some(text) = bedrock_stream_event_to_text(&event) { + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": { + "content": text + }, + "finish_reason": serde_json::Value::Null + }] + }); + + yield Ok(Bytes::from(format!("data: {}\n\n", chunk))); + } + + if let Some(input_delta) = bedrock_stream_event_to_tool_delta(&event) { + let index = state.current_tool_index; + if let Some((_id, _name, ref mut args)) = state.tool_calls.get_mut(&index) { + args.push_str(&input_delta); + + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": index, + "function": { + "arguments": input_delta + } + }] + }, + "finish_reason": serde_json::Value::Null + }] + }); + + yield Ok(Bytes::from(format!("data: {}\n\n", chunk))); + } + } + + if bedrock_stream_event_is_block_stop(&event) { + state.current_tool_index += 1; + } + + if let aws_sdk_bedrockruntime::types::ConverseStreamOutput::MessageStop(stop) = &event { + let stop_reason = stop.stop_reason().as_str(); + let finish_reason = match stop_reason { + "end_turn" => "stop", + "max_tokens" => "length", + "tool_use" => "tool_calls", + "stop_sequence" => "stop", + "guardrail_intervened" | "content_filtered" => "content_filter", + _ => "stop", + }; + + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": finish_reason + }] + }); + + yield Ok(Bytes::from(format!("data: {}\n\n", chunk))); + } + } + Ok(None) => break, + Err(e) => { + yield Err(std::io::Error::new( + std::io::ErrorKind::Other, + e.to_string(), + )); + break; + } + } + } + + yield Ok(Bytes::from("data: [DONE]\n\n")); + } +} + +async fn handle_bedrock_sdk_non_streaming( + model: &str, + body: &[u8], + args: &ProxyBuildArgs<'_>, + region: &str, +) -> Result { + let openai_req: OpenAIRequest = serde_json::from_slice(body) + .map_err(|e| Error::internal_err(format!("Failed to parse OpenAI request: {}", e)))?; + + let bedrock_client = create_bedrock_client(args, region).await?; + let enable_prompt_caching = bedrock_model_supports_prompt_caching(model); + let (bedrock_messages, system_prompts) = + openai_messages_to_bedrock(&openai_req.messages, enable_prompt_caching)?; + let inference_config = create_inference_config(openai_req.temperature, openai_req.max_tokens); + let tool_config = build_tool_config_from_request( + openai_req.tools.as_deref(), + openai_req.tool_choice.as_ref(), + enable_prompt_caching, + )?; + + let mut request_builder = bedrock_client + .client() + .converse() + .model_id(model) + .set_messages(Some(bedrock_messages)); + + if !system_prompts.is_empty() { + request_builder = request_builder.set_system(Some(system_prompts)); + } + + if let Some(config) = inference_config { + request_builder = request_builder.inference_config(config); + } + + if let Some(config) = tool_config { + request_builder = request_builder.set_tool_config(Some(config)); + } + + tracing::debug!("Bedrock SDK non-streaming: sending converse request"); + let response = request_builder.send().await.map_err(|e| { + let error_msg = format!( + "Bedrock SDK non-streaming error: {}", + format_bedrock_error(&e) + ); + tracing::error!("Bedrock SDK non-streaming failed: {}", error_msg); + Error::internal_err(error_msg) + })?; + tracing::debug!( + "Bedrock SDK non-streaming: response received, stop_reason={}", + response.stop_reason().as_str() + ); + + let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); + let created = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + let stop_reason = response.stop_reason().as_str(); + let finish_reason = match stop_reason { + "end_turn" => "stop", + "max_tokens" => "length", + "tool_use" => "tool_calls", + "stop_sequence" => "stop", + "guardrail_intervened" | "content_filtered" => "content_filter", + _ => "stop", + }; + + let mut text_content = String::new(); + let mut tool_calls: Vec = Vec::new(); + + if let Some(aws_sdk_bedrockruntime::types::ConverseOutput::Message(message)) = response.output() + { + for block in message.content() { + match block { + aws_sdk_bedrockruntime::types::ContentBlock::Text(text) => { + text_content.push_str(text); + } + aws_sdk_bedrockruntime::types::ContentBlock::ToolUse(tool_use) => { + let input_json = document_to_json(tool_use.input()); + tool_calls.push(OpenAIToolCall { + id: tool_use.tool_use_id().to_string(), + function: OpenAIFunction { + name: tool_use.name().to_string(), + arguments: serde_json::to_string(&input_json).unwrap_or_default(), + }, + r#type: "function".to_string(), + extra_content: None, + }); + } + _ => {} + } + } + } + + let message = if !tool_calls.is_empty() { + serde_json::json!({ + "role": "assistant", + "content": if text_content.is_empty() { serde_json::Value::Null } else { serde_json::Value::String(text_content) }, + "tool_calls": tool_calls + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": text_content + }) + }; + + let usage = if let Some(usage_data) = response.usage() { + serde_json::json!({ + "prompt_tokens": usage_data.input_tokens(), + "completion_tokens": usage_data.output_tokens(), + "total_tokens": usage_data.total_tokens() + }) + } else { + serde_json::json!({ + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + }) + }; + + let openai_resp = serde_json::json!({ + "id": id, + "object": "chat.completion", + "created": created, + "model": model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": finish_reason + }], + "usage": usage + }); + + let body = serde_json::to_vec(&openai_resp) + .map_err(|e| Error::internal_err(format!("Failed to serialize OpenAI response: {}", e)))?; + + Ok(BedrockProxyResponse { + status_code: StatusCode::OK, + headers: json_response_headers(), + body: BedrockProxyResponseBody::Fixed(Bytes::from(body)), + }) +} + +fn json_response_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert("content-type", "application/json".parse().unwrap()); + headers +} + +fn event_stream_response_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert("content-type", "text/event-stream".parse().unwrap()); + headers.insert("cache-control", "no-cache".parse().unwrap()); + headers.insert("connection", "keep-alive".parse().unwrap()); + headers +} + +fn document_to_json(doc: &aws_smithy_types::Document) -> serde_json::Value { + match doc { + aws_smithy_types::Document::Object(map) => { + let mut json_map = serde_json::Map::new(); + for (key, value) in map { + json_map.insert(key.clone(), document_to_json(value)); + } + serde_json::Value::Object(json_map) + } + aws_smithy_types::Document::Array(values) => { + serde_json::Value::Array(values.iter().map(document_to_json).collect()) + } + aws_smithy_types::Document::Number(number) => match number { + aws_smithy_types::Number::PosInt(number) => serde_json::Value::Number((*number).into()), + aws_smithy_types::Number::NegInt(number) => serde_json::Value::Number((*number).into()), + aws_smithy_types::Number::Float(number) => serde_json::json!(*number), + }, + aws_smithy_types::Document::String(value) => serde_json::Value::String(value.clone()), + aws_smithy_types::Document::Bool(value) => serde_json::Value::Bool(*value), + aws_smithy_types::Document::Null => serde_json::Value::Null, + } +} // ============================================================================ // Query Builder @@ -256,3 +966,60 @@ impl BedrockQueryBuilder { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn determine_auth_config_prioritizes_bearer_token() { + let config = determine_auth_config( + Some("bearer-token"), + Some("AKIA123"), + Some("secret"), + Some("session-token"), + ); + + match config { + BedrockAuthConfig::BearerToken(token) => assert_eq!(token, "bearer-token"), + _ => panic!("expected bearer token auth config"), + } + } + + #[test] + fn determine_auth_config_uses_iam_with_optional_session_token() { + let config = + determine_auth_config(None, Some("AKIA123"), Some("secret"), Some("session-token")); + + match config { + BedrockAuthConfig::IamCredentials { + access_key_id, + secret_access_key, + session_token, + } => { + assert_eq!(access_key_id, "AKIA123"); + assert_eq!(secret_access_key, "secret"); + assert_eq!(session_token.as_deref(), Some("session-token")); + } + _ => panic!("expected IAM auth config"), + } + } + + #[test] + fn determine_auth_config_treats_empty_session_token_as_none() { + let config = determine_auth_config(None, Some("AKIA123"), Some("secret"), Some("")); + + match config { + BedrockAuthConfig::IamCredentials { session_token, .. } => { + assert!(session_token.is_none()); + } + _ => panic!("expected IAM auth config"), + } + } + + #[test] + fn determine_auth_config_falls_back_to_environment() { + let config = determine_auth_config(None, Some("AKIA123"), None, Some("session-token")); + assert!(matches!(config, BedrockAuthConfig::Environment)); + } +} diff --git a/backend/windmill-ai/src/providers/google_ai.rs b/backend/windmill-ai/src/providers/google_ai.rs index b6a6295d88..57abae5182 100644 --- a/backend/windmill-ai/src/providers/google_ai.rs +++ b/backend/windmill-ai/src/providers/google_ai.rs @@ -408,6 +408,8 @@ fn build_google_ai_model_endpoint( action: &str, is_vertex: bool, ) -> String { + let model = model.strip_prefix("models/").unwrap_or(model); + if is_vertex { format!("{}/{}:{}", base_url, model, action) } else { @@ -416,6 +418,10 @@ fn build_google_ai_model_endpoint( } fn add_google_ai_auth_header(headers: &mut Vec<(String, String)>, api_key: &str, is_vertex: bool) { + // Native Google AI proxy intentionally does not apply AI_HTTP_HEADERS or + // resource custom headers yet. Gemini/Vertex header semantics are + // provider-specific; keep this limited to required auth headers until + // explicit custom-header support is designed. if is_vertex { headers.push(("Authorization".to_string(), format!("Bearer {}", api_key))); } else { @@ -749,6 +755,32 @@ mod tests { assert!(body["contents"].is_array()); } + #[test] + fn builds_standard_google_ai_endpoint_from_model_resource_name() { + assert_eq!( + build_google_ai_model_endpoint( + "https://generativelanguage.googleapis.com/v1beta", + "models/gemini-2.0-flash", + "generateContent", + false, + ), + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent" + ); + } + + #[test] + fn builds_vertex_google_ai_endpoint_from_model_resource_name() { + assert_eq!( + build_google_ai_model_endpoint( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models", + "models/gemini-2.0-flash", + "streamGenerateContent", + true, + ), + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent" + ); + } + #[test] fn builds_vertex_google_ai_streaming_proxy_request() { let credentials = credentials( diff --git a/backend/windmill-api-auth/src/lib.rs b/backend/windmill-api-auth/src/lib.rs index b9bc2e2417..99fa677cee 100644 --- a/backend/windmill-api-auth/src/lib.rs +++ b/backend/windmill-api-auth/src/lib.rs @@ -235,6 +235,50 @@ where Ok(()) } +/// Returns a predicate that checks whether `path` is within the token's +/// scope for `{domain}:{action}:{path}`. For tokens without scope +/// restrictions (no scopes at all, or only `if_jobs:filter_tags:*` scopes), +/// the predicate always returns `true`. +/// +/// Pre-parses the token's scopes once so the returned closure can cheaply +/// filter large listings without re-parsing on each call. +pub fn build_scope_path_predicate( + authed: &ApiAuthed, + domain: &str, + action: &str, +) -> impl Fn(&str) -> bool { + // Mirror check_scopes semantics: a token is "scope-restricted" iff it has + // at least one non-`if_jobs:filter_tags:` scope. Unparseable scopes still + // count as restrictive — they just match nothing. + let (is_scoped_token, parsed): (bool, Vec) = match authed.scopes.as_ref() { + Some(scopes) => { + let mut is_scoped = false; + let parsed = scopes + .iter() + .filter(|s| !s.starts_with("if_jobs:filter_tags:")) + .inspect(|_| is_scoped = true) + .filter_map(|s| ScopeDefinition::from_scope_string(s).ok()) + .collect(); + (is_scoped, parsed) + } + None => (false, Vec::new()), + }; + let domain = domain.to_string(); + let action = action.to_string(); + + move |path: &str| -> bool { + if !is_scoped_token { + return true; + } + let required = + match ScopeDefinition::from_scope_string(&format!("{}:{}:{}", domain, action, path)) { + Ok(r) => r, + Err(_) => return false, + }; + parsed.iter().any(|s| s.includes(&required)) + } +} + pub async fn require_devops_role(db: &DB, email: &str) -> error::Result<()> { let is_devops = is_devops_email(db, email).await?; @@ -803,3 +847,65 @@ pub fn require_path_read_access_for_preview( ))), } } + +#[cfg(test)] +mod tests { + use super::*; + + fn authed_with_scopes(scopes: Option>) -> ApiAuthed { + ApiAuthed { + scopes: scopes.map(|v| v.into_iter().map(String::from).collect()), + ..Default::default() + } + } + + #[test] + fn predicate_no_scopes_allows_all() { + let authed = authed_with_scopes(None); + let allowed = build_scope_path_predicate(&authed, "resources", "read"); + assert!(allowed("u/alice/anything")); + assert!(allowed("u/bob/other")); + } + + #[test] + fn predicate_tag_filter_only_allows_all() { + let authed = authed_with_scopes(Some(vec!["if_jobs:filter_tags:default"])); + let allowed = build_scope_path_predicate(&authed, "resources", "read"); + assert!(allowed("u/alice/foo")); + } + + #[test] + fn predicate_single_resource_scope_filters_others() { + // Regression test for WIN-1981: a token scoped to one resource must + // not match unrelated paths in listings (e.g. /resources/list_search). + let authed = authed_with_scopes(Some(vec!["resources:read:u/alice/allowed_resource"])); + let allowed = build_scope_path_predicate(&authed, "resources", "read"); + assert!(allowed("u/alice/allowed_resource")); + assert!(!allowed("u/alice/other_resource")); + assert!(!allowed("u/bob/foo")); + } + + #[test] + fn predicate_wildcard_scope_matches_subtree() { + let authed = authed_with_scopes(Some(vec!["resources:read:f/team/*"])); + let allowed = build_scope_path_predicate(&authed, "resources", "read"); + assert!(allowed("f/team/db")); + assert!(allowed("f/team/sub/nested")); + assert!(!allowed("f/other/db")); + } + + #[test] + fn predicate_wrong_domain_is_rejected() { + let authed = authed_with_scopes(Some(vec!["variables:read:u/alice/secret"])); + let allowed = build_scope_path_predicate(&authed, "resources", "read"); + assert!(!allowed("u/alice/secret")); + } + + #[test] + fn predicate_write_implies_read() { + let authed = authed_with_scopes(Some(vec!["resources:write:u/alice/foo"])); + let allowed = build_scope_path_predicate(&authed, "resources", "read"); + assert!(allowed("u/alice/foo")); + assert!(!allowed("u/alice/bar")); + } +} diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml index 2af5e17406..fbe6fb23c8 100644 --- a/backend/windmill-api/Cargo.toml +++ b/backend/windmill-api/Cargo.toml @@ -40,7 +40,7 @@ gcp_trigger = ["dep:windmill-trigger-gcp", "windmill-store/gcp_trigger"] azure_trigger = ["dep:windmill-trigger-azure", "windmill-store/azure_trigger"] cloud = ["windmill-common/cloud", "windmill-api-auth/cloud", "windmill-store/cloud", "windmill-api-workspaces/cloud"] mcp = ["dep:windmill-mcp", "windmill-mcp/server", "windmill-mcp/auth", "windmill-api-auth/mcp", "windmill-store/mcp"] -bedrock = ["windmill-ai/bedrock", "dep:aws-sdk-bedrock", "dep:aws-sdk-bedrockruntime", "dep:aws-config", "dep:aws-credential-types", "dep:aws-smithy-types"] +bedrock = ["windmill-ai/bedrock"] python = ["windmill-dep-map/python", "dep:windmill-parser-py", "dep:windmill-parser-py-imports", "windmill-api-scripts/python", "windmill-api-configs/python", "windmill-api-agent-workers?/python", "windmill-trigger/python", "windmill-common/python"] no_auth = ["windmill-api-auth/no_auth", "windmill-store/no_auth", "windmill-api-users/no_auth"] quickjs = ["windmill-jseval/quickjs"] @@ -172,11 +172,6 @@ rustls = { workspace = true } aws-sigv4 = { workspace = true, optional = true } aws-sdk-config = { workspace = true, optional = true } -aws-config = { workspace = true, optional = true } -aws-credential-types = { workspace = true, optional = true } -aws-sdk-bedrock = { workspace = true, optional = true } -aws-sdk-bedrockruntime = { workspace = true, optional = true } -aws-smithy-types = { workspace = true, optional = true } async-trait.workspace = true eventsource-stream.workspace = true windmill-jseval.workspace = true diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index e7bc089d7c..ef8c610f1f 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -1,5 +1,3 @@ -#[cfg(feature = "bedrock")] -use crate::bedrock; use crate::db::{ApiAuthed, DB}; use crate::utils::check_scopes; @@ -20,6 +18,10 @@ use windmill_ai::ai_cache::current_instance_ai_config_revision; use windmill_ai::ai_providers::{ empty_string_as_none, AIPlatform, AIProvider, ProviderConfig, ProviderModel, }; +#[cfg(feature = "bedrock")] +use windmill_ai::providers::bedrock::{ + handle_bedrock_proxy, BedrockProxyResponse, BedrockProxyResponseBody, +}; use windmill_ai::providers::{ create_proxy_query_builder, google_ai::{ @@ -540,6 +542,18 @@ fn google_ai_proxy_response_to_body( (response.status_code, response.headers, body) } +#[cfg(feature = "bedrock")] +fn bedrock_proxy_response_to_body( + response: BedrockProxyResponse, +) -> (http::StatusCode, HeaderMap, axum::body::Body) { + let body = match response.body { + BedrockProxyResponseBody::Fixed(body) => axum::body::Body::from(body), + BedrockProxyResponseBody::Stream(stream) => axum::body::Body::from_stream(stream), + }; + + (response.status_code, response.headers, body) +} + pub(crate) fn inject_keepalives( upstream: S, interval: Duration, @@ -885,95 +899,31 @@ async fn proxy( // Handle Bedrock-specific logic when the feature is enabled #[cfg(feature = "bedrock")] - { - // Extract model and streaming flag for Bedrock transformation (only for POST requests) - let (model, is_streaming) = if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) - && method == Method::POST - { - #[derive(Deserialize, Debug)] - struct BedrockRequest { - model: String, - #[serde(default)] - stream: bool, - } - let parsed: BedrockRequest = serde_json::from_slice(&body) - .map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?; - (Some(parsed.model), parsed.stream) - } else { - (None, false) - }; + if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) { + let mut tx = db.begin().await?; + audit_log( + &mut *tx, + &authed, + "ai.request", + ActionKind::Execute, + &w_id, + Some(&authed.email), + Some([("ai_config_path", &format!("{:?}", ai_path)[..])].into()), + ) + .await?; + tx.commit().await?; - // For Bedrock requests, use the SDK-based approach - if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) { - let region = request_config - .region - .as_deref() - .unwrap_or(windmill_ai::ai_providers::USE_ENV_REGION); + let credentials = request_config.into_provider_credentials(provider.clone()); + let response = handle_bedrock_proxy(&ProxyBuildArgs { + method: &method, + path: &ai_path, + headers: &headers, + body: &body, + credentials: &credentials, + }) + .await?; - // Audit log before making the SDK request - let mut tx = db.begin().await?; - audit_log( - &mut *tx, - &authed, - "ai.request", - ActionKind::Execute, - &w_id, - Some(&authed.email), - Some([("ai_config_path", &format!("{:?}", ai_path)[..])].into()), - ) - .await?; - tx.commit().await?; - - // Handle GET requests for control plane operations - if method == Method::GET { - if ai_path == "foundation-models" { - return bedrock::list_foundation_models( - request_config.api_key.as_deref(), - request_config.aws_access_key_id.as_deref(), - request_config.aws_secret_access_key.as_deref(), - request_config.aws_session_token.as_deref(), - region, - ) - .await; - } else if ai_path == "inference-profiles" { - return bedrock::list_inference_profiles( - request_config.api_key.as_deref(), - request_config.aws_access_key_id.as_deref(), - request_config.aws_secret_access_key.as_deref(), - request_config.aws_session_token.as_deref(), - region, - ) - .await; - } - } - - // Handle POST requests for inference - if method == Method::POST && model.is_some() { - if is_streaming { - return bedrock::handle_bedrock_sdk_streaming( - model.as_ref().unwrap(), - &body, - request_config.api_key.as_deref(), - request_config.aws_access_key_id.as_deref(), - request_config.aws_secret_access_key.as_deref(), - request_config.aws_session_token.as_deref(), - region, - ) - .await; - } else { - return bedrock::handle_bedrock_sdk_non_streaming( - model.as_ref().unwrap(), - &body, - request_config.api_key.as_deref(), - request_config.aws_access_key_id.as_deref(), - request_config.aws_secret_access_key.as_deref(), - request_config.aws_session_token.as_deref(), - region, - ) - .await; - } - } - } + return Ok(bedrock_proxy_response_to_body(response)); } // When bedrock feature is disabled, return error for Bedrock provider diff --git a/backend/windmill-api/src/bedrock.rs b/backend/windmill-api/src/bedrock.rs deleted file mode 100644 index cc84fda1e8..0000000000 --- a/backend/windmill-api/src/bedrock.rs +++ /dev/null @@ -1,873 +0,0 @@ -//! AWS Bedrock SDK-based operations for the AI chat proxy. -//! -//! This module provides SDK-based request handling for Bedrock: -//! -//! ## Inference (Runtime SDK): -//! - `handle_bedrock_sdk_streaming`: Uses BedrockClient for streaming requests -//! - `handle_bedrock_sdk_non_streaming`: Uses BedrockClient for non-streaming requests -//! - `sdk_stream_to_sse`: Converts SDK ConverseStream events to SSE format -//! -//! ## Control Plane (Bedrock SDK): -//! - `list_foundation_models`: Lists available foundation models -//! - `list_inference_profiles`: Lists inference profiles -//! -//! Shared AWS SDK code is available in `windmill_common::ai_bedrock`, including: -//! - `BedrockClient`: SDK wrapper with bearer token and IAM auth -//! - Stream event parsing functions -//! - Helper utilities - -use axum::body::Bytes; -use serde::Deserialize; -use windmill_ai::ai_bedrock::build_tool_config; -use windmill_ai::ai_bedrock::{ - bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text, - bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, format_bedrock_error, - BedrockClient, -}; -use windmill_ai::ai_types::{ - OpenAIFunction, OpenAIMessage, OpenAIToolCall, ToolDef, ToolDefFunction, -}; -use windmill_common::error::{Error, Result}; - -// ============================================================================ -// Shared Request Types for SDK-Based Handlers -// ============================================================================ - -/// OpenAI-format request body for Bedrock SDK handlers -#[derive(Deserialize, Debug)] -struct OpenAIRequest { - messages: Vec, - #[serde(default)] - tools: Option>, - #[serde(default)] - tool_choice: Option, - #[serde(default)] - max_tokens: Option, - #[serde(default)] - temperature: Option, -} - -#[derive(Deserialize, Debug)] -struct OpenAIToolDef { - #[serde(default)] - #[allow(dead_code)] - r#type: Option, - function: OpenAIToolFunction, -} - -#[derive(Deserialize, Debug)] -struct OpenAIToolFunction { - name: String, - #[serde(default)] - description: Option, - #[serde(default)] - parameters: Option, -} - -// ============================================================================ -// Shared Helper Functions for SDK-Based Handlers -// ============================================================================ - -/// Authentication configuration for Bedrock clients -enum BedrockAuthConfig { - BearerToken(String), - IamCredentials { - access_key_id: String, - secret_access_key: String, - session_token: Option, - }, - Environment, -} - -/// Determine auth configuration with priority: bearer token → IAM credentials → environment -fn determine_auth_config( - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, -) -> BedrockAuthConfig { - if let Some(key) = api_key.filter(|k| !k.is_empty()) { - BedrockAuthConfig::BearerToken(key.to_string()) - } else if let (Some(access_key_id), Some(secret_access_key)) = ( - aws_access_key_id.filter(|s| !s.is_empty()), - aws_secret_access_key.filter(|s| !s.is_empty()), - ) { - BedrockAuthConfig::IamCredentials { - access_key_id: access_key_id.to_string(), - secret_access_key: secret_access_key.to_string(), - session_token: aws_session_token - .filter(|token| !token.is_empty()) - .map(str::to_string), - } - } else { - BedrockAuthConfig::Environment - } -} - -/// Create a BedrockClient with auth priority: bearer token → IAM credentials → environment -async fn create_bedrock_client( - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, - region: &str, -) -> Result { - match determine_auth_config( - api_key, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - ) { - BedrockAuthConfig::BearerToken(key) => BedrockClient::from_bearer_token(key, region).await, - BedrockAuthConfig::IamCredentials { access_key_id, secret_access_key, session_token } => { - BedrockClient::from_credentials(access_key_id, secret_access_key, session_token, region) - .await - } - BedrockAuthConfig::Environment => BedrockClient::from_env(region).await, - } -} - -/// Convert OpenAIToolDef array to tool configuration for Bedrock SDK -fn build_tool_config_from_request( - tools: Option<&[OpenAIToolDef]>, - tool_choice: Option<&serde_json::Value>, - enable_prompt_caching: bool, -) -> Result> { - if let Some(tools) = tools { - let tool_defs: Vec = tools - .iter() - .map(|t| ToolDef { - r#type: "function".to_string(), - function: ToolDefFunction { - name: t.function.name.clone(), - description: t.function.description.clone(), - parameters: Box::from( - serde_json::value::RawValue::from_string( - serde_json::to_string( - &t.function - .parameters - .clone() - .unwrap_or(serde_json::json!({})), - ) - .unwrap_or_default(), - ) - .unwrap_or_else(|_| { - serde_json::value::RawValue::from_string("{}".to_string()).unwrap() - }), - ), - }, - }) - .collect(); - - // Determine if we should force tool use based on tool_choice - let force_tool_use = tool_choice - .map(|tc| tc == "required" || tc.as_str() == Some("required")) - .unwrap_or(false); - - build_tool_config(Some(&tool_defs), force_tool_use, enable_prompt_caching) - } else { - Ok(None) - } -} - -// ============================================================================ -// Control Plane Operations (using aws-sdk-bedrock) -// ============================================================================ - -/// Create a Bedrock control plane client with auth priority: bearer token → IAM credentials → environment -async fn create_bedrock_control_client( - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, - region: &str, -) -> Result { - use aws_config::BehaviorVersion; - use windmill_ai::ai_bedrock::BearerTokenProvider; - - let region_provider = aws_sdk_bedrock::config::Region::new(region.to_string()); - - match determine_auth_config( - api_key, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - ) { - BedrockAuthConfig::BearerToken(key) => { - let config = aws_sdk_bedrock::config::Builder::new() - .region(region_provider) - .behavior_version(BehaviorVersion::latest()) - .token_provider(BearerTokenProvider::new(key)) - .build(); - Ok(aws_sdk_bedrock::Client::from_conf(config)) - } - BedrockAuthConfig::IamCredentials { access_key_id, secret_access_key, session_token } => { - let credentials = aws_credential_types::Credentials::new( - access_key_id, - secret_access_key, - session_token, - None, - "windmill", - ); - let config = aws_sdk_bedrock::config::Builder::new() - .region(region_provider) - .behavior_version(BehaviorVersion::latest()) - .credentials_provider(credentials) - .build(); - Ok(aws_sdk_bedrock::Client::from_conf(config)) - } - BedrockAuthConfig::Environment => { - let config = aws_config::defaults(BehaviorVersion::latest()) - .region(region_provider) - .load() - .await; - Ok(aws_sdk_bedrock::Client::new(&config)) - } - } -} - -/// List foundation models using the Bedrock SDK -pub async fn list_foundation_models( - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, - region: &str, -) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { - let client = create_bedrock_control_client( - api_key, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - region, - ) - .await?; - - let response = client - .list_foundation_models() - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to list foundation models: {}", e)))?; - - // Convert to JSON response - let models: Vec = response - .model_summaries() - .iter() - .map(|m| { - serde_json::json!({ - "modelId": m.model_id(), - "modelName": m.model_name(), - "providerName": m.provider_name(), - "modelArn": m.model_arn(), - "inputModalities": m.input_modalities().iter().map(|i| i.as_str()).collect::>(), - "outputModalities": m.output_modalities().iter().map(|o| o.as_str()).collect::>(), - "responseStreamingSupported": m.response_streaming_supported(), - "inferenceTypesSupported": m.inference_types_supported().iter().map(|i| i.as_str()).collect::>(), - }) - }) - .collect(); - - let body = serde_json::json!({ "modelSummaries": models }); - let body_bytes = serde_json::to_vec(&body) - .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; - - let mut headers = http::HeaderMap::new(); - headers.insert("content-type", "application/json".parse().unwrap()); - - Ok(( - http::StatusCode::OK, - headers, - axum::body::Body::from(body_bytes), - )) -} - -/// List inference profiles using the Bedrock SDK -pub async fn list_inference_profiles( - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, - region: &str, -) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { - let client = create_bedrock_control_client( - api_key, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - region, - ) - .await?; - - let response = - client.list_inference_profiles().send().await.map_err(|e| { - Error::internal_err(format!("Failed to list inference profiles: {}", e)) - })?; - - // Convert to JSON response - let profiles: Vec = response - .inference_profile_summaries() - .iter() - .map(|p| { - serde_json::json!({ - "inferenceProfileId": p.inference_profile_id(), - "inferenceProfileName": p.inference_profile_name(), - "inferenceProfileArn": p.inference_profile_arn(), - "description": p.description(), - "status": p.status().as_str(), - "type": p.r#type().as_str(), - }) - }) - .collect(); - - let body = serde_json::json!({ "inferenceProfileSummaries": profiles }); - let body_bytes = serde_json::to_vec(&body) - .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; - - let mut headers = http::HeaderMap::new(); - headers.insert("content-type", "application/json".parse().unwrap()); - - Ok(( - http::StatusCode::OK, - headers, - axum::body::Body::from(body_bytes), - )) -} - -// ============================================================================ -// Inference Operations (using aws-sdk-bedrockruntime) -// ============================================================================ - -/// Handle Bedrock streaming request using the AWS SDK. -/// -/// This function uses the shared BedrockClient to make streaming requests -/// and converts the SDK stream events to SSE format for the proxy response. -/// -/// Auth priority: bearer token → IAM credentials → environment credentials -pub async fn handle_bedrock_sdk_streaming( - model: &str, - body: &Bytes, - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, - region: &str, -) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { - let openai_req: OpenAIRequest = serde_json::from_slice(body) - .map_err(|e| Error::internal_err(format!("Failed to parse OpenAI request: {}", e)))?; - - // Create Bedrock client using shared helper - let bedrock_client = create_bedrock_client( - api_key, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - region, - ) - .await?; - - // Convert messages using shared conversion - let enable_prompt_caching = - windmill_ai::ai_bedrock::bedrock_model_supports_prompt_caching(model); - let (bedrock_messages, system_prompts) = - windmill_ai::ai_bedrock::openai_messages_to_bedrock( - &openai_req.messages, - enable_prompt_caching, - )?; - - // Build inference configuration - let inference_config = windmill_ai::ai_bedrock::create_inference_config( - openai_req.temperature, - openai_req.max_tokens, - ); - - // Convert tools using shared helper - let tool_config = build_tool_config_from_request( - openai_req.tools.as_deref(), - openai_req.tool_choice.as_ref(), - enable_prompt_caching, - )?; - - // Build the SDK request - let mut request_builder = bedrock_client - .client() - .converse_stream() - .model_id(model) - .set_messages(Some(bedrock_messages)); - - if !system_prompts.is_empty() { - request_builder = request_builder.set_system(Some(system_prompts)); - } - - if let Some(config) = inference_config { - request_builder = request_builder.inference_config(config); - } - - if let Some(config) = tool_config { - request_builder = request_builder.set_tool_config(Some(config)); - } - - // Send the request and get the stream - tracing::debug!("Bedrock SDK streaming: sending converse_stream request"); - let stream_output = request_builder.send().await.map_err(|e| { - let error_msg = format!("Bedrock SDK streaming error: {}", format_bedrock_error(&e)); - tracing::error!("Bedrock SDK streaming failed: {}", error_msg); - Error::internal_err(error_msg) - })?; - tracing::debug!("Bedrock SDK streaming: stream established successfully"); - - // Convert SDK stream to SSE (pass the inner stream, not the full output) - let sse_stream = sdk_stream_to_sse(stream_output.stream, model.to_string()); - - // Build response headers - let mut response_headers = http::HeaderMap::new(); - response_headers.insert("content-type", "text/event-stream".parse().unwrap()); - response_headers.insert("cache-control", "no-cache".parse().unwrap()); - response_headers.insert("connection", "keep-alive".parse().unwrap()); - - Ok(( - http::StatusCode::OK, - response_headers, - axum::body::Body::from_stream(sse_stream), - )) -} - -/// Convert AWS SDK ConverseStream events to SSE format. -/// -/// Uses shared stream parsing functions from windmill_common::ai_bedrock -/// to extract text deltas and tool calls from the SDK stream events. -pub fn sdk_stream_to_sse( - stream: aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver< - aws_sdk_bedrockruntime::types::ConverseStreamOutput, - aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError, - >, - model: String, -) -> impl futures::Stream> + Send { - use std::collections::HashMap; - - let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); - let created = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs(); - - // State to track partial tool calls - struct StreamState { - id: String, - model: String, - created: u64, - tool_calls: HashMap, // index -> (id, name, args) - current_tool_index: usize, - } - - let state = std::sync::Arc::new(tokio::sync::Mutex::new(StreamState { - id: id.clone(), - model: model.clone(), - created, - tool_calls: HashMap::new(), - current_tool_index: 0, - })); - - async_stream::stream! { - let mut stream = stream; - let state = state.clone(); - - loop { - match stream.recv().await { - Ok(Some(event)) => { - let mut state = state.lock().await; - - // Handle tool use start - if let Some(tool_call) = bedrock_stream_event_to_tool_start(&event) { - let index = state.current_tool_index; - state.tool_calls.insert( - index, - (tool_call.id.clone(), tool_call.name.clone(), String::new()), - ); - - // Send initial tool call chunk - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "id": tool_call.id, - "type": "function", - "function": { - "name": tool_call.name, - "arguments": "" - } - }] - }, - "finish_reason": serde_json::Value::Null - }] - }); - - yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); - } - - // Handle text delta - if let Some(text) = bedrock_stream_event_to_text(&event) { - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": { - "content": text - }, - "finish_reason": serde_json::Value::Null - }] - }); - - yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); - } - - // Handle tool use input delta - if let Some(input_delta) = bedrock_stream_event_to_tool_delta(&event) { - let index = state.current_tool_index; - if let Some((_id, _name, ref mut args)) = state.tool_calls.get_mut(&index) { - args.push_str(&input_delta); - - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "function": { - "arguments": input_delta - } - }] - }, - "finish_reason": serde_json::Value::Null - }] - }); - - yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); - } - } - - // Handle content block stop - if bedrock_stream_event_is_block_stop(&event) { - state.current_tool_index += 1; - } - - // Handle message stop - if let aws_sdk_bedrockruntime::types::ConverseStreamOutput::MessageStop(stop) = &event { - let stop_reason = stop.stop_reason().as_str(); - let finish_reason = match stop_reason { - "end_turn" => "stop", - "max_tokens" => "length", - "tool_use" => "tool_calls", - "stop_sequence" => "stop", - "guardrail_intervened" | "content_filtered" => "content_filter", - _ => "stop", - }; - - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": {}, - "finish_reason": finish_reason - }] - }); - - yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); - } - } - Ok(None) => break, - Err(e) => { - yield Err(std::io::Error::new( - std::io::ErrorKind::Other, - e.to_string(), - )); - break; - } - } - } - - // Send [DONE] at the end - yield Ok(bytes::Bytes::from("data: [DONE]\n\n")); - } -} - -/// Handle non-streaming Bedrock request using the AWS SDK. -/// -/// Auth priority: bearer token → IAM credentials → environment credentials -pub async fn handle_bedrock_sdk_non_streaming( - model: &str, - body: &Bytes, - api_key: Option<&str>, - aws_access_key_id: Option<&str>, - aws_secret_access_key: Option<&str>, - aws_session_token: Option<&str>, - region: &str, -) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { - let openai_req: OpenAIRequest = serde_json::from_slice(body) - .map_err(|e| Error::internal_err(format!("Failed to parse OpenAI request: {}", e)))?; - - // Create Bedrock client using shared helper - let bedrock_client = create_bedrock_client( - api_key, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - region, - ) - .await?; - - // Convert messages using shared conversion - let enable_prompt_caching = - windmill_ai::ai_bedrock::bedrock_model_supports_prompt_caching(model); - let (bedrock_messages, system_prompts) = - windmill_ai::ai_bedrock::openai_messages_to_bedrock( - &openai_req.messages, - enable_prompt_caching, - )?; - - // Build inference configuration - let inference_config = windmill_ai::ai_bedrock::create_inference_config( - openai_req.temperature, - openai_req.max_tokens, - ); - - // Convert tools using shared helper - let tool_config = build_tool_config_from_request( - openai_req.tools.as_deref(), - openai_req.tool_choice.as_ref(), - enable_prompt_caching, - )?; - - // Build the SDK request (non-streaming) - let mut request_builder = bedrock_client - .client() - .converse() - .model_id(model) - .set_messages(Some(bedrock_messages)); - - if !system_prompts.is_empty() { - request_builder = request_builder.set_system(Some(system_prompts)); - } - - if let Some(config) = inference_config { - request_builder = request_builder.inference_config(config); - } - - if let Some(config) = tool_config { - request_builder = request_builder.set_tool_config(Some(config)); - } - - // Send the request - tracing::debug!("Bedrock SDK non-streaming: sending converse request"); - let response = request_builder.send().await.map_err(|e| { - let error_msg = format!( - "Bedrock SDK non-streaming error: {}", - format_bedrock_error(&e) - ); - tracing::error!("Bedrock SDK non-streaming failed: {}", error_msg); - Error::internal_err(error_msg) - })?; - tracing::debug!( - "Bedrock SDK non-streaming: response received, stop_reason={}", - response.stop_reason().as_str() - ); - - // Convert response to OpenAI format - let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); - let created = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs(); - - // Extract stop reason - let stop_reason = response.stop_reason().as_str(); - let finish_reason = match stop_reason { - "end_turn" => "stop", - "max_tokens" => "length", - "tool_use" => "tool_calls", - "stop_sequence" => "stop", - "guardrail_intervened" | "content_filtered" => "content_filter", - _ => "stop", - }; - - // Extract message content - let mut text_content = String::new(); - let mut tool_calls: Vec = Vec::new(); - - if let Some(output) = response.output() { - if let aws_sdk_bedrockruntime::types::ConverseOutput::Message(message) = output { - for block in message.content() { - match block { - aws_sdk_bedrockruntime::types::ContentBlock::Text(text) => { - text_content.push_str(text); - } - aws_sdk_bedrockruntime::types::ContentBlock::ToolUse(tool_use) => { - // Convert Document back to JSON string - let input_json = document_to_json(tool_use.input()); - tool_calls.push(OpenAIToolCall { - id: tool_use.tool_use_id().to_string(), - function: OpenAIFunction { - name: tool_use.name().to_string(), - arguments: serde_json::to_string(&input_json).unwrap_or_default(), - }, - r#type: "function".to_string(), - extra_content: None, - }); - } - _ => {} - } - } - } - } - - // Build the message - let message = if !tool_calls.is_empty() { - serde_json::json!({ - "role": "assistant", - "content": if text_content.is_empty() { serde_json::Value::Null } else { serde_json::Value::String(text_content) }, - "tool_calls": tool_calls - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": text_content - }) - }; - - // Extract usage information - let usage = if let Some(usage_data) = response.usage() { - serde_json::json!({ - "prompt_tokens": usage_data.input_tokens(), - "completion_tokens": usage_data.output_tokens(), - "total_tokens": usage_data.total_tokens() - }) - } else { - serde_json::json!({ - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - }) - }; - - // Build OpenAI-format response - let openai_resp = serde_json::json!({ - "id": id, - "object": "chat.completion", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": finish_reason - }], - "usage": usage - }); - - let response_body = serde_json::to_vec(&openai_resp) - .map_err(|e| Error::internal_err(format!("Failed to serialize OpenAI response: {}", e)))?; - - let mut response_headers = http::HeaderMap::new(); - response_headers.insert("content-type", "application/json".parse().unwrap()); - - Ok(( - http::StatusCode::OK, - response_headers, - axum::body::Body::from(response_body), - )) -} - -/// Convert AWS Smithy Document to serde_json::Value -fn document_to_json(doc: &aws_smithy_types::Document) -> serde_json::Value { - match doc { - aws_smithy_types::Document::Object(map) => { - let mut json_map = serde_json::Map::new(); - for (k, v) in map { - json_map.insert(k.clone(), document_to_json(v)); - } - serde_json::Value::Object(json_map) - } - aws_smithy_types::Document::Array(arr) => { - serde_json::Value::Array(arr.iter().map(document_to_json).collect()) - } - aws_smithy_types::Document::Number(num) => match num { - aws_smithy_types::Number::PosInt(n) => serde_json::Value::Number((*n).into()), - aws_smithy_types::Number::NegInt(n) => serde_json::Value::Number((*n).into()), - aws_smithy_types::Number::Float(f) => serde_json::json!(*f), - }, - aws_smithy_types::Document::String(s) => serde_json::Value::String(s.clone()), - aws_smithy_types::Document::Bool(b) => serde_json::Value::Bool(*b), - aws_smithy_types::Document::Null => serde_json::Value::Null, - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn determine_auth_config_prioritizes_bearer_token() { - let config = determine_auth_config( - Some("bearer-token"), - Some("AKIA123"), - Some("secret"), - Some("session-token"), - ); - - match config { - BedrockAuthConfig::BearerToken(token) => assert_eq!(token, "bearer-token"), - _ => panic!("expected bearer token auth config"), - } - } - - #[test] - fn determine_auth_config_uses_iam_with_optional_session_token() { - let config = - determine_auth_config(None, Some("AKIA123"), Some("secret"), Some("session-token")); - - match config { - BedrockAuthConfig::IamCredentials { - access_key_id, - secret_access_key, - session_token, - } => { - assert_eq!(access_key_id, "AKIA123"); - assert_eq!(secret_access_key, "secret"); - assert_eq!(session_token.as_deref(), Some("session-token")); - } - _ => panic!("expected IAM auth config"), - } - } - - #[test] - fn determine_auth_config_treats_empty_session_token_as_none() { - let config = determine_auth_config(None, Some("AKIA123"), Some("secret"), Some("")); - - match config { - BedrockAuthConfig::IamCredentials { session_token, .. } => { - assert!(session_token.is_none()); - } - _ => panic!("expected IAM auth config"), - } - } - - #[test] - fn determine_auth_config_falls_back_to_environment() { - let config = determine_auth_config(None, Some("AKIA123"), None, Some("session-token")); - assert!(matches!(config, BedrockAuthConfig::Environment)); - } -} diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index e7e4508971..8b2b9d8128 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -73,8 +73,6 @@ pub mod auth; #[cfg(all(feature = "private", feature = "parquet"))] pub mod azure_proxy_ee; mod azure_proxy_oss; -#[cfg(feature = "bedrock")] -mod bedrock; mod capture; mod concurrency_groups; mod db; diff --git a/backend/windmill-store/src/resources.rs b/backend/windmill-store/src/resources.rs index 5c17308bbf..3e5292e805 100644 --- a/backend/windmill-store/src/resources.rs +++ b/backend/windmill-store/src/resources.rs @@ -10,8 +10,8 @@ use std::collections::HashMap; use std::net::IpAddr; use windmill_api_auth::{ - check_scopes, maybe_refresh_folders, require_owner_of_path, require_super_admin, ApiAuthed, - Tokened, + build_scope_path_predicate, check_scopes, maybe_refresh_folders, require_owner_of_path, + require_super_admin, ApiAuthed, Tokened, }; use windmill_common::db::DB; use windmill_common::workspaces::{check_deploy_rules, RuleCheckResult}; @@ -194,6 +194,7 @@ async fn list_names( Extension(user_db): Extension, ) -> JsonResult> { let mut tx = user_db.begin(&authed).await?; + let allowed = build_scope_path_predicate(&authed, "resources", "read"); let rows = sqlx::query!( "SELECT value->>'name' as name, path from resource WHERE resource_type = $1 AND workspace_id = $2", rt, @@ -203,6 +204,7 @@ async fn list_names( .await? .into_iter() .filter_map(|x| x.name.map(|name| NamePath { name, path: x.path })) + .filter(|np| allowed(&np.path)) .collect::>(); tx.commit().await?; Ok(Json(rows)) @@ -225,6 +227,7 @@ async fn list_search_resources( #[cfg(not(feature = "enterprise"))] let n = 3; + let allowed = build_scope_path_predicate(&authed, "resources", "read"); let rows = sqlx::query_as!( SearchResource, "SELECT path, value from resource WHERE workspace_id = $1 LIMIT $2", @@ -234,6 +237,7 @@ async fn list_search_resources( .fetch_all(&mut *tx) .await? .into_iter() + .filter(|r| allowed(&r.path)) .collect::>(); tx.commit().await?; Ok(Json(rows)) @@ -338,9 +342,13 @@ async fn list_resources( let sql = sqlb.sql().map_err(|e| Error::internal_err(e.to_string()))?; let mut tx = user_db.begin(&authed).await?; + let allowed = build_scope_path_predicate(&authed, "resources", "read"); let rows = sqlx::query_as::<_, ListableResource>(&sql) .fetch_all(&mut *tx) - .await?; + .await? + .into_iter() + .filter(|r| allowed(&r.path)) + .collect::>(); tx.commit().await?; diff --git a/backend/windmill-store/src/variables.rs b/backend/windmill-store/src/variables.rs index 4c2a0ed670..2d767b393b 100644 --- a/backend/windmill-store/src/variables.rs +++ b/backend/windmill-store/src/variables.rs @@ -6,7 +6,10 @@ * LICENSE-AGPL for a copy of the license. */ -use windmill_api_auth::{check_scopes, maybe_refresh_folders, require_owner_of_path, ApiAuthed}; +use windmill_api_auth::{ + build_scope_path_predicate, check_scopes, maybe_refresh_folders, require_owner_of_path, + ApiAuthed, +}; use windmill_common::db::DB; use windmill_common::workspaces::{check_deploy_rules, RuleCheckResult}; @@ -188,9 +191,13 @@ async fn list_variables( let sql = sqlb.sql().map_err(|e| Error::internal_err(e.to_string()))?; let mut tx = user_db.begin(&authed).await?; + let allowed = build_scope_path_predicate(&authed, "variables", "read"); let rows = sqlx::query_as::<_, ListableVariable>(&sql) .fetch_all(&mut *tx) - .await?; + .await? + .into_iter() + .filter(|r| allowed(&r.path)) + .collect::>(); tx.commit().await?; Ok(Json(rows)) diff --git a/docs/windmill-ai-refactor-plan.md b/docs/windmill-ai-refactor-plan.md index ea2b9f451e..b69841d665 100644 --- a/docs/windmill-ai-refactor-plan.md +++ b/docs/windmill-ai-refactor-plan.md @@ -5,10 +5,10 @@ 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 (`ai.rs`, `google.rs`, `bedrock.rs`) with its own request building for Google/Bedrock, plus `AIRequestConfig::prepare_request` for auth/URL handling +- **windmill-api** — chat proxy routes (`ai.rs`), audit logging, caching, and DB-backed credential resolution through `AIRequestConfig` - **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. Both the API proxy and worker agent use `QueryBuilder` for every provider — no more duplicate logic. +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 @@ -25,7 +25,7 @@ windmill-common does **NOT** re-export from windmill-ai (would be circular). All ## Reviewer Note: Keep API Proxy Unification Split -The crate boundary, shared utilities, SSE parsers, image handling, and worker provider implementations are now in `windmill-ai`. The remaining duplication is the API proxy path: `AIRequestConfig::prepare_request`, `windmill-api/src/google.rs`, and `windmill-api/src/bedrock.rs` still own API-specific request transformation. +The crate boundary, shared utilities, SSE parsers, image handling, worker provider implementations, and provider-specific API proxy transformations are now in `windmill-ai`. The remaining duplication is credential shape and resolution: `windmill-api` still resolves DB-backed proxy credentials through `AIRequestConfig`, while worker agent execution still receives `ProviderWithResource`. 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. @@ -70,7 +70,7 @@ Follow-up status: Anthropic/Vertex proxy handling has since moved into `windmill-ai`, and the dead `AIRequestConfig::prepare_request` fallback has been removed. -## Current Phase PR: Proxy Execution Mode + Google AI Proxy Migration +## 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 @@ -101,6 +101,62 @@ Validation: - `cargo test -p windmill-api maps_request_config_to_provider_credentials` - `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. + +## Current Phase PR: 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. + ## Step-by-Step Plan Each step produces a compiling, working backend. @@ -200,9 +256,10 @@ Move `AI_HTTP_HEADERS` lazy_static (currently duplicated in `windmill-api/src/ai --- -### Step 8: Add proxy support to QueryBuilder — API uses QueryBuilder for all providers +### Step 8: Add API proxy execution support to windmill-ai ✅ -This is the key unification step. Add a new method to the `QueryBuilder` trait: +This is the key proxy unification step. HTTP-forwarding providers use +`QueryBuilder::build_proxy_request`: ```rust /// Build a request from a raw OpenAI-format proxy request. @@ -237,19 +294,21 @@ pub struct ProxyRequest { **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**: Convert OpenAI format → Gemini format (using existing `ai_google` functions). Replaces `windmill-api/src/google.rs`. -- **Bedrock**: Convert OpenAI format → Bedrock SDK calls. Replaces `windmill-api/src/bedrock.rs`. +- **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. Call `query_builder.build_proxy_request(&proxy_args)` → `ProxyRequest` -4. Send the request, return response with SSE keepalive injection +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 `GoogleAIQueryBuilder::build_proxy_request` -- `bedrock.rs` — replaced by `BedrockQueryBuilder::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 @@ -310,8 +369,8 @@ windmill-ai/src/ ├── 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 + build_proxy_request - ├── bedrock.rs # build_request + build_proxy_request (feature: bedrock) + ├── 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 ``` diff --git a/frontend/src/lib/components/FlowBuilder.svelte b/frontend/src/lib/components/FlowBuilder.svelte index 2dbdec772e..8c36de8b11 100644 --- a/frontend/src/lib/components/FlowBuilder.svelte +++ b/frontend/src/lib/components/FlowBuilder.svelte @@ -100,6 +100,7 @@ import { isRuleActive } from '$lib/workspaceProtectionRules.svelte' import { buildForkEditUrl } from '$lib/utils/editInFork' import { isCloudHosted } from '$lib/cloud' + import { UserDraft } from '$lib/userDraft.svelte' let { initialPath = $bindable(''), @@ -123,6 +124,7 @@ children, loadedFromHistoryFromUrl, noInitial = false, + liveEditorDraftStoragePath = undefined, onSaveInitial, onSaveDraft, onDeploy, @@ -588,6 +590,23 @@ const flowEditorDrawer = writable(undefined) const history = initHistory(untrack(() => flowStore).val) const pathStore = writable(untrack(() => pathStoreInit) ?? initialPath) + + $effect(() => { + if (liveEditorDraftStoragePath === undefined || !$workspaceStore) return + const workspace = $workspaceStore + UserDraft.setLiveEditorDraft({ + workspace, + itemKind: 'flow', + storagePath: liveEditorDraftStoragePath, + effectivePath: $pathStore + }) + return () => + UserDraft.clearLiveEditorDraft('flow', { + workspace, + storagePath: liveEditorDraftStoragePath + }) + }) + const captureOn = writable(false) const showCaptureHint = writable(undefined) const flowInputEditorStateStore = writable({ diff --git a/frontend/src/lib/components/Path.svelte b/frontend/src/lib/components/Path.svelte index 8306c8a410..b4bd67632e 100644 --- a/frontend/src/lib/components/Path.svelte +++ b/frontend/src/lib/components/Path.svelte @@ -5,7 +5,7 @@ diff --git a/frontend/src/lib/components/sessions/appDraftCodec.ts b/frontend/src/lib/components/sessions/appDraftCodec.ts index 683d075781..77eb881e03 100644 --- a/frontend/src/lib/components/sessions/appDraftCodec.ts +++ b/frontend/src/lib/components/sessions/appDraftCodec.ts @@ -1,10 +1,21 @@ import type { RawAppData } from '$lib/components/raw_apps/dataTableRefUtils' -import type { AppDraftValue } from '$lib/components/copilot/chat/global/draftStore.svelte' + +// The raw-app draft shape stored under `UserDraft` — matches the +// regular `/apps_raw/edit` route's UserDraft handle exactly. The chat's +// `userDraftAdapter.saveGlobalAppDraft` writes through the same shape, so +// session previews and the chat round-trip identically. +export type RawAppDraft = { + files: Record + runnables: Record + data: RawAppData + summary: string + policy?: any + custom_path?: string +} // The shape `runtime.rawApp.val` actually holds (see SessionRuntime in -// sessionRuntime.svelte.ts lines 74-84). Slightly flatter than the AI's -// `AppDraftValue`: `path` is metadata not present on the AI side, and -// `summary` is required here. +// sessionRuntime.svelte.ts). Adds `path` (a key, not a draft field) and +// makes `policy` required for the editor's live binding. export type RuntimeRawApp = { summary: string path: string @@ -14,28 +25,27 @@ export type RuntimeRawApp = { policy: any } -// Strip runtime metadata (just `path` for raw apps) and project into the -// AI-facing `AppDraftValue` envelope. -export function rawAppToDraftValue(raw: RuntimeRawApp): AppDraftValue { +// Strip runtime-only metadata (just `path`, the storage key) when persisting +// to UserDraft. +export function runtimeRawAppToDraft(raw: RuntimeRawApp): RawAppDraft { return { summary: raw.summary, files: raw.files, runnables: raw.runnables, data: raw.data, policy: raw.policy - // custom_path is read-only and not held on the runtime. } } -// Overlay an AI-produced draft onto an existing runtime raw app, -// preserving metadata fields (path) that don't live in `AppDraftValue`. -export function applyDraftValueToRawApp(raw: RuntimeRawApp, dv: AppDraftValue): RuntimeRawApp { +// Overlay a UserDraft-stored raw-app draft onto an existing runtime raw app, +// preserving the runtime-only `path` field. +export function applyDraftToRuntimeRawApp(raw: RuntimeRawApp, dv: RawAppDraft): RuntimeRawApp { return { ...raw, - summary: dv.summary ?? raw.summary, + summary: dv.summary, files: dv.files, runnables: dv.runnables, - data: (dv.data as RawAppData | undefined) ?? raw.data, + data: dv.data, policy: dv.policy ?? raw.policy } } diff --git a/frontend/src/lib/components/sessions/flowDraftCodec.ts b/frontend/src/lib/components/sessions/flowDraftCodec.ts deleted file mode 100644 index 66f4f2c617..0000000000 --- a/frontend/src/lib/components/sessions/flowDraftCodec.ts +++ /dev/null @@ -1,28 +0,0 @@ -import type { Flow } from '$lib/gen' -import type { FlowDraftValue } from '$lib/components/copilot/chat/global/draftStore.svelte' - -// Convert the editor's full `Flow` (carrying metadata like path, edited_by, -// edited_at, archived, etc.) into the slimmer `FlowDraftValue` shape the -// global AI chat's draft store uses. Metadata stays on the runtime side — -// the draft store only holds what the AI's tools need to round-trip the -// in-flight edit. -export function flowToDraftValue(flow: Flow): FlowDraftValue { - return { - value: flow.value, - schema: flow.schema ?? null, - groups: flow.value.groups ?? null - } -} - -// Overlay a draft from the store onto an existing `Flow`. Preserves the -// metadata fields that aren't in `FlowDraftValue` (path, edited_by, -// edited_at, archived, extra_perms, …). `groups` lives inside -// `FlowValue`, so it rides along on `dv.value` automatically; the -// sibling-key on `FlowDraftValue` is purely for the AI's tool I/O. -export function applyDraftValueToFlow(flow: Flow, dv: FlowDraftValue): Flow { - return { - ...flow, - value: dv.value, - schema: dv.schema ?? flow.schema - } -} diff --git a/frontend/src/lib/components/sessions/sessionRuntime.svelte.ts b/frontend/src/lib/components/sessions/sessionRuntime.svelte.ts index 153cf5a725..4a2418ce3a 100644 --- a/frontend/src/lib/components/sessions/sessionRuntime.svelte.ts +++ b/frontend/src/lib/components/sessions/sessionRuntime.svelte.ts @@ -28,13 +28,8 @@ import { type Session, type SessionTarget } from './sessionState.svelte' -import { - globalDraftStore, - type AppDraftValue, - type FlowDraftValue -} from '$lib/components/copilot/chat/global/draftStore.svelte' -import { applyDraftValueToFlow, flowToDraftValue } from './flowDraftCodec' -import { applyDraftValueToRawApp, rawAppToDraftValue } from './appDraftCodec' +import { UserDraft } from '$lib/userDraft.svelte' +import { applyDraftToRuntimeRawApp, runtimeRawAppToDraft, type RawAppDraft } from './appDraftCodec' import { setOpenPreviewHandler } from '$lib/components/copilot/chat/global/core' export interface SessionRuntime { @@ -211,21 +206,14 @@ function createRuntime(session: Session): SessionRuntime { loadingFlow = true notFound = false try { - // Draft first. globalDraftStore is the authoritative content - // source: the AI writes through it (write_flow / patch_flow_json - // / set_flow_module_code) and the editor's outbound $effect - // mirrors user edits back into it. If a draft exists we render - // from it, even when the path has never been deployed. - const aiDraft = globalDraftStore.getFlowDraft(workspace, path) - const draftValue = - aiDraft && - aiDraft.value && - typeof aiDraft.value === 'object' && - 'value' in (aiDraft.value as object) - ? (aiDraft.value as FlowDraftValue) - : undefined + // Draft first. UserDraft is the shared authoritative content + // source — the chat (write_flow / patch_flow_json / + // set_flow_module_code) and the editor's outbound $effect both + // write through it. If a draft exists we render from it, even + // when the path has never been deployed. + const aiDraft = UserDraft.get('flow', path, { workspace }) - if (draftValue) { + if (aiDraft) { // Best-effort fetch the backend baseline for the diff // drawer. Don't fail the load if the path doesn't exist // yet on the backend — draft-only flows are a valid state. @@ -235,18 +223,7 @@ function createRuntime(session: Session): SessionRuntime { } catch { savedFlow.val = undefined } - const skeleton: Flow = (savedFlow.val as Flow | undefined) ?? { - path, - summary: '', - value: { modules: [] }, - edited_by: '', - edited_at: '', - archived: false, - extra_perms: {}, - schema: emptySchema() - } - const flow = applyDraftValueToFlow(skeleton, draftValue) - await initFlow(flow, flowStore, flowStateStore) + await initFlow(aiDraft, flowStore, flowStateStore) loadedPath = path return } @@ -256,13 +233,7 @@ function createRuntime(session: Session): SessionRuntime { const result = await FlowService.getFlowByPathWithDraft({ workspace, path }) savedFlow.val = result const flow: Flow = (result.draft as Flow | undefined) ?? (result as Flow) - globalDraftStore.setDraft(workspace, { - type: 'flow', - path, - summary: flow.summary, - value: flowToDraftValue(flow), - isDraft: true - }) + UserDraft.save('flow', path, flow, { workspace }) await initFlow(flow, flowStore, flowStateStore) loadedPath = path } catch (err) { @@ -290,16 +261,14 @@ function createRuntime(session: Session): SessionRuntime { loadingScript = true notFoundScript = false try { - // Draft first. globalDraftStore is the authoritative content - // source: the AI writes through it (write_script / edit_script) - // and the editor's outbound $effect mirrors user edits back - // into it. If a draft exists we render from it, even when the - // path has never been deployed. - const aiDraft = globalDraftStore.getScriptDraft(workspace, path) - const draftContent = - aiDraft && typeof aiDraft.value === 'string' ? aiDraft.value : undefined + // Draft first. UserDraft is the shared authoritative content + // source — the chat (write_script / edit_script) and the + // editor's outbound $effect both write through it. If a draft + // exists we render from it, even when the path has never been + // deployed. + const aiDraft = UserDraft.get('script', path, { workspace }) - if (aiDraft && draftContent !== undefined) { + if (aiDraft && typeof aiDraft.content === 'string') { // Best-effort fetch the backend baseline for the diff // drawer + parent_hash. 404 means draft-only — leave // savedScript undefined and skip parent_hash. @@ -322,7 +291,7 @@ function createRuntime(session: Session): SessionRuntime { if (savedScript.val?.hash) { baseline.parent_hash = savedScript.val.hash } - baseline.content = draftContent + baseline.content = aiDraft.content if (aiDraft.language) baseline.language = aiDraft.language if (aiDraft.summary !== undefined) baseline.summary = aiDraft.summary scriptStore.val = baseline @@ -335,14 +304,7 @@ function createRuntime(session: Session): SessionRuntime { savedScript.val = result const baseline = (result.draft as NewScript | undefined) ?? (result as NewScript) baseline.parent_hash = result.hash - globalDraftStore.setDraft(workspace, { - type: 'script', - path, - language: baseline.language, - summary: baseline.summary, - value: baseline.content ?? '', - isDraft: true - }) + UserDraft.save('script', path, baseline, { workspace }) scriptStore.val = baseline loadedScriptPath = path } catch (err) { @@ -425,21 +387,14 @@ function createRuntime(session: Session): SessionRuntime { loadingRawApp = true notFoundRawApp = false try { - // Draft first. globalDraftStore is the authoritative content - // source: the AI writes through it (init_app / write_app_file - // / ...) and the editor's outbound $effect mirrors user edits - // back into it. If a draft exists we render from it, even - // when the path has never been deployed. - const aiDraft = globalDraftStore.getAppDraft(workspace, path) - const draftValue = - aiDraft && - aiDraft.value && - typeof aiDraft.value === 'object' && - 'files' in (aiDraft.value as object) - ? (aiDraft.value as AppDraftValue) - : undefined + // Draft first. UserDraft is the shared authoritative content + // source — the chat (init_app / write_app_file / ...) and the + // editor's outbound $effect both write through it. If a draft + // exists we render from it, even when the path has never been + // deployed. + const aiDraft = UserDraft.get('raw_app', path, { workspace }) - if (draftValue) { + if (aiDraft) { // Best-effort fetch the backend baseline for the diff // drawer. Don't fail the load if the path doesn't exist // yet on the backend — draft-only apps are a valid state. @@ -457,16 +412,16 @@ function createRuntime(session: Session): SessionRuntime { } catch { savedRawApp.val = undefined } - rawApp.val = applyDraftValueToRawApp( + rawApp.val = applyDraftToRuntimeRawApp( { files: {}, runnables: {}, data: { ...DEFAULT_DATA }, policy: undefined, - summary: draftValue.summary ?? '', + summary: aiDraft.summary ?? '', path }, - draftValue + aiDraft ) loadedRawAppPath = path return @@ -510,13 +465,7 @@ function createRuntime(session: Session): SessionRuntime { summary: result.summary ?? '', path: result.path } - globalDraftStore.setDraft(workspace, { - type: 'app', - path, - summary: runtimeValue.summary, - value: rawAppToDraftValue(runtimeValue), - isDraft: true - }) + UserDraft.save('raw_app', path, runtimeRawAppToDraft(runtimeValue), { workspace }) rawApp.val = runtimeValue loadedRawAppPath = path } catch (err) { diff --git a/frontend/src/lib/userDraft.svelte.ts b/frontend/src/lib/userDraft.svelte.ts index 3d83952513..03db43dc1a 100644 --- a/frontend/src/lib/userDraft.svelte.ts +++ b/frontend/src/lib/userDraft.svelte.ts @@ -130,7 +130,26 @@ export type UserDraftEntry = { live: boolean } +export type LiveEditorDraft = { + workspace: string + itemKind: UserDraftItemKind + storagePath: string + effectivePath?: string +} + +export type LiveEditorDraftSpec = { + itemKind: UserDraftItemKind + storagePath: string + effectivePath?: string + workspace?: string +} + +export type ClearLiveEditorDraftOptions = UserDraftOptions & { + storagePath?: string +} + const entries = new Map() +const liveEditorDrafts = new Map() function resolveWorkspace(opts?: UserDraftOptions): string { const ws = opts?.workspace ?? get(workspaceStore) @@ -230,6 +249,10 @@ function localStorageKey(workspace: string, itemKind: UserDraftItemKind, path: s return `userdraft/w/${workspace}/${itemKind}/${path}` } +function liveEditorDraftKey(workspace: string, itemKind: UserDraftItemKind): string { + return `${workspace}/${itemKind}` +} + function parseLocalStorageKey( key: string, workspace: string, @@ -331,10 +354,13 @@ export const UserDraft = { const mk = mapKey(ws, itemKind, path) const entry = entries.get(mk) if (entry) { - // Notify observers; preserve existing rev metadata. `untrack`ed - // read — see `set draft` below for why. + // Static writes are external mutations. Update live observers and + // force the storage slot to match, even if the live entry still has + // its initial-write skip armed. const current = untrack(() => entry.state.val as StoredDraft | undefined) - entry.state.val = wrap(value, extractMeta(current)) + const meta = extractMeta(current) + entry.state.setWithoutPersist(wrap(value, meta)) + persistDirect(localStorageKey(ws, itemKind, path), value, meta) return } // No live handle: preserve any persisted meta so the staleness @@ -361,10 +387,10 @@ export const UserDraft = { const mk = mapKey(ws, itemKind, path) const entry = entries.get(mk) if (entry) { - entry.state.val = wrap(value, meta) // Static writes represent explicit external draft mutations. A // freshly acquired live entry may still have the initial-write skip // armed, so force the storage slot to match the live value. + entry.state.setWithoutPersist(wrap(value, meta)) persistDirect(localStorageKey(ws, itemKind, path), value, meta) return } @@ -401,9 +427,9 @@ export const UserDraft = { const mk = mapKey(ws, itemKind, path) const entry = entries.get(mk) if (entry) { - return unwrap(entry.state.val as StoredDraft | undefined) + return snapshotDraftValue(unwrap(entry.state.val as StoredDraft | undefined)) } - return unwrap(readPersisted(localStorageKey(ws, itemKind, path))) + return snapshotDraftValue(unwrap(readPersisted(localStorageKey(ws, itemKind, path)))) }, /** @@ -523,6 +549,34 @@ export const UserDraft = { return Array.from(out.values()) }, + setLiveEditorDraft(spec: LiveEditorDraftSpec): void { + const ws = resolveWorkspace({ workspace: spec.workspace }) + liveEditorDrafts.set(liveEditorDraftKey(ws, spec.itemKind), { + workspace: ws, + itemKind: spec.itemKind, + storagePath: spec.storagePath, + effectivePath: spec.effectivePath || undefined + }) + }, + + getLiveEditorDraft( + itemKind: UserDraftItemKind, + opts?: UserDraftOptions + ): LiveEditorDraft | undefined { + const ws = resolveWorkspace(opts) + const draft = liveEditorDrafts.get(liveEditorDraftKey(ws, itemKind)) + return draft ? { ...draft } : undefined + }, + + clearLiveEditorDraft(itemKind: UserDraftItemKind, opts?: ClearLiveEditorDraftOptions): void { + const ws = resolveWorkspace(opts) + const key = liveEditorDraftKey(ws, itemKind) + const draft = liveEditorDrafts.get(key) + if (!draft) return + if (opts?.storagePath !== undefined && draft.storagePath !== opts.storagePath) return + liveEditorDrafts.delete(key) + }, + /** * Like `remove`, but also resets any live handle's `draft` to * `fallback` in-memory (so reactive readers see it immediately) and @@ -802,4 +856,5 @@ export function gcUserDrafts(maxAgeMs: number = USER_DRAFT_GC_MAX_AGE_MS): void /** Test-only: clear all in-memory entries. */ export function __resetUserDraftForTesting(): void { entries.clear() + liveEditorDrafts.clear() } diff --git a/frontend/src/lib/userDraft.test.ts b/frontend/src/lib/userDraft.test.ts index f703953f11..e2dd97418f 100644 --- a/frontend/src/lib/userDraft.test.ts +++ b/frontend/src/lib/userDraft.test.ts @@ -18,6 +18,7 @@ vi.mock('svelte', async (importOriginal) => { const { UserDraft, normalizeForCompare, localDraftDiffers, __resetUserDraftForTesting } = await import('./userDraft.svelte') const { workspaceStore } = await import('./stores') +const { deleteGlobalDraft } = await import('./components/copilot/chat/global/userDraftAdapter') function flushDestroyCallbacks(): void { const callbacks = onDestroyCallbacks.splice(0, onDestroyCallbacks.length) @@ -119,6 +120,78 @@ describe('UserDraft.save / get / remove (no observers)', () => { }) }) +describe('UserDraft live editor draft registry', () => { + it('stores the live editor storage path and effective path per workspace and kind', () => { + UserDraft.setLiveEditorDraft({ + itemKind: 'script', + storagePath: '', + effectivePath: 'u/me/generated_script' + }) + + expect(UserDraft.getLiveEditorDraft('script')).toEqual({ + workspace: 'test_ws', + itemKind: 'script', + storagePath: '', + effectivePath: 'u/me/generated_script' + }) + }) + + it('keeps live editor registrations isolated by workspace', () => { + UserDraft.setLiveEditorDraft({ + workspace: 'ws_a', + itemKind: 'flow', + storagePath: '', + effectivePath: 'u/me/a' + }) + UserDraft.setLiveEditorDraft({ + workspace: 'ws_b', + itemKind: 'flow', + storagePath: '', + effectivePath: 'u/me/b' + }) + + expect(UserDraft.getLiveEditorDraft('flow', { workspace: 'ws_a' })?.effectivePath).toBe( + 'u/me/a' + ) + expect(UserDraft.getLiveEditorDraft('flow', { workspace: 'ws_b' })?.effectivePath).toBe( + 'u/me/b' + ) + }) + + it('clears only the matching live editor storage path when provided', () => { + UserDraft.setLiveEditorDraft({ + itemKind: 'raw_app', + storagePath: '', + effectivePath: 'u/me/live_app' + }) + + UserDraft.clearLiveEditorDraft('raw_app', { storagePath: 'u/me/other' }) + expect(UserDraft.getLiveEditorDraft('raw_app')).toBeDefined() + + UserDraft.clearLiveEditorDraft('raw_app', { storagePath: '' }) + expect(UserDraft.getLiveEditorDraft('raw_app')).toBeUndefined() + }) + + it('can remove persisted global draft storage without blanking the live editor', () => { + const draft = { path: 'u/me/live_script', content: 'export async function main() {}' } + localStorage.setItem('userdraft/w/test_ws/script/', wrapped(draft)) + const handle = UserDraft.use('script', '') + UserDraft.setLiveEditorDraft({ + itemKind: 'script', + storagePath: '', + effectivePath: 'u/me/live_script' + }) + + deleteGlobalDraft('test_ws', 'script', 'u/me/live_script', undefined, { + preserveLiveDraft: true + }) + flushPersist() + + expect(handle.draft).toEqual(draft) + expect(localStorage.getItem('userdraft/w/test_ws/script/')).toBeNull() + }) +}) + describe('UserDraft.use() — observer sync', () => { it('loads the existing localStorage value on first use', () => { localStorage.setItem('userdraft/w/test_ws/flow/u/me/loaded', wrapped('preloaded')) @@ -138,24 +211,29 @@ describe('UserDraft.use() — observer sync', () => { expect(a.draft).toBe(99) }) - it('save() propagates to live use() handles (in-memory)', () => { + it('save() propagates to live use() handles and persists immediately', () => { const handle = UserDraft.use('flow', 'u/me/observed') expect(handle.draft).toBeUndefined() - // First write through a live entry is treated as the "initial value" - // (saveInitialValue=false) and is NOT persisted — observers still see it. UserDraft.save('flow', 'u/me/observed', 7) expect(handle.draft).toBe(7) - flushPersist() - expect(localStorage.getItem('userdraft/w/test_ws/flow/u/me/observed')).toBeNull() + expect(storedShape('userdraft/w/test_ws/flow/u/me/observed')).toBe(wrapped(7)) - // Subsequent writes persist. UserDraft.save('flow', 'u/me/observed', 9) expect(handle.draft).toBe(9) - flushPersist() expect(storedShape('userdraft/w/test_ws/flow/u/me/observed')).toBe(wrapped(9)) }) + it('get() returns a cloneable snapshot of live handle values', () => { + const handle = UserDraft.use<{ path: string; nested: { value: number } }>('script', '') + handle.draft = { path: 'u/me/live', nested: { value: 1 } } + + const draft = UserDraft.get<{ path: string; nested: { value: number } }>('script', '') + expect(draft).toEqual({ path: 'u/me/live', nested: { value: 1 } }) + expect(draft).not.toBe(handle.draft) + expect(() => structuredClone(draft)).not.toThrow() + }) + it('remove() clears localStorage without touching the in-memory handle', () => { // Seed localStorage so the live handle initialises from it. localStorage.setItem('userdraft/w/test_ws/flow/u/me/removed', wrapped(1)) @@ -411,6 +489,28 @@ describe('UserDraft — rev metadata for staleness checks', () => { ) }) + it('UserDraft.save persists immediately when a live handle exists', () => { + const handle = UserDraft.use('flow', 'u/me/live-save') + + UserDraft.save('flow', 'u/me/live-save', 'external') + + expect(handle.draft).toBe('external') + expect(storedShape('userdraft/w/test_ws/flow/u/me/live-save')).toBe(wrapped('external')) + }) + + it('UserDraft.save preserves live rev metadata while forcing persistence', () => { + const handle = UserDraft.use('flow', 'u/me/live-save-meta') + handle.setDraftAndMeta('baseline', { remoteRev: 5 }) + expect(localStorage.getItem('userdraft/w/test_ws/flow/u/me/live-save-meta')).toBeNull() + + UserDraft.save('flow', 'u/me/live-save-meta', 'external') + + expect(handle.draft).toBe('external') + expect(storedShape('userdraft/w/test_ws/flow/u/me/live-save-meta')).toBe( + JSON.stringify({ value: 'external', remoteRev: 5 }) + ) + }) + it('handle.meta is empty for a draft persisted without rev (forward compat with older entries)', () => { localStorage.setItem( 'userdraft/w/test_ws/flow/u/me/legacy', @@ -503,8 +603,7 @@ describe('UserDraft.use() — reference counting & cleanup', () => { UserDraft.save('flow', 'u/me/ref', 2) expect(a.draft).toBe(2) - // Now persisted (second write after the baseline). - flushPersist() + // External save() calls persist immediately, even with a live handle. expect(storedShape('userdraft/w/test_ws/flow/u/me/ref')).toBe(wrapped(2)) // Releasing the second handle drops the entry; subsequent save() @@ -878,9 +977,7 @@ describe('UserDraft.list / clear / setDraftAndMeta', () => { ) flushPersist() - expect(storedShape(key)).toBe( - wrapped({ path: 'f/rewrite-after-clear', content: 'new' }) - ) + expect(storedShape(key)).toBe(wrapped({ path: 'f/rewrite-after-clear', content: 'new' })) }) it('list hides persisted drafts when a live handle has cleared the value', () => { diff --git a/frontend/src/routes/(root)/(logged)/apps_raw/add/+page.svelte b/frontend/src/routes/(root)/(logged)/apps_raw/add/+page.svelte index fb68f09f1b..716000215f 100644 --- a/frontend/src/routes/(root)/(logged)/apps_raw/add/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/apps_raw/add/+page.svelte @@ -79,6 +79,8 @@ runnables: Record data: RawAppData summary: string + policy?: Policy + custom_path?: string }>('raw_app', '') // Restore the persisted autosave so a plain reload of /apps_raw/add // resumes the last session. Captured once; the $effect below mirrors @@ -114,13 +116,15 @@ let summary = $state(restoredDraft?.summary ?? '') let files: Record = $state(restoredDraft?.files ?? react19Template) - let policy: Policy = $state({ - on_behalf_of: $userStore?.username.includes('@') - ? $userStore?.username - : `u/${$userStore?.username}`, - on_behalf_of_email: $userStore?.email, - execution_mode: 'publisher' - }) + let policy: Policy = $state( + restoredDraft?.policy ?? { + on_behalf_of: $userStore?.username.includes('@') + ? $userStore?.username + : `u/${$userStore?.username}`, + on_behalf_of_email: $userStore?.email, + execution_mode: 'publisher' + } + ) let runnables: Record = $state(restoredDraft?.runnables ?? defaultRunnables) /** Data configuration including tables and creation policy */ @@ -133,13 +137,14 @@ readFieldsRecursively(files) readFieldsRecursively(runnables) readFieldsRecursively(data) + readFieldsRecursively(policy) void summary untrack(() => { if (firstMirror) { firstMirror = false draftHandle.setDraftAndMeta(undefined, {}) } - draftHandle.draft = { files, runnables, data, summary } + draftHandle.draft = { files, runnables, data, summary, policy } }) }) @@ -150,11 +155,12 @@ const d = draftHandle.draft if (d == null) return untrack(() => { - if (localDraftDiffers(d, { files, runnables, data, summary })) { + if (localDraftDiffers(d, { files, runnables, data, summary, policy })) { files = d.files runnables = d.runnables data = d.data summary = d.summary + if (d.policy !== undefined) policy = d.policy } }) }) @@ -666,6 +672,7 @@ bind:data {policy} path={''} + liveEditorDraftStoragePath="" bind:summary newApp /> diff --git a/frontend/src/routes/(root)/(logged)/apps_raw/edit/[...path]/+page.svelte b/frontend/src/routes/(root)/(logged)/apps_raw/edit/[...path]/+page.svelte index 93417aa3f4..84ce23e7a5 100644 --- a/frontend/src/routes/(root)/(logged)/apps_raw/edit/[...path]/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/apps_raw/edit/[...path]/+page.svelte @@ -32,6 +32,8 @@ runnables: Record data: RawAppData summary: string + policy?: any + custom_path?: string } let files: Record | undefined = $state(undefined) @@ -105,25 +107,45 @@ // Persist the bundle whenever any of the four pieces of state changes. $effect(() => { - if (!files) return - readFieldsRecursively(files) + const currentFiles = files + if (!currentFiles) return + readFieldsRecursively(currentFiles) readFieldsRecursively(runnables) readFieldsRecursively(data) + readFieldsRecursively(policy) void summary - draftHandle.draft = { files, runnables, data, summary } + draftHandle.draft = { + files: currentFiles, + runnables, + data, + summary, + policy, + custom_path: savedApp?.custom_path + } }) // Reflect an external UserDraft.save into the form. Idempotent; the // `!files` guard skips the reload window so it doesn't fight loadApp. $effect(() => { const d = draftHandle.draft - if (d == null || !files) return + const currentFiles = files + if (d == null || !currentFiles) return untrack(() => { - if (localDraftDiffers(d, { files, runnables, data, summary })) { + if ( + localDraftDiffers(d, { + files: currentFiles, + runnables, + data, + summary, + policy, + custom_path: savedApp?.custom_path + }) + ) { files = d.files runnables = d.runnables data = d.data summary = d.summary + if (d.policy !== undefined) policy = d.policy } }) }) @@ -192,7 +214,9 @@ (backendSource.value?.datatables ? { ...DEFAULT_DATA, tables: backendSource.value.datatables } : { ...DEFAULT_DATA }), - summary: backendSource.summary ?? '' + summary: backendSource.summary ?? '', + policy: backendSource.policy ?? app_w_draft.policy, + custom_path: backendSource.custom_path ?? app_w_draft.custom_path } if ( @@ -240,7 +264,7 @@ runnables = localDraft.runnables data = localDraft.data summary = localDraft.summary - policy = app_w_draft.policy + policy = localDraft.policy ?? app_w_draft.policy newPath = app_w_draft.path files = localDraft.files } else { @@ -370,6 +394,7 @@ bind:summary {newPath} path={page.params.path ?? ''} + liveEditorDraftStoragePath={path} {policy} bind:savedApp {diffDrawer} diff --git a/frontend/src/routes/(root)/(logged)/flows/add/+page.svelte b/frontend/src/routes/(root)/(logged)/flows/add/+page.svelte index cbeb9becb5..0af3ea5b3c 100644 --- a/frontend/src/routes/(root)/(logged)/flows/add/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/flows/add/+page.svelte @@ -200,6 +200,7 @@ onNavigate={(item) => goto(editPathFor(item))} {initialPath} {pathStoreInit} + liveEditorDraftStoragePath="" bind:this={flowBuilder} newFlow {initialArgs} diff --git a/frontend/src/routes/(root)/(logged)/flows/edit/[...path]/+page.svelte b/frontend/src/routes/(root)/(logged)/flows/edit/[...path]/+page.svelte index c33d6b6278..7d3d3c14b4 100644 --- a/frontend/src/routes/(root)/(logged)/flows/edit/[...path]/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/flows/edit/[...path]/+page.svelte @@ -367,6 +367,7 @@ {flowStore} {flowStateStore} initialPath={page.params.path ?? ''} + liveEditorDraftStoragePath={flowDraftPath} newFlow={false} {selectedId} {initialArgs} diff --git a/frontend/src/routes/(root)/(logged)/global_drafts/+page.svelte b/frontend/src/routes/(root)/(logged)/global_drafts/+page.svelte index 4b7f8e8d0c..d626c084bc 100644 --- a/frontend/src/routes/(root)/(logged)/global_drafts/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/global_drafts/+page.svelte @@ -1,9 +1,11 @@ @@ -41,9 +65,9 @@
-

Global AI drafts

+

Global local drafts

- Dev-only inspector for the in-memory global draft store. + Dev-only inspector for global local drafts.