From b411d19e76403c78ac75d4dc7c70d30e941da5fb Mon Sep 17 00:00:00 2001 From: natefleming Date: Mon, 31 Aug 2026 13:51:28 -0400 Subject: [PATCH 1/2] feat(agents): add Continue (continue.dev) as a ucode harness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Configure and launch Continue's `cn` CLI (and, via the shared ~/.continue/config.yaml, the VS Code/JetBrains extensions) against the Databricks AI Gateway, following the existing baked-token harness pattern. - New `agents/continue_dev.py` (module name avoids the `continue` keyword; tool name stays "continue"): OpenAI-compatible provider at the MLflow chat-completions gateway, chat/edit/apply roles, discovered Claude/GPT models, baked bearer + background refresh + OAUTH_TOKEN env — the same strategy as OpenCode/Copilot/Gemini/Pi. Writes modern config.yaml. - MCP parity: write/remove_mcp_server_config register the shared `ucode mcp-proxy` stdio bridge into config.yaml's mcpServers list, wired into mcp.py's client registry and register/remove dispatch (same pattern as Cursor/OpenCode/Copilot). - Registration across the usual surfaces: _MODULES/aliases/dispatch/ availability/discovery in agents/__init__.py, discovery consumers and want_claude/want_codex in cli.py, build_continue_base_url in databricks.py. - YAML I/O helpers (read_yaml_safe/write_yaml_file) in config_io.py; add pyyaml dependency. - Docs: README agent list, usage, --agents names, managed-files table, and MCP-capable tools list. New tests/test_agent_continue.py plus registry assertions in test_agents_init.py. Verified live on FEVM: configure discovers system.ai.claude-sonnet-5, writes config.yaml, validation round-trips through the gateway, and MCP register/ remove produce the documented schema without clobbering the models list. Co-authored-by: Isaac --- README.md | 10 +- pyproject.toml | 1 + src/ucode/agents/__init__.py | 14 +- src/ucode/agents/continue_dev.py | 250 +++++++++++++++++++++++++++++++ src/ucode/cli.py | 18 ++- src/ucode/config_io.py | 27 ++++ src/ucode/databricks.py | 8 + src/ucode/mcp.py | 12 +- tests/test_agent_continue.py | 246 ++++++++++++++++++++++++++++++ tests/test_agents_init.py | 36 ++++- uv.lock | 2 + 11 files changed, 612 insertions(+), 12 deletions(-) create mode 100644 src/ucode/agents/continue_dev.py create mode 100644 tests/test_agent_continue.py diff --git a/README.md b/README.md index 7030ef97..7e6efc18 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Unity AI Gateway Coding CLI (ucode) -`ucode` is a lightweight launcher for running Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, and Pi through Databricks. +`ucode` is a lightweight launcher for running Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, Pi, and Continue through Databricks. ## Requirements @@ -30,6 +30,7 @@ ucode gemini # Gemini CLI ucode opencode # OpenCode ucode copilot # GitHub Copilot CLI ucode pi # Pi +ucode continue # Continue (cn CLI + VS Code/JetBrains extensions) ucode cursor # Cursor Agent (MCP only — see below) ``` @@ -73,7 +74,9 @@ To configure specific tools without the picker, pass a comma-separated list: ucode configure --agents claude,codex ``` -Available agent names are `codex`, `claude`, `gemini`, `opencode`, `copilot`, and `pi`. `cursor` is also accepted (MCP-only — it registers Databricks MCP servers but configures no models). +Available agent names are `codex`, `claude`, `gemini`, `opencode`, `copilot`, `pi`, and `continue`. `cursor` is also accepted (MCP-only — it registers Databricks MCP servers but configures no models). + +`continue` writes `~/.continue/config.yaml`, which is shared by the `cn` CLI and the Continue VS Code/JetBrains extensions — so configuring it once wires up both. Its chat/edit/apply roles route to your Databricks Claude/GPT models through the Unity AI Gateway. Naming agents explicitly is treated as a request for all of them: if any one isn't available on the workspace, the run fails without configuring the others. Add `--skip-unavailable` to configure the available subset instead and skip the rest with a warning: @@ -111,7 +114,7 @@ ucode configure --profiles DEFAULT --agents claude,codex --use-pat --skip-valida ucode configure mcp ``` -Add Databricks MCP servers to installed MCP-capable tools: Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, and Cursor Agent. +Add Databricks MCP servers to installed MCP-capable tools: Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, Continue, and Cursor Agent. Options are shown in this order: - Discovered external MCP connections @@ -390,6 +393,7 @@ control the installation. | `~/.config/opencode/opencode.json` | OpenCode | | `~/.copilot/.env` | GitHub Copilot CLI | | `~/.pi/agent/models.json` | Pi | +| `~/.continue/config.yaml` | Continue (`cn` CLI + IDE extensions) | | `~/.cursor/mcp.json` | Cursor Agent (MCP servers only) | | `~/.ucode/managed-state.json` | The managed config — authored by `ucode setup` (admins) and refreshed from the workspace on launch | | `~/.ucode/managed-backups/` | Baseline backups for OS-managed files changed by ucode | diff --git a/pyproject.toml b/pyproject.toml index 52564f5c..b0e1ac97 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", "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 0cd4d2ce..0ee42a1a 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -39,7 +39,7 @@ spinner, ) -from . import claude, codex, copilot, gemini, opencode, pi +from . import claude, codex, continue_dev, copilot, gemini, opencode, pi from .args import explicit_model_arg_value as explicit_model_arg_value _MODULES = { @@ -49,6 +49,7 @@ "opencode": opencode, "copilot": copilot, "pi": pi, + "continue": continue_dev, } TOOL_SPECS: dict[str, ToolSpec] = {name: module.SPEC for name, module in _MODULES.items()} @@ -67,6 +68,9 @@ "opencode": "opencode", "copilot": "copilot", "pi": "pi", + "continue": "continue", + "continue-dev": "continue", + "cn": "continue", } DEFAULT_TOOL = "codex" @@ -103,7 +107,8 @@ 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, continue." ) return normalized @@ -350,6 +355,8 @@ def configure_tool( result = copilot.write_tool_config(state, model) elif tool == "pi": result = pi.write_tool_config(state, model) + elif tool == "continue": + result = continue_dev.write_tool_config(state, model) else: result = opencode.write_tool_config(state, model) # gemini/opencode/copilot/pi return (state, token); codex/claude return state @@ -380,6 +387,8 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: or bool(state.get("codex_models")) or bool(state.get("gemini_models")) ) + if tool == "continue": + return bool(state.get("claude_models")) or bool(state.get("codex_models")) return False @@ -390,6 +399,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: "gemini": ("gemini",), "copilot": ("claude", "codex"), "pi": ("claude", "codex", "gemini"), + "continue": ("claude", "codex"), } diff --git a/src/ucode/agents/continue_dev.py b/src/ucode/agents/continue_dev.py new file mode 100644 index 00000000..cfec0dfe --- /dev/null +++ b/src/ucode/agents/continue_dev.py @@ -0,0 +1,250 @@ +"""Continue agent: writes ~/.continue/config.yaml with a Databricks-backed model. + +Continue.dev's `cn` CLI and its VS Code/JetBrains extensions read the same +`~/.continue/config.yaml`, so the file ucode writes configures both. We point +Continue's OpenAI-compatible provider at the Databricks MLflow chat-completions +gateway (the same endpoint Copilot uses), which serves Claude and codex (gpt-5) +models behind one URL. `provider: openai` makes Continue append +`/chat/completions` to the configured `apiBase`. + +The gateway bearer token is baked into the config file (Continue has no +command-based auth refresh), so — like OpenCode/Copilot/Gemini — the token is +short-lived: `ucode continue` rewrites a fresh one on every launch and a +background thread refreshes it during the session. This is why Continue is not +in `GLOBAL_SETTINGS_AGENTS`: a bare `cn` launched outside ucode would eventually +see an expired token. +""" + +from __future__ import annotations + +import os +import signal +import subprocess +import threading +from pathlib import Path + +from ucode.config_io import ( + APP_DIR, + ToolSpec, + backup_existing_file, + read_yaml_safe, + write_yaml_file, +) +from ucode.databricks import ( + TOKEN_REFRESH_INTERVAL_SECONDS, + build_continue_base_url, + get_databricks_token, +) +from ucode.state import mark_tool_managed, save_state +from ucode.telemetry import agent_version, ucode_version + +CONTINUE_CONFIG_DIR = Path.home() / ".continue" +CONTINUE_CONFIG_PATH = CONTINUE_CONFIG_DIR / "config.yaml" +CONTINUE_BACKUP_PATH = APP_DIR / "continue-ucode-config.backup.yaml" + +# ucode's model entries in the shared `models:` list carry this name prefix so a +# rewrite can drop the stale ones without disturbing models the user added. +UCODE_MODEL_NAME_PREFIX = "Databricks (ucode)" + +SPEC: ToolSpec = { + "binary": "cn", + "package": "@continuedev/cli", + "display": "Continue", + "config_path": CONTINUE_CONFIG_PATH, + "backup_path": CONTINUE_BACKUP_PATH, +} + +# Informational: Continue's config is a list-shaped YAML document, so revert +# restores the backup (or deletes the ucode-created file) rather than pruning +# key paths. Recorded so `restore_file` sees the tool as managed. +MANAGED_KEYS: list[list[str]] = [["models"]] + + +def default_model(state: dict) -> str | None: + """Pick the best available Continue model. + + A managed config's ``continue_default_model`` wins outright. Otherwise prefer + Claude sonnet, then opus/haiku, then the first codex model — the same order + Copilot uses, since both draw from the shared chat-completions gateway. + """ + if isinstance(state.get("continue_default_model"), str): + return state.get("continue_default_model") + claude_models = state.get("claude_models") or {} + for family in ("sonnet", "opus", "haiku"): + if claude_models.get(family): + return claude_models[family] + codex_models = state.get("codex_models") or [] + if codex_models: + return codex_models[0] + return None + + +def _ucode_model_entry(model: str, token: str, workspace: str) -> dict: + """Build the ucode-managed `models:` entry for chat/edit/apply roles.""" + return { + "name": f"{UCODE_MODEL_NAME_PREFIX} {model}", + "provider": "openai", + "model": model, + "apiBase": build_continue_base_url(workspace), + "apiKey": token, + "roles": ["chat", "edit", "apply"], + "requestOptions": { + "headers": { + "User-Agent": f"ucode/{ucode_version()} continue/{agent_version('cn')}", + }, + }, + } + + +def _ensure_schema_header(doc: dict) -> None: + """Set Continue's required top-level keys, keeping the user's if present.""" + doc.setdefault("name", "ucode") + doc.setdefault("version", "0.0.1") + doc.setdefault("schema", "v1") + + +def render_config(model: str, token: str, workspace: str) -> dict: + """Return a complete Continue config document pinning the ucode model.""" + return { + "name": "ucode", + "version": "0.0.1", + "schema": "v1", + "models": [_ucode_model_entry(model, token, workspace)], + } + + +def write_tool_config( + state: dict, + model: str, + token: str | None = None, + *, + force_refresh: bool = False, +) -> tuple[dict, str]: + backup_existing_file(CONTINUE_CONFIG_PATH, CONTINUE_BACKUP_PATH) + if token is None: + token = get_databricks_token( + state["workspace"], state.get("profile"), force_refresh=force_refresh + ) + existing = read_yaml_safe(CONTINUE_CONFIG_PATH) + _ensure_schema_header(existing) + models = existing.get("models") + if not isinstance(models, list): + models = [] + # Drop any prior ucode entry so a rewrite replaces (not duplicates) it, while + # leaving the user's own models untouched. + models = [ + m + for m in models + if not (isinstance(m, dict) and str(m.get("name", "")).startswith(UCODE_MODEL_NAME_PREFIX)) + ] + models.append(_ucode_model_entry(model, token, state["workspace"])) + existing["models"] = models + write_yaml_file(CONTINUE_CONFIG_PATH, existing) + state = mark_tool_managed(state, "continue", MANAGED_KEYS) + save_state(state) + return state, token + + +def build_mcp_server_entry(name: str, argv: list[str]) -> dict: + # A local stdio MCP server: `command`/`args` run the `ucode mcp-proxy ...` + # bridge, which mints a fresh OAuth token per request — so Continue never + # speaks HTTP+bearer directly (same proxy pattern as Cursor/OpenCode). + # Continue's `mcpServers` is a list, so each entry carries its own `name`. + return { + "name": name, + "type": "stdio", + "command": argv[0], + "args": list(argv[1:]), + } + + +def write_mcp_server_config(name: str, argv: list[str]) -> bool: + backup_existing_file(CONTINUE_CONFIG_PATH, CONTINUE_BACKUP_PATH) + existing = read_yaml_safe(CONTINUE_CONFIG_PATH) + _ensure_schema_header(existing) + servers = existing.get("mcpServers") + if not isinstance(servers, list): + servers = [] + # `mcpServers` is a list keyed by `name`; drop any prior entry with this name + # so a re-register replaces it, leaving the user's own servers untouched. + removed = any(isinstance(s, dict) and s.get("name") == name for s in servers) + servers = [s for s in servers if not (isinstance(s, dict) and s.get("name") == name)] + servers.append(build_mcp_server_entry(name, argv)) + existing["mcpServers"] = servers + write_yaml_file(CONTINUE_CONFIG_PATH, existing) + return removed + + +def remove_mcp_server_config(name: str) -> bool: + existing = read_yaml_safe(CONTINUE_CONFIG_PATH) + servers = existing.get("mcpServers") + if not isinstance(servers, list): + return False + filtered = [s for s in servers if not (isinstance(s, dict) and s.get("name") == name)] + if len(filtered) == len(servers): + return False + existing["mcpServers"] = filtered + write_yaml_file(CONTINUE_CONFIG_PATH, existing) + return True + + +def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: + model = default_model(state) + if not model: + raise RuntimeError("No Continue 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(token: str) -> dict[str, str]: + env = os.environ.copy() + env["OAUTH_TOKEN"] = token + return env + + +def launch(state: dict, tool_args: list[str]) -> None: + """Launch `cn` with background token refresh (same pattern as OpenCode).""" + token = _refresh_token_once(state) + env = build_runtime_env(token) + + stop_event = threading.Event() + refresher = threading.Thread( + target=_refresh_forever, + args=(state, stop_event), + daemon=True, + ) + refresher.start() + + # `--config` forces ucode's config over any hub assistant the user last + # selected, so `ucode continue` always routes through the Databricks gateway. + argv = [SPEC["binary"], "--config", str(CONTINUE_CONFIG_PATH), *tool_args] + proc = subprocess.Popen(argv, 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]: + # `-p` runs headless (TTY-less); `--config` pins ucode's config for the probe. + return [ + binary, + "-p", + "say hi in 5 words or less", + "--config", + str(CONTINUE_CONFIG_PATH), + ] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a94eaa54..12d4c27f 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -142,8 +142,8 @@ from ucode.usage import usage as usage_report _DISCOVERY_CONSUMERS: dict[str, tuple[str, ...]] = { - "claude": ("claude", "opencode", "copilot", "pi"), - "codex": ("codex", "copilot", "pi"), + "claude": ("claude", "opencode", "copilot", "pi", "continue"), + "codex": ("codex", "copilot", "pi", "continue"), "gemini": ("gemini", "opencode", "pi"), "oss": ("opencode",), } @@ -638,10 +638,17 @@ def configure_shared_state( print_success("Unity AI Gateway detected") 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 "continue" 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 + want_codex = ( + fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools or "continue" 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 @@ -2565,7 +2572,8 @@ def configure( str | None, typer.Option( "--agent", - help="Configure only the named agent (e.g. claude, codex, gemini, opencode, copilot, pi).", + help="Configure only the named agent (e.g. claude, codex, gemini, opencode, copilot, " + "pi, continue).", ), ] = None, agents: Annotated[ diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index f67f3f32..764d8fb3 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -8,6 +8,7 @@ import tomlkit import tomlkit.exceptions +import yaml from ucode.ui import console @@ -177,6 +178,32 @@ def write_toml_file(path: Path, doc: tomlkit.TOMLDocument) -> None: raise RuntimeError(f"Failed to write config file: {path}") from exc +def read_yaml_safe(path: Path) -> dict: + # See read_json_safe: keep `path.exists()` inside the try so a PermissionError on a locked + # parent directory is treated as an empty document rather than propagating. + 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, doc: dict) -> None: + # `sort_keys=False` preserves the document's insertion order so the generated config reads + # top-down (name/version/schema before models) rather than alphabetized. + content = yaml.safe_dump(doc, sort_keys=False, default_flow_style=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 + + def parse_dotenv(path: Path) -> dict[str, str]: """Parse a simple KEY=VALUE / KEY="VALUE" .env file, preserving insertion order. diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 7b8e442f..89dd0e43 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3343,6 +3343,13 @@ def build_copilot_base_url(workspace: str) -> str: return f"{workspace}/ai-gateway/mlflow/v1" +def build_continue_base_url(workspace: str) -> str: + # Continue's `openai` provider appends `/chat/completions` to `apiBase`. The Databricks + # MLflow chat-completions gateway is OpenAI-compatible and serves Claude and codex (gpt-5) + # models behind one URL — the same endpoint Copilot uses. + return f"{workspace}/ai-gateway/mlflow/v1" + + def build_shared_base_urls(workspace: str) -> dict[str, str | dict[str, str]]: urls: dict[str, str | dict[str, str]] = { "codex": build_tool_base_url("codex", workspace), @@ -3351,5 +3358,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), + "continue": build_continue_base_url(workspace), } return urls diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 75e3699d..0b897c10 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 continue_dev, copilot, cursor, gemini, opencode from ucode.config_io import restore_file from ucode.databricks import ( apply_pat_environment, @@ -98,6 +98,11 @@ class _Back: "display": "Cursor", "list_command": "cursor-agent mcp list", }, + "continue": { + "binary": "cn", + "display": "Continue", + "list_command": "cn mcp list", + }, } SKILLS_MCP_KIND = "skills" SKILLS_MCP_SERVER_NAME = "databricks-skill-registry" @@ -330,6 +335,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 == "continue": + removed = continue_dev.write_mcp_server_config(name, argv) + return [MCP_USER_SCOPE] if removed else [] raise RuntimeError(f"Unsupported MCP client '{client}'.") @@ -346,6 +354,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 == "continue": + return [MCP_USER_SCOPE] if continue_dev.remove_mcp_server_config(name) else [] raise RuntimeError(f"Unsupported MCP client '{client}'.") diff --git a/tests/test_agent_continue.py b/tests/test_agent_continue.py new file mode 100644 index 00000000..f82f7077 --- /dev/null +++ b/tests/test_agent_continue.py @@ -0,0 +1,246 @@ +"""Tests for agents/continue_dev.py.""" + +from __future__ import annotations + +import yaml + +import ucode.config_io as config_io_mod +from ucode.agents import continue_dev + +WS = "https://example.databricks.com" + + +class TestContinueSpec: + def test_binary(self): + assert continue_dev.SPEC["binary"] == "cn" + + def test_package(self): + assert continue_dev.SPEC["package"] == "@continuedev/cli" + + def test_display(self): + assert continue_dev.SPEC["display"] == "Continue" + + def test_config_path_is_continue_config_yaml(self): + assert continue_dev.SPEC["config_path"].name == "config.yaml" + assert continue_dev.SPEC["config_path"].parent.name == ".continue" + + +class TestDefaultModel: + def test_prefers_claude_sonnet(self): + state = {"claude_models": {"sonnet": "s4", "opus": "o4", "haiku": "h4"}} + assert continue_dev.default_model(state) == "s4" + + def test_falls_back_to_opus(self): + assert continue_dev.default_model({"claude_models": {"opus": "o4"}}) == "o4" + + def test_falls_back_to_haiku(self): + assert continue_dev.default_model({"claude_models": {"haiku": "h4"}}) == "h4" + + def test_falls_back_to_codex_when_no_claude(self): + state = {"claude_models": {}, "codex_models": ["gpt-5", "gpt-4"]} + assert continue_dev.default_model(state) == "gpt-5" + + def test_returns_none_when_no_models(self): + assert continue_dev.default_model({}) is None + + def test_managed_default_wins(self): + state = {"continue_default_model": "pinned", "claude_models": {"sonnet": "s4"}} + assert continue_dev.default_model(state) == "pinned" + + +class TestRenderConfig: + def test_top_level_schema_fields(self): + doc = continue_dev.render_config("claude-sonnet-4-6", "tok", WS) + assert doc["name"] == "ucode" + assert doc["version"] == "0.0.1" + assert doc["schema"] == "v1" + + def test_single_model_entry(self): + doc = continue_dev.render_config("claude-sonnet-4-6", "tok", WS) + assert len(doc["models"]) == 1 + + def test_provider_is_openai(self): + entry = continue_dev.render_config("m", "tok", WS)["models"][0] + assert entry["provider"] == "openai" + + def test_api_base_is_mlflow_gateway(self): + entry = continue_dev.render_config("m", "tok", WS)["models"][0] + assert entry["apiBase"] == f"{WS}/ai-gateway/mlflow/v1" + + def test_model_id_verbatim(self): + entry = continue_dev.render_config("databricks-gpt-5", "tok", WS)["models"][0] + assert entry["model"] == "databricks-gpt-5" + + def test_api_key_is_token(self): + entry = continue_dev.render_config("m", "tok123", WS)["models"][0] + assert entry["apiKey"] == "tok123" + + def test_roles_are_chat_edit_apply(self): + entry = continue_dev.render_config("m", "tok", WS)["models"][0] + assert entry["roles"] == ["chat", "edit", "apply"] + + def test_name_carries_ucode_prefix(self): + entry = continue_dev.render_config("m", "tok", WS)["models"][0] + assert entry["name"].startswith(continue_dev.UCODE_MODEL_NAME_PREFIX) + + def test_user_agent_header_present(self): + entry = continue_dev.render_config("m", "tok", WS)["models"][0] + assert "User-Agent" in entry["requestOptions"]["headers"] + + +class TestWriteToolConfig: + def _patch_paths(self, tmp_path, monkeypatch): + config_file = tmp_path / "config.yaml" + backup_file = tmp_path / "continue-backup.yaml" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(continue_dev, "CONTINUE_CONFIG_PATH", config_file) + monkeypatch.setattr(continue_dev, "CONTINUE_BACKUP_PATH", backup_file) + monkeypatch.setattr(continue_dev, "get_databricks_token", lambda w, p, **k: "minted-token") + monkeypatch.setattr(continue_dev, "save_state", lambda s: None) + return config_file + + def test_writes_valid_yaml_with_ucode_model(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + state, token = continue_dev.write_tool_config( + {"workspace": WS, "profile": None}, "claude-sonnet-4-6" + ) + assert token == "minted-token" + doc = yaml.safe_load(config_file.read_text()) + assert doc["schema"] == "v1" + ucode_models = [ + m for m in doc["models"] if m["name"].startswith(continue_dev.UCODE_MODEL_NAME_PREFIX) + ] + assert len(ucode_models) == 1 + assert ucode_models[0]["model"] == "claude-sonnet-4-6" + assert ucode_models[0]["apiKey"] == "minted-token" + + def test_uses_passed_token_without_minting(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + continue_dev.write_tool_config({"workspace": WS, "profile": None}, "m", token="given") + doc = yaml.safe_load(config_file.read_text()) + assert doc["models"][0]["apiKey"] == "given" + + def test_rewrite_replaces_not_duplicates_ucode_entry(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + st = {"workspace": WS, "profile": None} + continue_dev.write_tool_config(st, "model-a", token="t1") + continue_dev.write_tool_config(st, "model-b", token="t2") + doc = yaml.safe_load(config_file.read_text()) + ucode_models = [ + m for m in doc["models"] if m["name"].startswith(continue_dev.UCODE_MODEL_NAME_PREFIX) + ] + assert len(ucode_models) == 1 + assert ucode_models[0]["model"] == "model-b" + + def test_preserves_user_authored_models(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + config_file.write_text( + yaml.safe_dump( + { + "name": "my-config", + "version": "1.0.0", + "schema": "v1", + "models": [{"name": "My Local Model", "provider": "ollama", "model": "llama"}], + } + ) + ) + continue_dev.write_tool_config({"workspace": WS, "profile": None}, "m", token="t") + doc = yaml.safe_load(config_file.read_text()) + names = [m["name"] for m in doc["models"]] + assert "My Local Model" in names + assert any(n.startswith(continue_dev.UCODE_MODEL_NAME_PREFIX) for n in names) + # The user's own top-level name is kept (setdefault does not clobber it). + assert doc["name"] == "my-config" + + def test_marks_tool_managed(self, tmp_path, monkeypatch): + self._patch_paths(tmp_path, monkeypatch) + state, _ = continue_dev.write_tool_config( + {"workspace": WS, "profile": None}, "m", token="t" + ) + assert state["managed_configs"]["continue"]["keys"] == [["models"]] + + +class TestMcpServerConfig: + _ARGV = [ + "/usr/local/bin/ucode", + "mcp-proxy", + "--url", + "https://ws/mcp/x", + "--host", + "https://ws", + ] + + def _patch_paths(self, tmp_path, monkeypatch): + config_file = tmp_path / "config.yaml" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(continue_dev, "CONTINUE_CONFIG_PATH", config_file) + monkeypatch.setattr(continue_dev, "CONTINUE_BACKUP_PATH", tmp_path / "continue-backup.yaml") + return config_file + + def test_builds_stdio_entry_from_proxy_argv(self): + entry = continue_dev.build_mcp_server_entry("databricks-slack", self._ARGV) + assert entry["name"] == "databricks-slack" + assert entry["type"] == "stdio" + assert entry["command"] == "/usr/local/bin/ucode" + assert entry["args"] == ["mcp-proxy", "--url", "https://ws/mcp/x", "--host", "https://ws"] + + def test_writes_server_without_clobbering_model(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + config_file.write_text( + yaml.safe_dump( + { + "name": "ucode", + "version": "0.0.1", + "schema": "v1", + "models": [ + {"name": "Databricks (ucode) m", "provider": "openai", "model": "m"} + ], + } + ) + ) + removed = continue_dev.write_mcp_server_config("databricks-slack", self._ARGV) + assert removed is False + doc = yaml.safe_load(config_file.read_text()) + assert [s["name"] for s in doc["mcpServers"]] == ["databricks-slack"] + # The model entry the writer produced is left intact. + assert doc["models"][0]["name"] == "Databricks (ucode) m" + + def test_rewrite_replaces_reports_removed(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + continue_dev.write_mcp_server_config("databricks-slack", self._ARGV) + removed = continue_dev.write_mcp_server_config("databricks-slack", self._ARGV) + assert removed is True + doc = yaml.safe_load(config_file.read_text()) + assert len(doc["mcpServers"]) == 1 + + def test_removes_without_clobbering_others(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + config_file.write_text( + yaml.safe_dump( + { + "schema": "v1", + "mcpServers": [ + {"name": "mine", "command": "x"}, + {"name": "databricks-slack", "command": "y"}, + ], + } + ) + ) + assert continue_dev.remove_mcp_server_config("databricks-slack") is True + doc = yaml.safe_load(config_file.read_text()) + assert [s["name"] for s in doc["mcpServers"]] == ["mine"] + + def test_remove_absent_returns_false(self, tmp_path, monkeypatch): + self._patch_paths(tmp_path, monkeypatch) + assert continue_dev.remove_mcp_server_config("nope") is False + + +class TestValidateCmd: + def test_starts_with_binary(self): + cmd = continue_dev.validate_cmd("cn") + assert cmd[0] == "cn" + + def test_has_headless_and_config_flags(self): + cmd = continue_dev.validate_cmd("cn") + assert "-p" in cmd + assert "--config" in cmd diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index a7692df0..3eb667e3 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -70,7 +70,15 @@ 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", + "continue", + } def test_each_spec_has_required_keys(self): required = {"binary", "package", "display", "config_path", "backup_path"} @@ -163,6 +171,9 @@ class TestNormalizeTool: ("opencode", "opencode"), ("copilot", "copilot"), ("pi", "pi"), + ("continue", "continue"), + ("continue-dev", "continue"), + ("cn", "continue"), ("CODEX", "codex"), (" Claude ", "claude"), ], @@ -218,6 +229,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_continue_available_with_claude(self): + assert check_gateway_endpoint({"claude_models": {"sonnet": "s4"}}, "continue") is True + + def test_continue_available_with_codex(self): + assert check_gateway_endpoint({"codex_models": ["m"]}, "continue") is True + + def test_continue_unavailable_with_only_gemini(self): + assert check_gateway_endpoint({"gemini_models": ["g"]}, "continue") is False + + def test_continue_unavailable_when_no_models(self): + assert check_gateway_endpoint({}, "continue") is False + class TestDefaultModelForTool: def test_codex_returns_none_without_a_configured_model(self): @@ -272,6 +295,17 @@ 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_continue_prefers_claude_sonnet(self): + state = {"claude_models": {"sonnet": "s4", "opus": "o4"}, "codex_models": ["c"]} + assert default_model_for_tool("continue", state) == "s4" + + def test_continue_falls_back_to_codex(self): + state = {"claude_models": {}, "codex_models": ["c1"]} + assert default_model_for_tool("continue", state) == "c1" + + def test_continue_returns_none_when_no_models(self): + assert default_model_for_tool("continue", {}) is None + class TestResolveLaunchModel: def test_codex_default_model_used_when_no_explicit(self): diff --git a/uv.lock b/uv.lock index bb61ef81..81954c2a 100644 --- a/uv.lock +++ b/uv.lock @@ -3160,6 +3160,7 @@ dependencies = [ { name = "databricks-sql-connector" }, { name = "httpx" }, { name = "mcp" }, + { name = "pyyaml" }, { name = "questionary" }, { name = "tomlkit" }, { name = "typer" }, @@ -3184,6 +3185,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" }, { name = "questionary", specifier = ">=2.0.0" }, { name = "tomlkit", specifier = ">=0.13.0" }, { name = "typer", specifier = ">=0.12.0" }, From a253e7feb0692dd1d7db48515da30d5a5a7026ec Mon Sep 17 00:00:00 2001 From: natefleming Date: Mon, 31 Aug 2026 14:07:37 -0400 Subject: [PATCH 2/2] fix(agents): address code-review findings on the Continue harness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Register the `ucode continue` launch command (was missing, so the documented entry point returned "No such command"); mirrors the other `ucode ` launchers. - Revert safety: Continue writes the user's shared ~/.continue/config.yaml (read by the IDE extensions), so `ucode revert` no longer whole-file restores it — it surgically strips only ucode's own model entries via continue_dev.revert_config(), mirroring Codex's shared-config handling. MCP entries are still removed by revert_mcp_configs. - DRY: build_copilot_base_url and build_continue_base_url now delegate to a shared _mlflow_chat_completions_base_url helper (was a duplicated body). - Remove dead render_config (only tests referenced it, risking drift); tests now assert on the production primitives _ucode_model_entry / _ensure_schema_header, and add revert_config coverage. - Factor the ucode-model-entry match into _is_ucode_model, shared by the write and revert paths. Verified: `ucode continue --help` resolves and a live headless `ucode continue -- -p ...` launches cn on system.ai.claude-sonnet-5 through the gateway. ruff + ty clean; 2171 unit tests pass (the 4 failures are the pre-existing e2e/PTY tests). Co-authored-by: Isaac --- src/ucode/agents/continue_dev.py | 48 +++++++++++---- src/ucode/cli.py | 20 ++++++ src/ucode/databricks.py | 21 ++++--- tests/test_agent_continue.py | 101 ++++++++++++++++++++++++------- 4 files changed, 146 insertions(+), 44 deletions(-) diff --git a/src/ucode/agents/continue_dev.py b/src/ucode/agents/continue_dev.py index cfec0dfe..00d41494 100644 --- a/src/ucode/agents/continue_dev.py +++ b/src/ucode/agents/continue_dev.py @@ -22,6 +22,7 @@ import subprocess import threading from pathlib import Path +from typing import cast from ucode.config_io import ( APP_DIR, @@ -103,14 +104,12 @@ def _ensure_schema_header(doc: dict) -> None: doc.setdefault("schema", "v1") -def render_config(model: str, token: str, workspace: str) -> dict: - """Return a complete Continue config document pinning the ucode model.""" - return { - "name": "ucode", - "version": "0.0.1", - "schema": "v1", - "models": [_ucode_model_entry(model, token, workspace)], - } +def _is_ucode_model(entry: object) -> bool: + """True for a ucode-managed `models:` entry (identified by its name prefix).""" + if not isinstance(entry, dict): + return False + name = cast("dict[str, object]", entry).get("name", "") + return isinstance(name, str) and name.startswith(UCODE_MODEL_NAME_PREFIX) def write_tool_config( @@ -132,11 +131,7 @@ def write_tool_config( models = [] # Drop any prior ucode entry so a rewrite replaces (not duplicates) it, while # leaving the user's own models untouched. - models = [ - m - for m in models - if not (isinstance(m, dict) and str(m.get("name", "")).startswith(UCODE_MODEL_NAME_PREFIX)) - ] + models = [m for m in models if not _is_ucode_model(m)] models.append(_ucode_model_entry(model, token, state["workspace"])) existing["models"] = models write_yaml_file(CONTINUE_CONFIG_PATH, existing) @@ -188,6 +183,33 @@ def remove_mcp_server_config(name: str) -> bool: return True +def revert_config() -> bool: + """Surgically remove ucode's model entries from the shared ``config.yaml``. + + Continue's config is the user's real IDE config — unlike the other agents' + ucode-owned/isolated files — so a whole-file restore would clobber models the + user added since ucode first ran. Strip only ucode's own entries instead + (mirroring Codex's ``revert_legacy_shared_config``); the MCP servers ucode + registered are removed separately by ``mcp.revert_mcp_configs``. Returns True + if anything was removed. + """ + if not CONTINUE_CONFIG_PATH.exists(): + return False + existing = read_yaml_safe(CONTINUE_CONFIG_PATH) + models = existing.get("models") + if not isinstance(models, list): + return False + kept = [m for m in models if not _is_ucode_model(m)] + if len(kept) == len(models): + return False + if kept: + existing["models"] = kept + else: + existing.pop("models", None) + write_yaml_file(CONTINUE_CONFIG_PATH, existing) + return True + + def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: model = default_model(state) if not model: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 12d4c27f..8a3469c5 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -16,6 +16,7 @@ configure_selected_tools, configure_single_tool, configure_tool, + continue_dev, ensure_bootstrap_dependencies, ensure_provider_state, explicit_model_arg_value, @@ -1118,7 +1119,13 @@ def revert() -> int: spec["config_path"], spec["backup_path"], bool(managed_configs.get(tool)) ) for tool, spec in TOOL_SPECS.items() + if tool != "continue" } + # Continue writes the user's shared ~/.continue/config.yaml (read by the IDE + # extensions), so a whole-file restore would clobber models the user added + # since ucode first ran — strip only ucode's own entries instead. Its MCP + # servers are removed above by revert_mcp_configs. + results["continue"] = continue_dev.revert_config() pi_settings_restored = restore_file( PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, bool(managed_configs.get("pi")) ) @@ -2529,6 +2536,19 @@ def pi_cmd( _launch_tool("pi", ctx, skip_preflight=skip_preflight) +@app.command( + "continue", context_settings={"allow_extra_args": True, "ignore_unknown_options": True} +) +def continue_cmd( + ctx: typer.Context, + skip_preflight: SkipPreflightOption = False, + skip_managed_config: SkipManagedConfigOption = False, +) -> None: + """Launch Continue via Databricks.""" + _disable_managed_config_if_requested(skip_managed_config) + _launch_tool("continue", 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/databricks.py b/src/ucode/databricks.py index 89dd0e43..966342bf 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3335,19 +3335,22 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: } -def build_copilot_base_url(workspace: str) -> str: - # Copilot CLI's `openai` provider appends `/chat/completions` to the - # configured base URL. The Databricks MLflow chat-completions gateway is - # OpenAI-compatible and serves Claude, codex (gpt-5), and gemini models - # behind one URL. +def _mlflow_chat_completions_base_url(workspace: str) -> str: + # The Databricks MLflow chat-completions gateway: OpenAI-compatible, serving + # Claude, codex (gpt-5), and gemini models behind one URL. Shared by the tools + # whose `openai` provider appends `/chat/completions` to `apiBase` (Copilot, + # Continue), so the endpoint has a single definition. return f"{workspace}/ai-gateway/mlflow/v1" +def build_copilot_base_url(workspace: str) -> str: + return _mlflow_chat_completions_base_url(workspace) + + def build_continue_base_url(workspace: str) -> str: - # Continue's `openai` provider appends `/chat/completions` to `apiBase`. The Databricks - # MLflow chat-completions gateway is OpenAI-compatible and serves Claude and codex (gpt-5) - # models behind one URL — the same endpoint Copilot uses. - return f"{workspace}/ai-gateway/mlflow/v1" + # Continue's `openai` provider appends `/chat/completions` to `apiBase`; it + # shares Copilot's MLflow chat-completions endpoint. + return _mlflow_chat_completions_base_url(workspace) def build_shared_base_urls(workspace: str) -> dict[str, str | dict[str, str]]: diff --git a/tests/test_agent_continue.py b/tests/test_agent_continue.py index f82f7077..67cd871a 100644 --- a/tests/test_agent_continue.py +++ b/tests/test_agent_continue.py @@ -48,46 +48,49 @@ def test_managed_default_wins(self): assert continue_dev.default_model(state) == "pinned" -class TestRenderConfig: - def test_top_level_schema_fields(self): - doc = continue_dev.render_config("claude-sonnet-4-6", "tok", WS) - assert doc["name"] == "ucode" - assert doc["version"] == "0.0.1" - assert doc["schema"] == "v1" - - def test_single_model_entry(self): - doc = continue_dev.render_config("claude-sonnet-4-6", "tok", WS) - assert len(doc["models"]) == 1 - +class TestUcodeModelEntry: + # `_ucode_model_entry` is the primitive the production write path uses, so + # asserting on it directly (rather than a test-only builder) can't drift. def test_provider_is_openai(self): - entry = continue_dev.render_config("m", "tok", WS)["models"][0] - assert entry["provider"] == "openai" + assert continue_dev._ucode_model_entry("m", "tok", WS)["provider"] == "openai" def test_api_base_is_mlflow_gateway(self): - entry = continue_dev.render_config("m", "tok", WS)["models"][0] + entry = continue_dev._ucode_model_entry("m", "tok", WS) assert entry["apiBase"] == f"{WS}/ai-gateway/mlflow/v1" def test_model_id_verbatim(self): - entry = continue_dev.render_config("databricks-gpt-5", "tok", WS)["models"][0] - assert entry["model"] == "databricks-gpt-5" + assert continue_dev._ucode_model_entry("databricks-gpt-5", "tok", WS)["model"] == ( + "databricks-gpt-5" + ) def test_api_key_is_token(self): - entry = continue_dev.render_config("m", "tok123", WS)["models"][0] - assert entry["apiKey"] == "tok123" + assert continue_dev._ucode_model_entry("m", "tok123", WS)["apiKey"] == "tok123" def test_roles_are_chat_edit_apply(self): - entry = continue_dev.render_config("m", "tok", WS)["models"][0] - assert entry["roles"] == ["chat", "edit", "apply"] + assert continue_dev._ucode_model_entry("m", "tok", WS)["roles"] == ["chat", "edit", "apply"] def test_name_carries_ucode_prefix(self): - entry = continue_dev.render_config("m", "tok", WS)["models"][0] + entry = continue_dev._ucode_model_entry("m", "tok", WS) assert entry["name"].startswith(continue_dev.UCODE_MODEL_NAME_PREFIX) def test_user_agent_header_present(self): - entry = continue_dev.render_config("m", "tok", WS)["models"][0] + entry = continue_dev._ucode_model_entry("m", "tok", WS) assert "User-Agent" in entry["requestOptions"]["headers"] +class TestSchemaHeader: + def test_sets_required_keys_on_empty(self): + doc: dict = {} + continue_dev._ensure_schema_header(doc) + assert doc == {"name": "ucode", "version": "0.0.1", "schema": "v1"} + + def test_keeps_user_values(self): + doc = {"name": "mine", "version": "9.9", "schema": "v1"} + continue_dev._ensure_schema_header(doc) + assert doc["name"] == "mine" + assert doc["version"] == "9.9" + + class TestWriteToolConfig: def _patch_paths(self, tmp_path, monkeypatch): config_file = tmp_path / "config.yaml" @@ -235,6 +238,60 @@ def test_remove_absent_returns_false(self, tmp_path, monkeypatch): assert continue_dev.remove_mcp_server_config("nope") is False +class TestRevertConfig: + def _patch_paths(self, tmp_path, monkeypatch): + config_file = tmp_path / "config.yaml" + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(continue_dev, "CONTINUE_CONFIG_PATH", config_file) + monkeypatch.setattr(continue_dev, "CONTINUE_BACKUP_PATH", tmp_path / "continue-backup.yaml") + return config_file + + def test_strips_only_ucode_models_keeping_user_content(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + config_file.write_text( + yaml.safe_dump( + { + "name": "mine", + "schema": "v1", + "models": [ + {"name": "My Local Model", "provider": "ollama", "model": "llama"}, + {"name": f"{continue_dev.UCODE_MODEL_NAME_PREFIX} m", "model": "m"}, + ], + "mcpServers": [{"name": "mine-mcp", "command": "x"}], + } + ) + ) + assert continue_dev.revert_config() is True + doc = yaml.safe_load(config_file.read_text()) + # ucode's model is gone; the user's model, servers, and name survive. + assert [m["name"] for m in doc["models"]] == ["My Local Model"] + assert [s["name"] for s in doc["mcpServers"]] == ["mine-mcp"] + assert doc["name"] == "mine" + + def test_drops_models_key_when_only_ucode(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + config_file.write_text( + yaml.safe_dump( + { + "name": "ucode", + "schema": "v1", + "models": [{"name": f"{continue_dev.UCODE_MODEL_NAME_PREFIX} m", "model": "m"}], + } + ) + ) + assert continue_dev.revert_config() is True + assert "models" not in yaml.safe_load(config_file.read_text()) + + def test_returns_false_when_no_ucode_models(self, tmp_path, monkeypatch): + config_file = self._patch_paths(tmp_path, monkeypatch) + config_file.write_text(yaml.safe_dump({"schema": "v1", "models": [{"name": "Mine"}]})) + assert continue_dev.revert_config() is False + + def test_returns_false_when_no_config(self, tmp_path, monkeypatch): + self._patch_paths(tmp_path, monkeypatch) + assert continue_dev.revert_config() is False + + class TestValidateCmd: def test_starts_with_binary(self): cmd = continue_dev.validate_cmd("cn")