Sfoglia il codice sorgente

Support multiple Zed selections in TUI context (#25140)

Kit Langton 4 mesi fa
parent
commit
320527a3e4

+ 73 - 38
packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx

@@ -12,7 +12,7 @@ import { useRoute } from "@tui/context/route"
 import { useProject } from "@tui/context/project"
 import { useSync } from "@tui/context/sync"
 import { useEvent } from "@tui/context/event"
-import { useEditorContext, type EditorSelection } from "@tui/context/editor"
+import { editorSelectionKey, useEditorContext, type EditorSelection } from "@tui/context/editor"
 import { MessageID, PartID } from "@/session/schema"
 import { createStore, produce, unwrap } from "solid-js/store"
 import { useKeybind } from "@tui/context/keybind"
@@ -84,16 +84,30 @@ function fadeColor(color: RGBA, alpha: number) {
   return RGBA.fromValues(color.r, color.g, color.b, color.a * alpha)
 }
 
-function getEditorSelectionKey(selection: EditorSelection) {
-  return [
-    selection.filePath,
-    selection.text,
-    selection.source ?? "",
-    selection.selection.start.line,
-    selection.selection.start.character,
-    selection.selection.end.line,
-    selection.selection.end.character,
-  ].join("-")
+function hasEditorRangeSelection(selection: EditorSelection["ranges"][number]) {
+  return (
+    selection.selection.start.line !== selection.selection.end.line ||
+    selection.selection.start.character !== selection.selection.end.character
+  )
+}
+
+function getEditorRangeLabel(selection: EditorSelection["ranges"][number]) {
+  if (!hasEditorRangeSelection(selection)) return
+  if (selection.selection.start.line === selection.selection.end.line) return `#${selection.selection.start.line}`
+  return `#${selection.selection.start.line}-${selection.selection.end.line}`
+}
+
+function formatEditorContext(selection: EditorSelection) {
+  const selected = selection.ranges.filter(hasEditorRangeSelection)
+  if (selected.length === 0)
+    return `<system-reminder>Note: The user opened the file "${selection.filePath}". This may or may not be relevant to the current task.</system-reminder>\n`
+
+  const ranges = selected.map((range, index) => {
+    const prefix = selected.length > 1 ? `Selection ${index + 1}: ` : ""
+    return `Note: The user selected ${prefix}${getEditorRangeLabel(range)} from "${selection.filePath}". \`\`\`${range.text}\`\`\`\n\n`
+  })
+
+  return `<system-reminder>${ranges.join("\n")} This may or may not be relevant to the current task.</system-reminder>\n`
 }
 
 let stashed: { prompt: PromptInfo; cursor: number } | undefined
@@ -125,13 +139,21 @@ export function Prompt(props: PromptProps) {
   const list = createMemo(() => props.placeholders?.normal ?? [])
   const shell = createMemo(() => props.placeholders?.shell ?? [])
   const fileContextEnabled = createMemo(() => kv.get("file_context_enabled", true))
-  const editorPath = createMemo(() => (fileContextEnabled() ? editor.selection()?.filePath : undefined))
-  const editorSelectionLabel = createMemo(() => {
-    const selection = fileContextEnabled() ? editor.selection()?.selection : undefined
+  const [dismissedEditorSelectionKey, setDismissedEditorSelectionKey] = createSignal<string>()
+  const editorContext = createMemo(() => {
+    const selection = fileContextEnabled() ? editor.selection() : undefined
     if (!selection) return
-    if (selection.start.line === selection.end.line && selection.start.character === selection.end.character) return
-    if (selection.start.line === selection.end.line) return `#${selection.start.line}`
-    return `#${selection.start.line}-${selection.end.line}`
+    return editorSelectionKey(selection) === dismissedEditorSelectionKey() ? undefined : selection
+  })
+  const editorPath = createMemo(() => editorContext()?.filePath)
+  const editorSelectionLabel = createMemo(() => {
+    const ranges = editorContext()?.ranges
+    if (!ranges) return
+    const first = ranges.find(hasEditorRangeSelection) ?? ranges[0]
+    if (!first) return
+    return [getEditorRangeLabel(first), ranges.length > 1 ? `+${ranges.length - 1}` : undefined]
+      .filter(Boolean)
+      .join(" ")
   })
   const editorFileLabel = createMemo(() => {
     const value = editorPath()
@@ -147,6 +169,7 @@ export function Prompt(props: PromptProps) {
     if (!file) return
     return Locale.truncateMiddle(file, Math.max(12, Math.min(48, Math.floor(dimensions().width / 3))))
   })
+  const [editorContextHover, setEditorContextHover] = createSignal(false)
   let lastSubmittedEditorSelectionKey: string | undefined
   const [auto, setAuto] = createSignal<AutocompleteRef>()
   const currentProviderLabel = createMemo(() => local.model.parsed().provider)
@@ -163,6 +186,11 @@ export function Prompt(props: PromptProps) {
     }
   }
 
+  function dismissEditorContext() {
+    setDismissedEditorSelectionKey(editorSelectionKey(editorContext()))
+    editor.clearSelection()
+  }
+
   const textareaKeybindings = useTextareaKeybindings()
 
   const fileStyleId = syntax().getStyleId("extmark.file")!
@@ -292,6 +320,16 @@ export function Prompt(props: PromptProps) {
           dialog.clear()
         },
       },
+      {
+        title: "Remove editor context",
+        value: "prompt.editor_context.clear",
+        category: "Prompt",
+        enabled: Boolean(editorContext()),
+        onSelect: (dialog) => {
+          dismissEditorContext()
+          dialog.clear()
+        },
+      },
       {
         title: "Paste",
         value: "prompt.paste",
@@ -760,35 +798,21 @@ export function Prompt(props: PromptProps) {
     // Capture mode before it gets reset
     const currentMode = store.mode
     const variant = local.model.variant.current()
-    const editorSelection = fileContextEnabled() ? editor.selection() : undefined
-    const editorSelectionKey = editorSelection ? getEditorSelectionKey(editorSelection) : undefined
+    const editorSelection = editorContext()
+    const currentEditorSelectionKey = editorSelectionKey(editorSelection)
     const editorParts =
-      editorSelection && editorSelectionKey !== lastSubmittedEditorSelectionKey
+      editorSelection && currentEditorSelectionKey !== lastSubmittedEditorSelectionKey
         ? [
             {
               id: PartID.ascending(),
               type: "text" as const,
-              text: (() => {
-                const start = editorSelection.selection.start
-                const end = editorSelection.selection.end
-
-                let text = ""
-                if (start.line === end.line && start.character === end.character) {
-                  text = `Note: The user opened the file "${editorSelection.filePath}".`
-                } else if (start.line === end.line) {
-                  text = `Note: The user selected line ${start.line + 1} from "${editorSelection.filePath}". \`\`\`${editorSelection.text}\`\`\`\n\n`
-                } else {
-                  text = `Note: The user selected lines ${start.line + 1} to ${end.line + 1} from "${editorSelection.filePath}". \`\`\`${editorSelection.text}\`\`\`\n\n`
-                }
-
-                return `<system-reminder>${text} This may or may not be relevant to the current task.</system-reminder>\n`
-              })(),
+              text: formatEditorContext(editorSelection),
               synthetic: true,
               metadata: {
                 kind: "editor_context",
                 source: editorSelection.source ?? "editor",
                 filePath: editorSelection.filePath,
-                selection: editorSelection.selection,
+                ranges: editorSelection.ranges,
               },
             },
           ]
@@ -855,7 +879,7 @@ export function Prompt(props: PromptProps) {
           ],
         })
         .catch(() => {})
-      lastSubmittedEditorSelectionKey = editorSelectionKey
+      lastSubmittedEditorSelectionKey = currentEditorSelectionKey
     }
     history.append({
       ...store.prompt,
@@ -1406,7 +1430,18 @@ export function Prompt(props: PromptProps) {
           </Show>
           <Show when={status().type !== "retry"}>
             <box gap={2} flexDirection="row">
-              <Show when={editorFileLabelDisplay()}>{(file) => <text fg={theme.secondary}>{file()}</text>}</Show>
+              <Show when={editorFileLabelDisplay()}>
+                {(file) => (
+                  <text
+                    fg={theme.secondary}
+                    onMouseOver={() => setEditorContextHover(true)}
+                    onMouseOut={() => setEditorContextHover(false)}
+                    onMouseUp={dismissEditorContext}
+                  >
+                    {editorContextHover() ? `x ${file()}` : file()}
+                  </text>
+                )}
+              </Show>
               <Switch>
                 <Match when={store.mode === "normal"}>
                   <Switch>

+ 57 - 9
packages/opencode/src/cli/cmd/tui/context/editor-zed.ts

@@ -12,6 +12,9 @@ const ZedEditorRowSchema = z.object({
   workspace_paths: z.string().nullable(),
   timestamp: z.string(),
   buffer_path: z.string().nullable(),
+})
+
+const ZedSelectionRowSchema = z.object({
   selection_start: z.number().nullable(),
   selection_end: z.number().nullable(),
 })
@@ -24,6 +27,7 @@ const utf8 = new TextEncoder()
 
 type ZedEditorRow = z.infer<typeof ZedEditorRowSchema>
 type ZedActiveEditorRow = ZedEditorRow & { item_kind: "Editor"; editor_id: number }
+type ZedSelectionRow = z.infer<typeof ZedSelectionRowSchema>
 
 export type ZedSelectionResult =
   | { type: "selection"; selection: EditorSelection }
@@ -36,7 +40,21 @@ export async function resolveZedSelection(dbPath: string, cwd = process.cwd()):
 
   const row = active.row
   if (!row.buffer_path) return { type: "empty" }
-  if (row.selection_start == null || row.selection_end == null) return { type: "unavailable" }
+
+  const selections = queryZedEditorSelections(dbPath, row)
+  if (selections.type !== "selections") return selections
+  const byteRanges = selections.selections
+    .flatMap((selection) => {
+      if (selection.selection_start == null || selection.selection_end == null) return []
+      return [
+        {
+          start: Math.min(selection.selection_start, selection.selection_end),
+          end: Math.max(selection.selection_start, selection.selection_end),
+        },
+      ]
+    })
+    .sort((left, right) => left.start - right.start || left.end - right.end)
+  if (byteRanges.length === 0) return { type: "unavailable" }
 
   const contents = queryZedEditorContents(dbPath, row)
   const text =
@@ -47,16 +65,21 @@ export async function resolveZedSelection(dbPath: string, cwd = process.cwd()):
           .catch(() => undefined)
   if (text == null) return { type: "unavailable" }
 
-  const startOffset = utf8ByteOffsetToStringIndex(text, Math.min(row.selection_start, row.selection_end))
-  const endOffset = utf8ByteOffsetToStringIndex(text, Math.max(row.selection_start, row.selection_end))
+  const ranges = byteRanges.map((range) => {
+    const startOffset = utf8ByteOffsetToStringIndex(text, range.start)
+    const endOffset = utf8ByteOffsetToStringIndex(text, range.end)
+    return {
+      text: text.slice(startOffset, endOffset),
+      selection: offsetsToSelection(text, startOffset, endOffset),
+    }
+  })
 
   return {
     type: "selection",
     selection: {
-      text: text.slice(startOffset, endOffset),
       filePath: row.buffer_path,
       source: "zed",
-      selection: offsetsToSelection(text, startOffset, endOffset),
+      ranges,
     },
   }
 }
@@ -73,14 +96,11 @@ function queryZedActiveEditor(dbPath: string, cwd: string) {
           i.workspace_id as workspace_id,
           w.paths as workspace_paths,
           w.timestamp as timestamp,
-          e.buffer_path as buffer_path,
-          s.start as selection_start,
-          s.end as selection_end
+          e.buffer_path as buffer_path
         from items i
         join panes p on p.pane_id = i.pane_id and p.workspace_id = i.workspace_id
         join workspaces w on w.workspace_id = i.workspace_id
         left join editors e on e.item_id = i.item_id and e.workspace_id = i.workspace_id
-        left join editor_selections s on s.editor_id = e.item_id and s.workspace_id = e.workspace_id
         where i.active = 1 and p.active = 1
         order by w.timestamp desc`,
       )
@@ -108,6 +128,34 @@ function queryZedActiveEditor(dbPath: string, cwd: string) {
   }
 }
 
+function queryZedEditorSelections(dbPath: string, row: ZedActiveEditorRow) {
+  let db: Database | undefined
+  try {
+    db = new Database(dbPath, { readonly: true })
+    const raw = db
+      .query(
+        `select
+          start as selection_start,
+          end as selection_end
+        from editor_selections
+        where editor_id = $editorID and workspace_id = $workspaceID`,
+      )
+      .all({ $editorID: row.editor_id, $workspaceID: row.workspace_id })
+
+    const selections = raw.flatMap((selection) => {
+      const parsed = ZedSelectionRowSchema.safeParse(selection)
+      return parsed.success ? [parsed.data] : []
+    })
+
+    if (raw.length > 0 && selections.length === 0) return { type: "unavailable" as const }
+    return { type: "selections" as const, selections }
+  } catch {
+    return { type: "unavailable" as const }
+  } finally {
+    db?.close()
+  }
+}
+
 function queryZedEditorContents(dbPath: string, row: ZedActiveEditorRow) {
   let db: Database | undefined
   try {

+ 42 - 9
packages/opencode/src/cli/cmd/tui/context/editor.ts

@@ -28,16 +28,46 @@ const PositionSchema = z.object({
   character: z.number(),
 })
 
-const EditorSelectionSchema = z.object({
+const EditorSelectionRangeSchema = z.object({
   text: z.string(),
-  filePath: z.string(),
-  source: z.enum(["websocket", "zed"]).optional(),
   selection: z.object({
     start: PositionSchema,
     end: PositionSchema,
   }),
 })
 
+const EditorSelectionSchema = z
+  .union([
+    z.object({
+      filePath: z.string(),
+      source: z.enum(["websocket", "zed"]).optional(),
+      ranges: z.array(EditorSelectionRangeSchema).min(1),
+    }),
+    z.object({
+      text: z.string(),
+      filePath: z.string(),
+      source: z.enum(["websocket", "zed"]).optional(),
+      selection: z.object({
+        start: PositionSchema,
+        end: PositionSchema,
+      }),
+    }),
+  ])
+  .transform((value) =>
+    "ranges" in value
+      ? value
+      : {
+          filePath: value.filePath,
+          source: value.source,
+          ranges: [
+            {
+              text: value.text,
+              selection: value.selection,
+            },
+          ],
+        },
+  )
+
 const EditorMentionSchema = z.object({
   filePath: z.string(),
   lineStart: z.number(),
@@ -262,6 +292,7 @@ export const { use: useEditorContext, provider: EditorContextProvider } = create
         return store.selection
       },
       clearSelection() {
+        lastZedSelectionKey = undefined
         setStore("selection", undefined)
       },
       onMention(listener: (mention: EditorMention) => void) {
@@ -352,15 +383,17 @@ function readEditorLockFile(filePath: string): EditorLockFile | undefined {
   }
 }
 
-function editorSelectionKey(selection: EditorSelection | undefined) {
+export function editorSelectionKey(selection: EditorSelection | undefined) {
   if (!selection) return ""
   return [
     selection.filePath,
-    selection.selection.start.line,
-    selection.selection.start.character,
-    selection.selection.end.line,
-    selection.selection.end.character,
-    selection.text,
+    ...selection.ranges.flatMap((range) => [
+      range.selection.start.line,
+      range.selection.start.character,
+      range.selection.end.line,
+      range.selection.end.character,
+      range.text,
+    ]),
   ].join("\0")
 }
 

+ 97 - 28
packages/opencode/test/cli/tui/editor-context-zed.test.ts

@@ -10,6 +10,7 @@ type ZedFixtureOptions = {
   editor?: boolean
   selectionStart?: number | null
   selectionEnd?: number | null
+  selections?: Array<{ start: number | null; end: number | null }>
   contents?: string
 }
 
@@ -30,10 +31,16 @@ async function writeZedFixture(dir: string, options: ZedFixtureOptions = {}) {
   db.run("insert into items values (1, 1, 1, 1, ?)", [options.itemKind ?? "Editor"])
   if (options.editor !== false) {
     db.run("insert into editors values (1, 1, ?, ?)", [filePath, contents])
-    db.run("insert into editor_selections values (1, 1, ?, ?)", [
-      options.selectionStart === undefined ? 4 : options.selectionStart,
-      options.selectionEnd === undefined ? 7 : options.selectionEnd,
-    ])
+    ;(
+      options.selections ?? [
+        {
+          start: options.selectionStart === undefined ? 4 : options.selectionStart,
+          end: options.selectionEnd === undefined ? 7 : options.selectionEnd,
+        },
+      ]
+    ).forEach((selection) =>
+      db.run("insert into editor_selections values (1, 1, ?, ?)", [selection.start, selection.end]),
+    )
   }
   db.close()
 
@@ -66,13 +73,59 @@ test("resolveZedSelection returns active editor selection", async () => {
   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 },
+      ranges: [
+        {
+          text: "two",
+          selection: {
+            start: { line: 2, character: 1 },
+            end: { line: 2, character: 4 },
+          },
+        },
+      ],
+    },
+  })
+})
+
+test("resolveZedSelection returns all active editor selections sorted by offset", async () => {
+  await using tmp = await tmpdir()
+  const contents = "one\ntwo\nthree\nfour"
+  const fixture = await writeZedFixture(tmp.path, {
+    contents,
+    selections: [
+      {
+        start: utf8ByteOffset(contents, contents.indexOf("four")),
+        end: utf8ByteOffset(contents, contents.indexOf("four") + 4),
+      },
+      {
+        start: utf8ByteOffset(contents, contents.indexOf("two")),
+        end: utf8ByteOffset(contents, contents.indexOf("two") + 3),
       },
+    ],
+  })
+
+  expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({
+    type: "selection",
+    selection: {
+      filePath: fixture.filePath,
+      source: "zed",
+      ranges: [
+        {
+          text: "two",
+          selection: {
+            start: { line: 2, character: 1 },
+            end: { line: 2, character: 4 },
+          },
+        },
+        {
+          text: "four",
+          selection: {
+            start: { line: 4, character: 1 },
+            end: { line: 4, character: 5 },
+          },
+        },
+      ],
     },
   })
 })
@@ -90,13 +143,17 @@ test("resolveZedSelection converts Zed UTF-8 byte offsets to string offsets", as
   expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({
     type: "selection",
     selection: {
-      text: "TARGET",
       filePath: fixture.filePath,
       source: "zed",
-      selection: {
-        start: { line: 4, character: 1 },
-        end: { line: 4, character: 7 },
-      },
+      ranges: [
+        {
+          text: "TARGET",
+          selection: {
+            start: { line: 4, character: 1 },
+            end: { line: 4, character: 7 },
+          },
+        },
+      ],
     },
   })
 })
@@ -114,13 +171,17 @@ test("resolveZedSelection handles non-ASCII text inside the selected range", asy
   expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({
     type: "selection",
     selection: {
-      text: "выбор",
       filePath: fixture.filePath,
       source: "zed",
-      selection: {
-        start: { line: 3, character: 1 },
-        end: { line: 3, character: 6 },
-      },
+      ranges: [
+        {
+          text: "выбор",
+          selection: {
+            start: { line: 3, character: 1 },
+            end: { line: 3, character: 6 },
+          },
+        },
+      ],
     },
   })
 })
@@ -138,13 +199,17 @@ test("resolveZedSelection handles emoji before the selected range", async () =>
   expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({
     type: "selection",
     selection: {
-      text: "TARGET",
       filePath: fixture.filePath,
       source: "zed",
-      selection: {
-        start: { line: 2, character: 1 },
-        end: { line: 2, character: 7 },
-      },
+      ranges: [
+        {
+          text: "TARGET",
+          selection: {
+            start: { line: 2, character: 1 },
+            end: { line: 2, character: 7 },
+          },
+        },
+      ],
     },
   })
 })
@@ -162,13 +227,17 @@ test("resolveZedSelection handles reversed Zed byte offsets", async () => {
   expect(await resolveZedSelection(fixture.dbPath, tmp.path)).toEqual({
     type: "selection",
     selection: {
-      text: "TARGET",
       filePath: fixture.filePath,
       source: "zed",
-      selection: {
-        start: { line: 3, character: 1 },
-        end: { line: 3, character: 7 },
-      },
+      ranges: [
+        {
+          text: "TARGET",
+          selection: {
+            start: { line: 3, character: 1 },
+            end: { line: 3, character: 7 },
+          },
+        },
+      ],
     },
   })
 })

+ 9 - 5
packages/opencode/test/cli/tui/editor-context.test.tsx

@@ -190,13 +190,17 @@ test("useEditorContext resets selection when reconnecting", async () => {
     serverInfo: { name: "test", version: "0.0.0" },
   })
   expect(mounted.editor.selection()).toEqual({
-    text: "foo",
     filePath: path.join(startupDirectory, "file.ts"),
     source: "websocket",
-    selection: {
-      start: { line: 1, character: 1 },
-      end: { line: 1, character: 4 },
-    },
+    ranges: [
+      {
+        text: "foo",
+        selection: {
+          start: { line: 1, character: 1 },
+          end: { line: 1, character: 4 },
+        },
+      },
+    ],
   })
 
   mounted.editor.reconnect(startupDirectory)