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
'''

RETRY_FALLBACK_DISPATCH_BEFORE = '''        while retry_count < max_retries:
            # ── Nous Portal rate limit guard ──────────────────────
'''
RETRY_FALLBACK_DISPATCH_AFTER = '''        while retry_count < max_retries:
            # Fallback activation happens inside this retry loop. Dispatch a
            # newly selected Codex app-server route before the next iteration
            # tries to build OpenAI-compatible kwargs from the intentionally
            # absent bearer-token client.
            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,
                )

            # ── Nous Portal rate limit guard ──────────────────────
'''

STREAM_RECOVERY_BEFORE = '''                    # If we have prior messages, roll back to last complete state
                    if len(messages) > 1:
                        agent._vprint(f"{agent.log_prefix}   ⏪ Rolling back to last complete assistant turn")
                        rolled_back_messages = agent._get_messages_up_to_last_assistant(messages)

                        agent._cleanup_task_resources(effective_task_id)
                        agent._persist_session(messages, conversation_history)

                        return {
                            "final_response": "Response truncated due to output length limit",
                            "messages": rolled_back_messages,
                            "api_calls": api_call_count,
                            "completed": False,
                            "partial": True,
                            "error": "Response truncated due to output length limit"
                        }
                    else:
                        # First message was truncated - mark as failed
                        agent._flush_status_buffer()
                        agent._vprint(f"{agent.log_prefix}❌ First response truncated - cannot recover", force=True)
                        agent._persist_session(messages, conversation_history)
                        return {
                            "final_response": "First response truncated due to output length limit",
                            "messages": messages,
                            "api_calls": api_call_count,
                            "completed": False,
                            "failed": True,
                            "error": "First response truncated due to output length limit"
                        }
'''
STREAM_RECOVERY_AFTER = '''                    # A partial-stream stub is a transport failure, not proof that
                    # the model exhausted its output budget. Retry the unchanged
                    # turn through the routed alias so Switchyard can choose a
                    # healthy provider. Never append the incomplete response.
                    _is_transport_stub = (
                        getattr(response, "id", "") == PARTIAL_STREAM_STUB_ID
                    )
                    if _is_transport_stub and truncated_tool_call_retries < 4:
                        truncated_tool_call_retries += 1
                        agent._buffer_vprint(
                            f"⚠️  Provider stream ended before a complete response — "
                            f"rerouting ({truncated_tool_call_retries}/4)..."
                        )
                        agent._session_messages = messages
                        continue

                    _incomplete_error = (
                        "Provider stream repeatedly ended before a complete response"
                        if _is_transport_stub
                        else "Response truncated due to output length limit"
                    )
                    # If we have prior messages, roll back to last complete state.
                    if len(messages) > 1:
                        agent._vprint(f"{agent.log_prefix}   ⏪ Rolling back to last complete assistant turn")
                        rolled_back_messages = agent._get_messages_up_to_last_assistant(messages)

                        agent._cleanup_task_resources(effective_task_id)
                        agent._persist_session(messages, conversation_history)

                        return {
                            "final_response": _incomplete_error,
                            "messages": rolled_back_messages,
                            "api_calls": api_call_count,
                            "completed": False,
                            "partial": True,
                            "error": _incomplete_error,
                        }

                    agent._flush_status_buffer()
                    agent._vprint(
                        f"{agent.log_prefix}❌ First response did not complete",
                        force=True,
                    )
                    agent._persist_session(messages, conversation_history)
                    return {
                        "final_response": _incomplete_error,
                        "messages": messages,
                        "api_calls": api_call_count,
                        "completed": False,
                        "failed": True,
                        "error": _incomplete_error,
                    }
'''

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(
'''

AUXILIARY_TOKEN_BEFORE = '''    except Exception as exc:
        logger.debug("Could not read Codex auth for auxiliary client: %s", exc)
        return None
'''
AUXILIARY_TOKEN_AFTER = '''    except Exception as exc:
        logger.debug("Could not read Hermes Codex auth for auxiliary client: %s", exc)

        # Agent Hermes deliberately runs Codex through the authenticated CLI
        # app-server and does not duplicate those credentials into Hermes'
        # provider auth store. Auxiliary tasks still use Hermes' native Codex
        # Responses adapter, so read the current CLI access token without
        # copying or refreshing it here. The authenticated Codex broker and
        # Codex CLI coordinate through this canonical file; this auxiliary
        # path remains read-only and never creates a metered API-key lane.
        try:
            codex_home = os.environ.get("CODEX_HOME", "").strip()
            if not codex_home:
                codex_home = str(Path.home() / ".codex")
            auth_path = Path(codex_home).expanduser() / "auth.json"
            payload = json.loads(auth_path.read_text(encoding="utf-8"))
            tokens = payload.get("tokens") or {}
            access_token = tokens.get("access_token")
            if not isinstance(access_token, str) or not access_token.strip():
                return None

            # Match the native expiry check above. An expired CLI token is not
            # refreshed from this side channel because the broker owns the
            # serialized, atomic refresh operation for chat boundaries.
            try:
                import base64
                jwt_payload = access_token.split(".")[1]
                jwt_payload += "=" * (-len(jwt_payload) % 4)
                claims = json.loads(base64.urlsafe_b64decode(jwt_payload))
                expires_at = claims.get("exp", 0)
                if expires_at and time.time() > expires_at:
                    logger.debug("Codex CLI access token is expired; skipping auxiliary route")
                    return None
            except Exception:
                pass
            return access_token.strip()
        except Exception as cli_exc:
            logger.debug("Could not read Codex CLI auth for auxiliary client: %s", cli_exc)
            return None
'''


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")
    content = _replace_once(
        content,
        FALLBACK_DISPATCH_BEFORE,
        FALLBACK_DISPATCH_AFTER,
        "Codex outer-loop fallback dispatch",
    )
    content = _replace_once(
        content,
        RETRY_FALLBACK_DISPATCH_BEFORE,
        RETRY_FALLBACK_DISPATCH_AFTER,
        "Codex retry-loop fallback dispatch",
    )
    content = patch_stream_recovery_text(content)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(content, encoding="utf-8")


def patch_stream_recovery_text(content: str) -> str:
    """Retry transport stubs and keep output-limit reporting truthful."""
    return _replace_once(
        content,
        STREAM_RECOVERY_BEFORE,
        STREAM_RECOVERY_AFTER,
        "partial stream recovery",
    )


def patch_stream_recovery(source: Path, destination: Path) -> None:
    """Patch only shared stream recovery for non-Codex Hermes lanes."""
    content = source.read_text(encoding="utf-8")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        patch_stream_recovery_text(content),
        encoding="utf-8",
    )


def patch_auxiliary(source: Path, destination: Path) -> None:
    """Let auxiliary tasks reuse the current Codex CLI access token safely."""
    content = source.read_text(encoding="utf-8")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        _replace_once(
            content,
            AUXILIARY_TOKEN_BEFORE,
            AUXILIARY_TOKEN_AFTER,
            "Codex CLI auxiliary auth",
        ),
        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)
    parser.add_argument("auxiliary_source", type=Path)
    parser.add_argument("auxiliary_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)
    patch_auxiliary(args.auxiliary_source, args.auxiliary_destination)
    return 0


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