omnigent-ai--omnigent/omnigent/onboarding/ucode_state.py

"""Read ucode's state.json.

ucode (github.com/databricks/ucode) is a CLI that configures coding harnesses
to talk to Databricks Unity AI Gateway.  After ``ucode configure`` runs, it
writes ``~/.ucode/state.json`` with the workspace URL, available models,
base URLs, and per-agent auth/config snippets.

Omnigent reads this file to pick per-harness model defaults, base URLs,
and auth helpers instead of hardcoding them.  ucode is the source of
truth for these values.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path

_logger = logging.getLogger(__name__)

_STATE_PATH = Path.home() / ".ucode" / "state.json"


@dataclass
class UcodeAgentState:
    """Parsed per-agent entry from a ucode workspace state.

    :param model: Preferred model id, e.g. ``"databricks-gpt-5-5"``.
    :param base_url: Single gateway base URL for agents with one provider,
        e.g. ``"https://example.databricks.com/ai-gateway/codex/v1"``.
    :param base_urls: Provider-specific gateway base URLs for agents with
        multiple providers, e.g. ``{"claude": "https://example.databricks.com/..."}``.
    :param auth_command: Shell command that prints a bearer token, e.g.
        ``"databricks auth token --host https://example.databricks.com ..."``.
    :param auth_refresh_interval_ms: Refresh cadence in milliseconds,
        e.g. ``900000``.
    :param auth: Structured auth config for CLIs that support it, e.g.
        ``{"command": "sh", "args": ["-c", "..."], "timeout_ms": 5000}``.
    :param env: Environment values generated by ucode for this agent,
        e.g. ``{"ANTHROPIC_BASE_URL": "https://example.databricks.com/..."}``.
    """

    model: str | None = None
    base_url: str | None = None
    base_urls: dict[str, str] = field(default_factory=dict)
    auth_command: str | None = None
    auth_refresh_interval_ms: int | None = None
    auth: dict[str, object] = field(default_factory=dict)
    env: dict[str, str] = field(default_factory=dict)


@dataclass
class UcodeWorkspaceState:
    """Parsed per-workspace entry from ``~/.ucode/state.json``.

    :param workspace_url: Workspace URL for this entry, e.g.
        ``"https://example.databricks.com"``.
    :param claude_models: Mapping of tier to model id,
        e.g. ``{"opus": "databricks-claude-opus-4-7", "sonnet": "..."}``.
        An optional ``"sonnet_5"`` key pins Claude Code's one custom
        ``/model`` picker slot (see
        :data:`omnigent.claude_native._UCODE_CLAUDE_CUSTOM_TIER`) to the
        newer Sonnet generation, offered as an opt-in alongside the default
        ``"sonnet"`` tier, for workspaces that serve both side by side.
    :param codex_models: Ordered list of Codex model ids available on this
        workspace, e.g. ``["databricks-gpt-5-5"]``.
    :param base_urls: Mapping of tool name to base URL,
        e.g. ``{"claude": "https://example.databricks.com/ai-gateway/anthropic",
        "codex": "https://example.databricks.com/ai-gateway/codex/v1"}``.
    :param available_tools: Tools configured by ucode on this workspace,
        e.g. ``["claude", "codex"]``.
    :param agents: Per-agent reusable config generated by ucode,
        e.g. ``{"codex": UcodeAgentState(...)}``.
    """

    workspace_url: str
    claude_models: dict[str, str] = field(default_factory=dict)
    codex_models: list[str] = field(default_factory=list)
    base_urls: dict[str, str] = field(default_factory=dict)
    available_tools: list[str] = field(default_factory=list)
    agents: dict[str, UcodeAgentState] = field(default_factory=dict)

    @property
    def workspace_host(self) -> str:
        """Return the workspace URL recorded for this ucode entry.

        :returns: Scheme + netloc of the workspace, e.g.
            ``"https://example.databricks.com"``.
        """
        return self.workspace_url

    def agent(self, name: str) -> UcodeAgentState | None:
        """Return reusable ucode config for an agent.

        :param name: ucode agent name, e.g. ``"claude"`` or ``"codex"``.
        :returns: Parsed :class:`UcodeAgentState`, or ``None`` when the
            agent has no entry in state.json.
        """
        return self.agents.get(name)


def _parse_agent_state(raw: object) -> UcodeAgentState:
    """Parse one entry from ucode's ``agents`` mapping.

    :param raw: Raw JSON value for one agent entry.
    :returns: Parsed :class:`UcodeAgentState`. Invalid or missing fields
        are omitted rather than raising so unrelated state additions do
        not break Omnigent.
    """
    if not isinstance(raw, dict):
        return UcodeAgentState()
    base_urls_raw = raw.get("base_urls")
    auth_raw = raw.get("auth")
    env_raw = raw.get("env")
    refresh_raw = raw.get("auth_refresh_interval_ms")
    return UcodeAgentState(
        model=raw.get("model") if isinstance(raw.get("model"), str) else None,
        base_url=raw.get("base_url") if isinstance(raw.get("base_url"), str) else None,
        base_urls={
            str(key): str(value)
            for key, value in (base_urls_raw.items() if isinstance(base_urls_raw, dict) else [])
            if isinstance(value, str)
        },
        auth_command=(
            raw.get("auth_command") if isinstance(raw.get("auth_command"), str) else None
        ),
        auth_refresh_interval_ms=refresh_raw if isinstance(refresh_raw, int) else None,
        auth=dict(auth_raw) if isinstance(auth_raw, dict) else {},
        env={str(key): str(value) for key, value in env_raw.items()}
        if isinstance(env_raw, dict)
        else {},
    )


def read_ucode_state(workspace_url: str) -> UcodeWorkspaceState | None:
    """Return the ucode workspace state for *workspace_url*, or ``None``.

    Opens ``~/.ucode/state.json`` and returns the entry for *workspace_url*
    (trailing-slash-insensitive).  The reader is intentionally tolerant of
    state-version changes; it validates the keys Omnigent consumes rather
    than rejecting the whole file based on ``state_version``.

    :param workspace_url: The Databricks workspace URL to look up,
        e.g. ``"https://example.databricks.com"``.
    :returns: Parsed :class:`UcodeWorkspaceState`, or ``None``.
    """
    if not _STATE_PATH.exists():
        return None
    try:
        raw = json.loads(_STATE_PATH.read_text())
    except (json.JSONDecodeError, OSError) as exc:
        _logger.debug("Could not read %s: %s", _STATE_PATH, exc)
        return None

    if not isinstance(raw, dict):
        return None

    workspaces = raw.get("workspaces")
    if not isinstance(workspaces, dict):
        return None

    normalized = workspace_url.rstrip("/")
    ws_key: str | None = None
    ws_data: dict | None = None
    for key, value in workspaces.items():
        if key.rstrip("/") == normalized:
            ws_key = key.rstrip("/")
            ws_data = value
            break

    if ws_key is None or ws_data is None or not isinstance(ws_data, dict):
        return None

    claude_models_raw = ws_data.get("claude_models", {})
    claude_models = claude_models_raw if isinstance(claude_models_raw, dict) else {}

    codex_models_raw = ws_data.get("codex_models", [])
    codex_models = codex_models_raw if isinstance(codex_models_raw, list) else []

    base_urls_raw = ws_data.get("base_urls", {})
    base_urls = base_urls_raw if isinstance(base_urls_raw, dict) else {}

    available_tools_raw = ws_data.get("available_tools", [])
    available_tools = available_tools_raw if isinstance(available_tools_raw, list) else []

    agents_raw = ws_data.get("agents", {})
    agents = agents_raw if isinstance(agents_raw, dict) else {}

    return UcodeWorkspaceState(
        workspace_url=str(ws_data.get("workspace") or ws_key),
        claude_models={str(k): str(v) for k, v in claude_models.items()},
        codex_models=[str(m) for m in codex_models],
        base_urls={str(k): str(v) for k, v in base_urls.items()},
        available_tools=[str(t) for t in available_tools],
        agents={str(k): _parse_agent_state(v) for k, v in agents.items()},
    )


def read_current_ucode_state() -> UcodeWorkspaceState | None:
    """Return the current ucode workspace state, or ``None``.

    Uses ``current_workspace`` from ``~/.ucode/state.json`` when present.
    If older/newer state omits that key but contains exactly one workspace,
    reads that single entry. This keeps the reader tolerant of state shape
    changes while still confirming that ``ucode configure`` wrote usable
    harness configuration.

    :returns: Parsed :class:`UcodeWorkspaceState`, or ``None`` when no
        usable workspace entry exists.
    """
    if not _STATE_PATH.exists():
        return None
    try:
        raw = json.loads(_STATE_PATH.read_text())
    except (json.JSONDecodeError, OSError) as exc:
        _logger.debug("Could not read %s: %s", _STATE_PATH, exc)
        return None
    if not isinstance(raw, dict):
        return None
    workspaces = raw.get("workspaces")
    if not isinstance(workspaces, dict) or not workspaces:
        return None
    current_workspace = raw.get("current_workspace")
    if isinstance(current_workspace, str):
        return read_ucode_state(current_workspace)
    if len(workspaces) == 1:
        workspace_url = next(iter(workspaces))
        if isinstance(workspace_url, str):
            return read_ucode_state(workspace_url)
    return None