From 06753e8146332aa787857fc5cc41caa2b5f753f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc=20Cornell=C3=A0?= <marc@mcornella.com>
Date: Sat, 9 Mar 2024 18:22:35 +0100
Subject: [PATCH] fix(async): register the git prompt async handler correctly
 (#12267)

This fix conditionally registers the git prompt async handler only
if `git_prompt_info` is used anywhere in the prompt variables.

This is done in the proper order, so that the async request is
processed once the handler has been registered.

This fix also passes the return value of the previous command
to each of the async handlers, in case they are needed.
---
 lib/async_prompt.zsh | 12 +++++++++---
 lib/git.zsh          | 20 +++++++++++++++++++-
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/lib/async_prompt.zsh b/lib/async_prompt.zsh
index c6d03328b..384e49d33 100644
--- a/lib/async_prompt.zsh
+++ b/lib/async_prompt.zsh
@@ -5,8 +5,8 @@
 zmodload zsh/system
 
 # For now, async prompt function handlers are set up like so:
-# First, define the async function handler and add the function name
-# to the _omz_async_functions array:
+# First, define the async function handler and register the handler
+# with _omz_register_handler:
 #
 #  function _git_prompt_status_async {
 #    # Do some expensive operation that outputs to stdout
@@ -17,7 +17,7 @@ zmodload zsh/system
 # which will show the output of "$_OMZ_ASYNC_OUTPUT[handler_name]":
 #
 #  function git_prompt_status {
-#    echo -n $_OMZ_ASYNC_OUTPUT[_git_prompt_status]
+#    echo -n $_OMZ_ASYNC_OUTPUT[_git_prompt_status_async]
 #  }
 #
 #  RPROMPT='$(git_prompt_status)'
@@ -43,6 +43,7 @@ function _omz_register_handler {
 
 # Set up async handlers and callbacks
 function _omz_async_request {
+  local -i ret=$?
   typeset -gA _OMZ_ASYNC_FDS _OMZ_ASYNC_PIDS _OMZ_ASYNC_OUTPUT
 
   # executor runs a subshell for all async requests based on key
@@ -83,6 +84,8 @@ function _omz_async_request {
       builtin echo ${sysparams[pid]}
       # Store handler name for callback
       builtin echo $handler
+      # Set exit code for the handler if used
+      (exit $ret)
       # Run the async function handler
       $handler
     )
@@ -138,3 +141,6 @@ function _omz_async_callback() {
   _OMZ_ASYNC_FDS[$handler]=-1
   _OMZ_ASYNC_PIDS[$handler]=-1
 }
+
+autoload -Uz add-zsh-hook
+add-zsh-hook precmd _omz_async_request
diff --git a/lib/git.zsh b/lib/git.zsh
index 8fe999095..96df5589d 100644
--- a/lib/git.zsh
+++ b/lib/git.zsh
@@ -40,11 +40,29 @@ function _omz_git_prompt_status() {
 # Enable async prompt by default unless the setting is at false / no
 if zstyle -t ':omz:alpha:lib:git' async-prompt; then
   function git_prompt_info() {
-    _omz_register_handler _omz_git_prompt_status
     if [[ -n "$_OMZ_ASYNC_OUTPUT[_omz_git_prompt_status]" ]]; then
       echo -n "$_OMZ_ASYNC_OUTPUT[_omz_git_prompt_status]"
     fi
   }
+
+  # Conditionally register the async handler, only if it's needed in $PROMPT
+  # or any of the other prompt variables
+  function _defer_async_git_register() {
+    # Check if git_prompt_info is used in a prompt variable
+    case "${PS1}:${PS2}:${PS3}:${PS4}:${RPS1}:${RPS2}:${RPS3}:${RPS4}" in
+    *(\$\(git_prompt_info\)|\`git_prompt_info\`)*)
+      _omz_register_handler _omz_git_prompt_status
+      return
+      ;;
+    esac
+
+    add-zsh-hook -d precmd _defer_async_git_register
+    unset -f _defer_async_git_register
+  }
+
+  # Register the async handler first. This needs to be done before
+  # the async request prompt is run
+  precmd_functions=(_defer_async_git_register $precmd_functions)
 else
   function git_prompt_info() {
     _omz_git_prompt_status