diff --git a/scripts/test-builder-passkey-submit.mjs b/scripts/test-builder-passkey-submit.mjs new file mode 100644 index 000000000..8795c9781 --- /dev/null +++ b/scripts/test-builder-passkey-submit.mjs @@ -0,0 +1,122 @@ +import assert from 'node:assert/strict' +import { chromium } from 'playwright-core' + +// Run against the local dev server: node scripts/test-builder-passkey-submit.mjs +const origin = process.env.BUILDER_TEST_ORIGIN ?? 'http://127.0.0.1:3008' +assert.ok(['localhost', '127.0.0.1'].includes(new URL(origin).hostname)) +const browser = await chromium.launch({ channel: 'chrome', headless: true }) +try { + const page = await browser.newPage() + await page.addInitScript(() => { + class Passkey { + rawId = new Uint8Array([1]).buffer + getClientExtensionResults() { + return { + prf: { enabled: true, results: { first: new Uint8Array(32) } }, + } + } + } + window.passkeyCalls = 0 + Object.defineProperty(window, 'PublicKeyCredential', { value: Passkey }) + Object.defineProperty(navigator, 'credentials', { + value: { + async create() { + return new Passkey() + }, + get() { + window.passkeyCalls += 1 + return new Promise((resolve, reject) => { + window.cancelPasskey = () => + reject(new Error('Cancelled test passkey')) + window.unlockPasskey = () => resolve(new Passkey()) + }) + }, + }, + }) + }) + await page.goto(`${origin}/builder/ai`) + await page.evaluate(async () => { + const base = await crypto.subtle.importKey( + 'raw', + new Uint8Array(32), + 'HKDF', + false, + ['deriveKey'], + ) + const key = await crypto.subtle.deriveKey( + { + name: 'HKDF', + hash: 'SHA-256', + salt: new Uint8Array(0), + info: new TextEncoder().encode('byok:keyring:v1'), + }, + base, + { name: 'AES-GCM', length: 256 }, + false, + ['encrypt'], + ) + const iv = crypto.getRandomValues(new Uint8Array(12)) + const ciphertext = await crypto.subtle.encrypt( + { name: 'AES-GCM', iv }, + key, + new TextEncoder().encode( + JSON.stringify({ openai: 'sk-test-not-a-real-key' }), + ), + ) + await new Promise((resolve, reject) => { + const request = indexedDB.open( + 'tanstack-builder-ai:byok:v1:local-spike', + 1, + ) + request.onupgradeneeded = () => + request.result.createObjectStore('keyring', { keyPath: 'id' }) + request.onerror = () => reject(request.error) + request.onsuccess = () => { + const tx = request.result.transaction('keyring', 'readwrite') + tx.objectStore('keyring').put({ + id: 'default', + credentialId: new Uint8Array([1]).buffer, + salt: new Uint8Array(32).buffer, + iv: iv.buffer, + ciphertext, + preview: { openai: '-key' }, + }) + tx.oncomplete = () => { + request.result.close() + resolve() + } + tx.onerror = () => reject(tx.error) + } + }) + }) + await page.reload() + const composer = page.locator('#builder-ai-prompt') + await composer.waitFor({ timeout: 60_000 }) + await composer.pressSequentially('Keep this draft after a cancelled unlock', { + delay: 30, + }) + await page.getByRole('button', { name: 'Send message', exact: true }).click() + await page.waitForFunction(() => window.passkeyCalls === 1) + assert.equal(await composer.isDisabled(), true) + await page + .locator('form') + .filter({ has: composer }) + .evaluate((form) => { + form.dispatchEvent( + new Event('submit', { bubbles: true, cancelable: true }), + ) + }) + assert.equal(await page.evaluate(() => window.passkeyCalls), 1) + await page.evaluate(() => window.cancelPasskey()) + await page.waitForFunction( + () => !document.querySelector('#builder-ai-prompt').disabled, + ) + assert.equal( + await composer.inputValue(), + 'Keep this draft after a cancelled unlock', + ) + + console.log('Builder passkey submission checks passed') +} finally { + await browser.close() +} diff --git a/src/components/builder/BuilderAssistant.client.tsx b/src/components/builder/BuilderAssistant.client.tsx index aed28e8dd..ec195ac72 100644 --- a/src/components/builder/BuilderAssistant.client.tsx +++ b/src/components/builder/BuilderAssistant.client.tsx @@ -364,6 +364,8 @@ export const BuilderAssistant = React.forwardRef< const [queueAnnouncement, setQueueAnnouncement] = React.useState('') const [showLatest, setShowLatest] = React.useState(false) const abortRef = React.useRef(null) + const unlockingRef = React.useRef(false) + const [unlocking, setUnlocking] = React.useState(false) const onRunningChangeRef = React.useRef(onRunningChange) const abortIntentRef = React.useRef<'steer' | 'stop' | undefined>(undefined) const agentStreamingRef = React.useRef(false) @@ -412,6 +414,10 @@ export const BuilderAssistant = React.forwardRef< canUsePendingPromptRef.current = canUsePendingPrompt startPromptSequenceRef.current = startPromptSequence + React.useLayoutEffect(() => { + pendingSubmissionGenerationRef.current += 1 + }, [threadId, storageScope, credentialScope, selectedModel]) + React.useLayoutEffect(() => { if (credentialScopeRef.current === credentialScope) return credentialScopeRef.current = credentialScope @@ -660,7 +666,17 @@ export const BuilderAssistant = React.forwardRef< setHydratedThreadId(threadId) return } + }, [ + hydratedThreadId, + syncedMessages, + syncedProjectId, + syncedRuns, + syncedThreads, + threadId, + ]) + React.useEffect(() => { + if (syncedProjectId) return const generation = hydrationGenerationRef.current + 1 hydrationGenerationRef.current = generation setHydratedThreadId(undefined) @@ -691,16 +707,7 @@ export const BuilderAssistant = React.forwardRef< hydrationGenerationRef.current += 1 } } - }, [ - hydratedThreadId, - refreshThreads, - storageScope, - syncedMessages, - syncedProjectId, - syncedRuns, - syncedThreads, - threadId, - ]) + }, [refreshThreads, storageScope, syncedProjectId, threadId]) React.useEffect(() => { const currentProjectSync = projectSync @@ -1028,7 +1035,7 @@ export const BuilderAssistant = React.forwardRef< }, []) function selectModel(model: ModelChoice) { - if (running) return + if (running || unlockingRef.current) return didSelectConnectionRef.current = model.connection !== 'chatgpt' || Boolean(model.model) setSelectedModel(model) @@ -1325,7 +1332,10 @@ export const BuilderAssistant = React.forwardRef< const promptQueue = promptQueueRef.current promptQueue.enqueuePrompt(queuedPrompt) syncQueuedPrompts() - if (clearComposer) { + if ( + clearComposer && + promptValueRef.current.trim() === queuedPrompt.content + ) { promptValueRef.current = '' setPrompt('') setSendMode('queue') @@ -1410,11 +1420,50 @@ export const BuilderAssistant = React.forwardRef< !instruction || instruction.length > 10_000 || hydrating || - needsConnection + needsConnection || + unlockingRef.current ) { return false } + // Both the composer and preview comments enter here, directly from the + // user action, before persistence or sandbox work can expire activation. + if ( + selectedModel.connection === 'byok' && + !byokConnection.getClient(selectedModel.provider, { allowUnlock: false }) + ) { + const provider = selectedModel.provider + const generation = pendingSubmissionGenerationRef.current + unlockingRef.current = true + setUnlocking(true) + void unlockApiKey(provider).then(() => { + unlockingRef.current = false + if (mountedRef.current) setUnlocking(false) + if ( + !mountedRef.current || + generation !== pendingSubmissionGenerationRef.current + ) { + lifecycle?.onDiscarded?.() + return + } + if (!byokConnection.getClient(provider, { allowUnlock: false })) { + setError('Could not unlock the API key. Try again.') + lifecycle?.onDiscarded?.() + return + } + enqueueInstruction(instruction, mode, clearComposer, lifecycle) + }) + return true + } + return enqueueInstruction(instruction, mode, clearComposer, lifecycle) + } + + function enqueueInstruction( + instruction: string, + mode: BuilderAiSendMode, + clearComposer: boolean, + lifecycle?: BuilderAiPromptLifecycle, + ) { const promptQueue = promptQueueRef.current const claimed = promptQueue.claim() const queuedPrompt: BuilderAiQueuedPrompt = { @@ -1515,7 +1564,10 @@ export const BuilderAssistant = React.forwardRef< initialPrompt: BuilderAiQueuedPrompt, clearComposer: boolean, ) { - if (clearComposer) { + if ( + clearComposer && + promptValueRef.current.trim() === initialPrompt.content + ) { promptValueRef.current = '' setPrompt('') setSendMode('queue') @@ -2616,7 +2668,8 @@ export const BuilderAssistant = React.forwardRef< } } - const submitDisabled = hydrating || !prompt.trim() || needsConnection + const submitDisabled = + hydrating || unlocking || !prompt.trim() || needsConnection const stopLabel = queuedPrompts.length === 0 ? 'Stop response' @@ -2943,6 +2996,7 @@ export const BuilderAssistant = React.forwardRef< ref={promptRef} id="builder-ai-prompt" value={prompt} + disabled={unlocking} rows={1} maxLength={10_000} placeholder="Describe a builder change" @@ -2962,7 +3016,7 @@ export const BuilderAssistant = React.forwardRef<