omnigent-ai--omnigent/omnigent/inner/hook_scripts/codex_router_hook.py

"""Codex hook entrypoint for deterministic subagent routing.

Registered in the per-session private ``CODEX_HOME`` ``hooks.json``
generated by :mod:`omnigent.inner.codex_executor`. One subcommand:

- ``route-subagent`` (``PreToolUse``, matcher ``.*spawn_agent``) — asks
  the runner's ``route-subagent`` endpoint what model the spawn may use
  and rewrites / denies accordingly.

Stdlib-only and short-lived: ``PreToolUse`` blocks the spawn.

The spawn ``message`` arrives in the hook payload in plaintext and is the
only real routing signal — this codex's ``spawn_agent`` has no task-name
field, so without it every spawn scores the same placeholder and lands the
router's default arm. An explicit ``model`` in the spawn arguments is
forwarded as ``requested_model`` so the server can honor an in-family ask.

A rewrite echoes the rest of ``tool_input`` back untouched, but ``model``
and ``reasoning_effort`` are both rewritten: codex validates each against
its own model catalog *before* the request leaves the CLI, so a servable
catalog id or an out-of-ladder effort fails the spawn outright rather than
degrading it. See :mod:`omnigent.codex_model_vocabulary`.
"""

from __future__ import annotations

import sys
from collections.abc import Callable
from pathlib import Path
from typing import Any

from omnigent.codex_model_vocabulary import clamp_spawn_effort, codex_spawn_model
from omnigent.inner.hook_scripts.subagent_router import (
    run_route_subagent_main,
)

# Codex flattens MCP-ish tool names, so the spawn tool arrives as
# ``collaborationspawn_agent`` on 0.145.x. Both the hooks.json matcher
# and this check must stay suffix-based rather than a bare literal.
SPAWN_AGENT_TOOL_SUFFIX = "spawn_agent"
# Harness label sent to the endpoint when argv names none.
DEFAULT_HARNESS = "codex"

_LABEL = "omnigent codex router hook"
_MODULE = "omnigent.inner.hook_scripts.codex_router_hook"


def main(argv: list[str] | None = None) -> int:
    """
    Dispatch a codex router-hook subcommand.

    :param argv: Argv override excluding the program name. ``None`` reads
        :data:`sys.argv`.
    :returns: Process exit code. Always ``0`` — verdicts travel on stdout
        so a hook failure never wedges codex.
    """
    raw_argv = sys.argv[1:] if argv is None else argv
    command = raw_argv[0] if raw_argv else ""
    if command == "route-subagent":
        return run_route_subagent_main(
            raw_argv[1:],
            prog=_prog(command),
            harness=DEFAULT_HARNESS,
            label=_LABEL,
            **ROUTE_SEAMS,
        )
    print(f"{_LABEL}: unknown subcommand {command!r}", file=sys.stderr)
    return 0


def _prog(command: str) -> str:
    return f"python -m {_MODULE} {command}"


def is_spawn_agent_tool(tool_name: Any) -> bool:  # type: ignore[explicit-any]  # hook payloads are untrusted JSON
    """
    Report whether a hook payload names codex's subagent-spawn tool.

    :param tool_name: ``tool_name`` from the hook payload, e.g.
        ``"collaborationspawn_agent"``.
    :returns: ``True`` when the (flattened) name ends in
        ``spawn_agent``.
    """
    if not isinstance(tool_name, str):
        return False
    return tool_name.strip().lower().endswith(SPAWN_AGENT_TOOL_SUFFIX)


def spawn_model_translator(
    bridge_dir: str | Path | None,
) -> Callable[[str], str | None]:
    """
    Build the model translator for codex's spawn tool.

    Codex's ``spawn_agent`` validates ``model`` against its own catalog
    before the request leaves the CLI, so a servable catalog id has to be
    spelled in codex's slugs — see :mod:`omnigent.codex_model_vocabulary`.

    :param bridge_dir: Unused; the vocabulary is a property of the CLI, not
        of the session (accepted so this matches the translator-factory
        seam).
    :returns: Callable mapping a servable id to a slug codex accepts.
    """
    del bridge_dir
    return codex_spawn_model


def finalize_spawn_input(
    output: dict[str, Any] | None,  # type: ignore[explicit-any]  # hook output JSON
    tool_input: dict[str, Any] | None = None,  # type: ignore[explicit-any]  # hook payload JSON
) -> dict[str, Any] | None:  # type: ignore[explicit-any]  # hook output JSON
    """
    Clamp the rewritten spawn's effort, then announce the routed model.

    Two things only a rewrite needs. The effort clamp comes first because
    codex refuses a model/effort pairing outside the model's ladder, so a
    session default of ``xhigh`` would fail a GLM spawn the router just
    approved. The notice comes second because a rewrite is otherwise
    invisible: codex reports no model change of its own, so the top-level
    ``systemMessage`` (alongside, not inside, ``hookSpecificOutput``) is the
    only place the decision shows up. When the spawn named a model of its
    own and the router picked another, the notice says so — otherwise the
    parent reads its own ask back and never learns it was substituted.

    :param output: Hook output from ``decision_to_hook_output``, or
        ``None`` for "no opinion".
    :param tool_input: The spawn's incoming ``tool_input``, whose ``model``
        is the ask the routed model replaced.
    :returns: *output* with the effort clamped and a ``systemMessage``,
        when it rewrote the spawn's model; otherwise *output* unchanged.
    """
    if not isinstance(output, dict):
        return output
    hook_output = output.get("hookSpecificOutput")
    if not isinstance(hook_output, dict) or hook_output.get("permissionDecision") != "allow":
        return output
    updated_input = hook_output.get("updatedInput")
    if not isinstance(updated_input, dict):
        return output
    model = updated_input.get("model")
    if not isinstance(model, str) or not model:
        return output
    effort = updated_input.get("reasoning_effort")
    clamped = clamp_spawn_effort(effort if isinstance(effort, str) and effort else None, model)
    if clamped != effort and clamped is not None:
        hook_output = {
            **hook_output,
            "updatedInput": {**updated_input, "reasoning_effort": clamped},
        }
        output = {**output, "hookSpecificOutput": hook_output}
    return {**output, "systemMessage": _routing_notice(tool_input, model)}


def _routing_notice(
    tool_input: dict[str, Any] | None,  # type: ignore[explicit-any]  # hook payload JSON
    model: str,
) -> str:
    """
    Word the TUI notice for a rewritten spawn.

    :param tool_input: The spawn's incoming ``tool_input``, when available.
    :param model: The slug the spawn now runs on.
    :returns: The ``systemMessage`` text.
    """
    asked = (tool_input or {}).get("model")
    if isinstance(asked, str) and asked and asked != model:
        return f"Using Smart Routing. Requested {asked}; routing to {model}."
    return f"Using Smart Routing. Routing to {model}."


def _payload_model(payload: dict[str, Any]) -> str | None:  # type: ignore[explicit-any]  # hook payloads are untrusted JSON
    """
    Extract the parent session's model from a hook payload.

    :param payload: Codex hook payload.
    :returns: Model id, or ``None`` when the payload carries none.
    """
    model = payload.get("model")
    return model if isinstance(model, str) and model else None


#: How codex differs from the claude-native default: its own spawn-tool
#: name, its task-name keys, the ``message`` prompt key (delivered plaintext
#: in hook payloads — measured, not encrypted as once assumed), its own
#: spawn-model vocabulary, and the effort clamp plus TUI routing notice a
#: rewrite needs. ``requested_model_resolver_factory`` stays unset: codex
#: spells the ask in slugs the server's bare-id comparison already lands, so
#: the ask is forwarded verbatim.
ROUTE_SEAMS: dict[str, Any] = {  # type: ignore[explicit-any]  # route_pre_tool_use seams
    "tool_matcher": is_spawn_agent_tool,
    "task_keys": ("task_name", "agent_name"),
    "include_prompt": True,
    "prompt_keys": ("message",),
    "parent_model_resolver": _payload_model,
    "model_translator_factory": spawn_model_translator,
    "requested_model_resolver_factory": None,
    "post_process": finalize_spawn_input,
}


if __name__ == "__main__":
    raise SystemExit(main())