fix(core): better state handling of editor context (#25911)

packages/opencode/test/cli/tui/editor-context.test.tsx

@@ -59,6 +59,39 @@ function createWebSocketImpl(...sockets: FakeWebSocket[]) {
   } as unknown as typeof WebSocket
 }
 
+function sendSelection(socket: FakeWebSocket, filePath: string, text = "foo") {
+  socket.message(
+    JSON.stringify({
+      jsonrpc: "2.0",
+      method: "selection_changed",
+      params: {
+        text,
+        filePath,
+        selection: {
+          start: { line: 1, character: 1 },
+          end: { line: 1, character: 4 },
+        },
+      },
+    }),
+  )
+}
+
+function expectedSelection(filePath: string, text = "foo") {
+  return {
+    filePath,
+    source: "websocket" as const,
+    ranges: [
+      {
+        text,
+        selection: {
+          start: { line: 1, character: 1 },
+          end: { line: 1, character: 4 },
+        },
+      },
+    ],
+  }
+}
+
 test("useEditorContext reconnect switches editor server by session directory", async () => {
   await using tmp = await tmpdir()
   const startupDirectory = path.join(tmp.path, "startup")
@@ -93,12 +126,18 @@ test("useEditorContext reconnect switches editor server by session directory", a
   await nextTick()
 
   expect(firstSocket.closed).toBeFalse()
+  sendSelection(firstSocket, path.join(startupDirectory, "file.ts"))
+
+  expect(mounted.editor.selection()).toEqual(expectedSelection(path.join(startupDirectory, "file.ts")))
+  expect(mounted.editor.labelState()).toBe("pending")
 
   mounted.editor.reconnect(sessionDirectory)
   await nextTick()
 
   expect(firstSocket.closed).toBeTrue()
   expect(secondSocket.closed).toBeFalse()
+  expect(mounted.editor.selection()).toBeUndefined()
+  expect(mounted.editor.labelState()).toBe("none")
 
   mounted.dispose()
 })
@@ -131,7 +170,7 @@ test("useEditorContext favors configured port over lock files", async () => {
   mounted.dispose()
 })
 
-test("useEditorContext resets selection when reconnecting", async () => {
+test("useEditorContext clears selection when reconnecting", async () => {
   await using tmp = await tmpdir()
   const startupDirectory = path.join(tmp.path, "startup")
   const ideDirectory = path.join(tmp.path, ".claude", "ide")
@@ -169,45 +208,66 @@ test("useEditorContext resets selection when reconnecting", async () => {
       },
     }),
   )
-  socket.message(
-    JSON.stringify({
-      jsonrpc: "2.0",
-      method: "selection_changed",
-      params: {
-        text: "foo",
-        filePath: path.join(startupDirectory, "file.ts"),
-        selection: {
-          start: { line: 1, character: 1 },
-          end: { line: 1, character: 4 },
-        },
-      },
-    }),
-  )
+  sendSelection(socket, path.join(startupDirectory, "file.ts"))
 
   expect(mounted.editor.connected()).toBeTrue()
   expect(mounted.editor.server()).toEqual({
     protocolVersion: "2025-11-25",
     serverInfo: { name: "test", version: "0.0.0" },
   })
-  expect(mounted.editor.selection()).toEqual({
-    filePath: path.join(startupDirectory, "file.ts"),
-    source: "websocket",
-    ranges: [
-      {
-        text: "foo",
-        selection: {
-          start: { line: 1, character: 1 },
-          end: { line: 1, character: 4 },
-        },
-      },
-    ],
-  })
+  expect(mounted.editor.selection()).toEqual(expectedSelection(path.join(startupDirectory, "file.ts")))
+  expect(mounted.editor.labelState()).toBe("pending")
+  mounted.editor.markSelectionSent()
+  expect(mounted.editor.labelState()).toBe("sent")
 
   mounted.editor.reconnect(startupDirectory)
 
   expect(socket.closed).toBeFalse()
   expect(mounted.editor.connected()).toBeTrue()
   expect(mounted.editor.selection()).toBeUndefined()
+  expect(mounted.editor.labelState()).toBe("none")
+
+  mounted.dispose()
+})
+
+test("useEditorContext preserves selection for the next reconnect when requested", async () => {
+  await using tmp = await tmpdir()
+  const startupDirectory = path.join(tmp.path, "startup")
+  const ideDirectory = path.join(tmp.path, ".claude", "ide")
+  await mkdir(startupDirectory, { recursive: true })
+  await mkdir(ideDirectory, { recursive: true })
+  await writeFile(
+    path.join(ideDirectory, "3001.lock"),
+    JSON.stringify({
+      transport: "ws",
+      workspaceFolders: [startupDirectory],
+    }),
+  )
+
+  process.env.CLAUDE_CODE_SSE_PORT = undefined
+  process.env.OPENCODE_EDITOR_SSE_PORT = undefined
+  spyOn(process, "cwd").mockImplementation(() => startupDirectory)
+  spyOn(os, "homedir").mockImplementation(() => tmp.path)
+  const socket = new FakeWebSocket("ws://127.0.0.1:3001")
+
+  const mounted = mountEditorContext(createWebSocketImpl(socket))
+  await nextTick()
+
+  sendSelection(socket, path.join(startupDirectory, "file.ts"))
+  expect(mounted.editor.selection()).toEqual(expectedSelection(path.join(startupDirectory, "file.ts")))
+
+  mounted.editor.markSelectionSent()
+  mounted.editor.preserveSelectionFromNewSession()
+  mounted.editor.reconnect(startupDirectory)
+
+  expect(socket.closed).toBeFalse()
+  expect(mounted.editor.selection()).toEqual(expectedSelection(path.join(startupDirectory, "file.ts")))
+  expect(mounted.editor.labelState()).toBe("sent")
+
+  mounted.editor.reconnect(startupDirectory)
+
+  expect(mounted.editor.selection()).toBeUndefined()
+  expect(mounted.editor.labelState()).toBe("none")
 
   mounted.dispose()
 })