#!/bin/zsh
typeset -gH _autocomplete__last_widget=

.autocomplete.zle-flags() {
  emulate -L zsh
  setopt $_autocomplete__func_opts[@]

  # If no arg, assume same widget as last time.
  [[ -v 1 && -n $1 ]] &&
      typeset -gH _autocomplete__last_widget="$1"

  # Return 0 to continue, 1 to abort.
  case $_autocomplete__last_widget in
    ( *kill-*~vi-* )
      builtin zle -f kill
      return 0
    ;;
    ( *yank*~vi-* | vi-*put-*after | bracketed-paste )
      builtin zle -f yank
      return 1
    ;;
    ( vi-*put-*before )
      builtin zle -f yankbefore
      return 1
    ;;
    ( * )
      return 0
    ;;
  esac
}

.autocomplete.zle-flags "$@"