Просмотр исходного кода

feat(app): keep prompt state in tabs (#33566)

Brendan Allan 1 месяц назад
Родитель
Сommit
fed4f4d86b

+ 23 - 4
packages/app/src/components/titlebar.tsx

@@ -41,6 +41,8 @@ import { tabHref, useTabs } from "@/context/tabs"
 import "./titlebar.css"
 import { useServerSDK } from "@/context/server-sdk"
 import { Session } from "@opencode-ai/sdk/v2"
+import { base64Encode } from "@opencode-ai/core/util/encode"
+import { createTabPromptState } from "@/context/prompt"
 
 type TauriDesktopWindow = {
   startDragging?: () => Promise<void>
@@ -523,13 +525,18 @@ export function Titlebar(props: { update?: TitlebarUpdate }) {
                             )
                           }
 
+                          const sdk = createMemo(() => {
+                            const conn = server.list.find((s) => ServerConnection.key(s) === tab.server)
+                            if (!conn) return null
+                            const { sdk } = global.createServerCtx(conn)
+                            return sdk
+                          })
                           const [session] = createResource(
                             () => {
                               const id = tab.sessionId
-                              const conn = server.list.find((s) => ServerConnection.key(s) === tab.server)
-                              if (!conn) return null
-                              const { sdk } = global.createServerCtx(conn)
-                              return { id, sdk }
+                              const _sdk = sdk()
+                              if (!_sdk) return null
+                              return { id, sdk: _sdk }
                             },
                             ({ id, sdk }) =>
                               sdk.client.session
@@ -538,6 +545,18 @@ export function Titlebar(props: { update?: TitlebarUpdate }) {
                                 .catch(() => undefined),
                           )
 
+                          createEffect(() => {
+                            if (tab.type !== "session") return
+                            const _sdk = sdk()
+                            if (!_sdk) return
+                            const sess = session()
+                            if (!sess) return
+                            createTabPromptState(tabs, tab, _sdk.scope, {
+                              dir: base64Encode(sess.directory),
+                              id: sess.id,
+                            })
+                          })
+
                           return (
                             <>
                               {divider()}

+ 32 - 1
packages/app/src/context/prompt.tsx

@@ -8,6 +8,10 @@ import { Persist, persisted } from "@/utils/persist"
 import { useServerSDK } from "./server-sdk"
 import type { ServerScope } from "@/utils/server-scope"
 import { useSDK } from "./sdk"
+import { useTabs, type Tab } from "./tabs"
+import { useServer } from "./server"
+import { requireServerKey } from "@/utils/session-route"
+import { useSettings } from "./settings"
 
 interface PartBase {
   content: string
@@ -258,14 +262,23 @@ export function createPromptState() {
   }
 }
 
+export const createTabPromptState = (
+  tabs: ReturnType<typeof useTabs>,
+  tab: Tab,
+  ...args: Parameters<typeof createPromptSession>
+) => tabs.state(tab, "prompt", () => createPromptSession(...args))
+
 export const { use: usePrompt, provider: PromptProvider } = createSimpleContext({
   name: "Prompt",
   gate: false,
   init: () => {
-    const params = useParams()
+    const params = useParams<{ serverKey?: string; id?: string }>()
     const sdk = useSDK()
     const [search] = useSearchParams<{ draftId?: string }>()
     const serverSDK = useServerSDK()
+    const server = useServer()
+    const tabs = useTabs()
+    const settings = useSettings()
     const cache = new Map<string, PromptCacheEntry>()
 
     const disposeAll = () => {
@@ -288,7 +301,25 @@ export const { use: usePrompt, provider: PromptProvider } = createSimpleContext(
     }
 
     const owner = getOwner()
+    const tab = createMemo<Tab | undefined>(() => {
+      if (!settings.general.newLayoutDesigns()) return
+      if (search.draftId) {
+        return tabs.store.find((item) => item.type === "draft" && item.draftID === search.draftId)
+      }
+      if (!params.id) return
+      const serverKey = params.serverKey ? requireServerKey(params.serverKey) : server.key
+      return (
+        tabs.store.find(
+          (item) => item.type === "session" && item.server === serverKey && item.sessionId === params.id,
+        ) ?? { type: "session", server: serverKey, sessionId: params.id }
+      )
+    })
     const load = (scope: Scope) => {
+      const current = tab()
+      if (current) {
+        return createTabPromptState(tabs, current, serverSDK().scope, scope)
+      }
+
       const key = scopeKey(scope)
       const existing = cache.get(key)
       if (existing) {

+ 33 - 0
packages/app/src/context/tab-memory.ts

@@ -0,0 +1,33 @@
+import { createRoot, type Owner } from "solid-js"
+
+type Entry = {
+  value: unknown
+  dispose: VoidFunction
+}
+
+export function createTabMemory(owner: Owner | null) {
+  const entries = new Map<string, Map<string, Entry>>()
+
+  const remove = (key: string) => {
+    const state = entries.get(key)
+    if (!state) return
+    for (const entry of state.values()) entry.dispose()
+    entries.delete(key)
+  }
+
+  return {
+    ensure<T>(key: string, name: string, init: () => T) {
+      const state = entries.get(key) ?? new Map<string, Entry>()
+      if (!entries.has(key)) entries.set(key, state)
+      const existing = state.get(name)
+      if (existing) return existing.value as T
+      const entry = createRoot((dispose) => ({ value: init(), dispose }), owner)
+      state.set(name, entry)
+      return entry.value
+    },
+    remove,
+    dispose() {
+      for (const key of entries.keys()) remove(key)
+    },
+  }
+}

+ 24 - 0
packages/app/src/context/tabs.test.ts

@@ -0,0 +1,24 @@
+import { describe, expect, test } from "bun:test"
+import { createRoot, getOwner, onCleanup } from "solid-js"
+import { createTabMemory } from "./tab-memory"
+
+describe("tab memory", () => {
+  test("keeps state until its tab is removed", () => {
+    createRoot((dispose) => {
+      const memory = createTabMemory(getOwner())
+      let disposed = 0
+      const first = memory.ensure("tab", "prompt", () => {
+        onCleanup(() => disposed++)
+        return { value: "prompt" }
+      })
+
+      expect(memory.ensure("tab", "prompt", () => ({ value: "other" }))).toBe(first)
+      expect(memory.ensure("other", "prompt", () => ({ value: "other" }))).not.toBe(first)
+
+      memory.remove("tab")
+      expect(disposed).toBe(1)
+      expect(memory.ensure("tab", "prompt", () => ({ value: "new" }))).not.toBe(first)
+      dispose()
+    })
+  })
+})

+ 18 - 2
packages/app/src/context/tabs.tsx

@@ -3,12 +3,13 @@ import { createSimpleContext } from "@opencode-ai/ui/context"
 import { createStore, produce } from "solid-js/store"
 import { Persist, persisted, removePersisted, draftPersistedKeys } from "@/utils/persist"
 import { ServerConnection, useServer } from "./server"
-import { createEffect, startTransition } from "solid-js"
+import { createEffect, getOwner, onCleanup, startTransition } from "solid-js"
 import { useLocation, useNavigate, useParams } from "@solidjs/router"
 import { usePlatform } from "./platform"
 import { uuid } from "@/utils/uuid"
 import { SessionTabsRemovedDetail } from "@/components/titlebar-session-events"
 import { sessionHref } from "@/utils/session-route"
+import { createTabMemory } from "./tab-memory"
 
 export type SessionTab = {
   type: "session"
@@ -66,6 +67,7 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
     const params = useParams()
     const navigate = useNavigate()
     const location = useLocation()
+    const memory = createTabMemory(getOwner())
 
     const closing = new Set<string>()
     let recentWrite = 0
@@ -89,11 +91,18 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
       for (const key of draftPersistedKeys()) removePersisted(Persist.draft(draftID, key), platform)
     }
 
+    onCleanup(memory.dispose)
+
     createEffect(() => {
       if (!ready() || !recentReady()) return
       const servers = new Set(server.list.map(ServerConnection.key))
       const next = store.filter((tab) => servers.has(tab.server))
-      if (next.length !== store.length) setStore(() => next)
+      if (next.length !== store.length) {
+        for (const tab of store) {
+          if (!servers.has(tab.server)) memory.remove(tabKey(tab))
+        }
+        setStore(() => next)
+      }
       if (recent.key && !next.some((tab) => tabKey(tab) === recent.key)) setRecentKey(undefined)
     })
 
@@ -151,6 +160,7 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
           if (recent.key === `draft:${draftID}`) setRecentKey(tabKey(next))
           if (active) navigateTab(next)
         })
+        memory.remove(`draft:${draftID}`)
         removeDraftPersisted(draftID)
       },
       removeTab: (index: number) => {
@@ -170,12 +180,14 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
           if (nextTab) navigateTab(nextTab)
           else navigate("/")
         }).finally(() => closing.delete(key))
+        memory.remove(key)
         if (draftID) removeDraftPersisted(draftID)
       },
       removeServer(key: ServerConnection.Key) {
         const drafts = store.flatMap((tab) => (tab.type === "draft" && tab.server === key ? [tab.draftID] : []))
         const removed = store.filter((tab) => tab.server === key).map(tabKey)
         setStore((tabs) => tabs.filter((tab) => tab.server !== key))
+        for (const key of removed) memory.remove(key)
         if (recent.key && removed.includes(recent.key)) setRecentKey(undefined)
         for (const draftID of drafts) removeDraftPersisted(draftID)
         if (server.key === key) navigate("/")
@@ -227,6 +239,7 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
           )
           if (recent.key && removed.includes(recent.key)) setRecentKey(undefined)
         })
+        for (const key of removed) memory.remove(key)
       },
       select: navigateTab,
       remember(tab: Tab) {
@@ -246,6 +259,9 @@ export const { use: useTabs, provider: TabsProvider } = createSimpleContext({
         }
         navigate("/")
       },
+      state<T>(tab: Tab, name: string, init: () => T) {
+        return memory.ensure(tabKey(tab), name, init)
+      },
     }
 
     return { ...actions, store, ready, recentReady }

+ 2 - 0
packages/app/test-browser/prompt-persistence.test.ts

@@ -20,6 +20,8 @@ beforeAll(async () => {
   mock.module("@solidjs/router", () => ({
     useParams: () => ({}),
     useSearchParams: () => [{}],
+    useLocation: () => ({ pathname: "", query: {} }),
+    useNavigate: () => () => undefined,
   }))
   mock.module("@opencode-ai/ui/context", () => ({
     createSimpleContext: () => ({