Unable to redo after undo

#!/usr/bin/env python

from prompt_toolkit import PromptSession
from prompt_toolkit.key_binding import KeyBindings


def dyn_prompt(buf):
    def inner():
        return f"undo: {len(buf._undo_stack)} | redo: {len(buf._redo_stack)}"
    return inner


def main():
    kb = KeyBindings()

    @kb.add("escape", "u")
    def _(event):
        event.current_buffer.undo()

    @kb.add("escape", "r")
    def _(event):
        event.current_buffer.redo()

    print("Try to type some text, then try: Alt-u to undo | Alt-r to redo")
    s = PromptSession("> ", key_bindings=kb)
    s.rprompt = dyn_prompt(s.app.current_buffer)
    try:
        s.prompt()
    except (EOFError, KeyboardInterrupt):
        pass


if __name__ == "__main__":
    main()