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

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<Postgres>) -> 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<Postgres>) -> 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(())
+}

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<DB>,
@@ -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<crate::ai_free_tier_oss::FreeTierLease> = 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::<serde_json::Value>(&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::<AIResource>(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() =>

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}
 		/>
 		<TestConnection {resourceType} {args} />
+		{#if aiProvider}
+			<TestAiKey
+				{aiProvider}
+				workspace={workspace ?? $operatingWorkspace}
+				resourceValue={args}
+				model={$copilotInfo.aiModels.find((m) => m.provider === aiProvider)?.model}
+			/>
+		{/if}
 		{#if resourceType == 'postgresql'}
 			<Popover
 				floatingConfig={{

frontend/src/lib/components/ResourceForm.svelte

@@ -14,6 +14,9 @@
 	import FilesetEditor from './FilesetEditor.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 GfmMarkdown from './GfmMarkdown.svelte'
 	import TestTriggerConnection from './triggers/TestTriggerConnection.svelte'
 	import GitHubAppIntegration from './GitHubAppIntegration.svelte'
@@ -85,6 +88,7 @@
 	}: Props = $props()
 
 	let ws = $derived(workspace ?? $operatingWorkspace)
+	let aiProvider = $derived(aiResourceProvider(resource_type))
 
 	let rawCode: string | undefined = $state(undefined)
 	let textFileContent: string = $state('')
@@ -225,6 +229,14 @@
 				workspaceOverride={workspace}
 			/>
 		{/if}
+		{#if aiProvider}
+			<TestAiKey
+				{aiProvider}
+				workspace={ws}
+				resourceValue={args}
+				model={$copilotInfo.aiModels.find((m) => m.provider === aiProvider)?.model}
+			/>
+		{/if}
 		{#if resource_type === 'git_repository' && ws && (actingUser?.is_admin || actingUser?.is_super_admin)}
 			<GitHubAppIntegration
 				resourceType={resource_type}

frontend/src/lib/components/copilot/TestAIKey.svelte

@@ -2,15 +2,19 @@
 	import type { AIProvider } from '$lib/gen'
 	import { sendUserToast } from '$lib/toast'
 	import Button from '../common/button/Button.svelte'
-	import { testKey } from './lib'
+	import TextInput from '../text_input/TextInput.svelte'
+	import { AI_PROVIDERS, testKey } from './lib'
+	import { untrack } from 'svelte'
 
 	interface Props {
 		disabled?: boolean
 		apiKey?: string | undefined
 		workspace?: string | undefined
 		resourcePath?: string | undefined
+		/** An unsaved resource value to test as is, instead of a stored resource. */
+		resourceValue?: Record<string, any> | 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)
 </script>
 
-<Button
-	unifiedSize="md"
-	variant="default"
-	{disabled}
-	{loading}
-	on:click={async () => {
-		loading = true
-		try {
-			const abortController = new AbortController()
-			setTimeout(() => {
-				abortController.abort()
-			}, 10000)
+<div class="flex flex-row items-center gap-1">
+	{#if modelEditable}
+		<TextInput
+			bind:value={typedModel}
+			size="md"
+			inputProps={{ placeholder: 'Model to test', 'aria-label': 'Model to test' }}
+		/>
+	{/if}
+	<Button
+		unifiedSize="md"
+		variant="default"
+		disabled={disabled || !testedModel}
+		{loading}
+		onClick={async () => {
+			loading = true
+			try {
+				const abortController = new AbortController()
+				setTimeout(() => {
+					abortController.abort()
+				}, 10000)
 
-			await testKey({
-				apiKey,
-				workspace,
-				resourcePath,
-				messages: [
-					{
-						role: 'user',
-						content: "this is a test, simply reply with 'ok'"
-					}
-				],
-				abortController,
-				aiProvider,
-				model
-			})
-			sendUserToast('Valid key')
-		} catch (err) {
-			if (err.message === 'Request was aborted.') {
-				sendUserToast('Could not validate key within 10s', true)
-			} else {
-				sendUserToast(`Invalid key: ${err}`, true)
+				await testKey({
+					apiKey,
+					workspace,
+					resourcePath,
+					resourceValue,
+					messages: [
+						{
+							role: 'user',
+							content: "this is a test, simply reply with 'ok'"
+						}
+					],
+					abortController,
+					aiProvider,
+					model: testedModel
+				})
+				sendUserToast('Valid key')
+			} catch (err) {
+				if (err.message === 'Request was aborted.') {
+					sendUserToast('Could not validate key within 10s', true)
+				} else {
+					sendUserToast(`Invalid key: ${err}`, true)
+				}
+			} finally {
+				loading = false
 			}
-		} finally {
-			loading = false
-		}
-	}}
-	>Test key
-</Button>
+		}}
+		>Test key
+	</Button>
+</div>

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
+}

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<string, any>
 		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)

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<string, any>
+}
+
+/** The header naming the credential under test; a resource wins over a bare key. */
+export function testedCredentialHeaders({
+	apiKey,
+	resourcePath,
+	resourceValue
+}: TestedCredential): Record<string, string> {
+	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<string, any>
 	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<string, any>
 	maxTokensCap?: number
 }
 
@@ -647,6 +679,7 @@ function buildAnthropicProxyRequest({
 	apiKey,
 	workspace,
 	resourcePath,
+	resourceValue,
 	maxTokensCap
 }: Omit<AnthropicCompletionParams, 'abortController'>) {
 	const { system, messages: anthropicMessages } = convertOpenAIToAnthropicMessages(messages)
@@ -655,16 +688,11 @@ function buildAnthropicProxyRequest({
 	// resolves the right credentials and Anthropic URL.
 	const headers: Record<string, string> = {
 		'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<string, any> // 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)