From cae1eed1b18e6019a9b8789f5c9fa5a53353fe0a Mon Sep 17 00:00:00 2001
From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com>
Date: Mon, 27 Jul 2026 17:12:29 +0800
Subject: [PATCH] feat: persist user (provider, model) selection per notebook
 in localStorage

When the user picks a provider+model in the inline prompt, remember
that choice across panel close/reopen AND across JupyterLab restarts,
keyed by notebookPath. If the stored provider (or the model under it)
is no longer present in the current providers payload, fall back to
the default (first provider / default[pid] model) - never silently
apply a stale value that would result in an empty or disabled select.

- New module src/components/model_selection.ts with
  loadModelSelection / saveModelSelection / clearModelSelection.
  Best-effort: localStorage may be disabled (private mode), quota
  may be exhausted, value may be corrupt - all failure modes return
  null / no-op rather than throw.
- IOpenCodeInlinePromptOptions gains notebookPath?: string.
- Constructor resolves the initial selection: prefer stored+valid,
  else first provider / default[pid]. providerSelect.change and
  modelSelect.change listeners call _persistSelection() to keep
  storage in sync.
- _rebuildModelSelect now accepts a preferredModelId; if it exists
  in the current provider's models it's used, else the default.
- OpenCodeCellActions._showPrompt passes notebookPath =
  this._context?.notebookPath so the prompt can key the storage.

Tests
-----
- 10 unit tests in model_selection.spec.ts (load/save/clear +
  malformed JSON + per-notebook independence).
- 4 prompt behavior tests in opencode_cell_actions.spec.ts
  (valid stored applied, provider removed -> default, model removed
  -> default, change -> saved).
- 72 frontend jest (was 58), 52 backend pytest (unchanged),
  build green.

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/__tests__/model_selection.spec.ts       |  99 ++++++++++++++++
 src/__tests__/opencode_cell_actions.spec.ts | 122 ++++++++++++++++++++
 src/components/model_selection.ts           |  82 +++++++++++++
 src/components/opencode_cell_actions.ts     |   1 +
 src/components/opencode_inline_prompt.ts    |  96 ++++++++++++++-
 5 files changed, 395 insertions(+), 5 deletions(-)
 create mode 100644 src/__tests__/model_selection.spec.ts
 create mode 100644 src/components/model_selection.ts

diff --git a/src/__tests__/model_selection.spec.ts b/src/__tests__/model_selection.spec.ts
new file mode 100644
index 0000000..5afcd54
--- /dev/null
+++ b/src/__tests__/model_selection.spec.ts
@@ -0,0 +1,99 @@
+/**
+ * Tests for the localStorage-backed model selection persistence helper.
+ * jsdom provides a real (per-test-class) localStorage.
+ */
+import {
+  clearModelSelection,
+  loadModelSelection,
+  saveModelSelection
+} from '../components/model_selection';
+
+describe('loadModelSelection', () => {
+  beforeEach(() => {
+    localStorage.clear();
+  });
+
+  it('returns null when no entry exists', () => {
+    expect(loadModelSelection('foo.ipynb')).toBeNull();
+  });
+
+  it('returns null for an empty notebookPath (defensive)', () => {
+    expect(loadModelSelection('')).toBeNull();
+  });
+
+  it('returns the parsed entry when storage is valid', () => {
+    localStorage.setItem(
+      'opencode_bridge:model-selection:foo.ipynb',
+      JSON.stringify({ providerId: 'anthropic', modelId: 'claude-sonnet-4-20250514' })
+    );
+    expect(loadModelSelection('foo.ipynb')).toEqual({
+      providerId: 'anthropic',
+      modelId: 'claude-sonnet-4-20250514'
+    });
+  });
+
+  it('returns null for malformed JSON (does not throw)', () => {
+    localStorage.setItem(
+      'opencode_bridge:model-selection:foo.ipynb',
+      '{not json'
+    );
+    expect(loadModelSelection('foo.ipynb')).toBeNull();
+  });
+
+  it('returns null when the entry is missing required fields', () => {
+    localStorage.setItem(
+      'opencode_bridge:model-selection:foo.ipynb',
+      JSON.stringify({ providerId: 'anthropic' })
+    );
+    expect(loadModelSelection('foo.ipynb')).toBeNull();
+  });
+
+  it('keys entries independently per notebookPath', () => {
+    saveModelSelection('a.ipynb', { providerId: 'p1', modelId: 'm1' });
+    saveModelSelection('b.ipynb', { providerId: 'p2', modelId: 'm2' });
+    expect(loadModelSelection('a.ipynb')).toEqual({
+      providerId: 'p1',
+      modelId: 'm1'
+    });
+    expect(loadModelSelection('b.ipynb')).toEqual({
+      providerId: 'p2',
+      modelId: 'm2'
+    });
+  });
+});
+
+describe('saveModelSelection + clearModelSelection', () => {
+  beforeEach(() => {
+    localStorage.clear();
+  });
+
+  it('writes a JSON entry that can be re-read', () => {
+    saveModelSelection('foo.ipynb', { providerId: 'p1', modelId: 'm1' });
+    expect(loadModelSelection('foo.ipynb')).toEqual({
+      providerId: 'p1',
+      modelId: 'm1'
+    });
+  });
+
+  it('overwrites an existing entry', () => {
+    saveModelSelection('foo.ipynb', { providerId: 'p1', modelId: 'm1' });
+    saveModelSelection('foo.ipynb', { providerId: 'p2', modelId: 'm2' });
+    expect(loadModelSelection('foo.ipynb')).toEqual({
+      providerId: 'p2',
+      modelId: 'm2'
+    });
+  });
+
+  it('clearModelSelection removes the entry', () => {
+    saveModelSelection('foo.ipynb', { providerId: 'p1', modelId: 'm1' });
+    clearModelSelection('foo.ipynb');
+    expect(loadModelSelection('foo.ipynb')).toBeNull();
+  });
+
+  it('no-op for an empty notebookPath (defensive)', () => {
+    expect(() =>
+      saveModelSelection('', { providerId: 'p1', modelId: 'm1' })
+    ).not.toThrow();
+    expect(() => clearModelSelection('')).not.toThrow();
+  });
+});
diff --git a/src/__tests__/opencode_cell_actions.spec.ts b/src/__tests__/opencode_cell_actions.spec.ts
index 64c4413..6c139e9 100644
--- a/src/__tests__/opencode_cell_actions.spec.ts
+++ b/src/__tests__/opencode_cell_actions.spec.ts
@@ -103,6 +103,12 @@ beforeAll(() => {
   });
 });
 
+beforeEach(() => {
+  // Each test starts with a clean selection store so persistence tests
+  // are deterministic.
+  localStorage.clear();
+});
+
 // Mock the network-touching API module so jsdom tests don't try to hit
 // a real notebook server (and to keep the refresh-history logic decoupled
 // from the network). The cell_actions tests focus on widget behavior.
@@ -323,6 +329,122 @@ describe('OpenCodeCellActions', () => {
     setOpenCodeProviders(null);
   });
 
+  it('restores the user\'s previously-saved (provider, model) when still valid in the providers payload', () => {
+    // Seed localStorage with a prior selection that IS valid.
+    localStorage.setItem(
+      'opencode_bridge:model-selection:foo.ipynb',
+      JSON.stringify({ providerId: 'openai', modelId: 'gpt-5' })
+    );
+    setOpenCodeProviders({
+      providers: [
+        { id: 'anthropic', models: { 'claude-sonnet-4-20250514': { id: 'x' } } },
+        { id: 'openai', models: { 'gpt-5': { id: 'x' }, 'gpt-4': { id: 'x' } } }
+      ],
+      default: { anthropic: 'claude-sonnet-4-20250514' }
+    });
+    const cell = makeFakeCell('x = 1', [], 'foo.ipynb');
+    const actions = new OpenCodeCellActions(cell);
+    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
+    (actions as any).onAfterAttach({} as any);
+
+    (actions.node.querySelector('button') as HTMLButtonElement).click();
+    const prompt = (actions as any)._prompt;
+    const providerSel = prompt.node.querySelector(
+      '.opencode-provider-select'
+    ) as HTMLSelectElement;
+    const modelSel = prompt.node.querySelector(
+      '.opencode-model-select'
+    ) as HTMLSelectElement;
+    // Stored selection was applied, NOT the default (first provider).
+    expect(providerSel.value).toBe('openai');
+    expect(modelSel.value).toBe('gpt-5');
+    setOpenCodeProviders(null);
+  });
+
+  it('falls back to default when the stored provider is no longer present', () => {
+    localStorage.setItem(
+      'opencode_bridge:model-selection:foo.ipynb',
+      JSON.stringify({ providerId: 'removed-provider', modelId: 'm1' })
+    );
+    setOpenCodeProviders({
+      providers: [
+        { id: 'anthropic', models: { 'claude-sonnet-4-20250514': { id: 'x' } } }
+      ]
+    });
+    const cell = makeFakeCell('x = 1', [], 'foo.ipynb');
+    const actions = new OpenCodeCellActions(cell);
+    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
+    (actions as any).onAfterAttach({} as any);
+
+    (actions.node.querySelector('button') as HTMLButtonElement).click();
+    const providerSel = (actions as any)._prompt.node.querySelector(
+      '.opencode-provider-select'
+    ) as HTMLSelectElement;
+    // Provider is gone → fall back to first available.
+    expect(providerSel.value).toBe('anthropic');
+    setOpenCodeProviders(null);
+  });
+
+  it('falls back to default model when the stored model is no longer present under the same provider', () => {
+    localStorage.setItem(
+      'opencode_bridge:model-selection:foo.ipynb',
+      JSON.stringify({ providerId: 'openai', modelId: 'gpt-99-removed' })
+    );
+    setOpenCodeProviders({
+      providers: [
+        { id: 'openai', models: { 'gpt-5': { id: 'x' }, 'gpt-4': { id: 'x' } } }
+      ],
+      default: { openai: 'gpt-5' }
+    });
+    const cell = makeFakeCell('x = 1', [], 'foo.ipynb');
+    const actions = new OpenCodeCellActions(cell);
+    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
+    (actions as any).onAfterAttach({} as any);
+
+    (actions.node.querySelector('button') as HTMLButtonElement).click();
+    const modelSel = (actions as any)._prompt.node.querySelector(
+      '.opencode-model-select'
+    ) as HTMLSelectElement;
+    // Model is gone → fall back to default[openai] = gpt-5.
+    expect(modelSel.value).toBe('gpt-5');
+    setOpenCodeProviders(null);
+  });
+
+  it('saves the new selection to localStorage when the model select changes', () => {
+    setOpenCodeProviders({
+      providers: [
+        { id: 'anthropic', models: { 'claude-sonnet-4-20250514': { id: 'x' } } },
+        { id: 'openai', models: { 'gpt-5': { id: 'x' } } }
+      ]
+    });
+    const cell = makeFakeCell('x = 1', [], 'foo.ipynb');
+    const actions = new OpenCodeCellActions(cell);
+    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
+    (actions as any).onAfterAttach({} as any);
+
+    (actions.node.querySelector('button') as HTMLButtonElement).click();
+    const prompt = (actions as any)._prompt;
+    const providerSel = prompt.node.querySelector(
+      '.opencode-provider-select'
+    ) as HTMLSelectElement;
+    const modelSel = prompt.node.querySelector(
+      '.opencode-model-select'
+    ) as HTMLSelectElement;
+
+    // Switch provider to openai (whose model list includes gpt-5).
+    providerSel.value = 'openai';
+    providerSel.dispatchEvent(new Event('change'));
+    // Now change the model select to gpt-5.
+    modelSel.value = 'gpt-5';
+    modelSel.dispatchEvent(new Event('change'));
+
+    const stored = JSON.parse(
+      localStorage.getItem('opencode_bridge:model-selection:foo.ipynb') || 'null'
+    );
+    expect(stored).toEqual({ providerId: 'openai', modelId: 'gpt-5' });
+    setOpenCodeProviders(null);
+  });
+
   it('falls back to the first model when default[provider] is not in models', () => {
     setOpenCodeProviders({
       providers: [
diff --git a/src/components/model_selection.ts b/src/components/model_selection.ts
new file mode 100644
index 0000000..82f0912
--- /dev/null
+++ b/src/components/model_selection.ts
@@ -0,0 +1,82 @@
+/**
+ * Persistence of the user's (provider, model) selection in the inline
+ * prompt, keyed by notebook path. Stored in localStorage so the
+ * preference survives panel close/reopen and JupyterLab restart.
+ *
+ * The selection is ALWAYS validated against the current providers
+ * payload before being applied: if the stored provider (or the model
+ * under it) is no longer present, the caller falls back to its
+ * default-selection logic. We never silently apply a stale value that
+ * would result in an empty / disabled <select>.
+ *
+ * Storage is best-effort: localStorage may be disabled (private mode),
+ * the quota may be exhausted, or the value may be corrupt. All
+ * failure modes return null / no-op rather than throw.
+ */
+const STORAGE_PREFIX = 'opencode_bridge:model-selection:';
+
+export interface ModelSelection {
+  providerId: string;
+  modelId: string;
+}
+
+export function loadModelSelection(notebookPath: string): ModelSelection | null {
+  if (!notebookPath) {
+    return null;
+  }
+  try {
+    if (typeof localStorage === 'undefined') {
+      return null;
+    }
+    const raw = localStorage.getItem(STORAGE_PREFIX + notebookPath);
+    if (!raw) {
+      return null;
+    }
+    const parsed = JSON.parse(raw);
+    if (
+      parsed &&
+      typeof parsed === 'object' &&
+      typeof (parsed as ModelSelection).providerId === 'string' &&
+      typeof (parsed as ModelSelection).modelId === 'string'
+    ) {
+      return {
+        providerId: (parsed as ModelSelection).providerId,
+        modelId: (parsed as ModelSelection).modelId
+      };
+    }
+    return null;
+  } catch {
+    return null;
+  }
+}
+
+export function saveModelSelection(
+  notebookPath: string,
+  sel: ModelSelection
+): void {
+  if (!notebookPath) {
+    return;
+  }
+  try {
+    if (typeof localStorage === 'undefined') {
+      return;
+    }
+    localStorage.setItem(STORAGE_PREFIX + notebookPath, JSON.stringify(sel));
+  } catch {
+    // Ignore quota exceeded / disabled storage.
+  }
+}
+
+export function clearModelSelection(notebookPath: string): void {
+  if (!notebookPath) {
+    return;
+  }
+  try {
+    if (typeof localStorage === 'undefined') {
+      return;
+    }
+    localStorage.removeItem(STORAGE_PREFIX + notebookPath);
+  } catch {
+    // ignore
+  }
+}
diff --git a/src/components/opencode_cell_actions.ts b/src/components/opencode_cell_actions.ts
index 916c7cb..cc070ed 100644
--- a/src/components/opencode_cell_actions.ts
+++ b/src/components/opencode_cell_actions.ts
@@ -127,6 +127,7 @@ export class OpenCodeCellActions extends Widget {
     const prompt = new OpenCodeInlinePrompt(this._cell, {
       disabled: !this._context || this._status === 'loading',
       providers: _providers,
+      notebookPath: this._context?.notebookPath,
       onSubmit: (text: string, providerId?: string, modelId?: string) => {
         void this._onSubmit(text, providerId, modelId);
       },
diff --git a/src/components/opencode_inline_prompt.ts b/src/components/opencode_inline_prompt.ts
index b29ba49..f04a155 100644
--- a/src/components/opencode_inline_prompt.ts
+++ b/src/components/opencode_inline_prompt.ts
@@ -30,12 +30,23 @@ import type {
   OpenCodeMessage,
   OpenCodeProvidersResponse
 } from '../types';
+import {
+  loadModelSelection,
+  saveModelSelection
+} from './model_selection';
 
 export interface IOpenCodeInlinePromptOptions {
   onSubmit: (text: string, providerId?: string, modelId?: string) => void;
   onCancel: () => void;
   disabled: boolean;
   providers: OpenCodeProvidersResponse | null;
+  /**
+   * Notebook path used to key the persisted (provider, model) selection
+   * in localStorage. If omitted (or the stored selection is no longer
+   * valid in the current `providers` payload), the prompt falls back
+   * to the default selection (first provider / `default[pid]` model).
+   */
+  notebookPath?: string;
   /**
    * Respond to a `permission.asked` event. Returns a promise that resolves
    * to true on success, false on failure (the prompt updates the UI in
@@ -143,6 +154,12 @@ export class OpenCodeInlinePrompt extends Widget {
     this._providerModels = flat.providers;
     this._defaultMap = flat.defaultMap;
 
+    // Resolve the initial selection: prefer a previously-saved
+    // (providerId, modelId) IF it's still valid in the current
+    // providers payload; otherwise fall back to the default
+    // (first provider / `default[pid]` model).
+    const stored = this._resolveStoredSelection(options.notebookPath);
+
     if (this._providerModels.length > 0) {
       this._providerSelect = document.createElement('select');
       this._providerSelect.className = 'opencode-provider-select';
@@ -152,16 +169,25 @@ export class OpenCodeInlinePrompt extends Widget {
         o.textContent = p.name ? `${p.id} (${p.name})` : p.id;
         this._providerSelect.appendChild(o);
       }
-      this._providerSelect.value = this._providerModels[0].id;
+      this._providerSelect.value = stored
+        ? stored.providerId
+        : this._providerModels[0].id;
       this._providerSelect.addEventListener('change', () => {
         this._rebuildModelSelect(this._providerSelect!.value);
+        this._persistSelection();
       });
     }
 
     if (this._providerModels.length > 0) {
       this._modelSelect = document.createElement('select');
       this._modelSelect.className = 'opencode-model-select';
-      this._rebuildModelSelect(this._providerSelect?.value);
+      this._rebuildModelSelect(
+        this._providerSelect?.value,
+        stored ? stored.modelId : undefined
+      );
+      this._modelSelect.addEventListener('change', () => {
+        this._persistSelection();
+      });
     }
 
     this._textarea = document.createElement('textarea');
@@ -230,7 +256,10 @@ export class OpenCodeInlinePrompt extends Widget {
     this.node.appendChild(actions);
   }
 
-  private _rebuildModelSelect(providerId: string | undefined): void {
+  private _rebuildModelSelect(
+    providerId: string | undefined,
+    preferredModelId?: string
+  ): void {
     if (!this._modelSelect) {
       return;
     }
@@ -258,8 +287,65 @@ export class OpenCodeInlinePrompt extends Widget {
       this._modelSelect.appendChild(o);
     }
     this._modelSelect.disabled = false;
-    this._modelSelect.value =
-      pickDefaultModelId(providerId, provider.models, this._defaultMap) ?? keys[0];
+    // Honor a preferred model (e.g. from a stored selection), but only
+    // if it actually exists in this provider's model list. Otherwise
+    // fall back to the default.
+    if (
+      preferredModelId &&
+      Object.prototype.hasOwnProperty.call(provider.models, preferredModelId)
+    ) {
+      this._modelSelect.value = preferredModelId;
+    } else {
+      this._modelSelect.value =
+        pickDefaultModelId(providerId, provider.models, this._defaultMap) ??
+        keys[0];
+    }
+  }
+
+  /**
+   * Read the stored (providerId, modelId) for this notebook and
+   * validate it against the current providers payload. Returns the
+   * stored selection only when BOTH the provider and the model still
+   * exist; otherwise returns null and the caller uses the default.
+   */
+  private _resolveStoredSelection(
+    notebookPath: string | undefined
+  ): { providerId: string; modelId: string } | null {
+    if (!notebookPath) {
+      return null;
+    }
+    const sel = loadModelSelection(notebookPath);
+    if (!sel) {
+      return null;
+    }
+    const provider = this._providerModels.find(p => p.id === sel.providerId);
+    if (!provider) {
+      // The provider was removed since the user last picked.
+      return null;
+    }
+    if (!Object.prototype.hasOwnProperty.call(provider.models, sel.modelId)) {
+      // The model was removed (or renamed) under a still-valid provider.
+      return null;
+    }
+    return sel;
+  }
+
+  /**
+   * Persist the current (providerSelect.value, modelSelect.value) for
+   * this notebook. No-op if the prompt has no notebookPath, no
+   * selects, or the values are empty.
+   */
+  private _persistSelection(): void {
+    const notebookPath = this._options.notebookPath;
+    if (!notebookPath || !this._providerSelect || !this._modelSelect) {
+      return;
+    }
+    const providerId = this._providerSelect.value;
+    const modelId = this._modelSelect.value;
+    if (!providerId || !modelId) {
+      return;
+    }
+    saveModelSelection(notebookPath, { providerId, modelId });
   }
 
   setDisabled(disabled: boolean): void {