add managed CodingAgentConfig read client, normalize, and local persist

src/ucode/databricks.py

@@ -1374,6 +1374,31 @@ def discover_model_services(
     return claude_models, codex_models, gemini_models, oss_models, None
 
 
+# --- Managed coding-agent config (admin-authored, developer-read) -----------
+
+# The workspace-admin authors a CodingAgentConfig via the AI Gateway; developers read it
+# (non-admin) through the List endpoint and apply it locally.
+_CODING_AGENT_CONFIGS_API_PATH = "/api/ai-gateway/v2/coding-agent-configs"
+
+
+def fetch_managed_coding_agent_configs(workspace: str, token: str) -> tuple[list[dict], str | None]:
+    """List the workspace's managed CodingAgentConfig(s) via the AI Gateway."""
+    hostname = workspace_hostname(workspace)
+    url = f"https://{hostname}{_CODING_AGENT_CONFIGS_API_PATH}"
+    payload, reason = _http_get_json(url, token, timeout=30)
+    if reason is not None:
+        return [], reason
+    if isinstance(payload, dict):
+        configs = payload.get("coding_agent_configs") or []
+    elif isinstance(payload, list):
+        configs = payload
+    else:
+        return [], "coding-agent-configs listing returned an unexpected response shape"
+    if not isinstance(configs, list):
+        return [], "coding-agent-configs listing returned an unexpected response shape"
+    return [c for c in configs if isinstance(c, dict)], None
+
+
 # --- MCP services (parallel to model services) -----------------------------
 
 

src/ucode/managed_config.py

@@ -0,0 +1,317 @@
+"""Admin-authored managed coding-agent config: fetch, normalize, and local persistence.
+
+An org admin authors a ``CodingAgentConfig`` through the Databricks AI Gateway; developers read it
+(non-admin) and ``ucode`` applies it locally. This module owns the developer-read half:
+
+- fetching the raw manifest (via :func:`ucode.databricks.fetch_managed_coding_agent_configs`),
+- normalizing the proto-JSON into a stable internal dict keyed by ucode's own tool names, and
+- persisting it to ``~/.ucode/managed-state.json`` (0600) so launches can reconcile against it.
+
+Reconciliation against the local ``state.json`` and applying the manifest to agents live in later
+changes; this module deliberately stops at "read + normalize + persist".
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import cast
+
+import ucode.config_io as config_io
+from ucode.databricks import fetch_managed_coding_agent_configs
+
+MANAGED_STATE_PATH = config_io.APP_DIR / "managed-state.json"
+
+# Shown to a developer when their workspace has no admin-defined managed config yet — the normal
+# case, not an error. Kept here so the CLI (which surfaces it) uses one consistent message.
+NO_MANAGED_CONFIG_MESSAGE = "No coding-agent config has been set up by your workspace admin yet."
+
+# CodingAgent proto enum -> ucode tool name. Anything unrecognized (e.g. a newer agent this ucode
+# build doesn't know) is dropped during normalization rather than guessed at.
+_AGENT_ENUM_TO_TOOL: dict[str, str] = {
+    "CODING_AGENT_CLAUDE_CODE": "claude",
+    "CODING_AGENT_CODEX": "codex",
+    "CODING_AGENT_GEMINI": "gemini",
+    "CODING_AGENT_COPILOT": "copilot",
+    "CODING_AGENT_PI": "pi",
+    "CODING_AGENT_OPENCODE": "opencode",
+}
+
+# McpServerType proto enum -> ucode's short type tag. Mirrors the selection prefixes in ``mcp.py``;
+# the actual name->URL resolution happens there when the manifest is applied (a later change).
+_MCP_TYPE_ENUM_TO_TAG: dict[str, str] = {
+    "MCP_SERVER_TYPE_UC_SERVICE": "mcp-service",
+    "MCP_SERVER_TYPE_EXTERNAL": "external",
+    "MCP_SERVER_TYPE_GENIE": "genie-space",
+    "MCP_SERVER_TYPE_VECTOR_SEARCH": "vector-search",
+    "MCP_SERVER_TYPE_UC_FUNCTIONS": "uc-functions",
+    "MCP_SERVER_TYPE_DATABRICKS_APP": "app",
+    "MCP_SERVER_TYPE_DATABRICKS_SQL": "sql",
+}
+
+
+def _as_dict(value: object) -> dict[str, object]:
+    """Return ``value`` as a ``dict[str, object]`` when it is a dict, else an empty dict.
+
+    Centralizes the isinstance-narrowing so downstream ``.get`` calls type-check (a bare
+    ``isinstance(x, dict)`` narrows to ``dict[Never, Never]``, which rejects string keys)."""
+    return cast("dict[str, object]", value) if isinstance(value, dict) else {}
+
+
+def _str(value: object) -> str | None:
+    """Return a non-empty stripped string, or None."""
+    if isinstance(value, str):
+        stripped = value.strip()
+        return stripped or None
+    return None
+
+
+def _str_list(value: object) -> list[str]:
+    if not isinstance(value, list):
+        return []
+    out: list[str] = []
+    for item in value:
+        s = _str(item)
+        if s:
+            out.append(s)
+    return out
+
+
+def _normalize_model_config(model_config: object) -> dict | None:
+    """Normalize an ``AgentModelConfig`` oneof into ``{model_provider_service?, default_model?,
+    models}``.
+
+    The proto is a oneof over per-agent variants (claude/codex/opencode/pi/gemini/copilot). We
+    don't care which variant tag it is here — the enclosing agent already tells us — so we read the
+    common fields. Claude's ``models`` is a dict of family slots; the rest are a flat list. Returns
+    None when there's no usable model config.
+    """
+    mc = _as_dict(model_config)
+    if not mc:
+        return None
+    # Unwrap the oneof: take whichever single variant sub-dict is present.
+    variant = next((_as_dict(v) for v in mc.values() if isinstance(v, dict)), None)
+    if not variant:
+        return None
+    result: dict = {}
+    mps = _str(variant.get("model_provider_service"))
+    if mps:
+        result["model_provider_service"] = mps
+    default_model = _str(variant.get("default_model"))
+    if default_model:
+        result["default_model"] = default_model
+    models = variant.get("models")
+    if isinstance(models, dict):
+        # Claude family slots (default_opus_model, default_sonnet_model, ...).
+        slots = {k: _str(v) for k, v in _as_dict(models).items() if _str(v)}
+        if slots:
+            result["models"] = slots
+    else:
+        model_list = _str_list(models)
+        if model_list:
+            result["models"] = model_list
+    return result or None
+
+
+def _normalize_enabled_agent(entry: object) -> tuple[str, dict] | None:
+    """Normalize one ``EnabledAgent`` into ``(tool, agent_config)``, or None if unusable.
+
+    Drops entries whose agent enum is unset/unknown to this ucode build.
+    """
+    entry_dict = _as_dict(entry)
+    if not entry_dict:
+        return None
+    tool = _AGENT_ENUM_TO_TOOL.get(_str(entry_dict.get("agent")) or "")
+    if tool is None:
+        return None
+    config_in = _as_dict(entry_dict.get("config"))
+    agent_config: dict = {}
+    if isinstance(config_in.get("use_as_global_settings"), bool):
+        agent_config["use_as_global_settings"] = config_in["use_as_global_settings"]
+    headers = config_in.get("custom_headers")
+    if isinstance(headers, dict):
+        clean = {
+            k: v for k, v in _as_dict(headers).items() if isinstance(k, str) and isinstance(v, str)
+        }
+        if clean:
+            agent_config["custom_headers"] = clean
+    tracing_table = _tracing_table(config_in.get("tracing_config"))
+    if tracing_table:
+        agent_config["tracing_table"] = tracing_table
+    model_config = _normalize_model_config(config_in.get("model_config"))
+    if model_config is not None:
+        agent_config["model_config"] = model_config
+    return tool, agent_config
+
+
+def _tracing_table(tracing: object) -> str | None:
+    """Extract ``TracingConfig.table`` (a UC table FQN), or None."""
+    return _str(_as_dict(tracing).get("table"))
+
+
+def _normalize_mcp_servers(value: object) -> list[dict]:
+    if not isinstance(value, list):
+        return []
+    out: list[dict] = []
+    for entry in value:
+        entry_dict = _as_dict(entry)
+        name = _str(entry_dict.get("name"))
+        tag = _MCP_TYPE_ENUM_TO_TAG.get(_str(entry_dict.get("type")) or "")
+        if name and tag:
+            out.append({"name": name, "type": tag})
+    return out
+
+
+def _normalize_budget_policy(value: object) -> dict | None:
+    bp = _as_dict(value)
+    if not bp:
+        return None
+    policy: dict = {}
+    display_name = _str(bp.get("display_name"))
+    if display_name:
+        policy["display_name"] = display_name
+    budget_id = _str(bp.get("budget_id"))
+    if budget_id:
+        policy["budget_id"] = budget_id
+    tiers: list[dict] = []
+    raw_tiers = bp.get("tiers")
+    for tier in raw_tiers if isinstance(raw_tiers, list) else []:
+        tier_dict = _as_dict(tier)
+        pct = tier_dict.get("spending_percentage")
+        if not isinstance(pct, (int, float)) or isinstance(pct, bool):
+            continue
+        tier_out: dict = {"spending_percentage": float(pct)}
+        agent = _AGENT_ENUM_TO_TOOL.get(_str(tier_dict.get("default_agent")) or "")
+        if agent:
+            tier_out["default_agent"] = agent
+        model = _str(tier_dict.get("default_model"))
+        if model:
+            tier_out["default_model"] = model
+        tiers.append(tier_out)
+    if tiers:
+        policy["tiers"] = tiers
+    return policy or None
+
+
+def normalize_managed_config(raw: dict) -> dict:
+    """Normalize a raw ``CodingAgentConfig`` proto-JSON dict into ucode's internal shape.
+
+    The internal shape uses ucode's own tool names and short MCP type tags so downstream reconcile
+    and apply code never touches proto enum spellings. Unknown agents / MCP types are dropped.
+    """
+    raw = _as_dict(raw)
+    result: dict = {}
+    name = _str(raw.get("name"))
+    if name:
+        result["name"] = name
+    default_agent = _AGENT_ENUM_TO_TOOL.get(_str(raw.get("default_agent")) or "")
+    if default_agent:
+        result["default_agent"] = default_agent
+    enabled_agents: dict[str, dict] = {}
+    raw_agents = raw.get("enabled_agents")
+    for entry in raw_agents if isinstance(raw_agents, list) else []:
+        normalized = _normalize_enabled_agent(entry)
+        if normalized is not None:
+            tool, agent_config = normalized
+            enabled_agents[tool] = agent_config
+    if enabled_agents:
+        result["enabled_agents"] = enabled_agents
+    mcp_servers = _normalize_mcp_servers(raw.get("mcp_servers"))
+    if mcp_servers:
+        result["mcp_servers"] = mcp_servers
+    skill_names = _str_list(_as_dict(raw.get("skills")).get("names"))
+    if skill_names:
+        result["skills"] = {"names": skill_names}
+    tracing_table = _tracing_table(raw.get("tracing"))
+    if tracing_table:
+        result["tracing_table"] = tracing_table
+    budget_policy = _normalize_budget_policy(raw.get("budget_policy"))
+    if budget_policy is not None:
+        result["budget_policy"] = budget_policy
+    return result
+
+
+def get_managed_config(workspace: str, token: str) -> tuple[dict | None, str | None]:
+    """Fetch and normalize the workspace's managed config.
+
+    Returns ``(config, reason)``:
+    - ``(config, None)`` — the normalized manifest for the workspace's single config;
+    - ``(None, None)`` — no managed config is defined for the workspace (not an error);
+    - ``(None, reason)`` — the read failed; ``reason`` says why.
+
+    "No config defined" arrives two ways depending on the backend: an empty listing (HTTP 200
+    with no configs) or a NOT_FOUND (HTTP 404). Both are the normal, non-error case for a workspace
+    whose admin hasn't set one up, so both collapse to ``(None, None)``.
+
+    v0 stores at most one config per workspace, so the first entry is the workspace's config.
+    """
+    configs, reason = fetch_managed_coding_agent_configs(workspace, token)
+    if reason is not None:
+        # A NOT_FOUND means the admin hasn't defined a config for this workspace — not a failure.
+        if _is_not_found(reason):
+            return None, None
+        return None, reason
+    if not configs:
+        return None, None
+    return normalize_managed_config(configs[0]), None
+
+
+def _is_not_found(reason: str) -> bool:
+    """True when a read failure reason indicates the config simply doesn't exist yet.
+
+    ``_http_get_json`` formats failures as ``HTTP <code> <text>[: <body>]``; a NOT_FOUND surfaces
+    as an ``HTTP 404`` there (and the API's error body carries ``NOT_FOUND``)."""
+    lowered = reason.lower()
+    return "http 404" in lowered or "not_found" in lowered
+
+
+def save_managed_state(workspace: str, config: dict) -> None:
+    """Persist the normalized managed config to ``~/.ucode/managed-state.json`` at mode 0600.
+
+    The file is org-authored, not developer-editable — 0600 keeps it readable/writable only by the
+    user (a light guard; hard enforcement / sudo ownership is a separate concern). No-op in dry-run.
+    """
+    if config_io.is_dry_run():
+        return
+    payload = {"workspace": workspace, "config": config}
+    config_io.ensure_parent_dir(MANAGED_STATE_PATH)
+    try:
+        MANAGED_STATE_PATH.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+    except OSError as exc:
+        raise RuntimeError(f"Failed to write managed state file: {MANAGED_STATE_PATH}") from exc
+    _restrict_permissions(MANAGED_STATE_PATH)
+
+
+def _restrict_permissions(path: Path) -> None:
+    """Best-effort chmod 0600. No-op where unsupported (e.g. Windows), where the effective
+    read-only guarantee is left to a later change."""
+    try:
+        os.chmod(path, 0o600)
+    except (OSError, NotImplementedError):
+        pass
+
+
+def load_managed_state(workspace: str | None) -> dict | None:
+    """Load the persisted managed config for ``workspace``, or None if absent/mismatched.
+
+    Returns the normalized config dict (the ``config`` field), only when the stored file is for the
+    same workspace — so a stale file from another workspace is ignored rather than misapplied.
+    """
+    if not workspace:
+        return None
+    data = config_io.read_json_safe(MANAGED_STATE_PATH)
+    if data.get("workspace") != workspace:
+        return None
+    config = data.get("config")
+    return config if isinstance(config, dict) else None
+
+
+def delete_managed_state() -> None:
+    """Remove the managed-state file, if any. No-op in dry-run."""
+    if config_io.is_dry_run():
+        return
+    try:
+        MANAGED_STATE_PATH.unlink(missing_ok=True)
+    except OSError as exc:
+        raise RuntimeError(f"Failed to remove managed state file: {MANAGED_STATE_PATH}") from exc