fix(tui): stabilize Zed editor context polling (#24656)

4 files changed, 128 insertions, 21 deletions

diff --git a/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx b/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx
index fc0286892..05afd3369 100644
--- a/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx
+++ b/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx
@@ -764,6 +764,12 @@ export function Prompt(props: PromptProps) {
               return `Note: The user selected lines ${start.line} to ${end.line} from "${editorSelection.filePath}": ${editorSelection.text}`
             })(),
             synthetic: true,
+            metadata: {
+              kind: "editor_context",
+              source: editorSelection.source ?? "editor",
+              filePath: editorSelection.filePath,
+              selection: editorSelection.selection,
+            },
           },
         ]
       : []
diff --git a/packages/opencode/src/cli/cmd/tui/context/editor-zed.ts b/packages/opencode/src/cli/cmd/tui/context/editor-zed.ts
index aee819a8b..bae88f28d 100644
--- a/packages/opencode/src/cli/cmd/tui/context/editor-zed.ts
+++ b/packages/opencode/src/cli/cmd/tui/context/editor-zed.ts
@@ -21,24 +21,37 @@ const ZedEditorContentsSchema = z.object({
 
 type ZedEditorRow = z.infer<typeof ZedEditorRowSchema>
 
-export async function resolveZedSelection(dbPath: string): Promise<EditorSelection | undefined> {
-  const row = queryZedActiveEditor(dbPath, process.cwd())
-  if (!row?.buffer_path || row.selection_start == null || row.selection_end == null) return
+export type ZedSelectionResult =
+  | { type: "selection"; selection: EditorSelection }
+  | { type: "empty" }
+  | { type: "unavailable" }
 
+export async function resolveZedSelection(dbPath: string, cwd = process.cwd()): Promise<ZedSelectionResult> {
+  const active = queryZedActiveEditor(dbPath, cwd)
+  if (active.type !== "row") return active
+
+  const row = active.row
+  if (!row.buffer_path) return { type: "empty" }
+  if (row.selection_start == null || row.selection_end == null) return { type: "unavailable" }
+
+  const contents = queryZedEditorContents(dbPath, row)
   const text =
-    queryZedEditorContents(dbPath, row) ??
-    (await Bun.file(row.buffer_path)
-      .text()
-      .catch(() => undefined))
-  if (text == null) return
+    contents.type === "contents" && contents.contents != null
+      ? contents.contents
+      : await Bun.file(row.buffer_path).text().catch(() => undefined)
+  if (text == null) return { type: "unavailable" }
 
   const startOffset = Math.min(row.selection_start, row.selection_end)
   const endOffset = Math.max(row.selection_start, row.selection_end)
 
   return {
-    text: text.slice(startOffset, endOffset),
-    filePath: row.buffer_path,
-    selection: offsetsToSelection(text, startOffset, endOffset),
+    type: "selection",
+    selection: {
+      text: text.slice(startOffset, endOffset),
+      filePath: row.buffer_path,
+      source: "zed",
+      selection: offsetsToSelection(text, startOffset, endOffset),
+    },
   }
 }
 
@@ -46,7 +59,7 @@ function queryZedActiveEditor(dbPath: string, cwd: string) {
   let db: Database | undefined
   try {
     db = new Database(dbPath, { readonly: true })
-    return db
+    const raw = db
       .query(
         `select
           e.item_id as editor_id,
@@ -65,15 +78,23 @@ function queryZedActiveEditor(dbPath: string, cwd: string) {
         order by w.timestamp desc`,
       )
       .all()
+
+    const rows = raw
       .flatMap((row) => {
         const parsed = ZedEditorRowSchema.safeParse(row)
         return parsed.success ? [parsed.data] : []
       })
+
+    if (raw.length > 0 && rows.length === 0) return { type: "unavailable" as const }
+
+    const row = rows
       .map((row) => ({ row, score: scoreZedWorkspace(row.workspace_paths, cwd) }))
       .filter((entry) => entry.score > 0)
       .sort((left, right) => right.score - left.score || right.row.timestamp.localeCompare(left.row.timestamp))[0]?.row
+    if (!row) return { type: "empty" as const }
+    return { type: "row" as const, row }
   } catch {
-    return
+    return { type: "unavailable" as const }
   } finally {
     db?.close()
   }
@@ -83,7 +104,7 @@ function queryZedEditorContents(dbPath: string, row: ZedEditorRow) {
   let db: Database | undefined
   try {
     db = new Database(dbPath, { readonly: true })
-    return ZedEditorContentsSchema.safeParse(
+    const parsed = ZedEditorContentsSchema.safeParse(
       db
         .query(
           `select contents
@@ -91,9 +112,11 @@ function queryZedEditorContents(dbPath: string, row: ZedEditorRow) {
         where item_id = $editorID and workspace_id = $workspaceID`,
         )
         .get({ $editorID: row.editor_id, $workspaceID: row.workspace_id }),
-    ).data?.contents
+    )
+    if (!parsed.success) return { type: "unavailable" as const }
+    return { type: "contents" as const, contents: parsed.data.contents }
   } catch {
-    return
+    return { type: "unavailable" as const }
   } finally {
     db?.close()
   }
diff --git a/packages/opencode/src/cli/cmd/tui/context/editor.ts b/packages/opencode/src/cli/cmd/tui/context/editor.ts
index 72b0785d6..7f2c46d08 100644
--- a/packages/opencode/src/cli/cmd/tui/context/editor.ts
+++ b/packages/opencode/src/cli/cmd/tui/context/editor.ts
@@ -31,6 +31,7 @@ const PositionSchema = z.object({
 const EditorSelectionSchema = z.object({
   text: z.string(),
   filePath: z.string(),
+  source: z.enum(["websocket", "zed"]).optional(),
   selection: z.object({
     start: PositionSchema,
     end: PositionSchema,
@@ -125,8 +126,10 @@ export const { use: useEditorContext, provider: EditorContextProvider } = create
             return
           }
           zedSelection ??= resolveZedSelection(dbPath)
-            .then((selection) => {
+            .then((result) => {
               if (closed || socket) return
+              if (result.type === "unavailable") return
+              const selection = result.type === "selection" ? result.selection : undefined
               const key = editorSelectionKey(selection)
               if (key !== lastZedSelectionKey) {
                 lastZedSelectionKey = key
@@ -135,8 +138,7 @@ export const { use: useEditorContext, provider: EditorContextProvider } = create
               }
             })
             .catch(() => {
-              if (closed || socket) return
-              setStore("status", "disabled")
+              // Keep the last known Zed selection for transient polling failures.
             })
             .finally(() => {
               zedSelection = undefined
@@ -171,7 +173,7 @@ export const { use: useEditorContext, provider: EditorContextProvider } = create
           const selection =
             message.method === "selection_changed" ? EditorSelectionSchema.safeParse(message.params) : undefined
           if (selection?.success) {
-            setStore("selection", selection.data)
+            setStore("selection", { ...selection.data, source: "websocket" })
             return
           }
 
diff --git a/packages/opencode/test/cli/tui/editor-context.test.ts b/packages/opencode/test/cli/tui/editor-context.test.ts
index c605029ca..60a1d5834 100644
--- a/packages/opencode/test/cli/tui/editor-context.test.ts
+++ b/packages/opencode/test/cli/tui/editor-context.test.ts
@@ -1,5 +1,41 @@
+import { Database } from "bun:sqlite"
+import path from "node:path"
 import { expect, test } from "bun:test"
-import { offsetToPosition } from "../../../src/cli/cmd/tui/context/editor-zed"
+import { offsetToPosition, resolveZedSelection } from "../../../src/cli/cmd/tui/context/editor-zed"
+import { tmpdir } from "../../fixture/fixture"
+
+type ZedFixtureOptions = {
+  workspacePaths?: string | null
+  selectionStart?: number | null
+  selectionEnd?: number | null
+}
+
+async function writeZedFixture(dir: string, options: ZedFixtureOptions = {}) {
+  const dbPath = path.join(dir, "zed.sqlite")
+  const filePath = path.join(dir, "file.ts")
+  await Bun.write(filePath, "one\ntwo\nthree")
+
+  const db = new Database(dbPath)
+  db.run("create table workspaces (workspace_id integer, paths text, timestamp text)")
+  db.run("create table panes (pane_id integer, workspace_id integer, active integer)")
+  db.run("create table items (item_id integer, workspace_id integer, pane_id integer, active integer, kind text)")
+  db.run("create table editors (item_id integer, workspace_id integer, buffer_path text, contents text)")
+  db.run("create table editor_selections (editor_id integer, workspace_id integer, start integer, end integer)")
+  db.run("insert into workspaces values (1, ?, ?)", [options.workspacePaths ?? JSON.stringify([dir]), "2026-04-27"])
+  db.run("insert into panes values (1, 1, 1)")
+  db.run("insert into items values (1, 1, 1, 1, 'Editor')")
+  db.run("insert into editors values (1, 1, ?, ?)", [filePath, "one\ntwo\nthree"])
+  db.run(
+    "insert into editor_selections values (1, 1, ?, ?)",
+    [
+      options.selectionStart === undefined ? 4 : options.selectionStart,
+      options.selectionEnd === undefined ? 7 : options.selectionEnd,
+    ],
+  )
+  db.close()
+
+  return { dbPath, filePath }
+}
 
 test("offsetToPosition converts Zed offsets to 1-based editor positions", () => {
   expect(offsetToPosition("one\ntwo\nthree", 0)).toEqual({ line: 1, character: 1 })
@@ -7,3 +43,43 @@ test("offsetToPosition converts Zed offsets to 1-based editor positions", () =>
   expect(offsetToPosition("one\ntwo\nthree", 6)).toEqual({ line: 2, character: 3 })
   expect(offsetToPosition("one\ntwo\nthree", 100)).toEqual({ line: 3, character: 6 })
 })
+
+test("resolveZedSelection returns active editor selection", async () => {
+  await using tmp = await tmpdir()
+  const fixture = await writeZedFixture(tmp.path)
+
+  expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({
+    type: "selection",
+    selection: {
+      text: "two",
+      filePath: fixture.filePath,
+      source: "zed",
+      selection: {
+        start: { line: 2, character: 1 },
+        end: { line: 2, character: 4 },
+      },
+    },
+  })
+})
+
+test("resolveZedSelection returns empty when no workspace matches", async () => {
+  await using tmp = await tmpdir()
+  const fixture = await writeZedFixture(tmp.path, {
+    workspacePaths: JSON.stringify([path.join(path.dirname(tmp.path), "other-workspace")]),
+  })
+
+  expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({ type: "empty" })
+})
+
+test("resolveZedSelection returns unavailable when the database cannot be queried", async () => {
+  await using tmp = await tmpdir()
+
+  expect(await resolveZedSelection(path.join(tmp.path, "missing.sqlite"), tmp.path)).toEqual({ type: "unavailable" })
+})
+
+test("resolveZedSelection returns unavailable when active selection is missing offsets", async () => {
+  await using tmp = await tmpdir()
+  const fixture = await writeZedFixture(tmp.path, { selectionStart: null, selectionEnd: null })
+
+  expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({ type: "unavailable" })
+})