From 9f3f8138d7771b6805f547b1c6a272d992a0b626 Mon Sep 17 00:00:00 2001 From: James Wade Date: Fri, 22 May 2026 08:22:09 -0400 Subject: [PATCH 1/2] Pin Pi defaultProvider/defaultModel in settings.json MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pi inherits the full environment from ucode, including any env vars that its pi-ai package recognizes as auth for a built-in provider (e.g. HF_TOKEN → huggingface). Pi's findInitialModel then picks the built-in before reaching our databricks-claude provider, routes to the wrong host, and validation 401s. Write settings.json alongside models.json so findInitialModel uses our provider/model directly. Fixes #68. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/ucode/agents/pi.py | 14 ++++++++++++++ tests/test_agent_pi.py | 26 ++++++++++++++++++++++---- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index aab1ec3..b8112e1 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -52,6 +52,7 @@ PI_UCODE_HOME = APP_DIR / "pi-home" PI_CONFIG_DIR = PI_UCODE_HOME / ".pi" / "agent" PI_CONFIG_PATH = PI_CONFIG_DIR / "models.json" +PI_SETTINGS_PATH = PI_CONFIG_DIR / "settings.json" PI_BACKUP_PATH = APP_DIR / "pi-models.backup.json" SPEC: ToolSpec = { @@ -184,11 +185,24 @@ def write_tool_config( providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) write_json_file(PI_CONFIG_PATH, merged) + _write_settings(overlay["model"]) state = mark_tool_managed(state, "pi", managed_keys) save_state(state) return state, token +def _write_settings(model_selector: str) -> None: + # Pin defaultProvider/defaultModel in settings.json so Pi doesn't fall + # through to an env-key-backed provider (e.g. HF_TOKEN exposing + # huggingface) in `findInitialModel` when no --model is passed. + provider, _, model_id = model_selector.partition("/") + if not model_id: + return + existing = read_json_safe(PI_SETTINGS_PATH) + merged = deep_merge_dict(existing, {"defaultProvider": provider, "defaultModel": model_id}) + write_json_file(PI_SETTINGS_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 {} diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 02ec6f1..11d11e7 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -267,9 +267,11 @@ def _setup(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) config_file = tmp_path / "models.json" backup_file = tmp_path / "pi-backup.json" + settings_file = tmp_path / "settings.json" monkeypatch.setattr(pi_mod, "PI_CONFIG_PATH", config_file) + monkeypatch.setattr(pi_mod, "PI_SETTINGS_PATH", settings_file) monkeypatch.setattr(pi_mod, "PI_BACKUP_PATH", backup_file) - return pi_mod, config_file + return pi_mod, config_file, settings_file def _state(self, **overrides) -> dict: state = { @@ -284,7 +286,7 @@ def _state(self, **overrides) -> dict: return state def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatch): - pi_mod, config_file = self._setup(tmp_path, monkeypatch) + pi_mod, config_file, _ = self._setup(tmp_path, monkeypatch) stale = { "providers": { @@ -312,7 +314,7 @@ 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.""" - pi_mod, config_file = self._setup(tmp_path, monkeypatch) + pi_mod, config_file, _ = self._setup(tmp_path, monkeypatch) config_file.write_text( json.dumps( @@ -339,7 +341,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): assert "databricks-claude" in written_providers def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch): - pi_mod, config_file = self._setup(tmp_path, monkeypatch) + pi_mod, config_file, _ = self._setup(tmp_path, monkeypatch) with ( patch("ucode.agents.pi.get_databricks_token", return_value="tok"), @@ -350,3 +352,19 @@ def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch written = json.loads(config_file.read_text()) assert written["model"] == "databricks-claude/claude-sonnet" assert written["providers"]["databricks-claude"]["apiKey"] == "tok" + + def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): + # Without this, Pi's `findInitialModel` can fall through to a built-in + # provider when an unrelated env var (e.g. HF_TOKEN) makes one look + # auth-configured. Pinning the default keeps Pi on our provider. + pi_mod, _, settings_file = self._setup(tmp_path, monkeypatch) + + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + settings = json.loads(settings_file.read_text()) + assert settings["defaultProvider"] == "databricks-claude" + assert settings["defaultModel"] == "claude-sonnet" From 83c7b275c4e6d9095f1b17137ef8cc469a700c75 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Thu, 28 May 2026 17:31:15 +0000 Subject: [PATCH 2/2] backup and restore pi settings.json --- src/ucode/agents/__init__.py | 4 +++ src/ucode/agents/pi.py | 2 ++ src/ucode/cli.py | 5 ++++ tests/test_agent_pi.py | 55 ++++++++++++++++++++++++++++++++---- 4 files changed, 61 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 14cfd1b..4770971 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -386,6 +386,7 @@ def validate_tool(tool: str) -> tuple[bool, str]: def validate_all_tools(state: dict) -> None: from rich.panel import Panel # local to avoid bumping module-level deps + from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import restore_file console.print() @@ -411,6 +412,9 @@ def validate_all_tools(state: dict) -> None: print_err(f"{spec['display']}: {err}") managed = bool(state.get("managed_configs", {}).get(tool)) restore_file(spec["config_path"], spec["backup_path"], managed) + # Rollback settings.json for Pi + if tool == "pi": + restore_file(PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, managed) available_tools.remove(tool) state["available_tools"] = available_tools save_state(state) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index b8112e1..e7c1760 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -54,6 +54,7 @@ PI_CONFIG_PATH = PI_CONFIG_DIR / "models.json" PI_SETTINGS_PATH = PI_CONFIG_DIR / "settings.json" PI_BACKUP_PATH = APP_DIR / "pi-models.backup.json" +PI_SETTINGS_BACKUP_PATH = APP_DIR / "pi-settings.backup.json" SPEC: ToolSpec = { "binary": "pi", @@ -198,6 +199,7 @@ def _write_settings(model_selector: str) -> None: provider, _, model_id = model_selector.partition("/") if not model_id: return + backup_existing_file(PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH) existing = read_json_safe(PI_SETTINGS_PATH) merged = deep_merge_dict(existing, {"defaultProvider": provider, "defaultModel": model_id}) write_json_file(PI_SETTINGS_PATH, merged) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1530609..47337e7 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -25,6 +25,7 @@ from ucode.agents import ( launch as launch_agent, ) +from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import restore_file, set_dry_run from ucode.databricks import ( build_shared_base_urls, @@ -398,12 +399,16 @@ def revert() -> int: ) for tool, spec in TOOL_SPECS.items() } + pi_settings_restored = restore_file( + PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, bool(managed_configs.get("pi")) + ) clear_state() print_heading("Revert") print_kv("Workspace", state.get("workspace") or "none") for tool, spec in TOOL_SPECS.items(): print_kv(f"{spec['display']} config", "restored" if results[tool] else "unchanged") + print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") for client, spec in MCP_CLIENTS.items(): print_kv( f"{spec['display']} MCP config", diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 11d11e7..0afc5fb 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from contextlib import nullcontext from unittest.mock import patch from ucode.agents import pi @@ -268,10 +269,12 @@ def _setup(self, tmp_path, monkeypatch): config_file = tmp_path / "models.json" backup_file = tmp_path / "pi-backup.json" settings_file = tmp_path / "settings.json" + settings_backup_file = tmp_path / "pi-settings-backup.json" monkeypatch.setattr(pi_mod, "PI_CONFIG_PATH", config_file) monkeypatch.setattr(pi_mod, "PI_SETTINGS_PATH", settings_file) monkeypatch.setattr(pi_mod, "PI_BACKUP_PATH", backup_file) - return pi_mod, config_file, settings_file + monkeypatch.setattr(pi_mod, "PI_SETTINGS_BACKUP_PATH", settings_backup_file) + return pi_mod, config_file, settings_file, settings_backup_file def _state(self, **overrides) -> dict: state = { @@ -286,7 +289,7 @@ def _state(self, **overrides) -> dict: return state def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatch): - pi_mod, config_file, _ = self._setup(tmp_path, monkeypatch) + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) stale = { "providers": { @@ -314,7 +317,7 @@ 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.""" - pi_mod, config_file, _ = self._setup(tmp_path, monkeypatch) + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) config_file.write_text( json.dumps( @@ -341,7 +344,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): assert "databricks-claude" in written_providers def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch): - pi_mod, config_file, _ = self._setup(tmp_path, monkeypatch) + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) with ( patch("ucode.agents.pi.get_databricks_token", return_value="tok"), @@ -357,7 +360,7 @@ def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): # Without this, Pi's `findInitialModel` can fall through to a built-in # provider when an unrelated env var (e.g. HF_TOKEN) makes one look # auth-configured. Pinning the default keeps Pi on our provider. - pi_mod, _, settings_file = self._setup(tmp_path, monkeypatch) + pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) with ( patch("ucode.agents.pi.get_databricks_token", return_value="tok"), @@ -368,3 +371,45 @@ def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): settings = json.loads(settings_file.read_text()) assert settings["defaultProvider"] == "databricks-claude" assert settings["defaultModel"] == "claude-sonnet" + + def test_pre_existing_settings_are_backed_up_before_first_write(self, tmp_path, monkeypatch): + pi_mod, _, settings_file, settings_backup_file = self._setup(tmp_path, monkeypatch) + + original = '{"theme": "Default Dark", "defaultProvider": "openai"}' + settings_file.parent.mkdir(parents=True, exist_ok=True) + settings_file.write_text(original, encoding="utf-8") + + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch("ucode.agents.pi.save_state"), + ): + pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + + assert settings_backup_file.read_text(encoding="utf-8") == original + # The on-disk settings still get the ucode pin applied via deep_merge. + merged = json.loads(settings_file.read_text()) + assert merged["defaultProvider"] == "databricks-claude" + assert merged["theme"] == "Default Dark" + + +class TestValidateAllToolsPiRollback: + def test_failed_pi_validation_rolls_back_settings(self, tmp_path, monkeypatch): + import ucode.agents as agents_mod + import ucode.agents.pi as pi_mod + + settings_file = tmp_path / "settings.json" + settings_file.write_text("{}", encoding="utf-8") + monkeypatch.setattr(pi_mod, "PI_SETTINGS_PATH", settings_file) + monkeypatch.setattr(pi_mod, "PI_SETTINGS_BACKUP_PATH", tmp_path / "settings.backup.json") + # Keep the generic models.json rollback off the user's real config dir. + monkeypatch.setitem(agents_mod.TOOL_SPECS["pi"], "config_path", tmp_path / "models.json") + monkeypatch.setitem( + agents_mod.TOOL_SPECS["pi"], "backup_path", tmp_path / "models.backup.json" + ) + 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": ["pi"], "managed_configs": {"pi": True}}) + + assert not settings_file.exists()