fix(builder): unlock BYOK passkey on the click, not mid-send

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

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<AbortController>(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<
               <div className="flex items-center justify-between gap-3 pl-1">
                 <ModelPicker
                   chatGptModels={chatGptModels}
-                  disabled={running}
+                  disabled={running || unlocking}
                   selected={selectedModel}
                   showChatGpt={supportsChatGptLogin}
                   onSelect={selectModel}