Add persistence to fetched / selected models

This has been a long time coming.

Author: aaronik

lua/gptmodels/store.lua

@@ -96,13 +96,29 @@ end
 ---@field get_job fun(self: Store): Job | nil
 ---@field clear_job fun(self: Store)
 ---@field get_models fun(self: Store, provider: Provider): string[]
----@field set_models fun(self: Store, provider: Provider, models: string[])
+---@field set_models fun(self: Store, provider: Provider, models: string[], skip_persistence?: boolean)
 ---@field get_model fun(self: Store): { provider: string, model: string }
 ---@field set_model fun(self: Store, provider: Provider, model: string)
 ---@field cycle_model_forward fun(self: Store)
 ---@field cycle_model_backward fun(self: Store)
 ---@field llm_model_strings fun(self: Store): string[]
 ---@field correct_potentially_missing_current_model fun(self: Store)
+---@field save_persisted_state fun(self: Store)
+---@field load_persisted_state fun(self: Store)
+
+-- Get the path to the persistence file
+---@return string
+local function get_persistence_file_path()
+  local data_path = vim.fn.stdpath("data")
+  return data_path .. "/gptmodels/state.json"
+end
+
+-- Get the directory for persistence files
+---@return string
+local function get_persistence_dir()
+  local data_path = vim.fn.stdpath("data")
+  return data_path .. "/gptmodels"
+end
 
 ---@return StrPane
 local function build_strpane()
@@ -140,6 +156,7 @@ local Store = {
   set_model = function(self, provider, model)
     self._llm_provider = provider
     self._llm_model = model
+    self:save_persisted_state()
   end,
 
   -- get all models for a provider
@@ -148,8 +165,11 @@ local Store = {
   end,
 
   -- set all models for a provider, overwriting previous values
-  set_models = function(self, provider, models)
+  set_models = function(self, provider, models, skip_persistence)
     self._llm_models[provider] = models
+    if not skip_persistence then
+      self:save_persisted_state()
+    end
   end,
 
   cycle_model_forward = function(self)
@@ -240,9 +260,10 @@ local Store = {
   clear = function(self)
     self.code:clear()
     self.chat:clear()
-    self:set_models("ollama", {})
-    self:set_models("openai", {})
-    -- TODO Need to clear default model as well?
+    self._llm_models.ollama = {}
+    self._llm_models.openai = {}
+    self._llm_provider = ""
+    self._llm_model = ""
   end,
 
   code = {
@@ -311,6 +332,78 @@ local Store = {
   clear_job = function(self)
     self._job = nil
   end,
+
+  -- Save current state to disk
+  save_persisted_state = function(self)
+    local persistence_dir = get_persistence_dir()
+    local persistence_file = get_persistence_file_path()
+
+    -- Create directory if it doesn't exist
+    vim.fn.mkdir(persistence_dir, "p")
+
+    -- Prepare data to save
+    local data = {
+      current_provider = self._llm_provider,
+      current_model = self._llm_model,
+      models = {
+        openai = self._llm_models.openai,
+        ollama = self._llm_models.ollama,
+      },
+    }
+
+    -- Write to file
+    local encoded = vim.json.encode(data)
+    local file = io.open(persistence_file, "w")
+    if file then
+      file:write(encoded)
+      file:close()
+    end
+  end,
+
+  -- Load persisted state from disk
+  load_persisted_state = function(self)
+    local persistence_file = get_persistence_file_path()
+
+    -- Check if file exists
+    if vim.fn.filereadable(persistence_file) == 0 then
+      return
+    end
+
+    -- Read and parse file
+    local file = io.open(persistence_file, "r")
+    if not file then
+      return
+    end
+
+    local content = file:read("*a")
+    file:close()
+
+    -- Handle empty file
+    if not content or content == "" then
+      return
+    end
+
+    -- Parse JSON safely
+    local ok, data = pcall(vim.json.decode, content)
+    if not ok or not data then
+      return
+    end
+
+    -- Restore state
+    if data.current_provider and data.current_model then
+      self._llm_provider = data.current_provider
+      self._llm_model = data.current_model
+    end
+
+    if data.models then
+      if data.models.openai then
+        self._llm_models.openai = data.models.openai
+      end
+      if data.models.ollama then
+        self._llm_models.ollama = data.models.ollama
+      end
+    end
+  end,
 }
 
 return Store

lua/gptmodels/windows/chat.lua

@@ -126,6 +126,10 @@ function M.build_and_mount(selection)
   Store.chat.chat.popup = chat
   Store.chat.input.popup = input
 
+  -- Load any persisted state first (model selection and cached models)
+  Store:load_persisted_state()
+  com.set_window_title(chat, WINDOW_TITLE_PREFIX .. com.model_display_name())
+
   -- Fetch all models so user can work with what they have on their system
   com.trigger_models_etl(function()
     local has_buf_and_win = chat.bufnr and chat.winid
@@ -141,7 +145,7 @@ function M.build_and_mount(selection)
 
     -- all providers, but especially openai, can have the etl finish after a window has been closed,
     -- if it opens then closes real fast
-    com.set_window_title(chat, "Chat w/ " .. com.model_display_name())
+    com.set_window_title(chat, WINDOW_TITLE_PREFIX .. com.model_display_name())
   end)
 
   -- Input window is text with no syntax

lua/gptmodels/windows/code.lua

@@ -171,7 +171,11 @@ function M.build_and_mount(selection)
   Store.code.left.popup = left
   Store.code.input.popup = input
 
-  -- Fetch all models so user can work with what they have on their system
+  -- Load any persisted state first (model selection and cached models)
+  Store:load_persisted_state()
+  com.set_window_title(right, com.model_display_name())
+
+  -- Fetch all models so user can work with what they have access to
   com.trigger_models_etl(function()
     local has_buf_and_win = right.bufnr and right.winid
     if not has_buf_and_win then

lua/gptmodels/windows/common.lua

@@ -214,27 +214,29 @@ end
 -- Triggers the fetching / saving of available models from the ollama and openai servers
 ---@param on_complete fun(): nil
 M.trigger_models_etl = function(on_complete)
+  local completed_providers = {}
+  local total_providers = 2
+
   ---@param err string | nil
   ---@param models string[] | nil
   ---@param provider Provider
   local function handle_models_fetch(err, models, provider)
     -- If there's an error fetching, assume we have no models
-    -- TODO We still need to inform the user somehow that their ollama models
-    -- fetching didn't work. Just not if we earlier detected a missing ollama
-    -- executable. Store.detected_missing_ollama_exe = true?
-    -- BUT DO WE? Maybe the models not appearing is sufficient feedback!
-    -- I think passing the err back is a good idea, because that can include
-    -- provider information
     if err or not models or #models == 0 then
-      Store:set_models(provider, {})
-      Store:correct_potentially_missing_current_model()
-      return on_complete()
+      Store:set_models(provider, {}, true) -- skip persistence during fetching
+    else
+      Store:set_models(provider, models, true) -- skip persistence during fetching
     end
 
-    Store:set_models(provider, models)
-    -- TODO Test that this gets called
-    Store:correct_potentially_missing_current_model()
-    on_complete()
+    -- Mark this provider as completed
+    completed_providers[provider] = true
+
+    -- Save state after ALL providers have completed
+    if vim.tbl_count(completed_providers) == total_providers then
+      -- Save the final state with all fetched models
+      Store:save_persisted_state()
+      on_complete()
+    end
   end
 
   -- Fetch models from ollama server

tests/gptmodels/persistence_spec.lua

@@ -0,0 +1,125 @@
+---@diagnostic disable: undefined-global
+
+local assert = require("luassert")
+local stub = require('luassert.stub')
+local Store = require("gptmodels.store")
+
+describe("Store | persistence", function()
+  local mock_data_path = "/tmp/nvim-test-data"
+  local vim_fn_stdpath_stub
+
+  before_each(function()
+    Store:clear()
+    -- Mock vim.fn.stdpath to return our test data path
+    vim_fn_stdpath_stub = stub(vim.fn, 'stdpath')
+    vim_fn_stdpath_stub.returns(mock_data_path)
+  end)
+
+  after_each(function()
+    vim_fn_stdpath_stub:revert()
+    -- Clean up test files
+    vim.fn.delete(mock_data_path .. "/gptmodels", "rf")
+  end)
+
+  describe("load_persisted_state", function()
+    it("loads previously saved model selection", function()
+      -- Setup: Save a model selection
+      Store:set_model("openai", "gpt-4o")
+      Store:save_persisted_state()
+
+      -- Clear the store and load persisted state
+      Store:clear()
+      Store:load_persisted_state()
+
+      local model_info = Store:get_model()
+      assert.equal("openai", model_info.provider)
+      assert.equal("gpt-4o", model_info.model)
+    end)
+
+    it("loads previously saved fetched models", function()
+      -- Setup: Save fetched models
+      Store:set_models("openai", {"gpt-4o", "gpt-4o-mini"})
+      Store:set_models("ollama", {"llama3.1:latest", "deepseek-v2:latest"})
+      Store:save_persisted_state()
+
+      -- Clear the store and load persisted state
+      Store:clear()
+      Store:load_persisted_state()
+
+      assert.same({"gpt-4o", "gpt-4o-mini"}, Store:get_models("openai"))
+      assert.same({"llama3.1:latest", "deepseek-v2:latest"}, Store:get_models("ollama"))
+    end)
+
+    it("handles missing persistence file gracefully", function()
+      -- Ensure no persistence file exists
+      vim.fn.delete(mock_data_path .. "/gptmodels", "rf")
+
+      Store:load_persisted_state()
+
+      -- Should not crash and should have empty state
+      local model_info = Store:get_model()
+      assert.equal("", model_info.provider)
+      assert.equal("", model_info.model)
+      assert.same({}, Store:get_models("openai"))
+      assert.same({}, Store:get_models("ollama"))
+    end)
+
+    it("handles corrupted persistence file gracefully", function()
+      -- Create corrupted file
+      vim.fn.mkdir(mock_data_path .. "/gptmodels", "p")
+      local file = io.open(mock_data_path .. "/gptmodels/state.json", "w")
+      file:write("invalid json {")
+      file:close()
+
+      Store:load_persisted_state()
+
+      -- Should not crash and should have empty state
+      local model_info = Store:get_model()
+      assert.equal("", model_info.provider)
+      assert.equal("", model_info.model)
+    end)
+  end)
+
+  describe("save_persisted_state", function()
+    it("creates data directory if it doesn't exist", function()
+      -- Ensure directory doesn't exist
+      vim.fn.delete(mock_data_path, "rf")
+
+      Store:set_model("openai", "gpt-4o")
+      Store:save_persisted_state()
+
+      -- Verify directory and file were created
+      assert.equal(1, vim.fn.isdirectory(mock_data_path .. "/gptmodels"))
+      assert.equal(1, vim.fn.filereadable(mock_data_path .. "/gptmodels/state.json"))
+    end)
+  end)
+
+  describe("auto persistence", function()
+    it("automatically saves state when model is changed", function()
+      Store:set_model("openai", "gpt-4o")
+
+      -- Should automatically save
+      local file = io.open(mock_data_path .. "/gptmodels/state.json", "r")
+      assert.is_not_nil(file)
+      local content = file:read("*a")
+      file:close()
+
+      local data = vim.json.decode(content)
+      assert.equal("openai", data.current_provider)
+      assert.equal("gpt-4o", data.current_model)
+    end)
+
+    it("automatically saves state when models are updated", function()
+      Store:set_models("ollama", {"llama3.1:latest", "deepseek-v2:latest"})
+
+      -- Should automatically save
+      local file = io.open(mock_data_path .. "/gptmodels/state.json", "r")
+      assert.is_not_nil(file)
+      local content = file:read("*a")
+      file:close()
+
+      local data = vim.json.decode(content)
+      assert.same({"llama3.1:latest", "deepseek-v2:latest"}, data.models.ollama)
+    end)
+  end)
+end)

tests/gptmodels/spec_helpers.lua

@@ -43,6 +43,10 @@ M.hook_reset_state = function()
     -- stubbing cmd.exec prevents the llm call from happening
     stub(cmd, "exec")
 
+    -- Mock persistence functions to prevent loading real saved data during tests
+    stub(Store, "load_persisted_state")
+    stub(Store, "save_persisted_state")
+
     Store:clear()
     snapshot = assert:snapshot()
   end)