3 files changed, 90 insertions, 8 deletions
diff --git a/lua/99/extensions/work/worker.lua b/lua/99/extensions/work/worker.lua
index be38aca..ce03ba7 100644
--- a/lua/99/extensions/work/worker.lua
+++ b/lua/99/extensions/work/worker.lua
@@ -19,6 +19,9 @@ local utils = require("99.utils")
 --- considered done
 local M = {}
 
+local DEFAULT_WORK_DESCRIPTION =
+  "Put in the description of the work you want to complete"
+
 --- @class _99.WorkOpts
 --- @docs included
 --- @field description string | nil
@@ -48,6 +51,14 @@ local function read_work_item()
   return contents
 end
 
+--- @return string | nil
+local function hydrate_current_work_item()
+  if M.current_work_item == nil then
+    M.current_work_item = read_work_item()
+  end
+  return M.current_work_item or DEFAULT_WORK_DESCRIPTION
+end
+
 --- @param success boolean
 ---@param result string
 local function set_work_item_cb(success, result)
@@ -66,8 +77,7 @@ local function set_work_item_cb(success, result)
 end
 
 function M.update_work()
-  local work = M.current_work_item
-    or "Put in the description of the work you want to complete"
+  local work = hydrate_current_work_item()
   Window.capture_input(" Work ", {
     cb = set_work_item_cb,
     content = vim.split(work, "\n"),
@@ -79,11 +89,12 @@ function M.set_work(opts)
   opts = opts or {}
   local description = opts.description
   if description then
-    M.current_work_item = description
+    set_work_item_cb(true, description)
   else
+    local work = hydrate_current_work_item()
     Window.capture_input(" Work ", {
       cb = set_work_item_cb,
-      content = { M.current_work_item or "Put in the description of the work you want to complete" },
+      content = { work },
     })
   end
 
@@ -132,9 +143,7 @@ end
 
 function M.work()
   local _99 = require("99")
-  if M.current_work_item == nil then
-    M.current_work_item = read_work_item()
-  end
+  hydrate_current_work_item()
 
   assert(
     M.current_work_item,
diff --git a/lua/99/prompt.lua b/lua/99/prompt.lua
index a03b183..e832a34 100644
--- a/lua/99/prompt.lua
+++ b/lua/99/prompt.lua
@@ -222,10 +222,21 @@ function Prompt:_observer(obs)
   }
 end
 
+local allowed_context_types = {
+  "visual",
+  "search",
+  "tutorial",
+  "vibe",
+}
 --- @return boolean
 function Prompt:valid()
   local t = self.data.type
-  return t == "visual" or t == "search" or t == "tutorial"
+  for _, allowed in ipairs(allowed_context_types) do
+    if t == allowed then
+      return true
+    end
+  end
+  return false
 end
 
 --- @param observer _99.Providers.Observer?
diff --git a/lua/99/test/worker_spec.lua b/lua/99/test/worker_spec.lua
new file mode 100644
index 0000000..02e78b0
--- /dev/null
+++ b/lua/99/test/worker_spec.lua
@@ -0,0 +1,62 @@
+-- luacheck: globals describe it assert before_each after_each
+local _99 = require("99")
+local Window = require("99.window")
+local Worker = require("99.extensions.work.worker")
+local test_utils = require("99.test.test_utils")
+local utils = require("99.utils")
+local eq = assert.are.same
+
+describe("worker", function()
+  local previous_capture_input
+  local captured_content
+  local tmp_dir
+
+  before_each(function()
+    tmp_dir = vim.fn.tempname()
+    vim.fn.mkdir(tmp_dir, "p")
+
+    local provider = test_utils.TestProvider.new()
+    _99.setup(
+      test_utils.get_test_setup_options({ tmp_dir = tmp_dir }, provider)
+    )
+
+    Worker.current_work_item = nil
+    Worker.last_work_search = nil
+
+    captured_content = nil
+    previous_capture_input = Window.capture_input
+    Window.capture_input = function(_, opts)
+      captured_content = opts.content
+    end
+  end)
+
+  after_each(function()
+    Window.capture_input = previous_capture_input
+    Worker.current_work_item = nil
+    Worker.last_work_search = nil
+
+    if tmp_dir then
+      vim.fn.delete(tmp_dir, "rf")
+    end
+  end)
+
+  it("set_work preloads existing persisted work", function()
+    local work_path = utils.named_tmp_file(tmp_dir, "work-item")
+    local file = assert(io.open(work_path, "w"))
+    file:write("fix flaky tests")
+    file:close()
+
+    Worker.set_work()
+
+    eq({ "fix flaky tests" }, captured_content)
+  end)
+
+  it("set_work shows default when no persisted work exists", function()
+    Worker.set_work()
+
+    eq(
+      { "Put in the description of the work you want to complete" },
+      captured_content
+    )
+  end)
+end)