perf(app): fix composer lag via buffered blob draft storage (#40207)

packages/desktop/src/main/draft-store.test.ts

@@ -0,0 +1,16 @@
+import { expect, test } from "bun:test"
+import { createDesktopDraftStore } from "./draft-store"
+
+test("flushes the latest buffered draft and stores blobs", () => {
+  const store = createDesktopDraftStore(":memory:")
+  store.set("prompt", "first")
+  store.set("prompt", "latest")
+  expect(store.get("prompt")).toBe("latest")
+  store.flush()
+  expect(store.get("prompt")).toBe("latest")
+
+  const bytes = new TextEncoder().encode("image")
+  const id = store.putBlob(bytes)
+  expect(store.getBlob(id)).toEqual(bytes)
+  store.close()
+})

packages/desktop/src/main/draft-store.ts

@@ -0,0 +1,82 @@
+import { createHash } from "node:crypto"
+import { DatabaseSync } from "node:sqlite"
+import { eq } from "drizzle-orm"
+import { drizzle } from "drizzle-orm/node-sqlite"
+import { blob, sqliteTable, text } from "drizzle-orm/sqlite-core"
+
+const documents = sqliteTable("document", {
+  key: text().primaryKey(),
+  value: text().notNull(),
+})
+const blobs = sqliteTable("blob", {
+  id: text().primaryKey(),
+  data: blob({ mode: "buffer" }).notNull(),
+})
+
+export function createDesktopDraftStore(filename: string) {
+  const native = new DatabaseSync(filename)
+  native.exec(
+    "PRAGMA journal_mode=WAL; CREATE TABLE IF NOT EXISTS document (key TEXT PRIMARY KEY, value TEXT NOT NULL); CREATE TABLE IF NOT EXISTS blob (id TEXT PRIMARY KEY, data BLOB NOT NULL);",
+  )
+  const db = drizzle({ client: native })
+  const used = new Set<string>()
+  db.select({ value: documents.value })
+    .from(documents)
+    .all()
+    .forEach(({ value }) =>
+      JSON.parse(value, (_key, item) => {
+        if (item?.blob && typeof item.blob.id === "string") used.add(item.blob.id)
+        return item
+      }),
+    )
+  db.select({ id: blobs.id })
+    .from(blobs)
+    .all()
+    .filter(({ id }) => !used.has(id))
+    .forEach(({ id }) => db.delete(blobs).where(eq(blobs.id, id)).run())
+  const pending = new Map<string, string | null>()
+  let timer: ReturnType<typeof setTimeout> | undefined
+  const flush = () => {
+    if (timer) clearTimeout(timer)
+    timer = undefined
+    const writes = [...pending]
+    pending.clear()
+    db.transaction((tx) => {
+      writes.forEach(([key, value]) => {
+        if (value === null) tx.delete(documents).where(eq(documents.key, key)).run()
+        else
+          tx.insert(documents)
+            .values({ key, value })
+            .onConflictDoUpdate({ target: documents.key, set: { value } })
+            .run()
+      })
+    })
+  }
+  const schedule = () => {
+    if (!timer) timer = setTimeout(flush, 500)
+  }
+  return {
+    get: (key: string) =>
+      pending.has(key)
+        ? (pending.get(key) ?? null)
+        : (db.select({ value: documents.value }).from(documents).where(eq(documents.key, key)).get()?.value ?? null),
+    set(key: string, value: string | null) {
+      pending.set(key, value)
+      schedule()
+    },
+    putBlob(data: Uint8Array) {
+      const id = createHash("sha256").update(data).digest("hex")
+      db.insert(blobs)
+        .values({ id, data: Buffer.from(data) })
+        .onConflictDoNothing()
+        .run()
+      return id
+    },
+    getBlob: (id: string) => db.select({ data: blobs.data }).from(blobs).where(eq(blobs.id, id)).get()?.data ?? null,
+    flush,
+    close() {
+      flush()
+      native.close()
+    },
+  }
+}

packages/desktop/src/main/index.ts

@@ -171,7 +171,7 @@ const main = Effect.gen(function* () {
     setAppQuitting()
     void stopSidecars().finally(() => {
       app.relaunch()
-      app.exit(0)
+      app.quit()
     })
   }
 
@@ -245,7 +245,7 @@ const main = Effect.gen(function* () {
   for (const signal of ["SIGINT", "SIGTERM"] as const) {
     process.on(signal, () => {
       setAppQuitting()
-      void stopSidecars().finally(() => app.exit(0))
+      void stopSidecars().finally(() => app.quit())
     })
   }
 

packages/desktop/src/main/ipc.ts

@@ -1,6 +1,6 @@
 import { execFile } from "node:child_process"
 import { stat } from "node:fs/promises"
-import { basename } from "node:path"
+import { basename, join } from "node:path"
 import { app, BrowserWindow, clipboard, dialog, ipcMain, shell } from "electron"
 import type { IpcMainEvent, IpcMainInvokeEvent } from "electron"
 import type { DesktopMenuAction } from "@opencode-ai/app/desktop-menu"
@@ -21,6 +21,7 @@ import {
 } from "./windows"
 import type { UpdaterController } from "./updater-controller"
 import { createUpdaterSubscriptions } from "./updater-subscriptions"
+import { createDesktopDraftStore } from "./draft-store"
 
 const pickerFilters = (ext?: string[]) => {
   if (!ext || ext.length === 0) return undefined
@@ -51,8 +52,12 @@ type Deps = {
 }
 
 export function registerIpcHandlers(deps: Deps) {
+  const drafts = createDesktopDraftStore(join(app.getPath("userData"), "drafts.sqlite"))
   const updaterSubscriptions = createUpdaterSubscriptions()
   app.once("will-quit", updaterSubscriptions.clear)
+  app.on("before-quit", () => drafts.flush())
+  app.once("will-quit", () => drafts.close())
+  app.on("browser-window-created", (_event, win) => win.on("session-end", () => drafts.flush()))
 
   ipcMain.handle("kill-sidecar", () => deps.killSidecar())
   ipcMain.handle("await-initialization", () => deps.awaitInitialization())
@@ -123,6 +128,14 @@ export function registerIpcHandlers(deps: Deps) {
     const store = getStore(name)
     return Object.keys(store.store).length
   })
+  ipcMain.handle("draft-get", (_event, key: string) => drafts.get(key))
+  ipcMain.handle("draft-set", (_event, key: string, value: string) => drafts.set(key, value))
+  ipcMain.handle("draft-delete", (_event, key: string) => drafts.set(key, null))
+  ipcMain.handle("draft-blob-put", (_event, data: ArrayBuffer) => drafts.putBlob(new Uint8Array(data)))
+  ipcMain.handle("draft-blob-get", (_event, id: string) => {
+    const data = drafts.getBlob(id)
+    return data ? data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength) : null
+  })
 
   ipcMain.handle(
     "open-directory-picker",

packages/desktop/src/preload/index.ts

@@ -73,6 +73,11 @@ const api: ElectronAPI = {
   storeClear: (name) => ipcRenderer.invoke("store-clear", name),
   storeKeys: (name) => ipcRenderer.invoke("store-keys", name),
   storeLength: (name) => ipcRenderer.invoke("store-length", name),
+  draftGet: (key) => ipcRenderer.invoke("draft-get", key),
+  draftSet: (key, value) => ipcRenderer.invoke("draft-set", key, value),
+  draftDelete: (key) => ipcRenderer.invoke("draft-delete", key),
+  draftBlobPut: (data) => ipcRenderer.invoke("draft-blob-put", data),
+  draftBlobGet: (id) => ipcRenderer.invoke("draft-blob-get", id),
 
   getWindowID: () => ipcRenderer.invoke("get-window-id"),
   onMenuCommand: (cb) => {

packages/desktop/src/preload/types.ts

@@ -63,6 +63,11 @@ export type ElectronAPI = {
   storeClear: (name: string) => Promise<void>
   storeKeys: (name: string) => Promise<string[]>
   storeLength: (name: string) => Promise<number>
+  draftGet: (key: string) => Promise<string | null>
+  draftSet: (key: string, value: string) => Promise<void>
+  draftDelete: (key: string) => Promise<void>
+  draftBlobPut: (data: ArrayBuffer) => Promise<string>
+  draftBlobGet: (id: string) => Promise<ArrayBuffer | null>
 
   getWindowID: () => Promise<string>
   onMenuCommand: (cb: (id: string) => void) => () => void

packages/desktop/src/renderer/index.tsx

@@ -9,6 +9,7 @@ import {
   type Locale,
   type Platform,
   PlatformProvider,
+  createDraftStore,
   ServerConnection,
   useCommand,
   useWslServers,
@@ -226,6 +227,13 @@ const createPlatform = (windowState: DesktopWindowState): Platform => {
     },
 
     storage,
+    draftStore: createDraftStore({
+      get: window.api.draftGet,
+      set: window.api.draftSet,
+      remove: window.api.draftDelete,
+      putBlob: (blob) => blob.arrayBuffer().then(window.api.draftBlobPut),
+      getBlob: (id) => window.api.draftBlobGet(id).then((data) => data && new Blob([data])),
+    }),
 
     updater: {
       state: updaterState,