Chore: Improve ai-code-implement-todo (#323)

tninja · Kang Tu · web-flow · commit 3d04acfef159 · 2026-04-28T10:47:57.000-07:00
* Reorder auto-test choices to prioritize no-test option

* Add org-mode TODO headline detection to detect-todo-info

* Route ask-question to implement-todo on TODO comments

* Default action to skip completing-read in implement-todo

* fix: Address PR review feedback and fix pre-existing test failures

Add explicit (require 'ai-code-change) to ai-code-discussion.el and
gate TODO routing on buffer-file-name to prevent errors in unsaved
buffers. Fix claude-code multiline test by mocking MCP agent launch,
and fix capf test by adding ai-code-git require and magit-git-lines
mock to test repo macro.

---------

Co-authored-by: Kang Tu &lt;kang_tu@apple.com&gt;
diff --git a/ai-code-change.el b/ai-code-change.el
@@ -124,18 +124,35 @@ returns that function's name.  Otherwise returns the result of `which-function`.
   "Detect TODO comment information at cursor or in selected region.
 REGION-ACTIVE indicates whether a region is selected.
 Returns (TEXT START-POS END-POS) if TODO found, nil otherwise."
+  ;; DONE: this function should include org-mode TODO headline detection as well, as it used inside ai-code-implement-todo.
   (let ((text (if region-active
                   (buffer-substring-no-properties (region-beginning) (region-end))
                 (thing-at-point 'line t))))
-    (when (and text comment-start)
-      (let* ((first-line (car (split-string text "\n")))
-             (comment-prefix-re (concat "^[ \t]*" (regexp-quote (string-trim-right comment-start)) "+[ \t]*")))
-        (when (string-match comment-prefix-re first-line)
-          (let ((rest (string-trim-left (substring first-line (match-end 0)))))
-            (when (string-prefix-p "TODO" rest)
-              (list text
-                    (if region-active (region-beginning) (line-beginning-position))
-                    (if region-active (region-end) (line-end-position))))))))))
+    (or
+     (when (and text comment-start)
+       (let* ((first-line (car (split-string text "\n")))
+              (comment-prefix-re (concat "^[ \t]*" (regexp-quote (string-trim-right comment-start)) "+[ \t]*")))
+         (when (string-match comment-prefix-re first-line)
+           (let ((rest (string-trim-left (substring first-line (match-end 0)))))
+             (when (string-prefix-p "TODO" rest)
+               (list text
+                     (if region-active (region-beginning) (line-beginning-position))
+                     (if region-active (region-end) (line-end-position))))))))
+     (when (and (not region-active)
+                (derived-mode-p 'org-mode)
+                text)
+       (save-excursion
+         (when (and (org-at-heading-p)
+                    (ignore-errors (org-back-to-heading t)))
+           (let ((heading-line (buffer-substring-no-properties
+                                (line-beginning-position)
+                                (line-end-position))))
+             (when (and (or (org-get-todo-state)
+                            (ai-code--implement-todo--org-todo-headline-p heading-line))
+                        (not (org-entry-is-done-p)))
+               (list heading-line
+                     (line-beginning-position)
+                     (line-end-position))))))))))
 
 (defun ai-code--generate-prompt-label (clipboard-context region-active function-name)
   "Generate appropriate prompt label based on context.
@@ -249,11 +266,12 @@ Argument ARG is the prefix argument."
     (let* ((region-active (region-active-p))
            (todo-info (ai-code--detect-todo-info region-active)))
       (if todo-info
-          (ai-code-implement-todo arg)
+          ;; DONE: in this code path, is it possible to default to Code Change in ai-code-implement-todo, no need to ask user to choose Code Change or Ask Question?
+          (ai-code-implement-todo arg "Code change")
         (ai-code--handle-regular-code-change arg region-active))))))
 
 ;;;###autoload
-(defun ai-code-implement-todo (arg)
+(defun ai-code-implement-todo (arg &optional default-action)
   "Generate prompt to implement TODO comments in current context.
 Implements code after TODO comments instead of replacing them in-place.
 With a prefix argument \\[universal-argument], append the clipboard
@@ -264,7 +282,8 @@ The input string will be prefixed with TODO: and insert to the current
 line, with proper indentation.  If cursor is inside a function, implement
 comments for that function.
 Otherwise implement comments for the entire current file.
-Argument ARG is the prefix argument."
+Argument ARG is the prefix argument.
+Optional DEFAULT-ACTION skips the action prompt when non-nil."
   ;; DONE: I want to implement the idea inside https://github.com/tninja/ai-code-interface.el/issues/316, it could to either code change or ask question, given user's input with completing-read selection. The difference of this org-mode section TODO, with the existing comment todo is, it won't replace the TODO section with implementation. It just use the section headline and content inside this section as part of prompt, and send to AI.
   (interactive "P")
   (if (not buffer-file-name)
@@ -274,7 +293,7 @@ Argument ARG is the prefix argument."
         (cl-return-from finalize nil))
       (when (ai-code--implement-todo--handle-blank-line)
         (cl-return-from finalize nil))
-      (ai-code--implement-todo--build-and-send-prompt arg))))
+      (ai-code--implement-todo--build-and-send-prompt arg default-action))))
 
 (defun ai-code--implement-todo--handle-done-line ()
   "Handle actions when current line is a DONE comment.
@@ -378,9 +397,10 @@ The plist contains `:heading-line', `:content', and `:line-number'."
   (string-match-p "^[[:space:]]*\\*+[[:space:]]+TODO:?\\(?:[[:space:]]\\|$\\)"
                   heading-line))
 
-(defun ai-code--implement-todo--build-and-send-prompt (arg)
+(defun ai-code--implement-todo--build-and-send-prompt (arg &optional default-action)
   "Build the TODO implementation prompt and insert it.
-ARG is the prefix argument for clipboard context."
+ARG is the prefix argument for clipboard context.
+Optional DEFAULT-ACTION skips the completing-read prompt when non-nil."
   ;; DONE: ask user with completing-read before build up prompt, candidate should be 1. Code change; 2. Ask question. Given selection, add suffix to them respectively to indicate AI to make code change, or do not make any code change
   (let* ((clipboard-context (when arg (ai-code--get-clipboard-text)))
          (current-line (string-trim (thing-at-point 'line t)))
@@ -423,9 +443,10 @@ ARG is the prefix argument for clipboard context."
               (user-error "Current line is not a TODO comment or Org TODO headline and cannot proceed with `ai-code-implement-todo'.  Please select a TODO comment (not DONE), an Org TODO headline, a region of comments, or activate on a blank line")))
          (_ (unless region-comment-block-p
               (user-error "Selected region must be a comment block")))
-         (action-intent (completing-read "Select action: "
-                                         '("Code change" "Ask question")
-                                         nil t))
+         (action-intent (or default-action
+                           (completing-read "Select action: "
+                                            '("Code change" "Ask question")
+                                            nil t)))
          (ask-question-p (string= action-intent "Ask question"))
          (prompt-label
           (cond
diff --git a/ai-code-discussion.el b/ai-code-discussion.el
@@ -13,6 +13,7 @@
 
 (require 'ai-code-input)
 (require 'ai-code-prompt-mode)
+(require 'ai-code-change)
 
 (declare-function ai-code-read-string "ai-code-input")
 (declare-function ai-code--insert-prompt "ai-code-prompt-mode")
@@ -102,13 +103,18 @@ Inserts the prompt into the AI prompt file and optionally sends to AI.
 
 Argument ARG is the prefix argument."
   (interactive "P")
-  (let ((clipboard-context (when arg (ai-code--get-clipboard-text))))
-    (cond
-     ;; Handle dired buffer
-     ((derived-mode-p 'dired-mode)
-      (ai-code--ask-question-dired clipboard-context))
-     ;; Handle regular file buffer
-     (t (ai-code--ask-question-file clipboard-context)))))
+  ;; DONE: similar to ai-code-code-change, when todo-info is available, call ai-code-implement-todo
+  (cond
+   ((derived-mode-p 'dired-mode)
+    (let ((clipboard-context (when arg (ai-code--get-clipboard-text))))
+      (ai-code--ask-question-dired clipboard-context)))
+   (t
+    (let ((todo-info (when buffer-file-name
+                       (ai-code--detect-todo-info (region-active-p)))))
+      (if todo-info
+          (ai-code-implement-todo arg "Ask question")
+        (let ((clipboard-context (when arg (ai-code--get-clipboard-text))))
+          (ai-code--ask-question-file clipboard-context)))))))
 
 (defun ai-code--ask-question-dired (clipboard-context)
   "Handle ask question for Dired buffer.
diff --git a/ai-code.el b/ai-code.el
@@ -182,9 +182,9 @@ See the later `defcustom' for user-facing documentation and default.")
 
 (defconst ai-code--auto-test-type-ask-choices
   '(("Run tests after code change" . test-after-change)
+    ("Do not write or run tests" . no-test)
     ("TDD Red + Green (write failing test, then make it pass)" . tdd)
-    ("TDD Red + Green + Blue (refactor after Green)" . tdd-with-refactoring)
-    ("Do not write or run tests" . no-test))
+    ("TDD Red + Green + Blue (refactor after Green)" . tdd-with-refactoring))
   "Resolve auto test suffix choices for `ask-me` mode.")
 
 (defconst ai-code--auto-test-type-persistent-choices
diff --git a/test/test_ai-code-change.el b/test/test_ai-code-change.el
@@ -581,6 +581,146 @@ is between the function definition and its body."
         (should (string-match-p "\\*\\* TODO: what is the most important verse in Bible"
                                 captured-prompt))))))
 
+(ert-deftest ai-code-test-detect-todo-info-org-todo-headline ()
+  "Test `ai-code--detect-todo-info' detects Org TODO headlines."
+  (with-temp-buffer
+    (require 'org)
+    (setq buffer-file-name "plan.org")
+    (insert "* TODO Build search feature\n")
+    (insert "Design the API first.\n")
+    (org-mode)
+    (goto-char (point-min))
+
+    (cl-letf (((symbol-function 'region-active-p) (lambda () nil)))
+      (let ((result (ai-code--detect-todo-info nil)))
+        (should result)
+        (should (stringp (nth 0 result)))
+        (should (string-match-p "TODO Build search feature" (nth 0 result)))
+        (should (integerp (nth 1 result)))
+        (should (integerp (nth 2 result)))))))
+
+(ert-deftest ai-code-test-detect-todo-info-org-done-headline-returns-nil ()
+  "Test `ai-code--detect-todo-info' returns nil for Org DONE headlines."
+  (with-temp-buffer
+    (require 'org)
+    (setq buffer-file-name "plan.org")
+    (insert "* DONE Completed task\n")
+    (org-mode)
+    (goto-char (point-min))
+
+    (cl-letf (((symbol-function 'region-active-p) (lambda () nil)))
+      (should-not (ai-code--detect-todo-info nil)))))
+
+(ert-deftest ai-code-test-detect-todo-info-org-non-todo-headline-returns-nil ()
+  "Test `ai-code--detect-todo-info' returns nil for non-TODO Org headlines."
+  (with-temp-buffer
+    (require 'org)
+    (setq buffer-file-name "notes.org")
+    (insert "* Regular heading\n")
+    (org-mode)
+    (goto-char (point-min))
+
+    (cl-letf (((symbol-function 'region-active-p) (lambda () nil)))
+      (should-not (ai-code--detect-todo-info nil)))))
+
+(ert-deftest ai-code-test-detect-todo-info-org-todo-colon-prefix ()
+  "Test `ai-code--detect-todo-info' detects `TODO:' prefixed Org headlines."
+  (with-temp-buffer
+    (require 'org)
+    (setq buffer-file-name "plan.org")
+    (insert "** TODO: Implement auth module\n")
+    (org-mode)
+    (goto-char (point-min))
+
+    (cl-letf (((symbol-function 'region-active-p) (lambda () nil)))
+      (let ((result (ai-code--detect-todo-info nil)))
+        (should result)
+        (should (string-match-p "TODO:" (nth 0 result)))))))
+
+(ert-deftest ai-code-test-implement-todo-default-action-skips-completing-read ()
+  "Test that passing default-action skips the completing-read prompt."
+  (with-temp-buffer
+    (setq buffer-file-name "test.el")
+    (setq-local comment-start ";")
+    (setq-local comment-end "")
+    (insert ";; TODO: implement feature\n")
+    (goto-char (point-min))
+
+    (let (captured-prompt (completing-read-called nil))
+      (cl-letf (((symbol-function 'completing-read)
+                 (lambda (&rest _args)
+                   (setq completing-read-called t)
+                   "Code change"))
+                ((symbol-function 'ai-code-read-string)
+                 (lambda (_label input) input))
+                ((symbol-function 'ai-code--get-clipboard-text) (lambda () nil))
+                ((symbol-function 'ai-code--get-context-files-string) (lambda () ""))
+                ((symbol-function 'ai-code--format-repo-context-info) (lambda () ""))
+                ((symbol-function 'ai-code--get-function-name-for-comment) (lambda () nil))
+                ((symbol-function 'which-function) (lambda () nil))
+                ((symbol-function 'region-active-p) (lambda () nil))
+                ((symbol-function 'ai-code--insert-prompt)
+                 (lambda (p) (setq captured-prompt p))))
+
+        (ai-code--implement-todo--build-and-send-prompt nil "Code change")
+
+        (should (stringp captured-prompt))
+        (should-not completing-read-called)
+        (should (string-match-p "Please implement code" captured-prompt))))))
+
+(ert-deftest ai-code-test-implement-todo-nil-default-action-prompts-user ()
+  "Test that nil default-action still prompts user with completing-read."
+  (with-temp-buffer
+    (setq buffer-file-name "test.el")
+    (setq-local comment-start ";")
+    (setq-local comment-end "")
+    (insert ";; TODO: implement feature\n")
+    (goto-char (point-min))
+
+    (let (captured-prompt (completing-read-called nil))
+      (cl-letf (((symbol-function 'completing-read)
+                 (lambda (_prompt candidates &rest _args)
+                   (setq completing-read-called t)
+                   (if (and (listp candidates)
+                            (member "Code change" candidates))
+                       "Code change"
+                     (if (listp candidates) (car candidates) ""))))
+                ((symbol-function 'ai-code-read-string)
+                 (lambda (_label input) input))
+                ((symbol-function 'ai-code--get-clipboard-text) (lambda () nil))
+                ((symbol-function 'ai-code--get-context-files-string) (lambda () ""))
+                ((symbol-function 'ai-code--format-repo-context-info) (lambda () ""))
+                ((symbol-function 'ai-code--get-function-name-for-comment) (lambda () nil))
+                ((symbol-function 'which-function) (lambda () nil))
+                ((symbol-function 'region-active-p) (lambda () nil))
+                ((symbol-function 'ai-code--insert-prompt)
+                 (lambda (p) (setq captured-prompt p))))
+
+        (ai-code--implement-todo--build-and-send-prompt nil nil)
+
+        (should (stringp captured-prompt))
+        (should completing-read-called)))))
+
+(ert-deftest ai-code-test-code-change-passes-code-change-action ()
+  "Test that `ai-code-code-change' passes \"Code change\" as default-action."
+  (with-temp-buffer
+    (setq buffer-file-name "test.el")
+    (setq-local comment-start ";")
+    (setq-local comment-end "")
+    (insert ";; TODO: implement feature\n")
+    (goto-char (point-min))
+
+    (let (captured-default-action)
+      (cl-letf (((symbol-function 'ai-code--get-clipboard-text) (lambda () nil))
+                ((symbol-function 'ai-code-implement-todo)
+                 (lambda (_arg &optional default-action)
+                   (setq captured-default-action default-action)))
+                ((symbol-function 'region-active-p) (lambda () nil)))
+
+        (ai-code-code-change nil)
+
+        (should (equal captured-default-action "Code change"))))))
+
 (provide 'test_ai-code-change)
 
 ;;; test_ai-code-change.el ends here
diff --git a/test/test_ai-code-claude-code.el b/test/test_ai-code-claude-code.el
@@ -25,6 +25,8 @@
               ((symbol-function 'ai-code-backends-infra--resolve-start-command)
                (lambda (&rest _args)
                  (list :command "claude")))
+              ((symbol-function 'ai-code-mcp-agent-prepare-launch)
+               (lambda (&rest _args) nil))
               ((symbol-function 'ai-code-backends-infra--toggle-or-create-session)
                (lambda (&rest args)
                  (cl-destructuring-bind
diff --git a/test/test_ai-code-discussion.el b/test/test_ai-code-discussion.el
diff --git a/test/test_ai-code-prompt-mode.el b/test/test_ai-code-prompt-mode.el
