atlas-iac/services/hermes/scripts/patch_codex_runtime.py

#!/usr/bin/env python3
"""Let Hermes use the authenticated Codex app-server with routed settings."""

from __future__ import annotations

import argparse
from pathlib import Path


PROVIDER_BEFORE = '''    if provider == "openai-codex":
        try:
            creds = resolve_codex_runtime_credentials()
'''
PROVIDER_AFTER = '''    if provider == "openai-codex":
        if (
            str(model_cfg.get("openai_runtime") or "").strip().lower()
            == "codex_app_server"
        ):
            return {
                "provider": "openai-codex",
                "api_mode": "codex_app_server",
                "base_url": DEFAULT_CODEX_BASE_URL,
                "api_key": "codex-cli-runtime",
                "source": "codex-cli",
                "requested_provider": requested_provider,
            }
        try:
            creds = resolve_codex_runtime_credentials()
'''

SESSION_SIGNATURE_BEFORE = '''    def run_turn(
        self,
        user_input: Any,
        *,
        turn_timeout: float = 600.0,
'''
SESSION_SIGNATURE_AFTER = '''    def run_turn(
        self,
        user_input: Any,
        *,
        model: Optional[str] = None,
        effort: Optional[str] = None,
        turn_timeout: float = 600.0,
'''

SESSION_REQUEST_BEFORE = '''            ts = self._client.request(
                "turn/start",
                {
                    "threadId": self._thread_id,
                    "input": [{"type": "text", "text": user_input_text}],
                },
                timeout=10,
            )
'''
SESSION_REQUEST_AFTER = '''            turn_params: dict[str, Any] = {
                "threadId": self._thread_id,
                "input": [{"type": "text", "text": user_input_text}],
                "approvalPolicy": "never",
                "sandboxPolicy": {"type": "dangerFullAccess"},
            }
            if model:
                turn_params["model"] = model
            if effort:
                turn_params["effort"] = effort
            ts = self._client.request(
                "turn/start",
                turn_params,
                timeout=10,
            )
'''

TURN_BEFORE = '''    try:
        turn = agent._codex_session.run_turn(user_input=user_message)
'''
TURN_AFTER = '''    reasoning = getattr(agent, "reasoning_config", None)
    effort = (
        str(reasoning.get("effort") or "").strip()
        if isinstance(reasoning, dict)
        else ""
    )
    try:
        turn = agent._codex_session.run_turn(
            user_input=user_message,
            model=str(getattr(agent, "model", "") or "").strip() or None,
            effort=effort or None,
        )
'''

FALLBACK_RESOLUTION_BEFORE = '''    # Use centralized router for client construction.
    # raw_codex=True because the main agent needs direct responses.stream()
    # access for Codex providers.
    try:
        from agent.auxiliary_client import resolve_provider_client
'''
FALLBACK_RESOLUTION_AFTER = '''    # The Codex subscription is exposed through the authenticated local CLI,
    # not as a bearer token that the OpenAI-compatible fallback resolver can
    # consume. Activate that runtime directly and let conversation_loop hand
    # the already-built turn context to codex app-server.
    try:
        if fb_provider == "openai-codex":
            from hermes_cli.runtime_provider import resolve_runtime_provider

            runtime = resolve_runtime_provider(
                requested=fb_provider,
                target_model=fb_model,
            )
            if runtime.get("api_mode") == "codex_app_server":
                try:
                    from hermes_cli.model_normalize import normalize_model_for_provider

                    fb_model = normalize_model_for_provider(fb_model, fb_provider)
                except Exception as _norm_err:
                    logger.warning(
                        "Could not normalize fallback model %r for provider %r: %s",
                        fb_model, fb_provider, _norm_err,
                    )

                old_model = agent.model
                agent._config_context_length = None
                agent.model = fb_model
                agent.provider = fb_provider
                agent.base_url = str(runtime.get("base_url") or "")
                agent.api_key = str(runtime.get("api_key") or "codex-cli-runtime")
                agent.api_mode = "codex_app_server"
                agent.client = None
                agent._anthropic_client = None
                agent._credential_pool = None
                agent._fallback_activated = True
                agent._codex_cross_provider_fallback = True
                if hasattr(agent, "_transport_cache"):
                    agent._transport_cache.clear()
                rewrite_prompt_model_identity(agent, fb_model, fb_provider)
                agent._buffer_status(
                    f"🔄 Primary model failed — switching to fallback: "
                    f"{fb_model} via {fb_provider}"
                )
                logger.info(
                    "Fallback activated through Codex app-server: %s → %s (%s)",
                    old_model, fb_model, fb_provider,
                )
                _reset_stale_streak(agent)
                return True

        from agent.auxiliary_client import resolve_provider_client
'''

FALLBACK_DISPATCH_BEFORE = '''    while (api_call_count < agent.max_iterations and agent.iteration_budget.remaining > 0) or agent._budget_grace_call:
        # Reset per-turn checkpoint dedup so each iteration can take one snapshot
'''
FALLBACK_DISPATCH_AFTER = '''    while (api_call_count < agent.max_iterations and agent.iteration_budget.remaining > 0) or agent._budget_grace_call:
        # A provider failure can activate the CLI-backed Codex runtime after
        # this function's initial runtime dispatch. Re-enter the native Codex
        # adapter here instead of attempting an OpenAI-compatible request with
        # credentials that deliberately do not exist in Hermes' auth store.
        if agent.api_mode == "codex_app_server":
            fallback_user_message = user_message
            if getattr(agent, "_codex_cross_provider_fallback", False):
                from agent.codex_runtime import build_cross_provider_codex_prompt

                fallback_user_message = build_cross_provider_codex_prompt(
                    messages,
                    user_message,
                )
                agent._codex_cross_provider_fallback = False
            return agent._run_codex_app_server_turn(
                user_message=fallback_user_message,
                original_user_message=original_user_message,
                messages=messages,
                effective_task_id=effective_task_id,
                should_review_memory=_should_review_memory,
            )

        # Reset per-turn checkpoint dedup so each iteration can take one snapshot
'''

FALLBACK_CONTEXT_BEFORE = '''def run_codex_app_server_turn(
'''
FALLBACK_CONTEXT_AFTER = '''def build_cross_provider_codex_prompt(
    messages: List[Dict[str, Any]],
    user_message: str,
    *,
    max_chars: int = 120_000,
) -> str:
    """Bridge recent user/assistant context into a new Codex fallback thread."""

    def _text(value: Any) -> str:
        if isinstance(value, str):
            return value
        if isinstance(value, list):
            parts = []
            for item in value:
                if isinstance(item, dict):
                    candidate = item.get("text") or item.get("content")
                    if isinstance(candidate, str):
                        parts.append(candidate)
            return "\\n".join(parts)
        return ""

    exchanges = []
    for message in messages:
        if not isinstance(message, dict):
            continue
        role = str(message.get("role") or "").strip().lower()
        if role not in {"user", "assistant"}:
            continue
        content = _text(message.get("content")).strip()
        if content:
            exchanges.append(f"{role.upper()}: {content}")

    transcript = "\\n\\n".join(exchanges)
    if len(transcript) > max_chars:
        transcript = "[Earlier conversation omitted for context size.]\\n\\n" + transcript[-max_chars:]
    if not transcript:
        transcript = f"USER: {user_message}"
    return (
        "Continue this Hermes conversation after the previously selected "
        "provider exhausted its capacity. Preserve the objective and answer "
        "the latest user request; do not restart or discard completed work.\\n\\n"
        f"<hermes_conversation>\\n{transcript}\\n</hermes_conversation>"
    )


def run_codex_app_server_turn(
'''


def _replace_once(content: str, before: str, after: str, label: str) -> str:
    """Apply one exact replacement and fail closed when upstream drifts."""
    if content.count(before) != 1:
        raise RuntimeError(f"Hermes {label} patch context changed")
    return content.replace(before, after, 1)


def patch_provider(source: Path, destination: Path) -> None:
    """Allow the explicit Codex app-server runtime without duplicate OAuth."""
    content = source.read_text(encoding="utf-8")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        _replace_once(
            content,
            PROVIDER_BEFORE,
            PROVIDER_AFTER,
            "runtime provider",
        ),
        encoding="utf-8",
    )


def patch_session(source: Path, destination: Path) -> None:
    """Pass the router's selected model and effort to each Codex turn."""
    content = source.read_text(encoding="utf-8")
    content = _replace_once(
        content,
        SESSION_SIGNATURE_BEFORE,
        SESSION_SIGNATURE_AFTER,
        "app-server session signature",
    )
    content = _replace_once(
        content,
        SESSION_REQUEST_BEFORE,
        SESSION_REQUEST_AFTER,
        "app-server turn request",
    )
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(content, encoding="utf-8")


def patch_turn(source: Path, destination: Path) -> None:
    """Forward route metadata and preserve context on Codex failover."""
    content = source.read_text(encoding="utf-8")
    content = _replace_once(
        content,
        FALLBACK_CONTEXT_BEFORE,
        FALLBACK_CONTEXT_AFTER,
        "Codex fallback context",
    )
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        _replace_once(content, TURN_BEFORE, TURN_AFTER, "Codex turn"),
        encoding="utf-8",
    )


def patch_fallback(source: Path, destination: Path) -> None:
    """Activate the CLI-backed Codex runtime in Hermes' fallback chain."""
    content = source.read_text(encoding="utf-8")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        _replace_once(
            content,
            FALLBACK_RESOLUTION_BEFORE,
            FALLBACK_RESOLUTION_AFTER,
            "Codex fallback activation",
        ),
        encoding="utf-8",
    )


def patch_loop(source: Path, destination: Path) -> None:
    """Dispatch a mid-turn Codex fallback through app-server."""
    content = source.read_text(encoding="utf-8")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        _replace_once(
            content,
            FALLBACK_DISPATCH_BEFORE,
            FALLBACK_DISPATCH_AFTER,
            "Codex fallback dispatch",
        ),
        encoding="utf-8",
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("provider_source", type=Path)
    parser.add_argument("provider_destination", type=Path)
    parser.add_argument("session_source", type=Path)
    parser.add_argument("session_destination", type=Path)
    parser.add_argument("turn_source", type=Path)
    parser.add_argument("turn_destination", type=Path)
    parser.add_argument("fallback_source", type=Path)
    parser.add_argument("fallback_destination", type=Path)
    parser.add_argument("loop_source", type=Path)
    parser.add_argument("loop_destination", type=Path)
    args = parser.parse_args()
    patch_provider(args.provider_source, args.provider_destination)
    patch_session(args.session_source, args.session_destination)
    patch_turn(args.turn_source, args.turn_destination)
    patch_fallback(args.fallback_source, args.fallback_destination)
    patch_loop(args.loop_source, args.loop_destination)
    return 0


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