ensuring state is kept between changes

10 files changed, 110 insertions, 36 deletions

diff --git a/lua/99/extensions/work/worker.lua b/lua/99/extensions/work/worker.lua
index 60825f3..eff66b9 100644
--- a/lua/99/extensions/work/worker.lua
+++ b/lua/99/extensions/work/worker.lua
@@ -108,6 +108,9 @@ end
 function M.craft_prompt(worker)
   return string.format(
     [[
+## Description
+%s
+
 ## You must complete the checklist
 [ ] - Inspect and understand all changed code
  [ ] - git diff
@@ -120,9 +123,6 @@ function M.craft_prompt(worker)
 [ ] - If you see bugs, also report those
 [ ] - if there are tests, run the tests
 
-<Description>
-%s
-</Description>
 ]],
     worker.current_work_item
   )
diff --git a/lua/99/ops/clean-up.lua b/lua/99/ops/clean-up.lua
index 56a7e11..9eea6da 100644
--- a/lua/99/ops/clean-up.lua
+++ b/lua/99/ops/clean-up.lua
@@ -7,10 +7,10 @@ local M = {}
 --- @field on_stderr? fun(line: string): nil
 --- @field on_start? fun(): nil
 
---- @param clean_up fun(): nil
+--- @param context _99.Prompt
 --- @param obs_or_fn _99.Providers.PartialObserver | _99.Providers.on_complete
 --- @return _99.Providers.Observer
-M.make_observer = function(clean_up, obs_or_fn)
+M.make_observer = function(context, obs_or_fn)
   --- @type _99.Providers.PartialObserver
   local obs = type(obs_or_fn) == "table" and obs_or_fn
     or {
@@ -23,8 +23,11 @@ M.make_observer = function(clean_up, obs_or_fn)
       end
     end,
     on_complete = function(status, res)
-      vim.schedule(clean_up)
-      obs.on_complete(status, res)
+      pcall(obs.on_complete, status, res)
+      vim.schedule(function()
+        context:stop()
+        context._99:sync()
+      end)
     end,
     on_stderr = function(line)
       if obs.on_stderr then
diff --git a/lua/99/ops/over-range.lua b/lua/99/ops/over-range.lua
index 3931d22..e8512b8 100644
--- a/lua/99/ops/over-range.lua
+++ b/lua/99/ops/over-range.lua
@@ -42,8 +42,6 @@ local function over_range(context, opts)
   local clean_up = make_clean_up(function()
     top_status:stop()
     bottom_status:stop()
-    context:clear_marks()
-    context:stop()
   end)
 
   local system_cmd = context._99.prompts.prompts.visual_selection(range)
@@ -55,7 +53,7 @@ local function over_range(context, opts)
 
   top_status:start()
   bottom_status:start()
-  context:start_request(make_observer(clean_up, {
+  context:start_request(make_observer(context, {
     on_complete = function(status, response)
       if status == "cancelled" then
         logger:debug("request cancelled for visual selection, removing marks")
@@ -90,6 +88,7 @@ local function over_range(context, opts)
         table.insert(lines, 1, "")
 
         new_range:replace_text(lines)
+        context._99:sync()
       end
     end,
     on_stdout = function(line)
diff --git a/lua/99/ops/search.lua b/lua/99/ops/search.lua
index 0b19bb4..d9fcd1e 100644
--- a/lua/99/ops/search.lua
+++ b/lua/99/ops/search.lua
@@ -2,7 +2,6 @@ local make_prompt = require("99.ops.make-prompt")
 local CleanUp = require("99.ops.clean-up")
 local QFixHelpers = require("99.ops.qfix-helpers")
 
-local make_clean_up = CleanUp.make_clean_up
 local make_observer = CleanUp.make_observer
 
 --- @param context _99.Prompt
@@ -31,16 +30,11 @@ local function search(context, opts)
   local logger = context.logger:set_area("search")
   logger:debug("search", "with opts", opts.additional_prompt)
 
-  local clean_up = make_clean_up(function()
-    context:stop()
-  end)
-
   local prompt, refs =
     make_prompt(context, context._99.prompts.prompts.semantic_search(), opts)
 
   context:add_prompt_content(prompt)
   context:add_references(refs)
-  context:add_clean_up(clean_up)
 
   --- TODO: part of the context request clean up there needs to be a refactoring of
   --- make observer... it really should just be within the context observer creation.
@@ -48,7 +42,7 @@ local function search(context, opts)
   --- once cleanup function wrapper.
   ---
   --- i think an interface, CleanUpI could be something that is worth it :)
-  context:start_request(make_observer(clean_up, function(status, response)
+  context:start_request(make_observer(context, function(status, response)
     if status == "cancelled" then
       logger:debug("request cancelled for search")
     elseif status == "failed" then
@@ -59,6 +53,7 @@ local function search(context, opts)
       )
     elseif status == "success" then
       create_search_locations(context, response)
+      context._99:sync()
     end
   end))
 end
diff --git a/lua/99/ops/tutorial.lua b/lua/99/ops/tutorial.lua
index 4b9a9b8..caf54d2 100644
--- a/lua/99/ops/tutorial.lua
+++ b/lua/99/ops/tutorial.lua
@@ -2,7 +2,6 @@ local CleanUp = require("99.ops.clean-up")
 local Window = require("99.window")
 local make_prompt = require("99.ops.make-prompt")
 
-local make_clean_up = CleanUp.make_clean_up
 local make_observer = CleanUp.make_observer
 
 --- @param context _99.Prompt
@@ -32,19 +31,13 @@ local function tutorial(context, opts)
   local logger = context.logger:set_area("tutorial")
   logger:debug("starting", "with opts", opts)
 
-  local clean_up = make_clean_up(function()
-    context:stop()
-  end)
-
   local prompt, refs =
     make_prompt(context, context._99.prompts.prompts.tutorial(), opts)
 
   context:add_references(refs)
   context:add_prompt_content(prompt)
-  context:add_clean_up(clean_up)
 
-  context:start_request(make_observer(clean_up, function(status, response)
-    vim.schedule(clean_up)
+  context:start_request(make_observer(context, function(status, response)
     if status == "cancelled" then
       logger:debug("cancelled")
     elseif status == "failed" then
@@ -55,6 +48,7 @@ local function tutorial(context, opts)
       )
     elseif status == "success" then
       open_tutorial(context, response)
+      context._99:sync()
     end
   end))
 end
diff --git a/lua/99/ops/vibe.lua b/lua/99/ops/vibe.lua
index 0e52307..8b53768 100644
--- a/lua/99/ops/vibe.lua
+++ b/lua/99/ops/vibe.lua
@@ -2,7 +2,6 @@ local make_prompt = require("99.ops.make-prompt")
 local CleanUp = require("99.ops.clean-up")
 local QFixHelpers = require("99.ops.qfix-helpers")
 
-local make_clean_up = CleanUp.make_clean_up
 local make_observer = CleanUp.make_observer
 
 --- @param context _99.Prompt
@@ -37,16 +36,11 @@ local function vibe(context, opts)
   local logger = context.logger:set_area("vibe")
   logger:debug("vibe", "with opts", opts.additional_prompt)
 
-  local clean_up = make_clean_up(function()
-    context:stop()
-  end)
-
   local prompt, refs =
     make_prompt(context, context._99.prompts.prompts.vibe(), opts)
 
   context:add_prompt_content(prompt)
   context:add_references(refs)
-  context:add_clean_up(clean_up)
 
   --- TODO: part of the context request clean up there needs to be a refactoring of
   --- make observer... it really should just be within the context observer creation.
@@ -54,7 +48,7 @@ local function vibe(context, opts)
   --- once cleanup function wrapper.
   ---
   --- i think an interface, CleanUpI could be something that is worth it :)
-  context:start_request(make_observer(clean_up, function(status, response)
+  context:start_request(make_observer(context, function(status, response)
     if status == "cancelled" then
       logger:debug("request cancelled for search")
     elseif status == "failed" then
@@ -65,6 +59,7 @@ local function vibe(context, opts)
       )
     elseif status == "success" then
       finish_vibe(context, response)
+      context._99:sync()
     end
   end))
 end
diff --git a/lua/99/prompt.lua b/lua/99/prompt.lua
index 96d4579..ca24ba7 100644
--- a/lua/99/prompt.lua
+++ b/lua/99/prompt.lua
@@ -31,7 +31,6 @@ local filetype_map = {
 --- @class _99.Prompt.Serialized
 --- @field data _99.Prompt.Data
 --- @field user_prompt string
-
 --- @class _99.Prompt.Data.Search
 --- @field type "search"
 --- @field qfix_items _99.Search.Result[]
@@ -119,6 +118,10 @@ function Prompt.deserialize(_99, data)
     operation = data.data.type,
     user_prompt = data.user_prompt,
     started_at = Time.now(),
+
+    --- we should only sync successful requests
+    state = "success",
+
     xid = get_id(),
   }, Prompt)
   assert(prompt:valid(), "prompt is not valid from data")
@@ -127,6 +130,7 @@ end
 
 --- @return _99.Prompt.Serialized
 function Prompt:serialize()
+  assert(self.state == "success", "you can only serialize successful prompts")
   return {
     data = self.data,
     user_prompt = self.user_prompt,
@@ -371,6 +375,8 @@ end
 
 function Prompt:stop()
   self:cancel()
+  self:clear_marks()
+
   for _, cb in ipairs(self.clean_ups) do
     cb()
   end
diff --git a/lua/99/state.lua b/lua/99/state.lua
index 7d94c07..6e38c5f 100644
--- a/lua/99/state.lua
+++ b/lua/99/state.lua
@@ -3,7 +3,7 @@ local Agents = require("99.extensions.agents")
 local Extensions = require("99.extensions")
 local Tracking = require("99.state.tracking")
 
-local _99_STATE_FILE = "99-state"
+local _99_STATE_FILE = "state"
 local function default_completion()
   return { source = nil, custom_rules = {} }
 end
diff --git a/lua/99/test/window_spec.lua b/lua/99/test/window_spec.lua
index c91c219..3f5c3c8 100644
--- a/lua/99/test/window_spec.lua
+++ b/lua/99/test/window_spec.lua
@@ -29,10 +29,47 @@ describe("Window", function()
     })
 
     eq(2, #Window.active_windows)
-    eq(2, vim.wo[win.win_id].scrolloff)
+    eq(1, vim.wo[win.win_id].scrolloff)
 
     local legend = Window.active_windows[2]
     local lines = vim.api.nvim_buf_get_lines(legend.buf_id, 0, -1, false)
-    eq({ "<CR>=submit", "q=cancel" }, lines)
+    eq({ " <CR>=submit q=cancel" }, lines)
+  end)
+
+  it("highlights keymap legend as warning=comment", function()
+    Window.capture_input("Prompt", {
+      cb = function() end,
+      keymap = {
+        q = "cancel",
+        ["<CR>"] = "submit",
+      },
+    })
+
+    local legend = Window.active_windows[2]
+    local legend_nsid = vim.api.nvim_get_namespaces()["99.window.legend"]
+    local extmarks = vim.api.nvim_buf_get_extmarks(
+      legend.buf_id,
+      legend_nsid,
+      0,
+      -1,
+      { details = true }
+    )
+
+    local highlights = {}
+    for _, extmark in ipairs(extmarks) do
+      table.insert(highlights, {
+        extmark[2],
+        extmark[3],
+        extmark[4].end_col,
+        extmark[4].hl_group,
+      })
+    end
+
+    eq({
+      { 0, 1, 5, "WarningMsg" },
+      { 0, 6, 12, "Comment" },
+      { 0, 13, 14, "WarningMsg" },
+      { 0, 15, 21, "Comment" },
+    }, highlights)
   end)
 end)
diff --git a/lua/99/window/init.lua b/lua/99/window/init.lua
index 4cc44f4..902f94f 100644
--- a/lua/99/window/init.lua
+++ b/lua/99/window/init.lua
@@ -13,6 +13,7 @@ local M = {
 }
 
 local nsid = vim.api.nvim_create_namespace("99.window.error")
+local legend_nsid = vim.api.nvim_create_namespace("99.window.legend")
 local win_valid = vim.api.nvim_win_is_valid
 local buf_valid = vim.api.nvim_buf_is_valid
 
@@ -403,6 +404,50 @@ local function create_window_legend(win, keymap)
   vim.bo[keymap_win.buf_id].modifiable = true
   vim.bo[keymap_win.buf_id].readonly = false
   vim.api.nvim_buf_set_lines(keymap_win.buf_id, 0, -1, false, lines)
+
+  vim.api.nvim_buf_clear_namespace(keymap_win.buf_id, legend_nsid, 0, -1)
+  for line_num, line in ipairs(lines) do
+    local start_col = 1
+    while true do
+      local found_start, found_end = string.find(line, "%S+=%S+", start_col)
+      if not found_start then
+        break
+      end
+
+      local legend = string.sub(line, found_start, found_end)
+      local separator = string.find(legend, "=", 1, true)
+      if separator and separator > 1 and separator < #legend then
+        local key_start = found_start - 1
+        local key_end = found_start + separator - 2
+        local value_start = found_start + separator - 1
+
+        vim.api.nvim_buf_set_extmark(
+          keymap_win.buf_id,
+          legend_nsid,
+          line_num - 1,
+          key_start,
+          {
+            end_col = key_end,
+            hl_group = "WarningMsg",
+          }
+        )
+
+        vim.api.nvim_buf_set_extmark(
+          keymap_win.buf_id,
+          legend_nsid,
+          line_num - 1,
+          value_start,
+          {
+            end_col = found_end,
+            hl_group = "Comment",
+          }
+        )
+      end
+
+      start_col = found_end + 1
+    end
+  end
+
   vim.bo[keymap_win.buf_id].modifiable = false
   vim.bo[keymap_win.buf_id].readonly = true