Explorar o código

fix(app): preserve model selection across draft remounts

Brendan Allan hai 1 mes
pai
achega
67225ae656

+ 53 - 34
packages/app/src/context/local.tsx

@@ -1,6 +1,6 @@
 import { createSimpleContext } from "@opencode-ai/ui/context"
 import { base64Encode } from "@opencode-ai/core/util/encode"
-import { useParams } from "@solidjs/router"
+import { useParams, useSearchParams } from "@solidjs/router"
 import { batch, createEffect, createMemo, startTransition } from "solid-js"
 import { createStore } from "solid-js/store"
 import { useModels } from "@/context/models"
@@ -11,21 +11,22 @@ import { useSDK } from "./sdk"
 import { useSync } from "./sync"
 import { useServerSDK } from "./server-sdk"
 import { ScopedKey, type ServerScope } from "@/utils/server-scope"
+import {
+  cloneModelSelectionState,
+  createModelSelectionState,
+  createTabModelSelectionState,
+  type ModelSelectionState,
+} from "./model-selection-state"
+import { useTabs } from "./tabs"
 
 export type ModelKey = { providerID: string; modelID: string; variant?: string }
 
-type State = {
-  agent?: string
-  model?: ModelKey
-  variant?: string | null
-}
-
 type Saved = {
-  session: Record<string, State | undefined>
+  session: Record<string, ModelSelectionState | undefined>
 }
 
 const WORKSPACE_KEY = "__workspace__"
-const handoff = new Map<string, State>()
+const handoff = new Map<string, ModelSelectionState>()
 
 const handoffKey = (scope: ServerScope, dir: string, id: string) => ScopedKey.from(scope, dir, id)
 
@@ -33,8 +34,8 @@ const migrate = (value: unknown) => {
   if (!value || typeof value !== "object") return { session: {} }
 
   const item = value as {
-    session?: Record<string, State | undefined>
-    pick?: Record<string, State | undefined>
+    session?: Record<string, ModelSelectionState | undefined>
+    pick?: Record<string, ModelSelectionState | undefined>
   }
 
   if (item.session && typeof item.session === "object") return { session: item.session }
@@ -45,14 +46,6 @@ const migrate = (value: unknown) => {
   }
 }
 
-const clone = (value: State | undefined) => {
-  if (!value) return
-  return {
-    ...value,
-    model: value.model ? { ...value.model } : undefined,
-  } satisfies State
-}
-
 export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
   name: "Local",
   init: () => {
@@ -62,6 +55,8 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
     const serverSDK = useServerSDK()
     const providers = useProviders(() => sdk().directory)
     const models = useModels()
+    const tabs = useTabs()
+    const [search] = useSearchParams<{ draftId?: string }>()
 
     const id = createMemo(() => params.id || undefined)
     const list = createMemo(() => sync().data.agent.filter((item) => item.mode !== "subagent" && !item.hidden))
@@ -79,8 +74,8 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
 
     const [store, setStore] = createStore<{
       current?: string
-      draft?: State
-      promoting?: State
+      draft?: ModelSelectionState
+      promoting?: ModelSelectionState
       last?: {
         type: "agent" | "model" | "variant"
         agent?: string
@@ -122,10 +117,33 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
       setStore("current", items[0]?.name)
     })
 
-    const scope = createMemo<State | undefined>(() => {
+    const draft = createMemo(() => {
+      const draftID = search.draftId
+      if (!draftID) return
+      const tab = tabs.store.find((item) => item.type === "draft" && item.draftID === draftID)
+      if (!tab || tab.type !== "draft") return
+      return createTabModelSelectionState(tabs, tab, () => {
+        const [state, setState] = persisted(
+          Persist.draft(tab.draftID, "model-selection"),
+          createStore<{ value?: ModelSelectionState }>({}),
+        )
+        return createModelSelectionState([state, setState])
+      })
+    })
+
+    const setDraft = (value: ModelSelectionState | undefined) => {
+      const state = draft()
+      if (state) {
+        state.set(value)
+        return
+      }
+      setStore("draft", cloneModelSelectionState(value))
+    }
+
+    const scope = createMemo<ModelSelectionState | undefined>(() => {
       const session = id()
-      if (!session) return store.draft ?? store.promoting
-      return saved.session[session] ?? handoff.get(handoffKey(serverSDK().scope, sdk().directory, session))
+      if (session) return saved.session[session] ?? handoff.get(handoffKey(serverSDK().scope, sdk().directory, session))
+      return draft()?.current() ?? store.draft ?? store.promoting
     })
 
     createEffect(() => {
@@ -141,7 +159,7 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
         return
       }
 
-      setSaved("session", session, clone(next))
+      setSaved("session", session, cloneModelSelectionState(next))
       handoff.delete(key)
       setStore("promoting", undefined)
     })
@@ -203,13 +221,13 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
             agent: item.name,
             model: item.model ?? prev?.model,
             variant: item.variant ?? prev?.variant,
-          } satisfies State
+          } satisfies ModelSelectionState
           const session = id()
           if (session) {
             setSaved("session", session, next)
             return
           }
-          setStore("draft", next)
+          setDraft(next)
         })
       },
       move(direction: 1 | -1) {
@@ -256,21 +274,21 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
         agent: agent.current()?.name,
         model: model ? { providerID: model.provider.id, modelID: model.id } : undefined,
         variant: selected(),
-      } satisfies State
+      } satisfies ModelSelectionState
     }
 
-    const write = (next: Partial<State>) => {
+    const write = (next: Partial<ModelSelectionState>) => {
       const state = {
         ...(scope() ?? { agent: agent.current()?.name }),
         ...next,
-      } satisfies State
+      } satisfies ModelSelectionState
 
       const session = id()
       if (session) {
         setSaved("session", session, state)
         return
       }
-      setStore("draft", state)
+      setDraft(state)
     }
 
     const recent = createMemo(() => models.recent.list().map(models.find).filter(Boolean))
@@ -376,10 +394,11 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
       agent,
       session: {
         reset() {
-          setStore({ draft: undefined, promoting: undefined })
+          setDraft(undefined)
+          setStore("promoting", undefined)
         },
         promote(dir: string, session: string) {
-          const next = clone(snapshot())
+          const next = cloneModelSelectionState(snapshot())
           if (!next) return
           const key = handoffKey(serverSDK().scope, dir, session)
           handoff.set(key, next)
@@ -389,7 +408,7 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
           }
 
           setStore("promoting", next)
-          setStore("draft", undefined)
+          setDraft(undefined)
         },
         restore(msg: { sessionID: string; agent: string; model: ModelKey }) {
           const session = id()

+ 39 - 0
packages/app/src/context/model-selection-state.test.ts

@@ -0,0 +1,39 @@
+import { describe, expect, test } from "bun:test"
+import { createRoot, getOwner } from "solid-js"
+import { createTabMemory } from "./tab-memory"
+import { createTabModelSelectionState } from "./model-selection-state"
+import type { DraftTab } from "./tabs"
+
+describe("model selection state", () => {
+  test("keeps model selection scoped to its draft tab across consumers", () => {
+    createRoot((dispose) => {
+      const memory = createTabMemory(getOwner())
+      const tabs = {
+        state<T>(tab: DraftTab, name: string, init: () => T) {
+          return memory.ensure(`draft:${tab.draftID}`, name, init)
+        },
+      }
+      const first = { type: "draft", draftID: "first", server: "server", directory: "/repo" } as DraftTab
+      const second = { type: "draft", draftID: "second", server: "server", directory: "/repo" } as DraftTab
+
+      createTabModelSelectionState(tabs, first).set({
+        agent: "build",
+        model: { providerID: "anthropic", modelID: "claude" },
+      })
+      createTabModelSelectionState(tabs, second).set({
+        agent: "build",
+        model: { providerID: "openai", modelID: "gpt" },
+      })
+
+      expect(createTabModelSelectionState(tabs, first).current()?.model).toEqual({
+        providerID: "anthropic",
+        modelID: "claude",
+      })
+      expect(createTabModelSelectionState(tabs, second).current()?.model).toEqual({
+        providerID: "openai",
+        modelID: "gpt",
+      })
+      dispose()
+    })
+  })
+})

+ 38 - 0
packages/app/src/context/model-selection-state.ts

@@ -0,0 +1,38 @@
+import { createStore, type SetStoreFunction, type Store } from "solid-js/store"
+import type { DraftTab, useTabs } from "./tabs"
+
+export type ModelSelectionState = {
+  agent?: string
+  model?: { providerID: string; modelID: string; variant?: string }
+  variant?: string | null
+}
+
+export function cloneModelSelectionState(value: ModelSelectionState | undefined) {
+  if (!value) return
+  return {
+    ...value,
+    model: value.model ? { ...value.model } : undefined,
+  } satisfies ModelSelectionState
+}
+
+type ModelSelectionStore = { value?: ModelSelectionState }
+
+export function createModelSelectionState(
+  state: [Store<ModelSelectionStore>, SetStoreFunction<ModelSelectionStore>] = createStore<ModelSelectionStore>({}),
+) {
+  const [store, setStore] = state
+  return {
+    current: () => store.value,
+    set(value: ModelSelectionState | undefined) {
+      setStore("value", cloneModelSelectionState(value))
+    },
+  }
+}
+
+export function createTabModelSelectionState(
+  tabs: Pick<ReturnType<typeof useTabs>, "state">,
+  tab: DraftTab,
+  init: () => ReturnType<typeof createModelSelectionState> = createModelSelectionState,
+) {
+  return tabs.state(tab, "model-selection", init)
+}