This codex plugin sets up completion for codex.
codex
To use it, add codex to the plugins array of your zshrc file:
plugins=(... codex)