setup_codex.py

#!/usr/bin/env python
"""Configure OpenAI Codex CLI with Databricks Model Serving.

Codex CLI is OpenAI's coding agent and uses the Responses API.
Databricks provides the Responses API through the workspace AI Gateway.

Config: ~/.codex/config.toml with custom model_providers for Databricks.
Auth: Bearer token via DATABRICKS_TOKEN environment variable.

Opt-out:
  Set ENABLE_CODEX=false in app.yaml to skip installation entirely.
"""
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

from cli_auth import _atomic_write_text
from gateway_models import codex_base_url, discover_model_catalog
from token_helper import resolve_databricks_token, write_token_helper
from utils import (
    adapt_instructions_file,
    ensure_https,
    get_npm_version,
    resolve_mlflow_experiment_id,
)

# Opt-out: allow operators to disable Codex bundling without removing the file.
if os.environ.get("ENABLE_CODEX", "true").strip().lower() in ("false", "0", "no"):
    print("ENABLE_CODEX=false — skipping Codex CLI setup")
    raise SystemExit(0)

# Set HOME if not properly set
if not os.environ.get("HOME") or os.environ["HOME"] == "/":
    os.environ["HOME"] = "/app/python/source_code"

home = Path(os.environ["HOME"])

host = os.environ.get("DATABRICKS_HOST", "")
# Use the same broker-aware credential resolution as Claude, Pi, and OpenCode
# so SP-authenticated boot works without requiring a pasted PAT.
token = resolve_databricks_token() or ""
requested_model = os.environ.get("CODEX_MODEL", "system.ai.gpt-5")

# 1. Install Codex CLI into ~/.local/bin (always, even without token)
local_bin = home / ".local" / "bin"
local_bin.mkdir(parents=True, exist_ok=True)
codex_bin = local_bin / "codex"

MAX_RETRIES = 3
RETRY_DELAY = 5  # seconds

if not codex_bin.exists():
    npm_prefix = str(home / ".local")
    codex_version = get_npm_version("@openai/codex")
    codex_pkg = f"@openai/codex@{codex_version}" if codex_version else "@openai/codex"

    for attempt in range(1, MAX_RETRIES + 1):
        print(f"Installing {codex_pkg} (attempt {attempt}/{MAX_RETRIES})...")
        result = subprocess.run(
            ["npm", "install", "-g", f"--prefix={npm_prefix}", codex_pkg],
            capture_output=True,
            text=True,
            env={**os.environ, "HOME": str(home)},
        )
        if result.returncode == 0 and codex_bin.exists():
            print(f"Codex CLI installed to {codex_bin}")
            break
        else:
            stderr = result.stderr.strip()
            print(f"Codex CLI install failed (attempt {attempt}/{MAX_RETRIES}, rc={result.returncode})")
            if stderr:
                print(f"  stderr: {stderr[:500]}")
            if result.stdout.strip():
                print(f"  stdout: {result.stdout.strip()[:500]}")
            if attempt < MAX_RETRIES:
                import time
                print(f"  Retrying in {RETRY_DELAY}s...")
                time.sleep(RETRY_DELAY)
            else:
                print(f"ERROR: Codex CLI installation failed after {MAX_RETRIES} attempts. "
                      f"Run manually: npm install -g --prefix=$HOME/.local @openai/codex")
else:
    print(f"Codex CLI already installed at {codex_bin}")

# 2. Skip auth config if no token (will be configured after PAT setup)
if not host or not token:
    print("Codex CLI installed — config will be set after PAT setup")
    exit(0)

# Strip trailing slash and ensure https:// prefix
host = ensure_https(host.rstrip("/"))

# Codex speaks the Responses API (wire_api="responses" below). Discover the
# system.ai model services that advertise the Responses dialect and select the
# newest compatible model when the configured request is not served here.
catalog = discover_model_catalog(host, token)
openai_models = catalog["openai"]
codex_model = requested_model
if openai_models:
    codex_model = requested_model if requested_model in openai_models else openai_models[0]
    if codex_model != requested_model:
        print(f"CODEX_MODEL={requested_model} not served here, using {codex_model}")

codex_url = codex_base_url(host)
print(f"Using workspace AI Gateway Responses API: {codex_url}")

# 3. Create ~/.codex directory and write config.toml
codex_dir = home / ".codex"
codex_dir.mkdir(exist_ok=True)

# Copy bundled Databricks model catalog into ~/.codex so it can be referenced
# by relative path in config.toml (codex resolves relatives against CODEX_HOME).
catalog_src = Path(__file__).parent / ".codex" / "databricks-models.json"
catalog_dst = codex_dir / "databricks-models.json"
if catalog_src.exists() and catalog_src.resolve() != catalog_dst.resolve():
    shutil.copyfile(catalog_src, catalog_dst)
    print(f"Codex model catalog copied: {catalog_dst}")

# Resolve a fresh credential per Codex process/request. This mirrors ucode's
# auth command and avoids persisting a short-lived broker bearer in ~/.codex.
helper_path = write_token_helper(codex_dir)
helper_python = os.environ.get("CODA_VENV_PYTHON") or sys.executable or "python3"
auth_config = (
    f"auth = {{ command = {json.dumps(helper_python)}, "
    f"args = [{json.dumps(str(helper_path))}], "
    "timeout_ms = 5000, refresh_interval_ms = 900000 }"
)
# Remove the static token file generated by older CoDA versions, but only when
# its marker proves that CoDA owns it.
legacy_env_path = codex_dir / ".env"
if legacy_env_path.exists() and legacy_env_path.read_text().startswith(
    "# Databricks token for Codex CLI"
):
    legacy_env_path.unlink()

# Optional: MLflow tracing notify hook (one switch enables Claude + Codex)
tracing_enabled = os.environ.get("MLFLOW_TRACING_ENABLED", "false").lower() == "true"
notify_line = ""
if tracing_enabled:
    notify_line = 'notify = ["mlflow-codex", "notify-hook"]\n'

# Codex CLI uses TOML config with custom model_providers
config_content = f"""# Databricks Model Serving Configuration for Codex CLI
# Generated by setup_codex.py

# Active model and provider
model = "{codex_model}"
model_provider = "databricks"
model_catalog_json = "databricks-models.json"

# Disable web_search - not supported by Databricks Responses API
web_search = "disabled"

{notify_line}
# Databricks custom provider
[model_providers.databricks]
name = "Databricks Model Serving"
base_url = "{codex_url}"
{auth_config}
wire_api = "responses"
"""

config_path = codex_dir / "config.toml"
config_path.write_text(config_content)
print(f"Codex CLI configured: {config_path}")

# 4. Codex invokes the helper configured above whenever it needs a bearer.
# No token is written to a terminal-visible .env file.

# MLflow tracing env vars (read by @mlflow/codex notify hook)
env_lines: list[str] = []
app_owner = os.environ.get("APP_OWNER", "")
app_name = os.environ.get("DATABRICKS_APP_NAME", "coding-agents")
experiment_name = f"/Users/{app_owner}/{app_name}" if app_owner else ""

if tracing_enabled and experiment_name:
    experiment_id = resolve_mlflow_experiment_id(host, token, experiment_name)

    # Install @mlflow/codex (provides the `mlflow-codex` binary used by the notify hook)
    mlflow_codex_bin = local_bin / "mlflow-codex"
    if not mlflow_codex_bin.exists():
        npm_prefix = str(home / ".local")
        print("Installing @mlflow/codex for MLflow tracing...")
        result = subprocess.run(
            ["npm", "install", "-g", f"--prefix={npm_prefix}", "@mlflow/codex"],
            capture_output=True, text=True,
            env={**os.environ, "HOME": str(home)},
        )
        if result.returncode == 0:
            print(f"@mlflow/codex installed to {mlflow_codex_bin}")
        else:
            print(f"WARNING: @mlflow/codex install failed (rc={result.returncode}): {result.stderr.strip()[:300]}")

    # Pass MLflow connection details via env (override mlflow-tracing.json).
    # DATABRICKS_HOST/TOKEN are inherited from the app process — no need to
    # re-write them here (and that avoids stale-token bugs on PAT rotation).
    env_lines.extend([
        "",
        "# MLflow tracing (enabled by MLFLOW_TRACING_ENABLED=true)",
        "MLFLOW_TRACKING_URI=databricks",
    ])
    if experiment_id:
        env_lines.append(f"MLFLOW_EXPERIMENT_ID={experiment_id}")
    else:
        env_lines.append(f"MLFLOW_EXPERIMENT_NAME={experiment_name}")

    # mlflow-tracing.json as a fallback for the notify hook when env isn't loaded
    tracing_cfg = {"trackingUri": "databricks"}
    if experiment_id:
        tracing_cfg["experimentId"] = experiment_id
    (codex_dir / "mlflow-tracing.json").write_text(json.dumps(tracing_cfg, indent=2))
    print(f"Codex MLflow tracing configured: experiment_id={experiment_id or 'unresolved'}")

if env_lines:
    env_path = codex_dir / ".env"
    _atomic_write_text(str(env_path), "\n".join(env_lines) + "\n")
    env_path.chmod(0o600)
    print(f"Codex CLI tracing env configured: {env_path}")

# 5. Copy Claude skills into ~/.agents/skills/ where Codex discovers them.
# Codex searches `$HOME/.agents/skills/` plus `.agents/skills/` walking up
# from cwd; both resolve to the same path on the deployed app since
# HOME == repo root, and the user-level lookup also covers local dev.
claude_skills_dir = home / ".claude" / "skills"
codex_skills_dir = home / ".agents" / "skills"
if claude_skills_dir.exists():
    codex_skills_dir.parent.mkdir(exist_ok=True)
    if codex_skills_dir.exists():
        shutil.rmtree(codex_skills_dir)
    shutil.copytree(claude_skills_dir, codex_skills_dir)
    print(f"Skills copied: {claude_skills_dir} -> {codex_skills_dir}")
else:
    print(f"No Claude skills found at {claude_skills_dir}, skipping copy")

# 6. Adapt CLAUDE.md to AGENTS.md for Codex
# Look for CLAUDE.md in common locations
claude_md_locations = [
    Path(__file__).parent / "CLAUDE.md",  # Same directory as setup script
    home / ".claude" / "CLAUDE.md",        # User's Claude config
    Path("/app/python/source_code/CLAUDE.md"),  # Databricks App location
]

claude_md_path = None
for loc in claude_md_locations:
    if loc.exists():
        claude_md_path = loc
        break

agents_path = codex_dir / "AGENTS.md"
adapt_instructions_file(
    source_path=claude_md_path or claude_md_locations[0],
    target_path=agents_path,
    new_header="# Codex Agent Instructions",
    cli_name="Codex",
)

print("\nCodex CLI ready! Usage:")
print("  codex                              # Start Codex CLI")
print("  codex 'explain this codebase'      # Run with prompt")
print(f"\nEndpoint: {codex_url}")
print(f"Model: {codex_model}")
print("Auth: Bearer token (Databricks token via OPENAI_API_KEY)")