From 8dcfec86ff69c2a066af20d605bbd0a62b04ec69 Mon Sep 17 00:00:00 2001 From: Zhu Lin Date: Mon, 7 Sep 2026 19:42:51 +0800 Subject: [PATCH 1/2] Add oh-my-pi (omp) as a supported coding agent Coexists with pi: isolated omp-home config, YAML models.yml, modelRoles.default pin in config.yml, MCP wiring, and tests. --- pyproject.toml | 1 + src/ucode/agents/__init__.py | 23 +- src/ucode/agents/omp.py | 326 +++++++++++++++++ src/ucode/cli.py | 40 +- src/ucode/config_io.py | 24 ++ src/ucode/databricks.py | 1 + src/ucode/mcp.py | 24 +- src/ucode/state.py | 7 + tests/test_agent_omp.py | 689 +++++++++++++++++++++++++++++++++++ tests/test_agents_init.py | 35 +- tests/test_cli.py | 3 +- tests/test_config_io.py | 32 ++ tests/test_databricks.py | 7 + tests/test_e2e.py | 2 +- tests/test_mcp.py | 43 ++- tests/test_state.py | 10 +- uv.lock | 2 + 17 files changed, 1246 insertions(+), 23 deletions(-) create mode 100644 src/ucode/agents/omp.py create mode 100644 tests/test_agent_omp.py diff --git a/pyproject.toml b/pyproject.toml index f69b8990..d79168ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ # request, via the SDK's stdio server + streamable-HTTP client. Works against # both mcp 1.x (httpx) and mcp 2.x (httpx2) — see mcp_proxy for the shared path. "mcp>=1.28.0", + "pyyaml>=6.0.0", "questionary>=2.0.0", "tomlkit>=0.13.0", "typer>=0.12.0", diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 8ea3e20d..53ba1312 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -40,7 +40,7 @@ spinner, ) -from . import claude, codex, copilot, gemini, opencode, pi +from . import claude, codex, copilot, gemini, omp, opencode, pi from .args import LaunchOptions as LaunchOptions from .args import explicit_model_arg_value as explicit_model_arg_value @@ -51,6 +51,7 @@ "opencode": opencode, "copilot": copilot, "pi": pi, + "omp": omp, } TOOL_SPECS: dict[str, ToolSpec] = {name: module.SPEC for name, module in _MODULES.items()} @@ -69,6 +70,8 @@ "opencode": "opencode", "copilot": "copilot", "pi": "pi", + "omp": "omp", + "oh-my-pi": "omp", } DEFAULT_TOOL = "codex" @@ -79,7 +82,7 @@ "codex": ["codex", "update"], } -# ucode tool -> `databricks aitools` agent id. gemini/pi aren't supported. +# ucode tool -> `databricks aitools` agent id. gemini/pi/omp aren't supported. AITOOLS_AGENT_TOKENS = { "claude": "claude-code", "codex": "codex", @@ -91,7 +94,7 @@ def install_databricks_ai_tools_for_agents(tools: list[str], state: dict) -> None: """Install Databricks AI Tools for supported agents. - Gemini and Pi have no ``aitools`` support and are dropped. + Gemini, Pi, and Oh My Pi have no ``aitools`` support and are dropped. """ if state.get("databricks_ai_tools_enabled", True) is False: return @@ -105,7 +108,7 @@ def normalize_tool(tool: str) -> str: normalized = TOOL_ALIASES.get(tool.strip().lower()) if not normalized: raise RuntimeError( - f"Unsupported tool '{tool}'. Use one of: codex, claude, gemini, opencode, copilot, pi." + f"Unsupported tool '{tool}'. Use one of: codex, claude, gemini, opencode, copilot, pi, omp." ) return normalized @@ -456,6 +459,8 @@ def configure_tool( result = copilot.write_tool_config(state, model) elif tool == "pi": result = pi.write_tool_config(state, model) + elif tool == "omp": + result = omp.write_tool_config(state, model) else: result = opencode.write_tool_config(state, model) # gemini/opencode/copilot/pi return (state, token); codex/claude return state @@ -492,6 +497,12 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: or bool(state.get("codex_models")) or bool(state.get("gemini_models")) ) + if tool == "omp": + return ( + bool(state.get("claude_models")) + or bool(state.get("codex_models")) + or bool(state.get("gemini_models")) + ) return False @@ -502,6 +513,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: "gemini": ("gemini",), "copilot": ("claude", "codex"), "pi": ("claude", "codex", "gemini"), + "omp": ("claude", "codex", "gemini"), } @@ -690,6 +702,7 @@ def provider_permission_error(tool: str, state: dict, err: str) -> str: def validate_all_tools(state: dict) -> None: from rich.panel import Panel # local to avoid bumping module-level deps + from ucode.agents.omp import OMP_CONFIG_BACKUP_PATH, OMP_CONFIG_PATH from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import restore_file @@ -723,6 +736,8 @@ def validate_all_tools(state: dict) -> None: # Rollback settings.json for Pi if tool == "pi": restore_file(PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, managed) + if tool == "omp": + restore_file(OMP_CONFIG_PATH, OMP_CONFIG_BACKUP_PATH, managed) available_tools.remove(tool) state["available_tools"] = available_tools save_state(state) diff --git a/src/ucode/agents/omp.py b/src/ucode/agents/omp.py new file mode 100644 index 00000000..a711f095 --- /dev/null +++ b/src/ucode/agents/omp.py @@ -0,0 +1,326 @@ +"""Oh-my-pi (omp) coding agent: writes a ucode-private models.yml with Databricks-backed providers. + +omp (https://omp.sh) is a multi-provider coding agentIAM-compatible with Pi's +provider model but configured in YAML. We register three providers in its +`models.yml`, each speaking the API dialect best suited to that family's +gateway path: + +- `databricks-claude` (api: anthropic-messages) → /ai-gateway/anthropic +- `databricks-openai` (api: openai-responses) → /ai-gateway/codex/v1 +- `databricks-gemini` (api: google-generative-ai) → /ai-gateway/gemini/v1beta + +Per-provider `compat` flags work around fields the gateway translators reject: + +- claude: `supportsEagerToolInputStreaming: false` — the Anthropic translator + rejects `tools[].eager_input_streaming` on the streaming + tools path that + omp uses for every request. With this flag omp omits the per-tool field and + sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header + instead, which the gateway accepts. + +The default model is pinned via `modelRoles.default` in omp's `config.yml` +(the role value omp's startup resolution consults before falling back to the +first available model), never in `models.yml` — whose root carries only +`providers` (unknown root keys fail omp's schema validation). + +OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via +omp today — they live behind /ai-gateway/mlflow/v1 with per-model +`max_tokens` caps that omp has no global way to honor without per-model +config we don't currently maintain. + +The bearer token is baked into the file and refreshed by a background thread +while the session runs (same pattern as OpenCode/Copilot). +""" + +from __future__ import annotations + +import os +import signal +import subprocess +import threading + +from ucode.config_io import ( + APP_DIR, + ToolSpec, + backup_existing_file, + deep_merge_dict, + read_json_safe, + read_yaml_safe, + write_json_file, + write_yaml_file, +) +from ucode.databricks import ( + TOKEN_REFRESH_INTERVAL_SECONDS, + build_pi_base_urls, + get_databricks_token, +) +from ucode.state import mark_tool_managed, save_state +from ucode.telemetry import agent_version, ucode_version + +from .args import LaunchOptions + +OMP_UCODE_HOME = APP_DIR / "omp-home" +OMP_AGENT_DIR = OMP_UCODE_HOME / ".omp" / "agent" +OMP_MODELS_PATH = OMP_AGENT_DIR / "models.yml" +OMP_CONFIG_PATH = OMP_AGENT_DIR / "config.yml" +OMP_MCP_PATH = OMP_AGENT_DIR / "mcp.json" +OMP_MODELS_BACKUP_PATH = APP_DIR / "omp-models.backup.yml" +OMP_CONFIG_BACKUP_PATH = APP_DIR / "omp-config.backup.yml" +OMP_MCP_BACKUP_PATH = APP_DIR / "omp-mcp.backup.json" + +SPEC: ToolSpec = { + "binary": "omp", + "package": "@oh-my-pi/pi-coding-agent", + "display": "Oh My Pi", + "config_path": OMP_MODELS_PATH, + "backup_path": OMP_MODELS_BACKUP_PATH, +} + +PROVIDER_NAMES = ( + "databricks-claude", + "databricks-openai", + "databricks-gemini", +) + +PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES] + +# Old provider names earlier ucode versions wrote; cleaned up on each write so +# users don't end up with stale entries pointing at routes that 400. +LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-oss") + + +def _resolve_model_selector( + model: str, + claude_models: dict[str, str], + codex_models: list[str], + gemini_models: list[str], +) -> str: + """Return an omp model selector in `/` form when possible.""" + for name in PROVIDER_NAMES: + if model.startswith(f"{name}/"): + return model + if model in claude_models.values(): + return f"databricks-claude/{model}" + if model in codex_models: + return f"databricks-openai/{model}" + if model in gemini_models: + return f"databricks-gemini/{model}" + return model + + +def render_overlay( + model: str, + token: str, + omp_base_urls: dict[str, str], + claude_models: dict[str, str], + codex_models: list[str], + gemini_models: list[str], +) -> tuple[dict, list[list[str]]]: + """Return (overlay, managed_key_paths) for omp's private agent config. + + The overlay carries only ``providers``: omp's ``models.yml`` schema + rejects unknown root keys, so the default-model selector is pinned in + ``config.yml`` by ``_write_default_model`` instead. + """ + providers: dict = {} + keys: list[list[str]] = [] + # omp expands header values that match an env var name. Our UA contains + # `/` and a space so it can never collide — safe to pass as a literal. + ua_headers = {"User-Agent": f"ucode/{ucode_version()} omp/{agent_version('omp')}"} + + claude_ids = sorted(set(claude_models.values())) + if claude_ids: + providers["databricks-claude"] = { + "baseUrl": omp_base_urls["claude"], + "api": "anthropic-messages", + "apiKey": token, + "authHeader": True, + # Gateway's Anthropic translator rejects per-tool + # `eager_input_streaming` on the streaming + tools path. omp sends + # the legacy beta header instead when this is false. + "compat": {"supportsEagerToolInputStreaming": False}, + "headers": ua_headers, + "models": [{"id": m} for m in claude_ids], + } + keys.append(["providers", "databricks-claude"]) + if codex_models: + providers["databricks-openai"] = { + "baseUrl": omp_base_urls["openai"], + "api": "openai-responses", + "apiKey": token, + "authHeader": True, + "headers": ua_headers, + "models": [{"id": m} for m in codex_models], + } + keys.append(["providers", "databricks-openai"]) + if gemini_models: + providers["databricks-gemini"] = { + "baseUrl": omp_base_urls["gemini"], + "api": "google-generative-ai", + "apiKey": token, + "authHeader": True, + "headers": ua_headers, + "models": [{"id": m} for m in gemini_models], + } + keys.append(["providers", "databricks-gemini"]) + overlay: dict = {} + if providers: + overlay["providers"] = providers + return overlay, keys + + +def write_tool_config( + state: dict, + model: str, + token: str | None = None, + *, + force_refresh: bool = False, +) -> tuple[dict, str]: + backup_existing_file(OMP_MODELS_PATH, OMP_MODELS_BACKUP_PATH) + if token is None: + token = get_databricks_token( + state["workspace"], state.get("profile"), force_refresh=force_refresh + ) + omp_base_urls = state.get("base_urls", {}).get("omp") or build_pi_base_urls(state["workspace"]) + claude_models = state.get("claude_models") or {} + codex_models = state.get("codex_models") or [] + gemini_models = state.get("gemini_models") or [] + overlay, managed_keys = render_overlay( + model, + token, + omp_base_urls, + claude_models, + codex_models, + gemini_models, + ) + existing = read_yaml_safe(OMP_MODELS_PATH) + providers = existing.get("providers") + if isinstance(providers, dict): + for stale in (*PROVIDER_NAMES, *LEGACY_PROVIDER_NAMES): + providers.pop(stale, None) + merged = deep_merge_dict(existing, overlay) + write_yaml_file(OMP_MODELS_PATH, merged) + _write_default_model(_resolve_model_selector(model, claude_models, codex_models, gemini_models)) + state = mark_tool_managed(state, "omp", managed_keys) + save_state(state) + return state, token + + +def _write_default_model(model_selector: str) -> None: + # Pin modelRoles.default in config.yml so omp starts on the Databricks + # model rather than falling through to the first available model (e.g. an + # env-key-backed provider) in its startup resolution order. + if "/" not in model_selector: + return + backup_existing_file(OMP_CONFIG_PATH, OMP_CONFIG_BACKUP_PATH) + existing = read_yaml_safe(OMP_CONFIG_PATH) + merged = deep_merge_dict(existing, {"modelRoles": {"default": model_selector}}) + write_yaml_file(OMP_CONFIG_PATH, merged) + + +def default_model(state: dict) -> str | None: + """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini.""" + claude_models = state.get("claude_models") or {} + for family in ("opus", "sonnet", "haiku"): + if claude_models.get(family): + return claude_models[family] + codex_models = state.get("codex_models") or [] + if codex_models: + return codex_models[0] + gemini_models = state.get("gemini_models") or [] + return gemini_models[0] if gemini_models else None + + +def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: + model = default_model(state) + if not model: + raise RuntimeError("No Oh My Pi model is available on this workspace.") + _, token = write_tool_config(state, model, force_refresh=force_refresh) + return token + + +def _refresh_forever(state: dict, stop_event: threading.Event) -> None: + while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS): + try: + _refresh_token_once(state, force_refresh=True) + except RuntimeError: + continue + + +def build_runtime_env() -> dict[str, str]: + # omp reads no token from the environment (auth is the baked models.yml + # apiKey); only redirect its agent dir into the ucode-private home. + env = os.environ.copy() + env["PI_CODING_AGENT_DIR"] = str(OMP_AGENT_DIR) + return env + + +def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: + _refresh_token_once(state) + env = build_runtime_env() + + stop_event = threading.Event() + refresher = threading.Thread( + target=_refresh_forever, + args=(state, stop_event), + daemon=True, + ) + refresher.start() + + proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() + finally: + stop_event.set() + refresher.join(timeout=1) + + raise SystemExit(returncode) + + +def validate_cmd(binary: str) -> list[str]: + return [binary, "--print", "say hi in 5 words or less"] + + +def validate_env(state: dict) -> dict[str, str]: + workspace = state.get("workspace") + if not workspace: + raise RuntimeError("No workspace configured.") + # Fetch a token to fail fast on bad auth; omp itself reads the baked file. + get_databricks_token(workspace, state.get("profile")) + return build_runtime_env() + + +def build_mcp_server_entry(argv: list[str]) -> dict: + # omp's stdioServer schema allows only command/args/env/cwd (stdio is the + # default transport when `command` is present without `url`), so the entry + # is exactly this — no `type`/`tools` keys like other clients use. + return { + "command": argv[0], + "args": list(argv[1:]), + } + + +def write_mcp_server_config(name: str, argv: list[str]) -> bool: + backup_existing_file(OMP_MCP_PATH, OMP_MCP_BACKUP_PATH) + existing = read_json_safe(OMP_MCP_PATH) + mcp_servers = existing.get("mcpServers") + if not isinstance(mcp_servers, dict): + mcp_servers = {} + removed = name in mcp_servers + mcp_servers[name] = build_mcp_server_entry(argv) + existing["mcpServers"] = mcp_servers + write_json_file(OMP_MCP_PATH, existing) + return removed + + +def remove_mcp_server_config(name: str) -> bool: + existing = read_json_safe(OMP_MCP_PATH) + mcp_servers = existing.get("mcpServers") + if not isinstance(mcp_servers, dict) or name not in mcp_servers: + return False + mcp_servers.pop(name) + existing["mcpServers"] = mcp_servers + write_json_file(OMP_MCP_PATH, existing) + return True diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 9857d78a..c2c1092e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -43,6 +43,7 @@ ) from ucode.agents.args import has_explicit_model_arg from ucode.agents.codex import revert_legacy_shared_config +from ucode.agents.omp import OMP_CONFIG_BACKUP_PATH, OMP_CONFIG_PATH from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import is_dry_run, restore_file, set_dry_run from ucode.databricks import ( @@ -150,9 +151,9 @@ from ucode.usage import usage as usage_report _DISCOVERY_CONSUMERS: dict[str, tuple[str, ...]] = { - "claude": ("claude", "opencode", "copilot", "pi"), - "codex": ("codex", "copilot", "pi"), - "gemini": ("gemini", "opencode", "pi"), + "claude": ("claude", "opencode", "copilot", "pi", "omp"), + "codex": ("codex", "copilot", "pi", "omp"), + "gemini": ("gemini", "opencode", "pi", "omp"), "oss": ("opencode",), } @@ -536,10 +537,19 @@ def configure_shared_state( print_warning(f"Model service: {model_service_probe.detail}") want_claude = ( - fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools + fetch_all + or "claude" in tools + or "opencode" in tools + or "copilot" in tools + or "pi" in tools + or "omp" in tools + ) + want_gemini = ( + fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools or "omp" in tools + ) + want_codex = ( + fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools or "omp" in tools ) - want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools - want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools # Codex smart routing can select OSS models such as GLM, so a Codex-only # configure must persist that discovered family too. want_oss = fetch_all or "opencode" in tools or "codex" in tools @@ -929,7 +939,8 @@ def status() -> int: and server.get("name") and server.get("kind") != SKILLS_MCP_KIND ] - print_kv("MCP list command", str(MCP_CLIENTS[tool]["list_command"])) + if MCP_CLIENTS[tool].get("list_command"): + print_kv("MCP list command", str(MCP_CLIENTS[tool]["list_command"])) print_kv( "MCP servers", ", ".join(tool_mcp_servers) if tool_mcp_servers else "none saved by ug", @@ -1012,6 +1023,9 @@ def revert() -> int: pi_settings_restored = restore_file( PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, bool(managed_configs.get("pi")) ) + omp_config_restored = restore_file( + OMP_CONFIG_PATH, OMP_CONFIG_BACKUP_PATH, bool(managed_configs.get("omp")) + ) # Older Codex (< 0.134.0) had ucode edit the shared ~/.codex/config.toml in # place; restoring the per-profile file above does not undo that. legacy_codex_stripped = revert_legacy_shared_config() @@ -1026,6 +1040,7 @@ def revert() -> int: print_kv("Claude Code OS-managed settings", claude_managed_result) print_kv("Codex OS-managed settings", codex_managed_result) print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") + print_kv("Oh My Pi config", "restored" if omp_config_restored else "unchanged") for client, spec in MCP_CLIENTS.items(): print_kv( f"{spec['display']} MCP config", @@ -2097,7 +2112,7 @@ def _launch_tool( f"{TOOL_SPECS[tool]['display']} may require one-time hook review. Open " "`/hooks` and trust the ug routing hooks if prompted." ) - if tool in ("gemini", "opencode", "copilot", "pi"): + if tool in ("gemini", "opencode", "copilot", "pi", "omp"): print_note( f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically " f"every 30 minutes while the session is running." @@ -2508,6 +2523,15 @@ def pi_cmd( _launch_tool("pi", ctx, skip_preflight=skip_preflight) +@app.command("omp", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) +def omp_cmd( + ctx: typer.Context, + skip_preflight: SkipPreflightOption = False, +) -> None: + """Launch Oh My Pi coding agent via Databricks.""" + _launch_tool("omp", ctx, skip_preflight=skip_preflight) + + @app.command("cursor", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def cursor_cmd(ctx: typer.Context) -> None: """Launch Cursor Agent. diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index f67f3f32..50629762 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -8,10 +8,34 @@ import tomlkit import tomlkit.exceptions +import yaml from ucode.ui import console +def read_yaml_safe(path: Path) -> dict: + # Same contract as read_json_safe: missing/unreadable/invalid → {}. + try: + if not path.exists(): + return {} + data = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError): + return {} + return data if isinstance(data, dict) else {} + + +def write_yaml_file(path: Path, payload: dict) -> None: + content = yaml.safe_dump(payload, sort_keys=False) + if _dry_run: + console.print(f"\n[bold]\\[dry run] {path}[/bold]\n{content}") + return + ensure_parent_dir(path) + try: + path.write_text(content, encoding="utf-8") + except OSError as exc: + raise RuntimeError(f"Failed to write config file: {path}") from exc + + class ToolSpec(TypedDict): binary: str package: str diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 49dba37f..994e0fd4 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3432,5 +3432,6 @@ def build_shared_base_urls(workspace: str) -> dict[str, str | dict[str, str]]: "opencode": build_opencode_base_urls(workspace), "copilot": build_copilot_base_url(workspace), "pi": build_pi_base_urls(workspace), + "omp": build_pi_base_urls(workspace), } return urls diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 59b4fb8a..6e48c55f 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -26,7 +26,7 @@ from questionary.question import Question from questionary.styles import merge_styles_default -from ucode.agents import copilot, cursor, gemini, opencode +from ucode.agents import copilot, cursor, gemini, omp, opencode from ucode.config_io import restore_file from ucode.databricks import ( apply_pat_environment, @@ -98,6 +98,12 @@ class _Back: "display": "Cursor", "list_command": "cursor-agent mcp list", }, + "omp": { + # No list_command: omp exposes no `mcp list` subcommand; status + # output skips the list row when the key is absent. + "binary": "omp", + "display": "Oh My Pi", + }, } SKILLS_MCP_KIND = "skills" SKILLS_MCP_SERVER_NAME = "databricks-skill-registry" @@ -330,6 +336,9 @@ def configure_client_mcp_server( if client == "cursor": removed = cursor.write_mcp_server_config(name, argv) return [MCP_USER_SCOPE] if removed else [] + if client == "omp": + removed = omp.write_mcp_server_config(name, argv) + return [MCP_USER_SCOPE] if removed else [] raise RuntimeError(f"Unsupported MCP client '{client}'.") @@ -346,6 +355,8 @@ def remove_client_mcp_server(client: str, name: str) -> list[str]: return [MCP_USER_SCOPE] if copilot.remove_mcp_server_config(name) else [] if client == "cursor": return [MCP_USER_SCOPE] if cursor.remove_mcp_server_config(name) else [] + if client == "omp": + return [MCP_USER_SCOPE] if omp.remove_mcp_server_config(name) else [] raise RuntimeError(f"Unsupported MCP client '{client}'.") @@ -376,6 +387,13 @@ def revert_mcp_configs(state: dict) -> dict[str, bool]: "copilot" in (server.get("clients") or []) for server in state.get("mcp_servers") or [] ), ) or results.get("copilot", False) + # omp stores MCP servers in a separate mcp.json like Copilot, so restore + # its original file after removing per-server entries above. + results["omp"] = restore_file( + omp.OMP_MCP_PATH, + omp.OMP_MCP_BACKUP_PATH, + any("omp" in (server.get("clients") or []) for server in state.get("mcp_servers") or []), + ) or results.get("omp", False) return results @@ -1687,7 +1705,7 @@ def setup_mcp_clients( if not installed_clients: raise RuntimeError( "No supported MCP clients are installed. Install Claude, Codex, Gemini, OpenCode, " - "or GitHub Copilot CLI." + "GitHub Copilot CLI, or Oh My Pi." ) clients = configured_mcp_clients(state, installed_clients) if agents is not None: @@ -1701,7 +1719,7 @@ def setup_mcp_clients( if not clients: raise RuntimeError( "No configured MCP-capable coding agents are installed. Run `ucode configure` " - "for Codex, Claude, Gemini, OpenCode, or GitHub Copilot CLI first." + "for Codex, Claude, Gemini, OpenCode, GitHub Copilot CLI, or Oh My Pi first." ) configured_tools = set(state.get("available_tools") or []) missing_clients = [ diff --git a/src/ucode/state.py b/src/ucode/state.py index 6344a7e5..a6eef4e5 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -182,6 +182,7 @@ def build_agent_state(state: dict) -> dict[str, dict]: ) codex_model = default_model_for_tool("codex", selection_state) pi_model = default_model_for_tool("pi", selection_state) + omp_model = default_model_for_tool("omp", selection_state) agents: dict[str, dict] = { "claude": { @@ -215,6 +216,12 @@ def build_agent_state(state: dict) -> dict[str, dict]: "auth_command": auth_command, "auth_refresh_interval_ms": AUTH_REFRESH_INTERVAL_MS, }, + "omp": { + "model": omp_model, + "base_urls": base_urls.get("omp") if isinstance(base_urls.get("omp"), dict) else {}, + "auth_command": auth_command, + "auth_refresh_interval_ms": AUTH_REFRESH_INTERVAL_MS, + }, } return { name: {key: value for key, value in config.items() if value is not None} diff --git a/tests/test_agent_omp.py b/tests/test_agent_omp.py new file mode 100644 index 00000000..3c379915 --- /dev/null +++ b/tests/test_agent_omp.py @@ -0,0 +1,689 @@ +"""Tests for agents/omp.py.""" + +from __future__ import annotations + +import re +from contextlib import nullcontext +from unittest.mock import patch + +import yaml + +from ucode.agents import LaunchOptions, omp + +WS = "https://example.databricks.com" + + +def _base_urls() -> dict[str, str]: + # Native API per family — see agents/omp.py docstring for path conventions. + return { + "claude": f"{WS}/ai-gateway/anthropic", + "openai": f"{WS}/ai-gateway/codex/v1", + "gemini": f"{WS}/ai-gateway/gemini/v1beta", + } + + +def _empty() -> dict: + """No-models input bundle for render_overlay.""" + return { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + } + + +def _overlay(model: str, token: str = "tok", **kwargs): + """Wrapper to call render_overlay with sensible defaults so tests stay terse.""" + bundle = {**_empty(), **kwargs} + return omp.render_overlay( + model, + token, + _base_urls(), + bundle["claude_models"], + bundle["codex_models"], + bundle["gemini_models"], + ) + + +class TestOmpSpec: + def test_binary(self): + assert omp.SPEC["binary"] == "omp" + + def test_package(self): + assert omp.SPEC["package"] == "@oh-my-pi/pi-coding-agent" + + def test_display(self): + assert omp.SPEC["display"] == "Oh My Pi" + + def test_config_path_is_models_yml_under_omp_agent_dir(self): + assert omp.SPEC["config_path"].name == "models.yml" + assert omp.SPEC["config_path"].parent.name == "agent" + assert omp.OMP_UCODE_HOME in omp.SPEC["config_path"].parents + # Isolated from pi's home so the two agents never share config. + assert "omp-home" in omp.SPEC["config_path"].parts + assert "pi-home" not in omp.SPEC["config_path"].parts + + +class TestRenderOverlayProviders: + def test_no_providers_when_no_models(self): + overlay, _ = _overlay("foo") + assert "providers" not in overlay + + def test_claude_provider_uses_anthropic_messages(self): + overlay, _ = _overlay("claude-sonnet", claude_models={"sonnet": "claude-sonnet"}) + provider = overlay["providers"]["databricks-claude"] + assert provider["api"] == "anthropic-messages" + assert provider["baseUrl"] == f"{WS}/ai-gateway/anthropic" + + def test_openai_provider_uses_openai_responses(self): + overlay, _ = _overlay("gpt-5", codex_models=["gpt-5"]) + provider = overlay["providers"]["databricks-openai"] + assert provider["api"] == "openai-responses" + assert provider["baseUrl"] == f"{WS}/ai-gateway/codex/v1" + + def test_gemini_provider_uses_google_generative_ai(self): + overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) + provider = overlay["providers"]["databricks-gemini"] + assert provider["api"] == "google-generative-ai" + assert provider["baseUrl"] == f"{WS}/ai-gateway/gemini/v1beta" + + def test_all_three_providers_when_all_present(self): + overlay, _ = _overlay( + "claude-sonnet", + claude_models={"sonnet": "claude-sonnet"}, + codex_models=["gpt-5"], + gemini_models=["gemini-2"], + ) + assert set(overlay["providers"].keys()) == { + "databricks-claude", + "databricks-openai", + "databricks-gemini", + } + + def test_overlay_carries_no_model_key(self): + # omp's models.yml schema rejects unknown root keys; the default model + # is pinned in config.yml instead (see TestWriteDefaultModel). + overlay, _ = _overlay( + "claude-sonnet", + claude_models={"sonnet": "claude-sonnet"}, + ) + assert "model" not in overlay + + +class TestRenderOverlayUserAgent: + def test_user_agent_set_on_all_three_providers(self, monkeypatch): + monkeypatch.setattr(omp, "ucode_version", lambda: "0.1.0") + monkeypatch.setattr(omp, "agent_version", lambda binary: "18.1.13") + overlay, _ = _overlay( + "claude-sonnet", + claude_models={"sonnet": "claude-sonnet"}, + codex_models=["gpt-5"], + gemini_models=["gemini-2"], + ) + expected = "ucode/0.1.0 omp/18.1.13" + for name in ("databricks-claude", "databricks-openai", "databricks-gemini"): + assert overlay["providers"][name]["headers"]["User-Agent"] == expected + + +class TestRenderOverlayCompatFlags: + def test_claude_disables_eager_tool_input_streaming(self): + # Gateway's Anthropic translator rejects per-tool + # `eager_input_streaming`; this flag makes omp send the legacy beta + # header instead. + overlay, _ = _overlay("claude-sonnet", claude_models={"sonnet": "claude-sonnet"}) + compat = overlay["providers"]["databricks-claude"]["compat"] + assert compat["supportsEagerToolInputStreaming"] is False + + def test_openai_and_gemini_have_no_compat_flags(self): + # Their gateway routes accept omp's request shape as-is. + overlay, _ = _overlay( + "gpt-5", + codex_models=["gpt-5"], + gemini_models=["gemini-2"], + ) + assert "compat" not in overlay["providers"]["databricks-openai"] + assert "compat" not in overlay["providers"]["databricks-gemini"] + + +class TestRenderOverlayAuthAndModels: + def test_token_in_api_key(self): + overlay, _ = _overlay( + "claude-sonnet", token="mytoken", claude_models={"sonnet": "claude-sonnet"} + ) + assert overlay["providers"]["databricks-claude"]["apiKey"] == "mytoken" + + def test_auth_header_flag_set_on_all_providers(self): + overlay, _ = _overlay( + "claude-sonnet", + claude_models={"sonnet": "claude-sonnet"}, + codex_models=["gpt-5"], + gemini_models=["gemini-2"], + ) + for name in ("databricks-claude", "databricks-openai", "databricks-gemini"): + assert overlay["providers"][name]["authHeader"] is True + + def test_claude_models_listed(self): + claude_models = {"opus": "claude-opus", "sonnet": "claude-sonnet"} + overlay, _ = _overlay("claude-sonnet", claude_models=claude_models) + ids = {m["id"] for m in overlay["providers"]["databricks-claude"]["models"]} + assert ids == {"claude-opus", "claude-sonnet"} + + def test_openai_models_listed(self): + overlay, _ = _overlay("gpt-5", codex_models=["gpt-5", "gpt-5-mini"]) + ids = {m["id"] for m in overlay["providers"]["databricks-openai"]["models"]} + assert ids == {"gpt-5", "gpt-5-mini"} + + def test_gemini_models_listed(self): + overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2", "gemini-2-pro"]) + ids = {m["id"] for m in overlay["providers"]["databricks-gemini"]["models"]} + assert ids == {"gemini-2", "gemini-2-pro"} + + +class TestRenderOverlayManagedKeys: + def test_managed_keys_exclude_model(self): + _, keys = _overlay("foo") + assert ["model"] not in keys + + def test_managed_keys_include_each_provider_present(self): + _, keys = _overlay( + "claude-sonnet", + claude_models={"sonnet": "claude-sonnet"}, + codex_models=["gpt-5"], + gemini_models=["gemini-2"], + ) + for name in ("databricks-claude", "databricks-openai", "databricks-gemini"): + assert ["providers", name] in keys + + +class TestResolveModelSelector: + def test_prefixes_claude_model(self): + selector = omp._resolve_model_selector("claude-sonnet", {"sonnet": "claude-sonnet"}, [], []) + assert selector == "databricks-claude/claude-sonnet" + + def test_prefixes_openai_model(self): + selector = omp._resolve_model_selector("gpt-5", {}, ["gpt-5"], []) + assert selector == "databricks-openai/gpt-5" + + def test_prefixes_gemini_model(self): + selector = omp._resolve_model_selector("gemini-2", {}, [], ["gemini-2"]) + assert selector == "databricks-gemini/gemini-2" + + def test_preserves_already_prefixed_model(self): + selector = omp._resolve_model_selector( + "databricks-claude/claude-sonnet", {"sonnet": "claude-sonnet"}, [], [] + ) + assert selector == "databricks-claude/claude-sonnet" + + def test_unknown_model_passes_through_unprefixed(self): + # Lets a user override to whatever omp accepts even if we + # didn't classify it. + assert omp._resolve_model_selector("custom/whatever", {}, [], []) == "custom/whatever" + + +class TestOmpDefaultModel: + def test_prefers_claude_opus(self): + state = {"claude_models": {"opus": "o4", "sonnet": "s4", "haiku": "h4"}} + assert omp.default_model(state) == "o4" + + def test_falls_back_to_sonnet(self): + state = {"claude_models": {"sonnet": "s4"}} + assert omp.default_model(state) == "s4" + + def test_falls_back_to_haiku(self): + state = {"claude_models": {"haiku": "h4"}} + assert omp.default_model(state) == "h4" + + def test_falls_back_to_codex(self): + state = {"claude_models": {}, "codex_models": ["gpt-5"]} + assert omp.default_model(state) == "gpt-5" + + def test_falls_back_to_gemini(self): + state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} + assert omp.default_model(state) == "gemini-2" + + def test_returns_none_when_empty(self): + assert omp.default_model({}) is None + assert ( + omp.default_model({"claude_models": {}, "codex_models": [], "gemini_models": []}) + is None + ) + + +class TestBuildRuntimeEnv: + def test_sets_private_agent_dir_without_replacing_home(self, monkeypatch): + monkeypatch.setenv("HOME", "/real-user-home") + + env = omp.build_runtime_env() + + assert env["PI_CODING_AGENT_DIR"] == str(omp.OMP_AGENT_DIR) + assert env["HOME"] == "/real-user-home" + + def test_sets_no_token_env(self, monkeypatch): + # omp never reads a token from the environment (no OAUTH_TOKEN-style + # hook anywhere in its packages); auth is the baked models.yml apiKey. + # ucode must not add one of its own — whatever the parent shell has + # is passed through untouched and ignored by omp. + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + env = omp.build_runtime_env() + assert "OAUTH_TOKEN" not in env + + +class TestOmpValidateCmd: + def test_starts_with_binary(self): + cmd = omp.validate_cmd("omp") + assert cmd[0] == "omp" + + def test_uses_print_flag(self): + # `--print` puts omp in non-interactive mode; without it the TUI hangs on stdin. + cmd = omp.validate_cmd("omp") + assert "--print" in cmd + + def test_has_prompt(self): + cmd = omp.validate_cmd("omp") + assert len(cmd) > 2 + + +class TestOmpValidateEnv: + def test_requires_workspace(self): + try: + omp.validate_env({}) + except RuntimeError as exc: + assert "workspace" in str(exc).lower() + else: # pragma: no cover - validate_env must raise + raise AssertionError("validate_env did not raise without a workspace") + + def test_fetches_token_to_fail_fast(self, monkeypatch): + monkeypatch.setattr(omp, "get_databricks_token", lambda workspace, profile: "tok") + env = omp.validate_env({"workspace": WS}) + assert env["PI_CODING_AGENT_DIR"] == str(omp.OMP_AGENT_DIR) + + +class TestWriteToolConfig: + def _setup(self, tmp_path, monkeypatch): + import ucode.agents.omp as omp_mod + import ucode.config_io as config_io_mod + + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + models_file = tmp_path / "models.yml" + models_backup = tmp_path / "models.backup.yml" + config_file = tmp_path / "config.yml" + config_backup = tmp_path / "config.backup.yml" + mcp_file = tmp_path / "mcp.json" + mcp_backup = tmp_path / "mcp.backup.json" + monkeypatch.setattr(omp_mod, "OMP_MODELS_PATH", models_file) + monkeypatch.setattr(omp_mod, "OMP_MODELS_BACKUP_PATH", models_backup) + monkeypatch.setattr(omp_mod, "OMP_CONFIG_PATH", config_file) + monkeypatch.setattr(omp_mod, "OMP_CONFIG_BACKUP_PATH", config_backup) + monkeypatch.setattr(omp_mod, "OMP_MCP_PATH", mcp_file) + monkeypatch.setattr(omp_mod, "OMP_MCP_BACKUP_PATH", mcp_backup) + return omp_mod, models_file, config_file, config_backup + + def _state(self, **overrides) -> dict: + state = { + "workspace": WS, + "base_urls": {"pi": _base_urls()}, + "claude_models": {"sonnet": "claude-sonnet"}, + "codex_models": [], + "gemini_models": [], + "managed_configs": {}, + } + state.update(overrides) + return state + + def _read_yaml(self, path): + return yaml.safe_load(path.read_text(encoding="utf-8")) + + def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatch): + omp_mod, models_file, _, _ = self._setup(tmp_path, monkeypatch) + + stale = { + "providers": { + "databricks-claude": {"old": True}, + "databricks-openai": {"old": True}, + "databricks-gemini": {"old": True}, + "user-provider": {"keep": True}, + } + } + models_file.write_text(yaml.safe_dump(stale), encoding="utf-8") + + with ( + patch("ucode.agents.omp.get_databricks_token", return_value="tok"), + patch("ucode.agents.omp.save_state"), + ): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + written = self._read_yaml(models_file) + providers = written.get("providers", {}) + assert providers.get("databricks-claude") != {"old": True} + assert "old" not in providers.get("databricks-claude", {}) + assert providers.get("user-provider") == {"keep": True} + + def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): + """Earlier ucode versions wrote `databricks-anthropic`, `databricks-codex`, + and `databricks-oss` providers. They must be stripped on the next write + so users don't end up with stale entries pointing at routes that 400.""" + omp_mod, models_file, _, _ = self._setup(tmp_path, monkeypatch) + + models_file.write_text( + yaml.safe_dump( + { + "providers": { + "databricks-anthropic": {"api": "anthropic-messages"}, + "databricks-codex": {"api": "openai-responses"}, + "databricks-oss": {"api": "openai-completions"}, + } + } + ), + encoding="utf-8", + ) + + with ( + patch("ucode.agents.omp.get_databricks_token", return_value="tok"), + patch("ucode.agents.omp.save_state"), + ): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + written_providers = self._read_yaml(models_file).get("providers", {}) + for legacy in ("databricks-anthropic", "databricks-codex", "databricks-oss"): + assert legacy not in written_providers + assert "databricks-claude" in written_providers + + def test_config_written_with_token_and_no_model_key(self, tmp_path, monkeypatch): + omp_mod, models_file, _, _ = self._setup(tmp_path, monkeypatch) + + with ( + patch("ucode.agents.omp.get_databricks_token", return_value="tok"), + patch("ucode.agents.omp.save_state"), + ): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + written = self._read_yaml(models_file) + assert written["providers"]["databricks-claude"]["apiKey"] == "tok" + # omp's models.yml schema rejects unknown root keys. + assert "model" not in written + + def test_config_pins_default_role_in_config_yml(self, tmp_path, monkeypatch): + # Without this, omp's startup resolution falls through to the first + # available model when an unrelated env var makes a built-in provider + # look auth-configured. Pinning the default role keeps omp on ours. + omp_mod, _, config_file, _ = self._setup(tmp_path, monkeypatch) + + with ( + patch("ucode.agents.omp.get_databricks_token", return_value="tok"), + patch("ucode.agents.omp.save_state"), + ): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + config = self._read_yaml(config_file) + assert config["modelRoles"]["default"] == "databricks-claude/claude-sonnet" + + def test_pre_existing_config_is_backed_up_before_first_write(self, tmp_path, monkeypatch): + omp_mod, _, config_file, config_backup = self._setup(tmp_path, monkeypatch) + + original = "modelRoles:\n default: some-other-provider/model\ntheme: titanium\n" + config_file.parent.mkdir(parents=True, exist_ok=True) + config_file.write_text(original, encoding="utf-8") + + with ( + patch("ucode.agents.omp.get_databricks_token", return_value="tok"), + patch("ucode.agents.omp.save_state"), + ): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + assert config_backup.read_text(encoding="utf-8") == original + # The on-disk config still gets the ucode pin applied via deep_merge. + merged = self._read_yaml(config_file) + assert merged["modelRoles"]["default"] == "databricks-claude/claude-sonnet" + assert merged["theme"] == "titanium" + + +class TestValidateAllToolsOmpRollback: + def test_failed_omp_validation_rolls_back_config(self, tmp_path, monkeypatch): + import ucode.agents as agents_mod + import ucode.agents.omp as omp_mod + + config_file = tmp_path / "config.yml" + config_file.write_text("modelRoles:\n default: x/y\n", encoding="utf-8") + monkeypatch.setattr(omp_mod, "OMP_CONFIG_PATH", config_file) + monkeypatch.setattr(omp_mod, "OMP_CONFIG_BACKUP_PATH", tmp_path / "config.backup.yml") + # Keep the generic models.yml rollback off the user's real config dir. + monkeypatch.setitem(agents_mod.TOOL_SPECS["omp"], "config_path", tmp_path / "models.yml") + monkeypatch.setitem( + agents_mod.TOOL_SPECS["omp"], "backup_path", tmp_path / "models.backup.yml" + ) + monkeypatch.setattr(agents_mod, "validate_tool", lambda tool: (False, "boom")) + monkeypatch.setattr(agents_mod, "save_state", lambda s: None) + monkeypatch.setattr(agents_mod, "spinner", lambda *_a, **_kw: nullcontext()) + + agents_mod.validate_all_tools( + {"available_tools": ["omp"], "managed_configs": {"omp": True}} + ) + + assert not config_file.exists() + + +class TestOmpMcpServerConfig: + def test_entry_is_exactly_command_and_args(self): + assert omp.build_mcp_server_entry(["ucode", "mcp-proxy", "--url", "https://x"]) == { + "command": "ucode", + "args": ["mcp-proxy", "--url", "https://x"], + } + + def test_ucode_server_names_match_omp_schema(self, tmp_path, monkeypatch): + # omp's mcp.json restricts names to ^[a-zA-Z0-9_.-]{1,100}$. + import ucode.agents.omp as omp_mod + + pattern = re.compile(r"^[a-zA-Z0-9_.-]{1,100}$") + mcp_file = tmp_path / "mcp.json" + monkeypatch.setattr(omp_mod, "OMP_MCP_PATH", mcp_file) + monkeypatch.setattr(omp_mod, "OMP_MCP_BACKUP_PATH", tmp_path / "mcp.backup.json") + + for name in ("databricks-sql", "databricks-skill-registry", "github-mcp"): + assert pattern.match(name), name + assert omp_mod.write_mcp_server_config(name, ["ucode", "mcp-proxy"]) is False + assert omp_mod.remove_mcp_server_config(name) is True + + def test_write_and_remove_round_trip(self, tmp_path, monkeypatch): + import json + + import ucode.agents.omp as omp_mod + + mcp_file = tmp_path / "mcp.json" + monkeypatch.setattr(omp_mod, "OMP_MCP_PATH", mcp_file) + monkeypatch.setattr(omp_mod, "OMP_MCP_BACKUP_PATH", tmp_path / "mcp.backup.json") + + assert omp_mod.write_mcp_server_config("github-mcp", ["ucode", "mcp-proxy"]) is False + written = json.loads(mcp_file.read_text(encoding="utf-8")) + assert written["mcpServers"]["github-mcp"] == { + "command": "ucode", + "args": ["mcp-proxy"], + } + assert omp_mod.write_mcp_server_config("github-mcp", ["ucode", "mcp-proxy"]) is True + assert omp_mod.remove_mcp_server_config("github-mcp") is True + assert omp_mod.remove_mcp_server_config("github-mcp") is False + + +class TestWriteToolConfigBaseUrls: + def _setup(self, tmp_path, monkeypatch): + import ucode.agents.omp as omp_mod + import ucode.config_io as config_io_mod + + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + models_file = tmp_path / "models.yml" + monkeypatch.setattr(omp_mod, "OMP_MODELS_PATH", models_file) + monkeypatch.setattr(omp_mod, "OMP_MODELS_BACKUP_PATH", tmp_path / "m.bak") + monkeypatch.setattr(omp_mod, "OMP_CONFIG_PATH", tmp_path / "config.yml") + monkeypatch.setattr(omp_mod, "OMP_CONFIG_BACKUP_PATH", tmp_path / "c.bak") + return omp_mod, models_file + + def _state(self, **overrides) -> dict: + state = { + "workspace": WS, + "claude_models": {"sonnet": "claude-sonnet"}, + "codex_models": [], + "gemini_models": [], + "managed_configs": {}, + } + state.update(overrides) + return state + + def test_prefers_omp_base_urls_from_state(self, tmp_path, monkeypatch): + import yaml as yaml_lib + + omp_mod, models_file = self._setup(tmp_path, monkeypatch) + state_urls = { + "claude": "https://state.example/anthropic", + "openai": "https://state.example/codex", + "gemini": "https://state.example/gemini", + } + + def _fail(_workspace): + raise AssertionError("should use state base_urls, not rebuild them") + + monkeypatch.setattr(omp_mod, "build_pi_base_urls", _fail) + with patch("ucode.agents.omp.save_state"): + omp_mod.write_tool_config( + self._state(base_urls={"omp": state_urls}), "claude-sonnet", token="tok" + ) + + written = yaml_lib.safe_load(models_file.read_text(encoding="utf-8")) + assert written["providers"]["databricks-claude"]["baseUrl"] == state_urls["claude"] + + def test_falls_back_to_building_urls_without_state_key(self, tmp_path, monkeypatch): + import yaml as yaml_lib + + omp_mod, models_file = self._setup(tmp_path, monkeypatch) + monkeypatch.setattr(omp_mod, "build_pi_base_urls", lambda workspace: _base_urls()) + with patch("ucode.agents.omp.save_state"): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + written = yaml_lib.safe_load(models_file.read_text(encoding="utf-8")) + assert written["providers"]["databricks-claude"]["baseUrl"] == f"{WS}/ai-gateway/anthropic" + + +class TestWriteToolConfigEdgeCases: + def _setup(self, tmp_path, monkeypatch): + import ucode.agents.omp as omp_mod + import ucode.config_io as config_io_mod + + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + models_file = tmp_path / "models.yml" + config_file = tmp_path / "config.yml" + monkeypatch.setattr(omp_mod, "OMP_MODELS_PATH", models_file) + monkeypatch.setattr(omp_mod, "OMP_MODELS_BACKUP_PATH", tmp_path / "m.bak") + monkeypatch.setattr(omp_mod, "OMP_CONFIG_PATH", config_file) + monkeypatch.setattr(omp_mod, "OMP_CONFIG_BACKUP_PATH", tmp_path / "c.bak") + return omp_mod, models_file, config_file + + def _state(self, **overrides) -> dict: + state = { + "workspace": WS, + "base_urls": {"omp": _base_urls()}, + "claude_models": {"sonnet": "claude-sonnet"}, + "codex_models": [], + "gemini_models": [], + "managed_configs": {}, + } + state.update(overrides) + return state + + def test_non_dict_providers_section_is_replaced(self, tmp_path, monkeypatch): + import yaml as yaml_lib + + omp_mod, models_file, _ = self._setup(tmp_path, monkeypatch) + models_file.write_text("providers:\n- not-a-mapping\n", encoding="utf-8") + + with patch("ucode.agents.omp.save_state"): + omp_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + written = yaml_lib.safe_load(models_file.read_text(encoding="utf-8")) + assert isinstance(written["providers"], dict) + assert "databricks-claude" in written["providers"] + + def test_passthrough_model_pins_verbatim_selector(self, tmp_path, monkeypatch): + import yaml as yaml_lib + + omp_mod, _, config_file = self._setup(tmp_path, monkeypatch) + + with patch("ucode.agents.omp.save_state"): + omp_mod.write_tool_config(self._state(), "custom/whatever", token="tok") + + # Unclassified selectors pass through like pi's `model` override and + # pin verbatim, so a user override still wins the startup default. + config = yaml_lib.safe_load(config_file.read_text(encoding="utf-8")) + assert config["modelRoles"]["default"] == "custom/whatever" + + def test_bare_model_name_writes_no_default_role(self, tmp_path, monkeypatch): + omp_mod, _, config_file = self._setup(tmp_path, monkeypatch) + + with patch("ucode.agents.omp.save_state"): + omp_mod.write_tool_config(self._state(), "somename", token="tok") + + assert not config_file.exists() + + def test_ignores_pi_managed_keys(self): + # omp has no managed-config support: pi's admin keys must not leak + # into omp's model choice even when both tools share a workspace. + state = { + "pi_default_model": "admin-chosen", + "pi_models": ["admin-allowlisted"], + "claude_models": {"sonnet": "discovered"}, + } + assert omp.default_model(state) == "discovered" + + +class TestLaunchAndValidateEnvEdgeCases: + def test_launch_raises_before_spawning_without_models(self, monkeypatch): + import ucode.agents.omp as omp_mod + + def _no_spawn(*args, **kwargs): + raise AssertionError("must not spawn a process without a model") + + monkeypatch.setattr(omp_mod.subprocess, "Popen", _no_spawn) + try: + omp_mod.launch({"workspace": WS}, [], options=LaunchOptions()) + except RuntimeError as exc: + assert "Oh My Pi" in str(exc) + else: # pragma: no cover - launch must raise + raise AssertionError("launch did not raise without models") + + def test_validate_env_propagates_token_errors(self, monkeypatch): + def _boom(workspace, profile=None): + raise RuntimeError("bad auth") + + monkeypatch.setattr(omp, "get_databricks_token", _boom) + try: + omp.validate_env({"workspace": WS}) + except RuntimeError as exc: + assert "bad auth" in str(exc) + else: # pragma: no cover - validate_env must raise + raise AssertionError("validate_env swallowed the token error") + + +class TestOmpMcpEdgeCases: + def _setup(self, tmp_path, monkeypatch): + import ucode.agents.omp as omp_mod + + mcp_file = tmp_path / "mcp.json" + monkeypatch.setattr(omp_mod, "OMP_MCP_PATH", mcp_file) + monkeypatch.setattr(omp_mod, "OMP_MCP_BACKUP_PATH", tmp_path / "mcp.backup.json") + return omp_mod, mcp_file + + def test_write_resets_non_dict_server_map(self, tmp_path, monkeypatch): + import json + + omp_mod, mcp_file = self._setup(tmp_path, monkeypatch) + mcp_file.write_text('{"mcpServers": ["not-a-mapping"]}', encoding="utf-8") + + assert omp_mod.write_mcp_server_config("github-mcp", ["ucode"]) is False + + written = json.loads(mcp_file.read_text(encoding="utf-8")) + assert written["mcpServers"] == {"github-mcp": {"command": "ucode", "args": []}} + + def test_remove_missing_file_returns_false(self, tmp_path, monkeypatch): + omp_mod, _ = self._setup(tmp_path, monkeypatch) + assert omp_mod.remove_mcp_server_config("github-mcp") is False + + def test_remove_non_dict_server_map_returns_false(self, tmp_path, monkeypatch): + omp_mod, mcp_file = self._setup(tmp_path, monkeypatch) + mcp_file.write_text('{"mcpServers": ["not-a-mapping"]}', encoding="utf-8") + assert omp_mod.remove_mcp_server_config("github-mcp") is False diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 9ff7d5ee..3bc985ca 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -71,7 +71,7 @@ def test_passthrough_for_unrelated_error(self): class TestToolSpecs: def test_all_tools_present(self): - assert set(TOOL_SPECS) == {"codex", "claude", "gemini", "opencode", "copilot", "pi"} + assert set(TOOL_SPECS) == {"codex", "claude", "gemini", "opencode", "copilot", "pi", "omp"} def test_each_spec_has_required_keys(self): required = {"binary", "package", "display", "config_path", "backup_path"} @@ -123,9 +123,9 @@ def _capture(self, monkeypatch): def test_maps_supported_tools_and_drops_others(self, monkeypatch): captured = self._capture(monkeypatch) - # Gemini and Pi aren't supported by `databricks aitools`, so they drop. + # Gemini, Pi, and Oh My Pi aren't supported by `databricks aitools`, so they drop. install_databricks_ai_tools_for_agents( - ["claude", "codex", "gemini", "pi"], {"profile": "prof"} + ["claude", "codex", "gemini", "pi", "omp"], {"profile": "prof"} ) assert captured == {"agents": ["claude-code", "codex"], "profile": "prof"} @@ -192,6 +192,8 @@ class TestNormalizeTool: ("opencode", "opencode"), ("copilot", "copilot"), ("pi", "pi"), + ("omp", "omp"), + ("oh-my-pi", "omp"), ("CODEX", "codex"), (" Claude ", "claude"), ], @@ -247,6 +249,18 @@ def test_pi_available_with_gemini(self): def test_pi_unavailable_when_no_models(self): assert check_gateway_endpoint({}, "pi") is False + def test_omp_available_with_claude(self): + assert check_gateway_endpoint({"claude_models": {"sonnet": "s4"}}, "omp") is True + + def test_omp_available_with_codex(self): + assert check_gateway_endpoint({"codex_models": ["m"]}, "omp") is True + + def test_omp_available_with_gemini(self): + assert check_gateway_endpoint({"gemini_models": ["gemini-2"]}, "omp") is True + + def test_omp_unavailable_when_no_models(self): + assert check_gateway_endpoint({}, "omp") is False + class TestDefaultModelForTool: def test_codex_returns_none_without_a_configured_model(self): @@ -301,6 +315,21 @@ def test_pi_falls_back_to_gemini(self): def test_pi_returns_none_when_no_models(self): assert default_model_for_tool("pi", {}) is None + def test_omp_prefers_claude_opus(self): + state = {"claude_models": {"opus": "o4", "sonnet": "s4"}, "codex_models": ["c"]} + assert default_model_for_tool("omp", state) == "o4" + + def test_omp_falls_back_to_codex(self): + state = {"claude_models": {}, "codex_models": ["c1"]} + assert default_model_for_tool("omp", state) == "c1" + + def test_omp_falls_back_to_gemini(self): + state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} + assert default_model_for_tool("omp", state) == "gemini-2" + + def test_omp_returns_none_when_no_models(self): + assert default_model_for_tool("omp", {}) is None + class TestResolveLaunchModel: def test_codex_default_model_used_when_no_explicit(self): diff --git a/tests/test_cli.py b/tests/test_cli.py index 18637d60..ffd880cf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1774,6 +1774,7 @@ def test_skipped_when_already_configured(self): ("opencode", "Launching OpenCode with Unity Gateway"), ("copilot", "Launching GitHub Copilot CLI with Unity Gateway"), ("pi", "Launching Pi with Unity Gateway"), + ("omp", "Launching Oh My Pi with Unity Gateway"), ], ) def test_launch_title(tool, expected): @@ -3255,7 +3256,7 @@ class TestSkipPreflightFlag: """`--skip-preflight` on a launch command threads through _launch_tool to configure_shared_state as skip_preflight.""" - LAUNCH_TOOLS = ["codex", "claude", "gemini", "opencode", "copilot", "pi"] + LAUNCH_TOOLS = ["codex", "claude", "gemini", "opencode", "copilot", "pi", "omp"] @staticmethod def _patches(cfg): diff --git a/tests/test_config_io.py b/tests/test_config_io.py index f2af8559..a738591c 100644 --- a/tests/test_config_io.py +++ b/tests/test_config_io.py @@ -17,12 +17,14 @@ prune_key_paths, read_json_safe, read_toml_safe, + read_yaml_safe, restore_file, set_dry_run, write_dotenv, write_json_file, write_text_file, write_toml_file, + write_yaml_file, ) @@ -172,6 +174,17 @@ def test_write_json_file_dry_run_no_write(self, tmp_path): write_json_file(p, {"a": 1}) assert not p.exists() + def test_write_yaml_file(self, tmp_path): + p = tmp_path / "out.yml" + write_yaml_file(p, {"a": 1}) + assert "a: 1" in p.read_text() + + def test_write_yaml_file_dry_run_no_write(self, tmp_path): + set_dry_run(True) + p = tmp_path / "out.yml" + write_yaml_file(p, {"a": 1}) + assert not p.exists() + def test_write_toml_file(self, tmp_path): p = tmp_path / "out.toml" doc = tomlkit.document() @@ -220,6 +233,25 @@ def test_read_json_safe_non_dict(self, tmp_path): p.write_text("[1, 2, 3]", encoding="utf-8") assert read_json_safe(p) == {} + def test_read_yaml_safe_missing_file(self, tmp_path): + result = read_yaml_safe(tmp_path / "missing.yml") + assert result == {} + + def test_read_yaml_safe_valid(self, tmp_path): + p = tmp_path / "data.yml" + p.write_text("x: 1\n", encoding="utf-8") + assert read_yaml_safe(p) == {"x": 1} + + def test_read_yaml_safe_invalid_yaml(self, tmp_path): + p = tmp_path / "bad.yml" + p.write_text("not: [unclosed\n bad indent: }", encoding="utf-8") + assert read_yaml_safe(p) == {} + + def test_read_yaml_safe_non_dict(self, tmp_path): + p = tmp_path / "arr.yml" + p.write_text("- 1\n- 2\n", encoding="utf-8") + assert read_yaml_safe(p) == {} + def test_read_toml_safe_missing_file(self, tmp_path): doc = read_toml_safe(tmp_path / "missing.toml") assert dict(doc) == {} diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 20301da3..1a565ccb 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -137,6 +137,13 @@ def test_contains_all_tools(self): assert "claude" in urls assert "gemini" in urls assert "opencode" in urls + assert "copilot" in urls + assert "pi" in urls + assert "omp" in urls + + def test_omp_reuses_pi_gateway_paths(self): + urls = build_shared_base_urls(WS) + assert urls["omp"] == urls["pi"] def test_opencode_is_dict(self): urls = build_shared_base_urls(WS) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 040182a9..fe211e44 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -225,7 +225,7 @@ def test_claude_url_contains_workspace(self, e2e_workspace): def test_shared_base_urls_all_tools(self, e2e_workspace): urls = build_shared_base_urls(e2e_workspace) - for tool in ("codex", "claude", "gemini", "opencode", "copilot", "pi"): + for tool in ("codex", "claude", "gemini", "opencode", "copilot", "pi", "omp"): assert tool in urls diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 5e1829bc..bcec53c0 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -299,6 +299,42 @@ def test_skipped_when_binary_not_installed(self): assert "cursor" not in clients +class TestOmpMcpClient: + def test_omp_registered_without_list_command(self): + # omp exposes no `mcp list` subcommand; status output skips the + # list row when the key is absent. + assert "omp" in mcp.MCP_CLIENTS + assert mcp.MCP_CLIENTS["omp"]["binary"] == "omp" + assert "list_command" not in mcp.MCP_CLIENTS["omp"] + + def test_configure_dispatches_proxy_argv_to_omp_writer(self, monkeypatch): + calls: list[tuple[str, list[str]]] = [] + monkeypatch.setattr( + mcp.omp, + "write_mcp_server_config", + lambda name, argv: calls.append((name, argv)) or False, + ) + + removed_scopes = mcp.configure_client_mcp_server("omp", "github", GH_URL, WS, "p") + + assert removed_scopes == [] + assert calls == [("github", _proxy_argv())] + + def test_configure_reports_user_scope_on_replace(self, monkeypatch): + monkeypatch.setattr(mcp.omp, "write_mcp_server_config", lambda name, argv: True) + assert mcp.configure_client_mcp_server("omp", "github", GH_URL, WS, "p") == [ + mcp.MCP_USER_SCOPE + ] + + def test_remove_dispatches_to_omp_writer(self, monkeypatch): + calls: list[str] = [] + monkeypatch.setattr( + mcp.omp, "remove_mcp_server_config", lambda name: calls.append(name) or True + ) + assert mcp.remove_client_mcp_server("omp", "github-mcp") == [mcp.MCP_USER_SCOPE] + assert calls == ["github-mcp"] + + class TestConfigureClientMcpServer: def test_configures_copilot_with_proxy_argv(self, monkeypatch): calls: list[tuple[str, list[str]]] = [] @@ -2634,7 +2670,7 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey "mcp_servers": [ { "name": "github-mcp", - "clients": ["claude", "codex", "gemini", "opencode", "copilot"], + "clients": ["claude", "codex", "gemini", "opencode", "copilot", "omp"], } ] } @@ -2646,9 +2682,11 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey ("gemini", "github-mcp"), ("opencode", "github-mcp"), ("copilot", "github-mcp"), + ("omp", "github-mcp"), ] assert restored == [ - (mcp.copilot.COPILOT_MCP_CONFIG_PATH, mcp.copilot.COPILOT_MCP_BACKUP_PATH, True) + (mcp.copilot.COPILOT_MCP_CONFIG_PATH, mcp.copilot.COPILOT_MCP_BACKUP_PATH, True), + (mcp.omp.OMP_MCP_PATH, mcp.omp.OMP_MCP_BACKUP_PATH, True), ] assert result == { "claude": True, @@ -2656,6 +2694,7 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey "gemini": True, "opencode": True, "copilot": True, + "omp": True, } def test_removes_skills_registry_across_its_clients(self, monkeypatch): diff --git a/tests/test_state.py b/tests/test_state.py index 36c8ce4f..f685de08 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -36,6 +36,11 @@ "openai": f"{FAKE_WS}/ai-gateway/codex/v1", "gemini": f"{FAKE_WS}/ai-gateway/gemini/v1beta", }, + "omp": { + "claude": f"{FAKE_WS}/ai-gateway/anthropic", + "openai": f"{FAKE_WS}/ai-gateway/codex/v1", + "gemini": f"{FAKE_WS}/ai-gateway/gemini/v1beta", + }, } @@ -125,6 +130,7 @@ def test_persists_codex_launcher_default_in_agent_state(self): assert persisted["codex_models"][0] == "system.ai.gpt-5" assert "model" not in persisted["agents"]["codex"] assert persisted["agents"]["pi"]["model"] == "system.ai.gpt-5" + assert persisted["agents"]["omp"]["model"] == "system.ai.gpt-5" def test_save_respects_dry_run(self): import ucode.config_io as config_io_mod @@ -229,6 +235,8 @@ def test_populates_agent_state_when_workspace_present(self): assert codex_auth["args"][0] == "auth-token" assert result["agents"]["pi"]["model"] == "claude-opus" assert result["agents"]["pi"]["base_urls"] == FAKE_URLS["pi"] + assert result["agents"]["omp"]["model"] == "claude-opus" + assert result["agents"]["omp"]["base_urls"] == FAKE_URLS["omp"] def test_normalizes_managed_configs_dict_entry(self): state = {"managed_configs": {"claude": {"keys": [["env", "X"]]}}} @@ -263,7 +271,7 @@ def test_use_pat_state_builds_pat_auth_command(self): ) # --use-pat threads through to the `ucode auth-token --use-pat` helper, # which resolves the static PAT internally on every platform. - for agent in ("claude", "codex", "pi"): + for agent in ("claude", "codex", "pi", "omp"): assert "--use-pat" in result[agent]["auth_command"] assert "--profile DEFAULT" in result[agent]["auth_command"] diff --git a/uv.lock b/uv.lock index 9ce5b4a6..81c61d1d 100644 --- a/uv.lock +++ b/uv.lock @@ -3361,6 +3361,7 @@ dependencies = [ { name = "databricks-sql-connector" }, { name = "httpx" }, { name = "mcp" }, + { name = "pyyaml" }, { name = "questionary" }, { name = "tomlkit" }, { name = "typer" }, @@ -3385,6 +3386,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.27.1" }, { name = "mcp", specifier = ">=1.28.0" }, { name = "mlflow", extras = ["databricks"], marker = "extra == 'tracing'", specifier = ">=3.4" }, + { name = "pyyaml", specifier = ">=6.0.0" }, { name = "questionary", specifier = ">=2.0.0" }, { name = "tomlkit", specifier = ">=0.13.0" }, { name = "typer", specifier = ">=0.12.0" }, From 1c980f878fdfad7d9883f9cf58ca962ae302373b Mon Sep 17 00:00:00 2001 From: Zhu Lin Date: Thu, 10 Sep 2026 05:23:06 +0800 Subject: [PATCH 2/2] Cover omp in managed setup, model families, and URL builder guard Exclude omp from the managed-setup picker (server proto has no omp variant; serialize would silently drop it), list its model families mirroring pi, and give build_tool_base_url an omp guard. --- src/ucode/databricks.py | 2 ++ src/ucode/managed_setup.py | 3 +++ src/ucode/managed_wizard.py | 4 +++- tests/test_databricks.py | 4 ++++ tests/test_managed_setup.py | 7 ++++++- tests/test_managed_wizard.py | 17 +++++++++++++++++ 6 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 994e0fd4..22583251 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3386,6 +3386,8 @@ def build_tool_base_url(tool: str, workspace: str) -> str: ) if tool == "pi": raise RuntimeError("Pi has multiple base URLs — use build_pi_base_urls() instead.") + if tool == "omp": + raise RuntimeError("Oh My Pi has multiple base URLs — use build_pi_base_urls() instead.") raise RuntimeError(f"Unsupported tool '{tool}'.") diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index 3c52f909..5ba03e4f 100644 --- a/src/ucode/managed_setup.py +++ b/src/ucode/managed_setup.py @@ -62,6 +62,9 @@ "opencode": ("claude", "codex", "gemini", "oss"), "pi": ("claude", "codex", "gemini", "oss"), "copilot": ("claude", "codex", "gemini", "oss"), + # omp mirrors pi (the multi-provider harness); it is still excluded from the + # setup picker (see SETUP_EXCLUDED_AGENTS) since the server has no omp variant. + "omp": ("claude", "codex", "gemini", "oss"), } diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index fa7ca124..7dda4535 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -98,7 +98,9 @@ # Agents not offered in `ug setup`'s picker, even when the workspace serves their models. # `ug gemini` still works as a launch target; it's just not part of the managed config authored # here. Serialize/validate keep supporting it, so a `--from-file` manifest can still name it. -SETUP_EXCLUDED_AGENTS = frozenset({"gemini"}) +# `omp` has no managed-config support at all (the server proto has no omp agent variant), so it +# is excluded too — `ug omp` remains a regular launch target. +SETUP_EXCLUDED_AGENTS = frozenset({"gemini", "omp"}) def _tracing_table_from_state(state: dict) -> str | None: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 1a565ccb..e120c487 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -117,6 +117,10 @@ def test_opencode_raises(self): with pytest.raises(RuntimeError, match="multiple base URLs"): build_tool_base_url("opencode", WS) + def test_omp_raises(self): + with pytest.raises(RuntimeError, match="multiple base URLs"): + build_tool_base_url("omp", WS) + def test_unsupported_tool_raises(self): with pytest.raises(RuntimeError, match="Unsupported"): build_tool_base_url("unknown", WS) diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 140bc3d9..eff45c65 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -326,13 +326,18 @@ def test_codex_sees_gpt_and_oss(self): ] def test_multi_provider_agents_see_everything(self): - for tool in ("opencode", "pi", "copilot"): + for tool in ("opencode", "pi", "copilot", "omp"): options = model_options_for_agent(tool, STATE) assert "system.ai.claude-opus-4-8" in options, tool assert "system.ai.gpt-5-6" in options, tool assert "system.ai.gemini-3-flash" in options, tool assert "system.ai.kimi-k2-6" in options, tool + def test_omp_matches_pi_options(self): + # omp is the same multi-provider harness as pi; the setup picker + # excludes it, but any residual path must offer pi-identical options. + assert model_options_for_agent("omp", STATE) == model_options_for_agent("pi", STATE) + def test_empty_state_yields_no_options(self): assert model_options_for_agent("claude", {}) == [] diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index a7173414..da713079 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2991,3 +2991,20 @@ def test_section_command_interrupt_exits_130(self): def _out(result) -> str: """CliRunner output with stderr folded in, since print_err writes to a stderr console.""" return result.output + (result.stderr if result.stderr_bytes else "") + + +class TestSetupExcludedAgents: + def test_omp_is_excluded_from_the_setup_picker(self): + # The server proto has no omp agent variant, so the managed-setup + # picker must never offer it even when the workspace serves omp's + # models (the picker's `available` filter honors this set). + assert "omp" in wizard.SETUP_EXCLUDED_AGENTS + assert "gemini" in wizard.SETUP_EXCLUDED_AGENTS + + def test_omp_would_otherwise_pass_the_gateway_gate(self): + # Guard against a vacuous exclusion: with models discovered, omp + # clears the gateway check, so the exclusion set is what keeps it + # out of the picker. + from ucode.agents import check_gateway_endpoint + + assert check_gateway_endpoint(STATE, "omp") is True