From 26be1d7da8be48c2e7d7a6b95d406513d649be4e Mon Sep 17 00:00:00 2001 From: Diego Imbert <70353967+diegoimbert@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:25:49 +0200 Subject: [PATCH] feat: add test key button to AI resource drawers (#11466) * feat: add test key button to AI resource drawers * fix: scope-check inline AI resource values and keep test model editable * fix: keep test model editable for unsaved resources, own-key provider check --- .../tests/ai_routes.rs | 74 ++++++++++++ backend/windmill-api/src/ai.rs | 40 +++++++ .../src/lib/components/ApiConnectForm.svelte | 13 +++ .../src/lib/components/ResourceForm.svelte | 12 ++ .../lib/components/copilot/TestAIKey.svelte | 105 +++++++++++------- .../components/copilot/aiResourceProvider.ts | 10 ++ .../copilot/chat/openai-responses.ts | 20 +--- frontend/src/lib/components/copilot/lib.ts | 69 +++++++----- 8 files changed, 263 insertions(+), 80 deletions(-) create mode 100644 frontend/src/lib/components/copilot/aiResourceProvider.ts diff --git a/backend/windmill-api-integration-tests/tests/ai_routes.rs b/backend/windmill-api-integration-tests/tests/ai_routes.rs index cfd93caef5..79152d82f7 100644 --- a/backend/windmill-api-integration-tests/tests/ai_routes.rs +++ b/backend/windmill-api-integration-tests/tests/ai_routes.rs @@ -232,3 +232,77 @@ async fn test_ai_proxy_x_resource_path_enforces_rls(db: Pool) -> anyho Ok(()) } + +/// An unsaved resource value sent in X-Resource-Value is what the upstream is called +/// with, and a scoped token cannot use it to read a variable outside its scopes. +#[sqlx::test(migrations = "../migrations", fixtures("base"))] +async fn test_ai_proxy_x_resource_value(db: Pool) -> anyhow::Result<()> { + use axum::{http::HeaderMap, http::StatusCode, routing::post, Json, Router}; + + initialize_tracing().await; + std::env::set_var("ALLOW_PRIVATE_AI_BASE_URLS", "true"); + let server = ApiServer::start(db.clone()).await?; + let port = server.addr.port(); + + // Only answers the key carried by the header, so another credential source fails. + let app = Router::new().fallback(post(|headers: HeaderMap| async move { + if headers.get("authorization").and_then(|v| v.to_str().ok()) == Some("Bearer sk-unsaved") { + Ok(Json(json!({ + "id": "chatcmpl-test", + "object": "chat.completion", + "choices": [{"message": {"role": "assistant", "content": "hello"}}] + }))) + } else { + Err(StatusCode::UNAUTHORIZED) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let base_url = format!("http://127.0.0.1:{}/v1", listener.local_addr()?.port()); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + + sqlx::query( + "INSERT INTO token (token_hash, token_prefix, token, email, label, super_admin, scopes) \ + VALUES (encode(sha256('SCOPED_AI_TOKEN'::bytea), 'hex'), 'SCOPED_AI_', 'SCOPED_AI_TOKEN', \ + 'test@windmill.dev', 'scoped', true, ARRAY['ai:write', 'resources:write'])", + ) + .execute(&db) + .await?; + + let send = |api_key: &'static str, token: &'static str| { + let value = json!({ "api_key": api_key, "base_url": base_url }).to_string(); + let encoded: String = value.bytes().map(|b| format!("%{b:02X}")).collect(); + authed_with( + client() + .post(format!( + "http://localhost:{port}/api/w/test-workspace/ai/proxy/chat/completions" + )) + .header("X-Provider", "openai") + .header("X-Resource-Value", encoded) + .json(&json!({ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}] + })), + token, + ) + .send() + }; + + let resp = send("sk-unsaved", "SECRET_TOKEN").await?; + assert_2xx( + resp.status().as_u16(), + &resp.text().await?, + "X-Resource-Value with an unsaved resource", + ); + + let resp = send("sk-other", "SECRET_TOKEN").await?; + assert!(resp.status().as_u16() >= 400, "a different key must reach the upstream as is"); + + let resp = send("$var:u/test-user/secret", "SCOPED_AI_TOKEN").await?; + let body = resp.text().await?; + assert!( + body.contains("variables:read:u/test-user/secret"), + "a token without variables:read must not resolve a $var: reference, got: {body}" + ); + + Ok(()) +} diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 3255b629da..855cd29987 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -1062,6 +1062,16 @@ where } } +/// Paths of every `$var:` reference found anywhere in a resource value. +fn var_references(value: &serde_json::Value) -> Vec<&str> { + match value { + serde_json::Value::String(s) => s.strip_prefix("$var:").into_iter().collect(), + serde_json::Value::Array(items) => items.iter().flat_map(var_references).collect(), + serde_json::Value::Object(fields) => fields.values().flat_map(var_references).collect(), + _ => vec![], + } +} + async fn global_proxy( authed: ApiAuthed, Extension(db): Extension, @@ -1220,7 +1230,37 @@ async fn proxy( // Set when serving the request through Windmill's free AI tier (the lent key). Holds // the per-user concurrency lock and drives response metering. let mut free_lease: Option = None; + // An unsaved resource value (percent-encoded JSON), so a resource can be tested + // before it is stored. `$var:` references resolve as the caller, and since the + // caller also picks the base URL they are sent to, a scoped token needs the same + // scopes as storing the resource and reading each variable would. + let inline_resource = headers + .get("X-Resource-Value") + .map(|v| { + let invalid = |e: String| Error::BadRequest(format!("Invalid X-Resource-Value: {e}")); + let decoded = + urlencoding::decode(v.to_str().unwrap_or("")).map_err(|e| invalid(e.to_string()))?; + let value = serde_json::from_str::(&decoded) + .map_err(|e| invalid(e.to_string()))?; + check_scopes(&authed, || "resources:write".to_string())?; + for path in var_references(&value) { + check_scopes(&authed, || format!("variables:read:{path}"))?; + } + serde_json::from_value::(value).map_err(|e| invalid(e.to_string())) + }) + .transpose()?; + let mut credentials = 'cred: { + if let Some(resource) = inline_resource { + break 'cred resolve_provider_credentials( + &provider, + &db, + &w_id, + resource, + Some(&authed), + ) + .await?; + } match workspace_cache { Some(request_cache) if !request_cache.is_expired() && forced_resource_path.is_none() => diff --git a/frontend/src/lib/components/ApiConnectForm.svelte b/frontend/src/lib/components/ApiConnectForm.svelte index 2e4e4b526c..37342912c7 100644 --- a/frontend/src/lib/components/ApiConnectForm.svelte +++ b/frontend/src/lib/components/ApiConnectForm.svelte @@ -5,6 +5,9 @@ import SchemaForm from './SchemaForm.svelte' import Toggle from './Toggle.svelte' import TestConnection from './TestConnection.svelte' + import TestAiKey from './copilot/TestAIKey.svelte' + import { aiResourceProvider } from './copilot/aiResourceProvider' + import { copilotInfo } from '$lib/aiStore' import Popover from './meltComponents/Popover.svelte' import Button from './common/button/Button.svelte' import { Loader2 } from 'lucide-svelte' @@ -51,6 +54,8 @@ onSynced = undefined }: Props = $props() + let aiProvider = $derived(aiResourceProvider(resourceType)) + let schema = $state(emptySchema()) let notFound = $state(false) @@ -180,6 +185,14 @@ isFileset={resourceTypeInfo?.is_fileset ?? false} /> + {#if aiProvider} + m.provider === aiProvider)?.model} + /> + {/if} {#if resourceType == 'postgresql'} {/if} + {#if aiProvider} + m.provider === aiProvider)?.model} + /> + {/if} {#if resource_type === 'git_repository' && ws && (actingUser?.is_admin || actingUser?.is_super_admin)} | undefined aiProvider: AIProvider - model: string + model?: string | undefined } let { @@ -18,50 +22,69 @@ apiKey = undefined, workspace = undefined, resourcePath = undefined, + resourceValue = undefined, aiProvider, - model + model = undefined }: Props = $props() let loading = $state(false) + // For an unsaved resource, or without a model from the caller, the model is only a + // guess: deployment-named providers (Azure, Bedrock, custom endpoints) reject a model + // the account lacks, which would read as an invalid key, so the guess stays editable. + let typedModel = $state( + untrack(() => model ?? AI_PROVIDERS[aiProvider]?.defaultModels[0] ?? '') + ) + let modelEditable = $derived(!!resourceValue || !model) + let testedModel = $derived(modelEditable ? typedModel.trim() : model) - + }} + >Test key + + diff --git a/frontend/src/lib/components/copilot/aiResourceProvider.ts b/frontend/src/lib/components/copilot/aiResourceProvider.ts new file mode 100644 index 0000000000..91ae926dfc --- /dev/null +++ b/frontend/src/lib/components/copilot/aiResourceProvider.ts @@ -0,0 +1,10 @@ +import type { AIProvider } from '$lib/gen' +import { AI_PROVIDERS } from './lib' + +/** The AI provider a resource of this type holds credentials for, if any. */ +export function aiResourceProvider(resourceType: string | undefined): AIProvider | undefined { + if (resourceType === 'openai_client_credentials_oauth') return 'openai' + return resourceType && Object.hasOwn(AI_PROVIDERS, resourceType) + ? (resourceType as AIProvider) + : undefined +} diff --git a/frontend/src/lib/components/copilot/chat/openai-responses.ts b/frontend/src/lib/components/copilot/chat/openai-responses.ts index 31d3b379ac..ebaa5b1c38 100644 --- a/frontend/src/lib/components/copilot/chat/openai-responses.ts +++ b/frontend/src/lib/components/copilot/chat/openai-responses.ts @@ -10,6 +10,8 @@ import { getAiProxyBaseURL, getProviderAndCompletionConfig, providerSupportsWebSearch, + testedCredentialHeaders, + usesGlobalAiProxy, workspaceAIClients } from '../lib' import { applyReasoningToConfig } from '../reasoningRegistry' @@ -624,6 +626,7 @@ export async function getNonStreamingOpenAIResponsesCompletion( apiKey?: string workspace?: string resourcePath?: string + resourceValue?: Record forceModelProvider?: AIProviderModel maxTokensCap?: number } @@ -645,23 +648,12 @@ export async function getNonStreamingOpenAIResponsesCompletion( } = { signal: abortController.signal, headers: { - 'X-Provider': provider + 'X-Provider': provider, + ...testedCredentialHeaders(options ?? {}) } } - if (options?.resourcePath) { - fetchOptions.headers = { - ...fetchOptions.headers, - 'X-Resource-Path': options.resourcePath - } - } else if (options?.apiKey) { - fetchOptions.headers = { - ...fetchOptions.headers, - 'X-API-Key': options.apiKey - } - } - - const openaiClient = options?.apiKey + const openaiClient = usesGlobalAiProxy(options ?? {}) ? createOpenAIProxyClient(getAiProxyBaseURL()) : options?.workspace ? workspaceAIClients.createOpenaiClient(options.workspace) diff --git a/frontend/src/lib/components/copilot/lib.ts b/frontend/src/lib/components/copilot/lib.ts index ba718bf760..40f1dee6a6 100644 --- a/frontend/src/lib/components/copilot/lib.ts +++ b/frontend/src/lib/components/copilot/lib.ts @@ -581,10 +581,39 @@ class WorkspacedAIClients { export const workspaceAIClients = new WorkspacedAIClients() +interface TestedCredential { + apiKey?: string + resourcePath?: string + resourceValue?: Record +} + +/** The header naming the credential under test; a resource wins over a bare key. */ +export function testedCredentialHeaders({ + apiKey, + resourcePath, + resourceValue +}: TestedCredential): Record { + if (resourceValue) { + // Header values must be ASCII, the resource may not be. + return { 'X-Resource-Value': encodeURIComponent(JSON.stringify(resourceValue)) } + } else if (resourcePath) { + return { 'X-Resource-Path': resourcePath } + } else if (apiKey) { + return { 'X-API-Key': apiKey } + } + return {} +} + +/** A bare key has no workspace resource behind it, so it goes through the global proxy. */ +export function usesGlobalAiProxy({ apiKey, resourcePath, resourceValue }: TestedCredential) { + return !!apiKey && !resourcePath && !resourceValue +} + export async function testKey({ apiKey, workspace, resourcePath, + resourceValue, model, abortController, messages, @@ -593,13 +622,14 @@ export async function testKey({ apiKey?: string workspace?: string resourcePath?: string + resourceValue?: Record model: string | undefined messages: ChatCompletionMessageParam[] abortController: AbortController aiProvider: AIProvider }) { - if (!apiKey && !resourcePath) { - throw new Error('API key or resource path is required') + if (!apiKey && !resourcePath && !resourceValue) { + throw new Error('API key, resource path or resource value is required') } const modelToTest = model ?? AI_PROVIDERS[aiProvider].defaultModels[0] @@ -617,6 +647,7 @@ export async function testKey({ apiKey, workspace, resourcePath, + resourceValue, forceModelProvider: { model: modelToTest, provider: aiProvider @@ -638,6 +669,7 @@ interface AnthropicCompletionParams { apiKey?: string workspace?: string resourcePath?: string + resourceValue?: Record maxTokensCap?: number } @@ -647,6 +679,7 @@ function buildAnthropicProxyRequest({ apiKey, workspace, resourcePath, + resourceValue, maxTokensCap }: Omit) { const { system, messages: anthropicMessages } = convertOpenAIToAnthropicMessages(messages) @@ -655,16 +688,11 @@ function buildAnthropicProxyRequest({ // resolves the right credentials and Anthropic URL. const headers: Record = { 'X-Provider': modelProvider.provider, - 'anthropic-version': '2023-06-01' + 'anthropic-version': '2023-06-01', + ...testedCredentialHeaders({ apiKey, resourcePath, resourceValue }) } - if (resourcePath) { - headers['X-Resource-Path'] = resourcePath - } else if (apiKey) { - headers['X-API-Key'] = apiKey - } - - const client = apiKey + const client = usesGlobalAiProxy({ apiKey, resourcePath, resourceValue }) ? createAnthropicProxyClient(getAiProxyBaseURL()) : workspace ? workspaceAIClients.createAnthropicClient(workspace) @@ -959,6 +987,7 @@ export async function getNonStreamingCompletion( options?: { apiKey?: string // testing API KEY using the global ai proxy resourcePath?: string // testing resource path passed as a header to the backend proxy + resourceValue?: Record // testing an unsaved resource value, same route workspace?: string // use a specific workspace proxy when testing a workspace resource forceModelProvider?: AIProviderModel maxTokensCap?: number // hard ceiling on output tokens (see METADATA_MAX_TOKENS) @@ -974,6 +1003,7 @@ export async function getNonStreamingCompletion( apiKey: options?.apiKey, workspace: options?.workspace, resourcePath: options?.resourcePath, + resourceValue: options?.resourceValue, maxTokensCap: options?.maxTokensCap }) } @@ -1009,22 +1039,11 @@ export async function getNonStreamingCompletion( 'X-Provider': provider } } - if (options?.resourcePath) { - fetchOptions.headers = { - ...fetchOptions.headers, - 'X-Resource-Path': options.resourcePath - } - } else if (options?.apiKey) { - if (provider === 'customai') { - throw new Error('Cannot test API key for Custom AI, only resource path is supported') - } - - fetchOptions.headers = { - ...fetchOptions.headers, - 'X-API-Key': options.apiKey - } + if (usesGlobalAiProxy(options ?? {}) && provider === 'customai') { + throw new Error('Cannot test API key for Custom AI, only resource path is supported') } - const openaiClient = options?.apiKey + fetchOptions.headers = { ...fetchOptions.headers, ...testedCredentialHeaders(options ?? {}) } + const openaiClient = usesGlobalAiProxy(options ?? {}) ? createOpenAIProxyClient(getAiProxyBaseURL()) : options?.workspace ? workspaceAIClients.createOpenaiClient(options.workspace)