fix(tui): keep model selection session scoped (#40913)

packages/tui/src/component/dialog-model.tsx

@@ -8,17 +8,21 @@ import * as fuzzysort from "fuzzysort"
 import { useConnected } from "./use-connected"
 import { useData } from "../context/data"
 import { modelPreferenceKey } from "../model-preference"
+import { useLocation } from "../context/location"
 
 export function DialogModel(props: { providerID?: string }) {
   const local = useLocal()
   const data = useData()
   const dialog = useDialog()
+  const location = useLocation()
   const [query, setQuery] = createSignal("")
   const favoritePriority = new Set(local.model.favorite().map(modelPreferenceKey))
 
   const connected = useConnected()
-  const providers = createMemo(() => new Map((data.location.provider.list() ?? []).map((item) => [item.id, item])))
-  const models = createMemo(() => data.location.model.list() ?? [])
+  const providers = createMemo(
+    () => new Map((data.location.provider.list(location.ref) ?? []).map((item) => [item.id, item])),
+  )
+  const models = createMemo(() => data.location.model.list(location.ref) ?? [])
 
   const showExtra = createMemo(() => connected() && !props.providerID)
 

packages/tui/src/component/prompt/index.tsx

@@ -327,10 +327,6 @@ export function Prompt(props: PromptProps) {
     if (!session) return
     const agent = session.agent && local.agent.list().find((agent) => agent.id === session.agent)
     if (agent && !args.agent) local.agent.set(agent.id)
-    if (session.model) {
-      local.model.set({ providerID: session.model.providerID, modelID: session.model.id })
-      local.model.variant.set(session.model.variant)
-    }
     syncedSessionID = sessionID
   })
 
@@ -943,15 +939,43 @@ export function Prompt(props: PromptProps) {
       await slash.command.run(slash.input)
       return true
     }
+    const inputText = expandTrackedPastedText(
+      store.prompt.text,
+      input.extmarks.getAllForTypeId(promptPartTypeId).flatMap((extmark) => {
+        const ref = store.extmarkToPart.get(extmark.id)
+        if (ref?.type !== "pasted") return []
+        const part = store.prompt.pasted[ref.index]
+        if (!part) return []
+        return [{ start: extmark.start, end: extmark.end, text: part.text }]
+      }),
+    )
+    const slashHead = parseSlashHead(inputText, /\s/)
+    const isSkill =
+      slashHead !== undefined &&
+      (data.location.skill.list(currentLocation.ref) ?? []).some(
+        (skill) => skill.slash === true && skill.id === slashHead.name,
+      )
+    const isCommand =
+      slashHead !== undefined &&
+      (data.location.command.list(currentLocation.ref) ?? []).some((command) => command.name === slashHead.name)
     const agent = local.agent.current()
     if (!agent) return false
-    const selectedModel = local.model.current()
-    if (!selectedModel) {
+    const selection = local.model.selection()
+    if (!selection) {
       void promptModelWarning()
       return false
     }
+    const usesModel = !props.sessionID || (store.mode !== "shell" && !isSkill)
+    if (usesModel && !local.model.available(selection)) {
+      toast.show({
+        title: "Model unavailable",
+        message: `${selection.providerID}/${selection.modelID} is not available in this session's location`,
+        variant: "warning",
+      })
+      return false
+    }
 
-    const variant = local.model.variant.current()
+    const variant = selection.variant
     let sessionID = props.sessionID
     let session = sessionID ? data.session.get(sessionID) : undefined
     let finishMoveProgress = false
@@ -969,8 +993,8 @@ export function Prompt(props: PromptProps) {
           location: directory ? { directory } : location,
           agent: agent.id,
           model: {
-            providerID: selectedModel.providerID,
-            id: selectedModel.modelID,
+            providerID: selection.providerID,
+            id: selection.modelID,
             variant,
           },
         })
@@ -990,17 +1014,6 @@ export function Prompt(props: PromptProps) {
       session = created
     }
 
-    const inputText = expandTrackedPastedText(
-      store.prompt.text,
-      input.extmarks.getAllForTypeId(promptPartTypeId).flatMap((extmark) => {
-        const ref = store.extmarkToPart.get(extmark.id)
-        if (ref?.type !== "pasted") return []
-        const part = store.prompt.pasted[ref.index]
-        if (!part) return []
-        return [{ start: extmark.start, end: extmark.end, text: part.text }]
-      }),
-    )
-
     // Capture mode before it gets reset
     const currentMode = store.mode
     const editorSelection = editorContext()
@@ -1013,43 +1026,30 @@ export function Prompt(props: PromptProps) {
         command: inputText,
       })
       setStore("mode", "normal")
-    } else if (
-      inputText.startsWith("/") &&
-      (data.location.command.list(currentLocation.current) ?? []).some(
-        (command) => command.name === inputText.split("\n")[0].split(" ")[0].slice(1),
-      )
-    ) {
+    } else if (slashHead && isCommand) {
       move.startSubmit()
-      // Parse command from first line, preserve multi-line content in arguments
-      const firstLineEnd = inputText.indexOf("\n")
-      const firstLine = firstLineEnd === -1 ? inputText : inputText.slice(0, firstLineEnd)
-      const [command, ...firstLineArgs] = firstLine.split(" ")
-      const restOfInput = firstLineEnd === -1 ? "" : inputText.slice(firstLineEnd + 1)
-      const args = firstLineArgs.join(" ") + (restOfInput ? "\n" + restOfInput : "")
+      const model = { providerID: selection.providerID, id: selection.modelID, variant }
+      const cancelCommit = local.model.trackSessionCommit(sessionID, model)
 
       void client.api.session
         .command({
           sessionID,
-          command: command.slice(1),
-          arguments: args,
+          command: slashHead.name,
+          arguments: slashHead.arguments,
           agent: agent.id,
-          model: { providerID: selectedModel.providerID, id: selectedModel.modelID, variant },
+          model,
           files: store.prompt.files,
           agents: store.prompt.agents,
         })
         .catch((error) => {
+          cancelCommit()
           toast.show({ title: "Failed to run command", message: errorMessage(error), variant: "error" })
         })
-    } else if (
-      inputText.startsWith("/") &&
-      (data.location.skill.list(currentLocation.current) ?? []).some(
-        (skill) => skill.slash === true && skill.id === inputText.split("\n")[0].split(" ")[0].slice(1),
-      )
-    ) {
+    } else if (isSkill) {
       move.startSubmit()
       void client.api.session.skill({
         sessionID,
-        skill: inputText.split("\n")[0].split(" ")[0].slice(1),
+        skill: slashHead!.name,
       })
     } else {
       move.startSubmit()
@@ -1061,13 +1061,15 @@ export function Prompt(props: PromptProps) {
         await client.api.session.switchAgent({ sessionID, agent: agent.id })
       }
       if (
-        session?.model?.providerID !== selectedModel.providerID ||
-        session.model.id !== selectedModel.modelID ||
+        session?.model?.providerID !== selection.providerID ||
+        session.model.id !== selection.modelID ||
         (session.model.variant ?? "default") !== (variant ?? "default")
       ) {
-        await client.api.session.switchModel({
-          sessionID,
-          model: { providerID: selectedModel.providerID, id: selectedModel.modelID, variant },
+        const model = { providerID: selection.providerID, id: selection.modelID, variant }
+        const cancelCommit = local.model.trackSessionCommit(sessionID, model)
+        await client.api.session.switchModel({ sessionID, model }).catch((error) => {
+          cancelCommit()
+          throw error
         })
       }
       if (session?.revert) {
@@ -1320,10 +1322,7 @@ export function Prompt(props: PromptProps) {
       return `Ask anything... "${list()[store.placeholder % list().length]}"`
     })()
     if (!value) return undefined
-    const width =
-      dimensions().width < 44
-        ? dimensions().width - 5
-        : Math.min(75, dimensions().width - 4) - 5
+    const width = dimensions().width < 44 ? dimensions().width - 5 : Math.min(75, dimensions().width - 4) - 5
     return Locale.takeWidth(value, Math.max(1, width)).trimEnd()
   })
   const locationLabel = createMemo(() => {

packages/tui/src/context/local.tsx

@@ -1,7 +1,7 @@
 import { createStore } from "solid-js/store"
 import { dedupeWith } from "effect/Array"
 import { createSimpleContext } from "./helper"
-import { batch, createMemo } from "solid-js"
+import { batch, createMemo, onCleanup } from "solid-js"
 import { useEvent } from "./event"
 import path from "path"
 import { useTuiPaths } from "./runtime"
@@ -22,6 +22,7 @@ import { useToast } from "../ui/toast"
 import { useRoute } from "./route"
 import { useData } from "./data"
 import { usePermission } from "./permission"
+import { useLocation } from "./location"
 
 export function parseModel(model: string) {
   const [providerID, ...rest] = model.split("/")
@@ -57,26 +58,29 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
     const args = useArgs()
     const event = useEvent()
     const permission = usePermission()
+    const location = useLocation()
+
+    const models = () => data.location.model.list(location.ref)
+    const providers = () => data.location.provider.list(location.ref)
 
     function isModelValid(model: ModelPreferenceModel) {
-      return !!data.location.model
-        .list()
-        ?.some((item) => item.providerID === model.providerID && item.id === model.modelID)
+      return !!models()?.some((item) => item.providerID === model.providerID && item.id === model.modelID)
     }
 
     function getFirstValidModel(...modelFns: (() => ModelPreferenceModel | undefined)[]) {
       for (const modelFn of modelFns) {
         const model = modelFn()
-        if (!model) continue
-        if (isModelValid(model)) return model
+        if (model && isModelValid(model)) return model
       }
     }
 
     function createAgent() {
       const agents = createMemo(() =>
-        (data.location.agent.list() ?? []).filter((agent) => agent.mode !== "subagent" && !agent.hidden),
+        (data.location.agent.list(location.ref) ?? []).filter((agent) => agent.mode !== "subagent" && !agent.hidden),
+      )
+      const visibleAgents = createMemo(() =>
+        (data.location.agent.list(location.ref) ?? []).filter((agent) => !agent.hidden),
       )
-      const visibleAgents = createMemo(() => (data.location.agent.list() ?? []).filter((agent) => !agent.hidden))
       const [agentStore, setAgentStore] = createStore({
         current: undefined as string | undefined,
       })
@@ -128,35 +132,40 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
     const agent = createAgent()
 
     function createModel() {
-      const [modelStore, setModelStore] = createStore<
-        ModelPreference & {
-          ready: boolean
-          model: Record<string, ModelPreferenceModel>
-        }
-      >({
+      type ModelSelection = ModelPreferenceModel & { variant?: string }
+      const [preferences, setPreferences] = createStore<ModelPreference & { ready: boolean }>({
         ready: false,
-        model: {},
         recent: [],
         favorite: [],
         variant: {},
       })
+      const [selectionState, setSelectionState] = createStore<{
+        newSessionModelByLocationAgent: Record<string, ModelPreferenceModel | undefined>
+        draftBySession: Record<string, ModelSelection | undefined>
+      }>({
+        newSessionModelByLocationAgent: {},
+        draftBySession: {},
+      })
 
       const repository = createModelPreferenceRepository(path.join(paths.state, "model.json"))
-      const state = {
+      const pendingSelectionCommits = new Map<string, string>()
+      const selectionKey = (value: ModelSelection) =>
+        `${modelPreferenceKey(value)}:${normalizeModelVariant(value.variant) ?? "default"}`
+      const saveState = {
         pending: false,
       }
 
-      function save() {
-        if (!modelStore.ready) {
-          state.pending = true
+      function savePreferences() {
+        if (!preferences.ready) {
+          saveState.pending = true
           return
         }
-        state.pending = false
+        saveState.pending = false
         void repository
           .patch({
-            recent: modelStore.recent,
-            favorite: modelStore.favorite,
-            variant: modelStore.variant,
+            recent: preferences.recent,
+            favorite: preferences.favorite,
+            variant: preferences.variant,
           })
           .catch(() => undefined)
       }
@@ -164,14 +173,14 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
       repository
         .load()
         .then((value) => {
-          setModelStore("recent", value.recent)
-          setModelStore("favorite", value.favorite)
-          setModelStore("variant", value.variant)
+          setPreferences("recent", value.recent)
+          setPreferences("favorite", value.favorite)
+          setPreferences("variant", value.variant)
         })
         .catch(() => {})
         .finally(() => {
-          setModelStore("ready", true)
-          if (state.pending) save()
+          setPreferences("ready", true)
+          if (saveState.pending) savePreferences()
         })
 
       const fallbackModel = createMemo(() => {
@@ -185,13 +194,13 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
           }
         }
 
-        for (const item of modelStore.recent) {
+        for (const item of preferences.recent) {
           if (isModelValid(item)) {
             return item
           }
         }
 
-        const model = data.location.model.list()?.[0]
+        const model = models()?.[0]
         if (!model) return undefined
         return {
           providerID: model.providerID,
@@ -199,30 +208,134 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
         }
       })
 
-      const currentModel = createMemo(() => {
+      const newSessionModel = createMemo(() => {
         const a = agent.current()
-        return (
-          getFirstValidModel(
-            () => a && modelStore.model[a.id],
-            () => a?.model && { providerID: a.model.providerID, modelID: a.model.id },
-            fallbackModel,
-          ) ?? undefined
+        return getFirstValidModel(
+          () => a && selectionState.newSessionModelByLocationAgent[locationAgentKey(a.id)],
+          () => a?.model && { providerID: a.model.providerID, modelID: a.model.id },
+          fallbackModel,
         )
       })
 
+      const currentSelection = createMemo<ModelSelection | undefined>(() => {
+        if (route.data.type === "session") return sessionSelection(route.data.sessionID)
+        const model = newSessionModel()
+        if (!model) return
+        return { ...model, variant: normalizeModelVariant(preferences.variant[modelPreferenceKey(model)]) }
+      })
+
+      const currentModel = createMemo(() => {
+        const selection = currentSelection()
+        if (!selection) return
+        return { providerID: selection.providerID, modelID: selection.modelID }
+      })
+
+      function locationAgentKey(agentID: string) {
+        const ref = location.ref ?? data.location.default()
+        return `${JSON.stringify([ref.directory, ref.workspaceID])}:${agentID}`
+      }
+
+      function durableSelection(sessionID: string): ModelSelection | undefined {
+        const model = data.session.get(sessionID)?.model
+        if (!model) return
+        return {
+          providerID: model.providerID,
+          modelID: model.id,
+          variant: normalizeModelVariant(model.variant),
+        }
+      }
+
+      function sessionSelection(sessionID: string) {
+        return selectionState.draftBySession[sessionID] ?? durableSelection(sessionID)
+      }
+
+      function setSessionDraft(sessionID: string, selection: ModelSelection) {
+        const durable = durableSelection(sessionID)
+        setSelectionState(
+          "draftBySession",
+          sessionID,
+          durable && selectionKey(durable) === selectionKey(selection) ? undefined : selection,
+        )
+      }
+
+      function selectModel(model: ModelPreferenceModel) {
+        if (route.data.type === "session") {
+          const sessionID = route.data.sessionID
+          const current = sessionSelection(sessionID)
+          const preferred = normalizeModelVariant(
+            current?.providerID === model.providerID && current.modelID === model.modelID
+              ? current.variant
+              : preferences.variant[modelPreferenceKey(model)],
+          )
+          const info = models()?.find((item) => item.providerID === model.providerID && item.id === model.modelID)
+          const variant = preferred && info?.variants?.some((item) => item.id === preferred) ? preferred : undefined
+          setSessionDraft(sessionID, { ...model, variant })
+          return true
+        }
+        const current = agent.current()
+        if (!current) return false
+        setSelectionState("newSessionModelByLocationAgent", locationAgentKey(current.id), model)
+        return true
+      }
+
+      onCleanup(
+        event.on("session.model.selected", (evt) => {
+          const expected = pendingSelectionCommits.get(evt.data.sessionID)
+          if (!expected) return
+          const committed = selectionKey({
+            providerID: evt.data.model.providerID,
+            modelID: evt.data.model.id,
+            variant: evt.data.model.variant,
+          })
+          if (committed !== expected) return
+          pendingSelectionCommits.delete(evt.data.sessionID)
+          const draft = selectionState.draftBySession[evt.data.sessionID]
+          if (draft && selectionKey(draft) === committed)
+            setSelectionState("draftBySession", evt.data.sessionID, undefined)
+        }),
+      )
+
+      onCleanup(
+        event.on("session.deleted", (evt) => {
+          pendingSelectionCommits.delete(evt.data.sessionID)
+          setSelectionState("draftBySession", evt.data.sessionID, undefined)
+        }),
+      )
+
       return {
         current: currentModel,
+        selection: currentSelection,
+        available(model = currentModel()) {
+          return model ? isModelValid(model) : false
+        },
+        trackSessionCommit(
+          sessionID: string,
+          value: {
+            providerID: string
+            id: string
+            variant?: string
+          },
+        ) {
+          const committed = selectionKey({ providerID: value.providerID, modelID: value.id, variant: value.variant })
+          pendingSelectionCommits.set(sessionID, committed)
+          return () => {
+            if (pendingSelectionCommits.get(sessionID) === committed) pendingSelectionCommits.delete(sessionID)
+          }
+        },
         get ready() {
-          return modelStore.ready
+          return preferences.ready
+        },
+        get catalogReady() {
+          return models() !== undefined
         },
         recent() {
-          return modelStore.recent
+          return preferences.recent
         },
         favorite() {
-          return modelStore.favorite
+          return preferences.favorite
         },
         parsed: createMemo(() => {
-          const value = currentModel()
+          const value = currentSelection()
           if (!value) {
             return {
               provider: "Connect a provider",
@@ -230,33 +343,28 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
               reasoning: false,
             }
           }
-          const provider = data.location.provider.list()?.find((item) => item.id === value.providerID)
-          const info = data.location.model
-            .list()
-            ?.find((item) => item.providerID === value.providerID && item.id === value.modelID)
+          const provider = providers()?.find((item) => item.id === value.providerID)
+          const info = models()?.find((item) => item.providerID === value.providerID && item.id === value.modelID)
           return {
             provider: provider?.name ?? value.providerID,
-            model: info?.name ?? value.modelID,
+            model: info?.name ?? `${value.modelID} (unavailable)`,
             reasoning: (info?.variants?.length ?? 0) !== 0,
           }
         }),
         cycle(direction: 1 | -1) {
-          const current = currentModel()
+          const current = currentSelection()
           if (!current) return
-          const recent = modelStore.recent
+          const recent = recentModels(current, preferences.recent).filter(isModelValid)
           const index = recent.findIndex((x) => x.providerID === current.providerID && x.modelID === current.modelID)
-          if (index === -1) return
-          let next = index + direction
+          let next = index === -1 ? (direction === 1 ? 0 : recent.length - 1) : index + direction
           if (next < 0) next = recent.length - 1
           if (next >= recent.length) next = 0
           const val = recent[next]
           if (!val) return
-          const a = agent.current()
-          if (!a) return
-          setModelStore("model", a.id, { ...val })
+          selectModel({ ...val })
         },
         cycleFavorite(direction: 1 | -1) {
-          const favorites = modelStore.favorite.filter((item) => isModelValid(item))
+          const favorites = preferences.favorite.filter((item) => isModelValid(item))
           if (!favorites.length) {
             toast.show({
               variant: "info",
@@ -265,7 +373,7 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
             })
             return
           }
-          const current = currentModel()
+          const current = currentSelection()
           let index = -1
           if (current) {
             index = favorites.findIndex((x) => x.providerID === current.providerID && x.modelID === current.modelID)
@@ -279,45 +387,39 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
           }
           const next = favorites[index]
           if (!next) return
-          const a = agent.current()
-          if (!a) return
-          setModelStore("model", a.id, { ...next })
-          setModelStore("recent", recentModels(next, modelStore.recent))
-          save()
+          if (!selectModel({ ...next })) return
+          setPreferences("recent", recentModels(next, preferences.recent))
+          savePreferences()
         },
         set(model: { providerID: string; modelID: string }, options?: { recent?: boolean }) {
           batch(() => {
             if (!isModelValid(model)) return
-            const a = agent.current()
-            if (!a) return
-            setModelStore("model", a.id, model)
+            if (!selectModel(model)) return
             if (options?.recent) {
-              setModelStore("recent", recentModels(model, modelStore.recent))
-              save()
+              setPreferences("recent", recentModels(model, preferences.recent))
+              savePreferences()
             }
           })
         },
         toggleFavorite(model: { providerID: string; modelID: string }) {
           batch(() => {
             if (!isModelValid(model)) return
-            const exists = modelStore.favorite.some(
+            const exists = preferences.favorite.some(
               (x) => x.providerID === model.providerID && x.modelID === model.modelID,
             )
             const next = exists
-              ? modelStore.favorite.filter((x) => x.providerID !== model.providerID || x.modelID !== model.modelID)
-              : [model, ...modelStore.favorite]
-            setModelStore(
+              ? preferences.favorite.filter((x) => x.providerID !== model.providerID || x.modelID !== model.modelID)
+              : [model, ...preferences.favorite]
+            setPreferences(
               "favorite",
               next.map((x) => ({ providerID: x.providerID, modelID: x.modelID })),
             )
-            save()
+            savePreferences()
           })
         },
         variant: {
           selected() {
-            const m = currentModel()
-            if (!m) return undefined
-            return normalizeModelVariant(modelStore.variant[modelPreferenceKey(m)])
+            return currentSelection()?.variant
           },
           current() {
             const v = this.selected()
@@ -325,18 +427,20 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({
             return undefined
           },
           list() {
-            const m = currentModel()
+            const m = currentSelection()
             if (!m) return []
-            const info = data.location.model
-              .list()
-              ?.find((item) => item.providerID === m.providerID && item.id === m.modelID)
+            const info = models()?.find((item) => item.providerID === m.providerID && item.id === m.modelID)
             return info?.variants?.map((variant) => variant.id) ?? []
           },
           set(value: string | undefined) {
-            const m = currentModel()
+            const m = currentSelection()
             if (!m) return
-            setModelStore("variant", modelPreferenceKey(m), normalizeModelVariant(value))
-            save()
+            if (route.data.type === "session") {
+              setSessionDraft(route.data.sessionID, { ...m, variant: normalizeModelVariant(value) })
+              return
+            }
+            setPreferences("variant", modelPreferenceKey(m), normalizeModelVariant(value))
+            savePreferences()
           },
           cycle() {
             const variants = this.list()

packages/tui/src/routes/session/index.tsx

@@ -361,7 +361,7 @@ export function Session() {
 
   createEffect(() => {
     const current = prompt()
-    if (sent || !current || !synced() || !local.model.ready) return
+    if (sent || !current || !synced() || !local.model.ready || !local.model.catalogReady) return
     if (!local.agent.current() || !local.model.current()) return
     if (!args.prompt || route.prompt?.text !== args.prompt || current.current.text !== args.prompt) return
     sent = true