moved range selection from visual into prompt visual

5 files changed, 39 insertions, 47 deletions

diff --git a/lua/99/init.lua b/lua/99/init.lua
index 7356d7b..500c6dc 100644
--- a/lua/99/init.lua
+++ b/lua/99/init.lua
@@ -217,15 +217,6 @@ local _99 = {
   FATAL = Level.FATAL,
 }
 
---- you can only set those marks after the visual selection is removed
-local function set_selection_marks()
-  vim.api.nvim_feedkeys(
-    vim.api.nvim_replace_termcodes("<Esc>", true, false, true),
-    "x",
-    false
-  )
-end
-
 --- @param cb fun(context: _99.Prompt, o: _99.ops.Opts?): nil
 --- @param name string
 --- @param context _99.Prompt
@@ -360,15 +351,10 @@ end
 function _99.visual(opts)
   opts = process_opts(opts)
   local context = Prompt.visual(_99_state)
-  local function perform_range()
-    set_selection_marks()
-    local range = Range.from_visual_selection()
-    ops.over_range(context, range, opts)
-  end
   if opts.additional_prompt then
-    perform_range()
+    ops.over_range(context, opts)
   else
-    capture_prompt(perform_range, "Visual", context, opts)
+    capture_prompt(ops.over_range, "Visual", context, opts)
   end
 end
 
diff --git a/lua/99/ops/over-range.lua b/lua/99/ops/over-range.lua
index 1738270..fe82cf2 100644
--- a/lua/99/ops/over-range.lua
+++ b/lua/99/ops/over-range.lua
@@ -11,17 +11,17 @@ local Range = geo.Range
 local Point = geo.Point
 
 --- @param context _99.Prompt
---- @param range _99.Range
 --- @param opts? _99.ops.Opts
-local function over_range(context, range, opts)
+local function over_range(context, opts)
   opts = opts or {}
   local logger = context.logger:set_area("visual")
 
+  local data = context:visual_data()
+  local range = data.range
   local top_mark = Mark.mark_above_range(range)
   local bottom_mark = Mark.mark_point(range.buffer, range.end_)
   context.marks.top_mark = top_mark
   context.marks.bottom_mark = bottom_mark
-  context.data.range = range
 
   logger:debug(
     "visual request start",
diff --git a/lua/99/prompt.lua b/lua/99/prompt.lua
index d4f5095..690c94d 100644
--- a/lua/99/prompt.lua
+++ b/lua/99/prompt.lua
@@ -7,6 +7,15 @@ local get_id = require("99.id")
 local Range = require("99.geo").Range
 local Time = require("99.time")
 
+--- you can only set those marks after the visual selection is removed
+local function set_selection_marks()
+  vim.api.nvim_feedkeys(
+    vim.api.nvim_replace_termcodes("<Esc>", true, false, true),
+    "x",
+    false
+  )
+end
+
 local filetype_map = {
   typescriptreact = "typescript",
 }
@@ -98,6 +107,9 @@ end
 function Prompt.visual(_99)
   _99:refresh_rules()
 
+  set_selection_marks()
+  local range = Range.from_visual_selection()
+
   local file_type = vim.bo[0].ft
   local buffer = vim.api.nvim_get_current_buf()
   file_type = filetype_map[file_type] or file_type
@@ -115,7 +127,7 @@ function Prompt.visual(_99)
     type = "visual",
     buffer = buffer,
     file_type = file_type,
-    range = Range.zero(),
+    range = range,
   }
   context.logger:debug("99 Request", "method", "visual")
 
diff --git a/lua/99/test/request_spec.lua b/lua/99/test/request_spec.lua
index 0565b51..975a570 100644
--- a/lua/99/test/request_spec.lua
+++ b/lua/99/test/request_spec.lua
@@ -1,6 +1,7 @@
 -- luacheck: globals describe it assert
 local _99 = require("99")
 local test_utils = require("99.test.test_utils")
+local Prompt = require("99.prompt")
 local eq = assert.are.same
 
 local content = {
@@ -13,9 +14,8 @@ describe("request test", function()
   it("should replace visual selection with AI response", function()
     local p = test_utils.test_setup(content, 2, 1, "lua")
     local state = _99.__get_state()
-    local Prompt = require("99.prompt")
 
-    local context = Prompt.visual(state)
+    local context = Prompt.search(state)
     context:finalize()
 
     local finished_called = false
diff --git a/lua/99/test/visual_spec.lua b/lua/99/test/visual_spec.lua
index 7f17458..498c47f 100644
--- a/lua/99/test/visual_spec.lua
+++ b/lua/99/test/visual_spec.lua
@@ -5,6 +5,8 @@ local eq = assert.are.same
 local Levels = require("99.logger.level")
 local Range = require("99.geo").Range
 local Point = require("99.geo").Point
+local visual_fn = require("99.ops.over-range")
+local Prompt = require("99.prompt")
 
 --- @param content string[]
 --- @param start_row number
@@ -43,17 +45,21 @@ local content = {
   "end",
 }
 
+--- @param context _99.Prompt
+local function visual_call_with_range(context, range)
+  context.data.range = range
+  visual_fn(context, {
+    additional_prompt = "test prompt",
+  })
+end
+
 describe("visual", function()
   it("should replace visual selection with AI response", function()
     local p, buffer, range = setup(content, 2, 1, 2, 23)
     local state = _99.__get_state()
-    local visual_fn = require("99.ops.over-range")
+    local context = Prompt.visual(state)
 
-    local context = require("99.prompt").visual(state)
-
-    visual_fn(context, range, {
-      additional_prompt = "test prompt",
-    })
+    visual_call_with_range(context, range)
 
     eq(1, state:active_request_count())
     eq(content, r(buffer))
@@ -80,12 +86,9 @@ describe("visual", function()
     }
     local p, buffer, range = setup(multi_line_content, 2, 1, 4, 17)
     local state = _99.__get_state()
-    local visual_fn = require("99.ops.over-range")
+    local context = Prompt.visual(state)
 
-    local context = require("99.prompt").visual(state)
-    visual_fn(context, range, {
-      additional_prompt = "test prompt",
-    })
+    visual_call_with_range(context, range)
 
     eq(1, state:active_request_count())
     eq(multi_line_content, r(buffer))
@@ -106,13 +109,10 @@ describe("visual", function()
 
   it("should cancel request when stop_all_requests is called", function()
     local p, buffer, range = setup(content, 2, 1, 2, 23)
-    local visual_fn = require("99.ops.over-range")
     local state = _99.__get_state()
-    local context = require("99.prompt").visual(state, 300)
+    local context = Prompt.visual(state)
 
-    visual_fn(context, range, {
-      additional_prompt = "test prompt",
-    })
+    visual_call_with_range(context, range)
 
     eq(content, r(buffer))
 
@@ -134,13 +134,10 @@ describe("visual", function()
 
   it("should handle error cases with graceful failures", function()
     local p, buffer, range = setup(content, 2, 1, 2, 23)
-    local visual_fn = require("99.ops.over-range")
     local state = _99.__get_state()
-    local context = require("99.prompt").visual(state)
+    local context = Prompt.visual(state)
 
-    visual_fn(context, range, {
-      additional_prompt = "test prompt",
-    })
+    visual_call_with_range(context, range)
 
     eq(content, r(buffer))
 
@@ -153,13 +150,10 @@ describe("visual", function()
 
   it("should handle cancelled status gracefully", function()
     local p, buffer, range = setup(content, 2, 1, 2, 23)
-    local visual_fn = require("99.ops.over-range")
     local state = _99.__get_state()
     local context = require("99.prompt").visual(state)
 
-    visual_fn(context, range, {
-      additional_prompt = "test prompt",
-    })
+    visual_call_with_range(context, range)
 
     eq(content, r(buffer))