From 14cc40194ed15557c78571f937491c62f8cebdb7 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Tue, 27 Oct 2026 17:07:06 +0000 Subject: [PATCH 01/18] Set up Copilot Chat in VS Code with ug configure --agents vscode `ug configure --agents vscode` configures the workspace the way Cursor's path does (no model picker), then: - finds VS Code's `code` command (or `code-insiders`), on PATH or where the installer puts it, since macOS doesn't add it to PATH by default; - packs the extension bundled with ug into a vsix with the standard library and installs it, unless that version is already installed; - writes unityGateway.workspace, unityGateway.models and unityGateway.authCommand into the user settings.json. Models come from discover_model_services: Claude, GPT and open models (not Gemini), GPT-5 and later on the Responses API, the rest on Chat Completions. The auth command is the one Claude Code already uses. The settings are edited with a new JSONC helper in config_io that changes only ug's top-level keys and keeps the user's comments and formatting. `ug revert` removes those keys and uninstalls the extension. Co-authored-by: Isaac --- src/ucode/agents/copilot.py | 5 + src/ucode/agents/vscode.py | 251 +++++++++++++++++++++++++++++++ src/ucode/cli.py | 18 ++- src/ucode/config_io.py | 176 ++++++++++++++++++++++ tests/conftest.py | 6 + tests/test_agent_vscode.py | 286 ++++++++++++++++++++++++++++++++++++ tests/test_config_io.py | 72 +++++++++ 7 files changed, 810 insertions(+), 4 deletions(-) create mode 100644 src/ucode/agents/vscode.py create mode 100644 tests/test_agent_vscode.py diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 5e7adfbe4..614f6c7a4 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -240,6 +240,11 @@ def _wire_from_name(name: str) -> str | None: return WIRE_COMPLETIONS if gpt else None +def uses_responses_api(model: str) -> bool: + """Whether Copilot calls ``model`` through the Responses API: GPT 5+ and o-series models.""" + return _wire_from_name(bare_model_name(model)) == WIRE_RESPONSES + + def copilot_catalog_model_id(model: str) -> str | None: """Copilot's catalog id for a Claude model, e.g. ``system.ai.claude-sonnet-5-5`` -> ``claude-sonnet-5.5``. diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py new file mode 100644 index 000000000..de48b35cf --- /dev/null +++ b/src/ucode/agents/vscode.py @@ -0,0 +1,251 @@ +"""VS Code: install ug's model extension so Copilot Chat can use the workspace's AI Gateway models. + +Configure-only, like Cursor (it stays out of `agents.__init__._MODULES`): `ug configure --agents +vscode` packs the extension that ships with ug into a vsix, installs it with VS Code's `code` +command, and writes the extension's settings into the user settings.json. `ug revert` undoes both. +""" + +from __future__ import annotations + +import json +import os +import shutil +import tempfile +import zipfile +from pathlib import Path +from typing import cast +from xml.sax.saxutils import escape + +from ucode import config_io +from ucode.agents.copilot import uses_responses_api +from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command, launch_token +from ucode.databricks import build_auth_shell_command, discover_model_services +from ucode.managed_files import OS, current_os +from ucode.os_compatibility import subprocess_cross_os +from ucode.os_compatibility.system_cross_os import windows_program_files_dir +from ucode.ui import print_kv, print_warning + +EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" +EXTENSION_ID = "databricks.unity-gateway" +WORKSPACE_SETTING = "unityGateway.workspace" +MODELS_SETTING = "unityGateway.models" +AUTH_COMMAND_SETTING = "unityGateway.authCommand" +SETTINGS = (WORKSPACE_SETTING, MODELS_SETTING, AUTH_COMMAND_SETTING) +# Each `code` command and the folder its user settings live in. +CODE_COMMANDS = (("code", "Code"), ("code-insiders", "Code - Insiders")) +_CODE_TIMEOUT_S = 120 + + +def _install_locations(command: str) -> list[Path]: + """Where VS Code's installer puts `command` when it isn't on PATH (by default on macOS).""" + insiders = command == "code-insiders" + system = current_os() + if system is OS.MACOS: + app = "Visual Studio Code - Insiders.app" if insiders else "Visual Studio Code.app" + return [ + root / app / "Contents/Resources/app/bin" / command + for root in (Path("/Applications"), Path.home() / "Applications") + ] + if system is OS.WINDOWS: + folder = "Microsoft VS Code Insiders" if insiders else "Microsoft VS Code" + roots = ( + [Path(os.environ["LOCALAPPDATA"]) / "Programs"] + if os.environ.get("LOCALAPPDATA") + else [] + ) + return [ + root / folder / "bin" / f"{command}.cmd" + for root in [*roots, windows_program_files_dir()] + ] + return [Path("/usr/bin") / command, Path("/snap/bin") / command] + + +def find_code() -> tuple[Path, str] | None: + """VS Code's `code` command (or Insiders') and its settings folder name, or None.""" + for command, product in CODE_COMMANDS: + if found := shutil.which(command): + return Path(found), product + for command, product in CODE_COMMANDS: + for path in _install_locations(command): + if path.is_file(): + return path, product + return None + + +def settings_path(product: str) -> Path: + system = current_os() + if system is OS.MACOS: + base = Path.home() / "Library" / "Application Support" + elif system is OS.WINDOWS: + base = Path(os.environ.get("APPDATA") or Path.home() / "AppData" / "Roaming") + else: + base = Path(os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config") + return base / product / "User" / "settings.json" + + +def _package() -> dict: + return json.loads((EXTENSION_DIR / "package.json").read_text(encoding="utf-8")) + + +def pack_vsix(directory: Path) -> Path: + """Zip the bundled extension into a vsix in ``directory``, with no Node or `vsce` needed.""" + package = _package() + manifest = f""" + + + + {escape(package["displayName"])} + {escape(package["description"])} + + + + + + + + + +""" + content_types = ( + '' + '' + '' + '' + '' + "" + ) + vsix = directory / f"{package['name']}-{package['version']}.vsix" + with zipfile.ZipFile(vsix, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("[Content_Types].xml", content_types) + archive.writestr("extension.vsixmanifest", manifest) + archive.write(EXTENSION_DIR / "package.json", "extension/package.json") + archive.write( + EXTENSION_DIR / package["main"], f"extension/{package['main'].removeprefix('./')}" + ) + return vsix + + +def _run_code(code: Path, *args: str): + return subprocess_cross_os.run( + [str(code), *args], capture_output=True, text=True, timeout=_CODE_TIMEOUT_S, check=False + ) + + +def installed_version(code: Path) -> str | None: + result = _run_code(code, "--list-extensions", "--show-versions") + for line in (result.stdout or "").splitlines(): + name, _, version = line.strip().partition("@") + if name.lower() == EXTENSION_ID: + return version + return None + + +def install_extension(code: Path) -> str: + """Install the bundled extension unless that version already is; returns what happened.""" + version = _package()["version"] + if installed_version(code) == version: + return f"{EXTENSION_ID} {version} already installed" + with tempfile.TemporaryDirectory(prefix="ug-vscode-") as directory: + vsix = pack_vsix(Path(directory)) + result = _run_code(code, "--install-extension", str(vsix), "--force") + if result.returncode != 0: + detail = (result.stderr or result.stdout or "").strip()[:300] + raise RuntimeError(f"`{code} --install-extension` failed: {detail}") + return f"installed {EXTENSION_ID} {version}" + + +def model_settings(workspace: str, token: str) -> tuple[list[dict[str, str]], str | None]: + """The workspace's Claude, GPT and open models, each with the API the extension calls it on.""" + # Gemini is left out: the gateway rejects the requests Copilot sends to it. + claude, gpt, _gemini, oss, reason = discover_model_services(workspace, token) + models = dict.fromkeys([*claude.values(), *gpt, *oss]) + return [ + {"id": model, "api": "responses" if uses_responses_api(model) else "chat_completions"} + for model in models + ], reason + + +def auth_command(state: dict) -> str: + workspace = state["workspace"] + custom_oauth = state.get("custom_oauth") + if isinstance(custom_oauth, dict): + return build_custom_auth_shell_command(workspace, cast(CustomOAuthConfig, custom_oauth)) + return build_auth_shell_command( + workspace, state.get("profile"), use_pat=bool(state.get("use_pat")) + ) + + +def _edit_settings(path: Path, values: dict[str, object], remove: tuple[str, ...] = ()) -> bool: + text = path.read_text(encoding="utf-8") if path.exists() else "" + try: + updated = config_io.update_jsonc_keys(text, values, remove) + except ValueError as exc: + print_warning( + f"Couldn't update {path}: {exc}. Fix the file, then rerun `ug configure --agents vscode`." + ) + return False + if updated != text: + config_io.write_text_file(path, updated) + return True + + +def configure(state: dict) -> bool: + """Write the extension's settings and install it. Returns False, after a warning, when VS + Code isn't found or the settings can't be updated.""" + found = find_code() + if found is None: + print_warning( + "VS Code's `code` command wasn't found. Install VS Code (on macOS, also run " + "\"Shell Command: Install 'code' command in PATH\"), then rerun `ug configure --agents vscode`." + ) + return False + code, product = found + workspace = state["workspace"] + models, reason = model_settings(workspace, launch_token(state, workspace)) + if not models: + print_warning(f"No models for VS Code: {reason}") + path = settings_path(product) + values = { + WORKSPACE_SETTING: workspace, + MODELS_SETTING: models, + AUTH_COMMAND_SETTING: auth_command(state), + } + if not _edit_settings(path, values): + return False + print_kv("VS Code settings", f"{len(models)} models in {path}") + if config_io.is_dry_run(): + return True + try: + print_kv("VS Code extension", install_extension(code)) + except RuntimeError as exc: + print_warning( + f"{exc} Install it from VS Code's Extensions view, or rerun `ug configure --agents vscode`." + ) + return False + return True + + +def revert() -> str: + """Remove ug's settings and uninstall the extension: "removed" or "unchanged".""" + changed = False + for _, product in CODE_COMMANDS: + path = settings_path(product) + if not path.exists(): + continue + try: + current = config_io.loads_jsonc(path.read_text(encoding="utf-8")) + except ValueError: + print_warning(f"Couldn't read {path}; remove its unityGateway settings by hand.") + continue + if isinstance(current, dict) and any(key in current for key in SETTINGS): + changed = _edit_settings(path, {}, SETTINGS) or changed + found = find_code() + if found is not None and installed_version(found[0]) is not None: + result = _run_code(found[0], "--uninstall-extension", EXTENSION_ID) + if result.returncode == 0: + changed = True + else: + print_warning( + f"Couldn't uninstall {EXTENSION_ID}: {(result.stderr or '').strip()[:300]}" + ) + return "removed" if changed else "unchanged" diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 127dc0106..2a280f93a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1580,6 +1580,9 @@ def revert() -> int: claude_managed_result = claude_agent.revert_managed_settings() codex_managed_result = codex_agent.revert_managed_config() workspace_policy_result = revert_machine_policy() + from ucode.agents import vscode + + vscode_result = vscode.revert() results: dict[str, bool] = { # Antigravity's skills.json is shared with its app and IDE extensions, so revert @@ -1606,6 +1609,7 @@ def revert() -> int: print_kv("Codex OS-managed settings", codex_managed_result) print_kv("Unity Gateway workspace policy", workspace_policy_result) print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") + print_kv("VS Code extension and settings", vscode_result) for client, spec in MCP_CLIENTS.items(): print_kv( f"{spec['display']} MCP config", @@ -4419,7 +4423,7 @@ def configure( str | None, typer.Option( "--agents", - help="Configure a comma-separated list of agents without prompting (e.g. claude,codex).", + help="Configure a comma-separated list of agents without prompting (e.g. claude,codex,vscode).", ), ] = None, workspace: Annotated[ @@ -4636,9 +4640,11 @@ def configure( # (which picks cursor up through MCP_ONLY_CLIENTS). If cursor is the # only agent, do a workspace-only configure so that a later `ug mcp # add` run has a current workspace to target. + # VS Code is configure-only too: it needs the workspace, then its own setup below. requested = [a.strip().lower() for a in agents.split(",") if a.strip()] wants_cursor = "cursor" in requested - model_agent_names = ",".join(a for a in requested if a != "cursor") + wants_vscode = "vscode" in requested + model_agent_names = ",".join(a for a in requested if a not in ("cursor", "vscode")) if model_agent_names: selected_tools = _parse_agents_option(model_agent_names) if workspace_entries is None: @@ -4652,8 +4658,8 @@ def configure( workspaces=workspace_entries, **skip_kwargs, ) - elif wants_cursor: - # Cursor-only: establish workspace state without the model picker. + elif wants_cursor or wants_vscode: + # Cursor/VS Code only: establish workspace state without the model picker. _configure_shared_workspace_states( workspace_entries or [_prompt_for_configuration(None)], tools=[], @@ -4664,6 +4670,10 @@ def configure( else: # Neither model agents nor cursor -> empty/invalid --agents list. _parse_agents_option(agents) + if wants_vscode: + from ucode.agents import vscode + + vscode.configure(load_state()) elif mcp is not None: # MCP-only: `--mcp` without --agent(s) (e.g. Cursor, which isn't a # model agent, or adding MCP servers to an already-configured setup). diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index c4ca3578c..1635b6a27 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -266,3 +266,179 @@ def parse_dotenv(path: Path) -> dict[str, str]: def write_dotenv(path: Path, env: dict[str, str]) -> None: content = "".join(f'{key}="{val}"\n' for key, val in env.items()) write_text_file(path, content) + + +# --- JSONC (JSON with comments and trailing commas, e.g. VS Code's settings.json) --------------- + +_JSONC_PUNCTUATION = frozenset("{}[]:,") + + +def _jsonc_tokens(text: str) -> list[tuple[str, int, int]]: + """``(token, start, end)`` for each token in ``text``, skipping whitespace and comments. + + Raises ValueError on an unterminated string or comment. + """ + tokens: list[tuple[str, int, int]] = [] + i, length = 0, len(text) + while i < length: + char = text[i] + if char.isspace(): + i += 1 + elif text.startswith("//", i): + newline = text.find("\n", i) + i = length if newline < 0 else newline + 1 + elif text.startswith("/*", i): + close = text.find("*/", i + 2) + if close < 0: + raise ValueError("unterminated comment") + i = close + 2 + elif char == '"': + end = i + 1 + while end < length and text[end] != '"': + end += 2 if text[end] == "\\" else 1 + if end >= length: + raise ValueError("unterminated string") + tokens.append((text[i : end + 1], i, end + 1)) + i = end + 1 + elif char in _JSONC_PUNCTUATION: + tokens.append((char, i, i + 1)) + i += 1 + else: + end = i + while end < length and not text[end].isspace() and text[end] not in _JSONC_PUNCTUATION: + if text.startswith(("//", "/*"), end): + break + end += 1 + tokens.append((text[i:end], i, end)) + i = end + return tokens + + +def loads_jsonc(text: str) -> object: + """Parse JSONC: comments and trailing commas are dropped, then it's read as JSON.""" + tokens = _jsonc_tokens(text) + kept = [ + token + for index, (token, _, _) in enumerate(tokens) + if not (token == "," and index + 1 < len(tokens) and tokens[index + 1][0] in "}]") + ] + return json.loads(" ".join(kept)) if kept else None + + +class _JsoncEntry(TypedDict): + key: str + key_start: int + value_start: int + value_end: int + comma_end: int | None + + +def _jsonc_top_level(text: str) -> tuple[int, int, list[_JsoncEntry]]: + """The top-level object's ``{`` and ``}`` offsets and its entries; ValueError if it isn't one.""" + tokens = _jsonc_tokens(text) + if not tokens or tokens[0][0] != "{": + raise ValueError("the top level isn't an object") + entries: list[_JsoncEntry] = [] + index = 1 + while index < len(tokens) and tokens[index][0] != "}": + key, key_start, _ = tokens[index] + if not key.startswith('"') or index + 2 >= len(tokens) or tokens[index + 1][0] != ":": + raise ValueError("expected a key and a value") + index += 2 + value_start = tokens[index][1] + depth = 0 + while True: + token = tokens[index][0] + depth += 1 if token in "{[" else -1 if token in "}]" else 0 + if depth <= 0: + break + index += 1 + if index >= len(tokens): + raise ValueError("unterminated value") + value_end = tokens[index][2] + index += 1 + comma_end = None + if index < len(tokens) and tokens[index][0] == ",": + comma_end = tokens[index][2] + index += 1 + entries.append( + { + "key": json.loads(key), + "key_start": key_start, + "value_start": value_start, + "value_end": value_end, + "comma_end": comma_end, + } + ) + if index >= len(tokens) or index != len(tokens) - 1: + raise ValueError("unexpected content after the top-level object") + return tokens[0][1], tokens[index][1], entries + + +def _line_start(text: str, offset: int) -> int: + """The start of ``offset``'s line when only whitespace precedes it there, else ``offset``.""" + start = text.rfind("\n", 0, offset) + 1 + return start if not text[start:offset].strip() else offset + + +def _line_end(text: str, offset: int) -> int: + newline = text.find("\n", offset) + return len(text) if newline < 0 else newline + + +def _remove_jsonc_entry(text: str, entries: list[_JsoncEntry], position: int) -> str: + entry = entries[position] + start = _line_start(text, entry["key_start"]) + end = entry["comma_end"] if entry["comma_end"] is not None else entry["value_end"] + rest = text[end : _line_end(text, end)].strip() + # A comment after the entry on its own line goes with it. + if start != entry["key_start"] and (not rest or rest.startswith("//")): + end = min(_line_end(text, end) + 1, len(text)) + text = text[:start] + text[end:] + previous_comma = entries[position - 1]["comma_end"] if position > 0 else None + if entry["comma_end"] is None and previous_comma is not None: + text = text[: previous_comma - 1] + text[previous_comma:] + return text + + +def update_jsonc_keys(text: str, values: dict[str, object], remove: tuple[str, ...] = ()) -> str: + """Set ``values`` and delete ``remove`` among the top-level keys of a JSONC object, leaving + every other byte (comments included) as it was. Raises ValueError if ``text`` isn't one.""" + if not _jsonc_tokens(text): + text = text + ("\n" if text.strip() else "") + "{}" + for key in remove: + _, _, entries = _jsonc_top_level(text) + position = next((i for i, entry in enumerate(entries) if entry["key"] == key), None) + if position is not None: + text = _remove_jsonc_entry(text, entries, position) + for key, value in values.items(): + open_brace, close_brace, entries = _jsonc_top_level(text) + first = entries[0]["key_start"] if entries else None + indent = text[_line_start(text, first) : first] if first is not None else "" + indent = indent if indent and not indent.strip() else " " + rendered = json.dumps(value, indent=indent).replace("\n", "\n" + indent) + existing = next((entry for entry in entries if entry["key"] == key), None) + if existing is not None: + text = text[: existing["value_start"]] + rendered + text[existing["value_end"] :] + continue + line = f"{indent}{json.dumps(key)}: {rendered}" + if not entries: + text = text[: open_brace + 1] + f"\n{line}\n" + text[open_brace + 1 :].lstrip(" \t") + continue + last = entries[-1] + if last["comma_end"] is None: + text = text[: last["value_end"]] + "," + text[last["value_end"] :] + close_brace += 1 + after = last["value_end"] + 1 + else: + after = last["comma_end"] + # After the last entry's line, so a comment on that line stays with it. + at = _line_end(text, after) + at = after if close_brace < at else at + text = text[:at] + f"\n{line}" + text[at:] + loaded = loads_jsonc(text) + if not isinstance(loaded, dict) or any( + loaded.get(key) != value for key, value in values.items() + ): + raise ValueError("the edited file doesn't hold the expected values") + return text diff --git a/tests/conftest.py b/tests/conftest.py index 64e173255..65b0dd3bb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -41,6 +41,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): from ucode.agents import antigravity as antigravity_mod from ucode.agents import claude as claude_mod from ucode.agents import codex as codex_mod + from ucode.agents import vscode as vscode_mod from ucode.agents.desktop import claude as desktop_mod state_dir = tmp_path / ".ucode" @@ -49,6 +50,11 @@ def _isolate_ucode_state(tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir) monkeypatch.setattr(codex_mod, "_otel_service_path", lambda: state_dir / "otel.service") monkeypatch.setattr(desktop_mod, "APP_DIR", state_dir) + # No test may run the developer's `code` or edit their VS Code settings. + monkeypatch.setattr(vscode_mod, "find_code", lambda: None) + monkeypatch.setattr( + vscode_mod, "settings_path", lambda product: tmp_path / product / "User" / "settings.json" + ) monkeypatch.setattr(desktop_auth_mod, "_AUTH_LOCK_PATH", state_dir / "claude-desktop-auth.lock") # Antigravity's customization root is shared with the developer's Antigravity app and CLI. agy_dir = tmp_path / ".gemini" / "config" diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py new file mode 100644 index 000000000..d50cd2dd4 --- /dev/null +++ b/tests/test_agent_vscode.py @@ -0,0 +1,286 @@ +"""Tests for `ug configure --agents vscode` and its revert.""" + +from __future__ import annotations + +import json +import subprocess +import zipfile +from pathlib import Path + +import pytest +from typer.testing import CliRunner + +from ucode import config_io +from ucode.agents import vscode +from ucode.cli import app +from ucode.managed_files import OS + +# The conftest fixture stubs these out for every test; the ones here exercise the real lookups. +REAL_FIND_CODE = vscode.find_code +REAL_SETTINGS_PATH = vscode.settings_path + +WS = "https://ws.cloud.databricks.com" +STATE = {"workspace": WS, "profile": "dev"} +USER_SETTINGS = """// My settings +{ + // Theme + "workbench.colorTheme": "Default Dark+", + "editor.fontSize": 14, // big +} +""" + + +class FakeCode: + """Stands in for VS Code's `code` command: records calls and tracks the installed version.""" + + def __init__(self, installed: str | None = None, fail: str | None = None) -> None: + self.installed = installed + self.fail = fail + self.calls: list[tuple[str, ...]] = [] + self.vsix_entries: list[str] = [] + + def __call__(self, code: Path, *args: str) -> subprocess.CompletedProcess: + self.calls.append(args) + if args[0] == self.fail: + return subprocess.CompletedProcess(args, 1, "", "boom") + if args[0] == "--list-extensions": + listed = f"{vscode.EXTENSION_ID}@{self.installed}\n" if self.installed else "" + return subprocess.CompletedProcess(args, 0, "ms-python.python@2026.1.0\n" + listed, "") + if args[0] == "--install-extension": + self.vsix_entries = zipfile.ZipFile(args[1]).namelist() + self.installed = json.loads((vscode.EXTENSION_DIR / "package.json").read_text())[ + "version" + ] + if args[0] == "--uninstall-extension": + self.installed = None + return subprocess.CompletedProcess(args, 0, "", "") + + +@pytest.fixture +def code(monkeypatch): + fake = FakeCode() + monkeypatch.setattr(vscode, "find_code", lambda: (Path("/fake/bin/code"), "Code")) + monkeypatch.setattr(vscode, "_run_code", fake) + monkeypatch.setattr(vscode, "launch_token", lambda state, workspace: "token") + monkeypatch.setattr( + vscode, + "discover_model_services", + lambda workspace, token: ( + {"sonnet": "system.ai.claude-sonnet-5-5"}, + ["system.ai.gpt-6-luna", "system.ai.gpt-oss-120b"], + ["system.ai.gemini-4-argon"], + ["system.ai.gpt-oss-120b", "system.ai.deepseek-v4-1-flash"], + None, + ), + ) + return fake + + +def _printed(capsys) -> str: + return " ".join(capsys.readouterr().out.split()) + + +def _settings() -> Path: + return vscode.settings_path("Code") + + +def _settings_values() -> dict: + return config_io.loads_jsonc(_settings().read_text(encoding="utf-8")) + + +class TestConfigure: + def test_merges_settings_keeping_the_users_own(self, code): + _settings().parent.mkdir(parents=True) + _settings().write_text(USER_SETTINGS, encoding="utf-8") + + assert vscode.configure(STATE) + + text = _settings().read_text(encoding="utf-8") + for kept in ("// My settings", "// Theme", '"editor.fontSize": 14, // big'): + assert kept in text + assert _settings_values() == { + "workbench.colorTheme": "Default Dark+", + "editor.fontSize": 14, + "unityGateway.workspace": WS, + "unityGateway.models": [ + {"id": "system.ai.claude-sonnet-5-5", "api": "chat_completions"}, + {"id": "system.ai.gpt-6-luna", "api": "responses"}, + {"id": "system.ai.gpt-oss-120b", "api": "chat_completions"}, + {"id": "system.ai.deepseek-v4-1-flash", "api": "chat_completions"}, + ], + "unityGateway.authCommand": vscode.build_auth_shell_command(WS, "dev"), + } + + def test_creates_missing_settings_and_installs_the_packed_extension(self, code): + assert vscode.configure(STATE) + + assert set(_settings_values()) == set(vscode.SETTINGS) + assert code.calls[-1][0] == "--install-extension" + assert code.calls[-1][2] == "--force" + assert code.vsix_entries == [ + "[Content_Types].xml", + "extension.vsixmanifest", + "extension/package.json", + "extension/out/extension.js", + ] + + def test_reruns_update_values_without_reinstalling_the_same_version(self, code): + vscode.configure(STATE) + code.calls.clear() + + assert vscode.configure({**STATE, "profile": "other"}) + + assert "--profile other" in _settings_values()["unityGateway.authCommand"] + assert [call[0] for call in code.calls] == ["--list-extensions"] + + def test_leaves_unreadable_settings_alone(self, code, capsys): + _settings().parent.mkdir(parents=True) + _settings().write_text('{"broken": ', encoding="utf-8") + + assert not vscode.configure(STATE) + + assert _settings().read_text(encoding="utf-8") == '{"broken": ' + assert code.calls == [] + assert "Couldn't update" in _printed(capsys) + + def test_warns_when_vscode_isnt_found(self, capsys): + assert not vscode.configure(STATE) + + assert "`code` command wasn't found" in _printed(capsys) + assert not _settings().exists() + + def test_reports_a_failed_install(self, code, capsys): + code.fail = "--install-extension" + + assert not vscode.configure(STATE) + + assert "--install-extension` failed: boom" in _printed(capsys) + + +class TestRevert: + def test_removes_only_ugs_settings_and_uninstalls(self, code): + _settings().parent.mkdir(parents=True) + _settings().write_text(USER_SETTINGS, encoding="utf-8") + vscode.configure(STATE) + + assert vscode.revert() == "removed" + + assert _settings_values() == { + "workbench.colorTheme": "Default Dark+", + "editor.fontSize": 14, + } + assert "// Theme" in _settings().read_text(encoding="utf-8") + assert code.calls[-1] == ("--uninstall-extension", vscode.EXTENSION_ID) + + def test_does_nothing_when_vscode_was_never_configured(self, code): + _settings().parent.mkdir(parents=True) + _settings().write_text(USER_SETTINGS, encoding="utf-8") + + assert vscode.revert() == "unchanged" + + assert _settings().read_text(encoding="utf-8") == USER_SETTINGS + assert [call[0] for call in code.calls] == ["--list-extensions"] + + +class TestFindCode: + @pytest.fixture(autouse=True) + def _real_lookup(self, monkeypatch): + monkeypatch.setattr(vscode, "find_code", REAL_FIND_CODE) + + def test_prefers_code_on_path(self, monkeypatch): + monkeypatch.setattr(vscode.shutil, "which", lambda name: f"/usr/local/bin/{name}") + + assert vscode.find_code() == (Path("/usr/local/bin/code"), "Code") + + def test_falls_back_to_insiders(self, monkeypatch): + monkeypatch.setattr( + vscode.shutil, + "which", + lambda name: "/opt/bin/code-insiders" if name == "code-insiders" else None, + ) + + assert vscode.find_code() == (Path("/opt/bin/code-insiders"), "Code - Insiders") + + def test_finds_the_macos_app_when_code_isnt_on_path(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) + monkeypatch.setattr(vscode.Path, "home", lambda: tmp_path) + app_code = tmp_path / "Applications/Visual Studio Code.app/Contents/Resources/app/bin/code" + app_code.parent.mkdir(parents=True) + app_code.touch() + + assert vscode.find_code() == (app_code, "Code") + + def test_finds_the_windows_user_install(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + monkeypatch.setattr(vscode, "current_os", lambda: OS.WINDOWS) + monkeypatch.setenv("LOCALAPPDATA", str(tmp_path)) + code_cmd = tmp_path / "Programs/Microsoft VS Code/bin/code.cmd" + code_cmd.parent.mkdir(parents=True) + code_cmd.touch() + + assert vscode.find_code() == (code_cmd, "Code") + + def test_none_when_vscode_isnt_installed(self, monkeypatch): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + monkeypatch.setattr(vscode, "_install_locations", lambda command: []) + + assert vscode.find_code() is None + + +class TestSettingsPath: + @pytest.mark.parametrize( + ("system", "env", "expected"), + [ + (OS.LINUX, {"XDG_CONFIG_HOME": "/xdg"}, "/xdg/Code/User/settings.json"), + (OS.WINDOWS, {"APPDATA": "/appdata"}, "/appdata/Code/User/settings.json"), + ], + ) + def test_user_settings_per_os(self, monkeypatch, system, env, expected): + monkeypatch.setattr(vscode, "current_os", lambda: system) + for name, value in env.items(): + monkeypatch.setenv(name, value) + + assert REAL_SETTINGS_PATH("Code") == Path(expected) + + def test_macos(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) + monkeypatch.setattr(vscode.Path, "home", lambda: tmp_path) + + assert ( + REAL_SETTINGS_PATH("Code - Insiders") + == tmp_path / "Library/Application Support/Code - Insiders/User/settings.json" + ) + + +class TestAuthCommand: + def test_pat_and_custom_oauth(self): + assert "--use-pat" in vscode.auth_command({**STATE, "use_pat": True}) + custom = { + "client_id": "c", + "redirect_url": "http://localhost:8020", + "scopes": "all-apis", + "profile": "custom", + } + assert "--profile custom" in vscode.auth_command({"workspace": WS, "custom_oauth": custom}) + + +def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch): + configured: list[dict] = [] + shared: list[dict] = [] + monkeypatch.setattr( + "ucode.cli._configure_shared_workspace_states", + lambda *args, **kwargs: shared.append(kwargs) or [], + ) + monkeypatch.setattr( + "ucode.cli.configure_workspace_command", + lambda **kwargs: pytest.fail("no model agents to configure"), + ) + monkeypatch.setattr("ucode.cli.load_state", lambda: STATE) + monkeypatch.setattr(vscode, "configure", lambda state: configured.append(state) or True) + + result = CliRunner().invoke(app, ["configure", "--agents", "vscode", "--workspace", WS]) + + assert result.exit_code == 0, result.output + assert shared[0]["tools"] == [] + assert configured == [STATE] diff --git a/tests/test_config_io.py b/tests/test_config_io.py index 1cb325ee4..6e11ca2b9 100644 --- a/tests/test_config_io.py +++ b/tests/test_config_io.py @@ -405,3 +405,75 @@ def test_partial_path_through_non_dict_is_noop(self): doc = {"model": "scalar"} assert prune_key_paths(doc, [["model", "nested"]]) is False assert doc == {"model": "scalar"} + + +class TestJsonc: + SETTINGS = """// Top comment +{ + // Theme + "workbench.colorTheme": "Default Dark+", /* inline */ + "url": "https://example.com/a//b", + "editor.fontSize": 14, // big +} +""" + + def test_loads_comments_and_trailing_commas(self): + assert config_io.loads_jsonc(self.SETTINGS) == { + "workbench.colorTheme": "Default Dark+", + "url": "https://example.com/a//b", + "editor.fontSize": 14, + } + assert config_io.loads_jsonc('{"a": [1, 2,], "b": "/* not a comment */",}') == { + "a": [1, 2], + "b": "/* not a comment */", + } + + def test_sets_keys_and_keeps_everything_else(self): + updated = config_io.update_jsonc_keys( + self.SETTINGS, {"editor.fontSize": 16, "new.key": {"x": [1]}} + ) + + assert config_io.loads_jsonc(updated)["editor.fontSize"] == 16 + assert config_io.loads_jsonc(updated)["new.key"] == {"x": [1]} + for kept in ( + "// Top comment", + "// Theme", + "/* inline */", + "// big", + '"url": "https://example.com/a//b"', + ): + assert kept in updated + assert ' "new.key": {\n "x": [\n 1\n ]\n }' in updated + + def test_removes_keys_with_their_line_comment(self): + updated = config_io.update_jsonc_keys( + self.SETTINGS, {}, remove=("editor.fontSize", "missing") + ) + + assert config_io.loads_jsonc(updated) == { + "workbench.colorTheme": "Default Dark+", + "url": "https://example.com/a//b", + } + assert "// big" not in updated + assert "// Theme" in updated + + def test_removing_the_last_key_drops_the_comma_before_it(self): + assert ( + config_io.update_jsonc_keys('{\n "a": 1,\n "b": 2\n}\n', {}, remove=("b",)) + == '{\n "a": 1\n}\n' + ) + assert config_io.update_jsonc_keys('{"a": 1}', {}, remove=("a",)) == "{}" + + @pytest.mark.parametrize("text", ["", "// only a comment\n", "{}"]) + def test_fills_an_empty_file(self, text): + updated = config_io.update_jsonc_keys(text, {"a": 1}) + + assert config_io.loads_jsonc(updated) == {"a": 1} + assert text.strip("{}") in updated + + @pytest.mark.parametrize( + "text", ['{"a": ', "[1, 2]", '{"a": 1} {"b": 2}', '{"a" 1}', '{"a": "x'] + ) + def test_rejects_what_isnt_one_object(self, text): + with pytest.raises(ValueError): + config_io.update_jsonc_keys(text, {"b": 2}) From 0b845b465c7707c0936b9b7595e1c3fe7bae9014 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 14:04:46 +0000 Subject: [PATCH 02/18] Pack every bundled extension file into the vsix The extension now ships its provider icon next to the bundle and its manifest. Pack every file under the bundled extension folder instead of naming the two, so a new asset doesn't need a change here. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 8 ++++---- tests/test_agent_vscode.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index de48b35cf..19980e333 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -111,6 +111,7 @@ def pack_vsix(directory: Path) -> Path: '' '' '' + '' '' "" ) @@ -118,10 +119,9 @@ def pack_vsix(directory: Path) -> Path: with zipfile.ZipFile(vsix, "w", zipfile.ZIP_DEFLATED) as archive: archive.writestr("[Content_Types].xml", content_types) archive.writestr("extension.vsixmanifest", manifest) - archive.write(EXTENSION_DIR / "package.json", "extension/package.json") - archive.write( - EXTENSION_DIR / package["main"], f"extension/{package['main'].removeprefix('./')}" - ) + for path in sorted(EXTENSION_DIR.rglob("*")): + if path.is_file(): + archive.write(path, f"extension/{path.relative_to(EXTENSION_DIR).as_posix()}") return vsix diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index d50cd2dd4..adf64d80f 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -120,8 +120,9 @@ def test_creates_missing_settings_and_installs_the_packed_extension(self, code): assert code.vsix_entries == [ "[Content_Types].xml", "extension.vsixmanifest", - "extension/package.json", + "extension/media/databricks.svg", "extension/out/extension.js", + "extension/package.json", ] def test_reruns_update_values_without_reinstalling_the_same_version(self, code): From 0343cf5ebe4c88760c11224cbf2e3323fc74fc49 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 14:10:01 +0000 Subject: [PATCH 03/18] Reinstall the VS Code extension on every configure Configure skipped the install when VS Code already had the same extension version. A ug upgrade can change the bundled extension without changing its version, and users then kept the old build. Install with --force every time. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 5 ++--- tests/test_agent_vscode.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 19980e333..e338e31d6 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -141,10 +141,9 @@ def installed_version(code: Path) -> str | None: def install_extension(code: Path) -> str: - """Install the bundled extension unless that version already is; returns what happened.""" + """Install the bundled extension over whatever is installed; returns what happened.""" + # Always reinstall: a ug upgrade can change the bundle without changing its version. version = _package()["version"] - if installed_version(code) == version: - return f"{EXTENSION_ID} {version} already installed" with tempfile.TemporaryDirectory(prefix="ug-vscode-") as directory: vsix = pack_vsix(Path(directory)) result = _run_code(code, "--install-extension", str(vsix), "--force") diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index adf64d80f..ca0925ffa 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -125,14 +125,14 @@ def test_creates_missing_settings_and_installs_the_packed_extension(self, code): "extension/package.json", ] - def test_reruns_update_values_without_reinstalling_the_same_version(self, code): + def test_reruns_update_values_and_reinstall_the_same_version(self, code): vscode.configure(STATE) code.calls.clear() assert vscode.configure({**STATE, "profile": "other"}) assert "--profile other" in _settings_values()["unityGateway.authCommand"] - assert [call[0] for call in code.calls] == ["--list-extensions"] + assert [call[0] for call in code.calls] == ["--install-extension"] def test_leaves_unreadable_settings_alone(self, code, capsys): _settings().parent.mkdir(parents=True) From ff4365937d5cf6614e46b5b96e0c00567e9ed20e Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 14:20:15 +0000 Subject: [PATCH 04/18] Write Claude's token limits into the VS Code model settings Each Claude entry in unityGateway.models now carries maxInputTokens and maxOutputTokens from the limits ug already gives Copilot CLI (copilot_model_limits), e.g. 936K / 32K for Opus 5.5 and 872K / 128K for Sonnet 5.5. Those limits hold on the Chat Completions route the extension uses: the gateway took an 861K-token prompt to claude-haiku-5-5 and a 128K max_tokens on Sonnet 5.5. GPT and open models get no entry, so they keep the extension's 128K / 32K. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 21 +++++++++++++++------ tests/test_agent_vscode.py | 7 ++++++- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index e338e31d6..18b7f9c12 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -17,7 +17,7 @@ from xml.sax.saxutils import escape from ucode import config_io -from ucode.agents.copilot import uses_responses_api +from ucode.agents.copilot import copilot_model_limits, uses_responses_api from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command, launch_token from ucode.databricks import build_auth_shell_command, discover_model_services from ucode.managed_files import OS, current_os @@ -153,15 +153,24 @@ def install_extension(code: Path) -> str: return f"installed {EXTENSION_ID} {version}" -def model_settings(workspace: str, token: str) -> tuple[list[dict[str, str]], str | None]: +def _model_entry(model: str) -> dict[str, object]: + entry: dict[str, object] = { + "id": model, + "api": "responses" if uses_responses_api(model) else "chat_completions", + } + # Claude's limits hold on Chat Completions too. Models without known limits get the + # extension's 128K-in, 32K-out default. + if limits := copilot_model_limits(model): + entry["maxInputTokens"], entry["maxOutputTokens"] = limits + return entry + + +def model_settings(workspace: str, token: str) -> tuple[list[dict[str, object]], str | None]: """The workspace's Claude, GPT and open models, each with the API the extension calls it on.""" # Gemini is left out: the gateway rejects the requests Copilot sends to it. claude, gpt, _gemini, oss, reason = discover_model_services(workspace, token) models = dict.fromkeys([*claude.values(), *gpt, *oss]) - return [ - {"id": model, "api": "responses" if uses_responses_api(model) else "chat_completions"} - for model in models - ], reason + return [_model_entry(model) for model in models], reason def auth_command(state: dict) -> str: diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index ca0925ffa..30c1ed7a5 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -103,7 +103,12 @@ def test_merges_settings_keeping_the_users_own(self, code): "editor.fontSize": 14, "unityGateway.workspace": WS, "unityGateway.models": [ - {"id": "system.ai.claude-sonnet-5-5", "api": "chat_completions"}, + { + "id": "system.ai.claude-sonnet-5-5", + "api": "chat_completions", + "maxInputTokens": 872_000, + "maxOutputTokens": 128_000, + }, {"id": "system.ai.gpt-6-luna", "api": "responses"}, {"id": "system.ai.gpt-oss-120b", "api": "chat_completions"}, {"id": "system.ai.deepseek-v4-1-flash", "api": "chat_completions"}, From 7bba2dbc9d5fa835c7540e40a49db330bc439741 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 14:23:51 +0000 Subject: [PATCH 05/18] Tell users to reload VS Code after configure installs the extension VS Code applies the settings right away, but a window that's already open keeps running the extension build it loaded, and configure now reinstalls on every run. Print a note to reload open windows. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 4 +++- tests/test_agent_vscode.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 18b7f9c12..841823045 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -23,7 +23,7 @@ from ucode.managed_files import OS, current_os from ucode.os_compatibility import subprocess_cross_os from ucode.os_compatibility.system_cross_os import windows_program_files_dir -from ucode.ui import print_kv, print_warning +from ucode.ui import print_kv, print_note, print_warning EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" EXTENSION_ID = "databricks.unity-gateway" @@ -230,6 +230,8 @@ def configure(state: dict) -> bool: f"{exc} Install it from VS Code's Extensions view, or rerun `ug configure --agents vscode`." ) return False + # Settings apply live, but an open window keeps running the extension build it loaded. + print_note("Reload open VS Code windows (Developer: Reload Window) to use the extension.") return True diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 30c1ed7a5..8bd809a7b 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -116,7 +116,7 @@ def test_merges_settings_keeping_the_users_own(self, code): "unityGateway.authCommand": vscode.build_auth_shell_command(WS, "dev"), } - def test_creates_missing_settings_and_installs_the_packed_extension(self, code): + def test_creates_missing_settings_and_installs_the_packed_extension(self, code, capsys): assert vscode.configure(STATE) assert set(_settings_values()) == set(vscode.SETTINGS) @@ -129,6 +129,7 @@ def test_creates_missing_settings_and_installs_the_packed_extension(self, code): "extension/out/extension.js", "extension/package.json", ] + assert "Developer: Reload Window" in _printed(capsys) def test_reruns_update_values_and_reinstall_the_same_version(self, code): vscode.configure(STATE) From 18c5e40e0d3efa7b4bd789b5f54162ef1c173fcd Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 14:36:17 +0000 Subject: [PATCH 06/18] Spell out the VS Code steps after configure installs the extension "Reload open VS Code windows (Developer: Reload Window)" assumed people know the Command Palette. Print the steps instead, with this OS's shortcuts: Next, in VS Code: 1. If VS Code is already open, reload each window: press Cmd+Shift+P, type "Reload Window", and press Enter. Or quit VS Code (Cmd+Q) and open it again. 2. Open Chat, click the model picker, and choose a Databricks model. The reload matters only for a window that's already open: in VS Code 1.143, a first install shows up in an open window within seconds, but a reinstall doesn't, and the window keeps running the build it loaded. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 23 ++++++++++++++++++++--- tests/test_agent_vscode.py | 23 ++++++++++++++++++++++- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 841823045..49b6c59e1 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -23,7 +23,7 @@ from ucode.managed_files import OS, current_os from ucode.os_compatibility import subprocess_cross_os from ucode.os_compatibility.system_cross_os import windows_program_files_dir -from ucode.ui import print_kv, print_note, print_warning +from ucode.ui import console, print_heading, print_kv, print_warning EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" EXTENSION_ID = "databricks.unity-gateway" @@ -197,6 +197,22 @@ def _edit_settings(path: Path, values: dict[str, object], remove: tuple[str, ... return True +def next_steps() -> list[str]: + """What to do in VS Code once the extension is installed, with this OS's shortcuts.""" + if current_os() is OS.MACOS: + palette, restart = "Cmd+Shift+P", "Or quit VS Code (Cmd+Q) and open it again." + else: + palette, restart = "Ctrl+Shift+P", "Or close VS Code and open it again." + return [ + # A first install shows up in an open window by itself, but a reinstall doesn't: the + # window keeps running the build it loaded until it reloads. + "1. If VS Code is already open, reload each window:", + f' press {palette}, type "Reload Window", and press Enter.', + f" {restart}", + "2. Open Chat, click the model picker, and choose a Databricks model.", + ] + + def configure(state: dict) -> bool: """Write the extension's settings and install it. Returns False, after a warning, when VS Code isn't found or the settings can't be updated.""" @@ -230,8 +246,9 @@ def configure(state: dict) -> bool: f"{exc} Install it from VS Code's Extensions view, or rerun `ug configure --agents vscode`." ) return False - # Settings apply live, but an open window keeps running the extension build it loaded. - print_note("Reload open VS Code windows (Developer: Reload Window) to use the extension.") + print_heading("Next, in VS Code:") + for line in next_steps(): + console.print(f" {line}", highlight=False) return True diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 8bd809a7b..b69305d10 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -129,7 +129,9 @@ def test_creates_missing_settings_and_installs_the_packed_extension(self, code, "extension/out/extension.js", "extension/package.json", ] - assert "Developer: Reload Window" in _printed(capsys) + assert "Next, in VS Code: 1. If VS Code is already open, reload each window" in _printed( + capsys + ) def test_reruns_update_values_and_reinstall_the_same_version(self, code): vscode.configure(STATE) @@ -260,6 +262,25 @@ def test_macos(self, monkeypatch, tmp_path): ) +class TestNextSteps: + @pytest.mark.parametrize( + ("system", "palette", "restart"), + [ + (OS.MACOS, "Cmd+Shift+P", "quit VS Code (Cmd+Q)"), + (OS.WINDOWS, "Ctrl+Shift+P", "close VS Code"), + (OS.LINUX, "Ctrl+Shift+P", "close VS Code"), + ], + ) + def test_uses_the_shortcuts_for_the_os(self, monkeypatch, system, palette, restart): + monkeypatch.setattr(vscode, "current_os", lambda: system) + + steps = " ".join(vscode.next_steps()) + + assert f'press {palette}, type "Reload Window"' in steps + assert restart in steps + assert "choose a Databricks model" in steps + + class TestAuthCommand: def test_pat_and_custom_oauth(self): assert "--use-pat" in vscode.auth_command({**STATE, "use_pat": True}) From 7f0c61cdc9945c44cac83e4fd2fe4bc1e77f9d41 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 14:43:38 +0000 Subject: [PATCH 07/18] Tell users to restart VS Code instead of reloading each window Reloading each window through the Command Palette was more to explain, and it misses any window the user doesn't reload. Restarting VS Code loads the new extension build in every window at once: Next, in VS Code: 1. If VS Code is open, quit it (Cmd+Q) and open it again. 2. Open Chat, click the model picker, and choose a Databricks model. Windows and Linux say "close every VS Code window and open it again". Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 19 +++++++++---------- tests/test_agent_vscode.py | 23 ++++++++++------------- 2 files changed, 19 insertions(+), 23 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 49b6c59e1..1fd87e880 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -198,17 +198,16 @@ def _edit_settings(path: Path, values: dict[str, object], remove: tuple[str, ... def next_steps() -> list[str]: - """What to do in VS Code once the extension is installed, with this OS's shortcuts.""" - if current_os() is OS.MACOS: - palette, restart = "Cmd+Shift+P", "Or quit VS Code (Cmd+Q) and open it again." - else: - palette, restart = "Ctrl+Shift+P", "Or close VS Code and open it again." + """What to do in VS Code once the extension is installed, for this OS.""" + restart = ( + "quit it (Cmd+Q) and open it again" + if current_os() is OS.MACOS + else "close every VS Code window and open it again" + ) return [ - # A first install shows up in an open window by itself, but a reinstall doesn't: the - # window keeps running the build it loaded until it reloads. - "1. If VS Code is already open, reload each window:", - f' press {palette}, type "Reload Window", and press Enter.', - f" {restart}", + # An open window keeps running the extension build it loaded, and configure reinstalls + # every time. A restart covers every open window at once. + f"1. If VS Code is open, {restart}.", "2. Open Chat, click the model picker, and choose a Databricks model.", ] diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index b69305d10..6a4b3dc66 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -129,9 +129,7 @@ def test_creates_missing_settings_and_installs_the_packed_extension(self, code, "extension/out/extension.js", "extension/package.json", ] - assert "Next, in VS Code: 1. If VS Code is already open, reload each window" in _printed( - capsys - ) + assert "Next, in VS Code: 1. If VS Code is open," in _printed(capsys) def test_reruns_update_values_and_reinstall_the_same_version(self, code): vscode.configure(STATE) @@ -264,21 +262,20 @@ def test_macos(self, monkeypatch, tmp_path): class TestNextSteps: @pytest.mark.parametrize( - ("system", "palette", "restart"), + ("system", "restart"), [ - (OS.MACOS, "Cmd+Shift+P", "quit VS Code (Cmd+Q)"), - (OS.WINDOWS, "Ctrl+Shift+P", "close VS Code"), - (OS.LINUX, "Ctrl+Shift+P", "close VS Code"), + (OS.MACOS, "If VS Code is open, quit it (Cmd+Q) and open it again."), + (OS.WINDOWS, "If VS Code is open, close every VS Code window and open it again."), + (OS.LINUX, "If VS Code is open, close every VS Code window and open it again."), ], ) - def test_uses_the_shortcuts_for_the_os(self, monkeypatch, system, palette, restart): + def test_says_how_to_restart_on_the_os(self, monkeypatch, system, restart): monkeypatch.setattr(vscode, "current_os", lambda: system) - steps = " ".join(vscode.next_steps()) - - assert f'press {palette}, type "Reload Window"' in steps - assert restart in steps - assert "choose a Databricks model" in steps + assert vscode.next_steps() == [ + f"1. {restart}", + "2. Open Chat, click the model picker, and choose a Databricks model.", + ] class TestAuthCommand: From 8b22c4cea754a2fbcc877585c8c2b0d9da26e2b8 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Wed, 28 Oct 2026 19:15:06 +0000 Subject: [PATCH 08/18] Offer only Claude and GPT models in VS Code Configure wrote every open model the workspace lists, but many of them don't work from VS Code. On e2-dogfood, five that the listing returns (deepseek-v4-flash-0731, deepseek-v4-pro-0813, glm-5-3, glm-5-3-flash, kimi-k3) fail with "Could not route request" on both Chat Completions and Responses, kimi-k3-neo returns 404, and gpt-oss-120b stops without calling a tool when Copilot sends its full tool list. Keep the models Copilot CLI offers: Claude and GPT, using the same rule (copilot_exclusion_reason), which leaves out gpt-oss and the other open models. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 19 ++++++++++++++----- tests/test_agent_vscode.py | 4 +--- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 1fd87e880..12db79017 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -17,7 +17,11 @@ from xml.sax.saxutils import escape from ucode import config_io -from ucode.agents.copilot import copilot_model_limits, uses_responses_api +from ucode.agents.copilot import ( + copilot_exclusion_reason, + copilot_model_limits, + uses_responses_api, +) from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command, launch_token from ucode.databricks import build_auth_shell_command, discover_model_services from ucode.managed_files import OS, current_os @@ -166,10 +170,15 @@ def _model_entry(model: str) -> dict[str, object]: def model_settings(workspace: str, token: str) -> tuple[list[dict[str, object]], str | None]: - """The workspace's Claude, GPT and open models, each with the API the extension calls it on.""" - # Gemini is left out: the gateway rejects the requests Copilot sends to it. - claude, gpt, _gemini, oss, reason = discover_model_services(workspace, token) - models = dict.fromkeys([*claude.values(), *gpt, *oss]) + """The workspace's Claude and GPT models, each with the API the extension calls it on.""" + # The models Copilot CLI offers: Gemini and open models, gpt-oss among them, fail on Copilot's + # requests or aren't served on every workspace. + claude, gpt, _gemini, _oss, reason = discover_model_services(workspace, token) + models = [ + model + for model in dict.fromkeys([*claude.values(), *gpt]) + if copilot_exclusion_reason(model) is None + ] return [_model_entry(model) for model in models], reason diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 6a4b3dc66..979144613 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -89,7 +89,7 @@ def _settings_values() -> dict: class TestConfigure: - def test_merges_settings_keeping_the_users_own(self, code): + def test_merges_settings_keeping_the_users_own_with_claude_and_gpt_only(self, code): _settings().parent.mkdir(parents=True) _settings().write_text(USER_SETTINGS, encoding="utf-8") @@ -110,8 +110,6 @@ def test_merges_settings_keeping_the_users_own(self, code): "maxOutputTokens": 128_000, }, {"id": "system.ai.gpt-6-luna", "api": "responses"}, - {"id": "system.ai.gpt-oss-120b", "api": "chat_completions"}, - {"id": "system.ai.deepseek-v4-1-flash", "api": "chat_completions"}, ], "unityGateway.authCommand": vscode.build_auth_shell_command(WS, "dev"), } From 7867ee48955c6f50e28812f00ce9291b88430df0 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 10:33:37 +0000 Subject: [PATCH 09/18] Fail a VS Code reply that doesn't end with a finished event Both stream readers stopped when the connection closed, so a stream cut off mid-answer (a dropped connection, a proxy timeout) reached Copilot Chat as a finished reply. Only an output-token stop was reported: a content-filter or any other early stop was ignored, and the user got a truncated answer with no reason. A reply now counts as done only after the gateway says the model finished: response.completed on Responses, finish_reason "stop" or "tool_calls" on Chat Completions. Any other stop reason fails with a message for it, and a stream that ends without a final event fails as cut off. Co-authored-by: Isaac --- src/ucode/vscode/extension/out/extension.js | 4 +-- vscode-extension/src/formats.ts | 34 ++++++++++++++++++--- vscode-extension/test/extension.test.ts | 4 +-- vscode-extension/test/formats.test.ts | 21 +++++++++++-- 4 files changed, 52 insertions(+), 11 deletions(-) diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 575c17279..4f32a40c5 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var U=Object.create;var f=Object.defineProperty;var F=Object.getOwnPropertyDescriptor;var B=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,H=Object.prototype.hasOwnProperty;var K=(t,n)=>{for(var e in n)f(t,e,{get:n[e],enumerable:!0})},M=(t,n,e,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let i of B(n))!H.call(t,i)&&i!==e&&f(t,i,{get:()=>n[i],enumerable:!(s=F(n,i))||s.enumerable});return t};var C=(t,n,e)=>(e=t!=null?U(W(t)):{},M(n||!t||!t.__esModule?f(e,"default",{value:t,enumerable:!0}):e,t)),Y=t=>M(f({},"__esModule",{value:!0}),t);var se={};K(se,{DatabricksModels:()=>v,activate:()=>oe,readSettings:()=>$,userAgent:()=>q});module.exports=Y(se);var l=C(require("vscode"));var b="0.1.0";var T=require("node:child_process"),X=5*6e4,Q=60*6e4,V=t=>new Promise((n,e)=>{(0,T.exec)(t,{timeout:6e4,windowsHide:!0},(s,i,o)=>{s?e(new Error(`Couldn't get a Databricks token: \`${t}\` failed: ${String(o).trim().slice(0,500)||s.message}`)):n(String(i).trim())})}),Z=t=>{try{let{exp:n}=JSON.parse(Buffer.from(t.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+Q},h=class{constructor(n,e=()=>{}){this.command=n;this.onFetch=e}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-X>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=V(n).then(e=>{if(!e)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:e,expiresAt:Z(e)},this.onFetch(this.cached.expiresAt),e}).finally(()=>{this.pending=void 0}),this.pending}};var c=C(require("vscode")),S=3,ee=1e3,y=t=>t instanceof c.LanguageModelTextPart?t.value:t instanceof c.LanguageModelDataPart&&!t.mimeType.startsWith("image/")?new TextDecoder().decode(t.data):t instanceof c.LanguageModelPromptTsxPart?JSON.stringify(t.value):"",_=t=>t instanceof c.LanguageModelDataPart&&t.mimeType.startsWith("image/")?`data:${t.mimeType};base64,${Buffer.from(t.data).toString("base64")}`:void 0,R=t=>t.inputSchema??{type:"object",properties:{}},L=(t,n,e)=>{let s;try{s=e.trim()?JSON.parse(e):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${e}`);return new c.LanguageModelToolCallPart(t,n,s)},P=()=>new Error("The model stopped at its output token limit before finishing the response."),A=(t,n)=>t.tools.length?{tools:t.tools.map(n),...t.toolRequired?{tool_choice:"required"}:{}}:{},E=t=>{if(typeof t=="string")return Math.ceil(t.length/4);let n=0,e=0;for(let s of t.content){let i=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let o of i)n+=o instanceof c.LanguageModelToolCallPart?o.name.length+JSON.stringify(o.input).length:y(o).length,e+=_(o)===void 0?0:1}return Math.ceil(n/4)+e*ee},O=t=>({results:t.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:t.filter(n=>n instanceof c.LanguageModelToolCallPart),text:t.map(y).join(""),images:t.map(_).filter(n=>n!==void 0)});function I(t){let n=t.messages.flatMap(e=>{let{results:s,calls:i,text:o,images:r}=O(e.content);if(Number(e.role)===S)return[{role:"system",content:o}];if(e.role===c.LanguageModelChatMessageRole.Assistant){let a=i.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return o||a.length?[{role:"assistant",content:o||null,...a.length?{tool_calls:a}:{}}]:[]}let d=[...o?[{type:"text",text:o}]:[],...r.map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...d.length?[{role:"user",content:d}]:[]]});return{model:t.model,messages:n,stream:!0,...A(t,e=>({type:"function",function:{name:e.name,description:e.description,parameters:R(e)}}))}}var te=t=>typeof t=="string"?t:(t??[]).filter(n=>n.type==="text").map(n=>n.text??"").join("");async function*j(t){let n=[];for await(let e of t)for(let s of e.choices??[]){let i=te(s.delta?.content);i&&(yield new c.LanguageModelTextPart(i));for(let o of s.delta?.tool_calls??[]){let r=n[o.index??0]??={id:"",name:"",args:""};r.id=o.id||r.id,r.name=o.function?.name||r.name,r.args+=o.function?.arguments??""}if(s.finish_reason==="length")throw P()}for(let e of n.filter(Boolean))yield L(e.id,e.name,e.args)}function D(t){let n=t.messages.flatMap(e=>{let{results:s,calls:i,text:o,images:r}=O(e.content);if(Number(e.role)===S)return[{role:"system",content:o}];if(e.role===c.LanguageModelChatMessageRole.Assistant)return[...o?[{role:"assistant",content:o}]:[],...i.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let d=[...o?[{type:"input_text",text:o}]:[],...r.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:a.content.map(y).join("")})),...d.length?[{role:"user",content:d}]:[]]});return{model:t.model,input:n,stream:!0,store:!1,...A(t,e=>({type:"function",name:e.name,description:e.description,parameters:R(e),strict:!1}))}}async function*G(t){for await(let n of t){let e=n;if(e.type==="response.output_text.delta"&&e.delta)yield new c.LanguageModelTextPart(e.delta);else if(e.type==="response.output_item.done"&&e.item?.type==="function_call")yield L(e.item.call_id??"",e.item.name??"",e.item.arguments??"");else{if(e.type==="response.incomplete"&&e.response?.incomplete_details?.reason==="max_output_tokens")throw P();if(e.type==="response.failed"||e.type==="error")throw new Error(e.response?.error?.message??e.message??"The model response failed.")}}}async function*k(t,n,e,s,i,o){let r=m=>fetch(`${t}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...i,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(e),signal:o}),d=await s.get(),a=await r(d);if(a.status===401&&(await a.body?.cancel(),d=await s.renew(d),a=await r(d)),a.status===401)throw await a.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${t}\`, then retry.`);if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let x=g.split(` -`);g=x.pop()??"";for(let J of x){let w=N(J);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=N(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var N=t=>t.startsWith("data:")?t.slice(5).trim():"";var ne=t=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let e=t[n];return typeof e=="number"&&Number.isInteger(e)&&e>0?[[n,e]]:[]})),$=t=>{let n=t("models");return{workspace:String(t("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(e=>typeof e?.id=="string"&&e.id!==""&&(e.api==="chat_completions"||e.api==="responses")).map(e=>({id:e.id,api:e.api,...ne(e)})),authCommand:String(t("authCommand")??"")}},q=t=>`ucode/${b} githubcopilotchat${typeof t=="string"?`/${t}`:""}`,v=class{constructor(n,e,s=()=>({})){this.settings=n;this.log=e;this.headers=s}settings;log;headers;changed=new l.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:e,authCommand:s}=this.settings();return!n||!s||e.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):e.map(({id:i,api:o,maxInputTokens:r,maxOutputTokens:d})=>({id:i,api:o,name:i,family:i.split(".").pop()??i,version:b,detail:"Databricks",maxInputTokens:r??128e3,maxOutputTokens:d??32e3,capabilities:{toolCalling:!0,imageInput:o==="responses"||i.includes("claude")}}))}async provideLanguageModelChatResponse(n,e,s,i,o){let{workspace:r,authCommand:d}=this.settings();this.tokens??=new h(d,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:e,tools:s.tools??[],toolRequired:s.toolMode===l.LanguageModelChatToolMode.Required},p=new AbortController,g=o.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?G(k(r,"responses",D(a),this.tokens,this.headers(),p.signal)):j(k(r,"chat/completions",I(a),this.tokens,this.headers(),p.signal));for await(let m of u)i.report(m)}catch(u){if(!o.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,e){return E(e)}};function oe(t){let n=l.window.createOutputChannel("Unity Gateway",{log:!0}),e=()=>{let o=l.workspace.getConfiguration("unityGateway");return $(r=>o.get(r))},s=()=>({"User-Agent":q(l.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),i=new v(e,o=>n.info(o),s);t.subscriptions.push(n,l.lm.registerLanguageModelChatProvider("databricks-unity-gateway",i),l.commands.registerCommand("unityGateway.refreshModels",()=>i.refresh()),l.workspace.onDidChangeConfiguration(o=>{o.affectsConfiguration("unityGateway")&&i.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); +"use strict";var F=Object.create;var f=Object.defineProperty;var B=Object.getOwnPropertyDescriptor;var H=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,K=Object.prototype.hasOwnProperty;var Y=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},M=(e,n,t,o)=>{if(n&&typeof n=="object"||typeof n=="function")for(let a of H(n))!K.call(e,a)&&a!==t&&f(e,a,{get:()=>n[a],enumerable:!(o=B(n,a))||o.enumerable});return e};var C=(e,n,t)=>(t=e!=null?F(W(e)):{},M(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),z=e=>M(f({},"__esModule",{value:!0}),e);var re={};Y(re,{DatabricksModels:()=>v,activate:()=>ae,readSettings:()=>q,userAgent:()=>J});module.exports=z(re);var d=C(require("vscode"));var b="0.1.0";var T=require("node:child_process"),Q=5*6e4,V=60*6e4,Z=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(o,a,s)=>{o?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(s).trim().slice(0,500)||o.message}`)):n(String(a).trim())})}),ee=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+V},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-Q>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=Z(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:ee(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var c=C(require("vscode")),_=3,te=1e3,y=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",S=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let o;try{o=t.trim()?JSON.parse(t):{}}catch{o=void 0}if(typeof o!="object"||o===null||Array.isArray(o))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new c.LanguageModelToolCallPart(e,n,o)},P=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),E=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},I=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let o of e.content){let a=o instanceof c.LanguageModelToolResultPart?o.content:[o];for(let s of a)n+=s instanceof c.LanguageModelToolCallPart?s.name.length+JSON.stringify(s.input).length:y(s).length,t+=S(s)===void 0?0:1}return Math.ceil(n/4)+t*te},O=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(S).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:o,calls:a,text:s,images:l}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===c.LanguageModelChatMessageRole.Assistant){let r=a.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return s||r.length?[{role:"assistant",content:s||null,...r.length?{tool_calls:r}:{}}]:[]}let i=[...s?[{type:"text",text:s}]:[],...l.map(r=>({type:"image_url",image_url:{url:r}}))];return[...o.map(r=>({role:"tool",tool_call_id:r.callId,content:r.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}var ne=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),oe=new Set(["stop","tool_calls","function_call"]);async function*j(e){let n=[],t;for await(let o of e)for(let a of o.choices??[]){let s=ne(a.delta?.content);s&&(yield new c.LanguageModelTextPart(s));for(let l of a.delta?.tool_calls??[]){let i=n[l.index??0]??={id:"",name:"",args:""};i.id=l.id||i.id,i.name=l.function?.name||i.name,i.args+=l.function?.arguments??""}if(a.finish_reason&&(t=a.finish_reason,!oe.has(t)))throw P(t)}if(!t)throw E();for(let o of n.filter(Boolean))yield L(o.id,o.name,o.args)}function N(e){let n=e.messages.flatMap(t=>{let{results:o,calls:a,text:s,images:l}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===c.LanguageModelChatMessageRole.Assistant)return[...s?[{role:"assistant",content:s}]:[],...a.map(r=>({type:"function_call",call_id:r.callId,name:r.name,arguments:JSON.stringify(r.input)}))];let i=[...s?[{type:"input_text",text:s}]:[],...l.map(r=>({type:"input_image",image_url:r,detail:"auto"}))];return[...o.map(r=>({type:"function_call_output",call_id:r.callId,output:r.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*$(e){let n=!1;for await(let t of e){let o=t;if(o.type==="response.output_text.delta"&&o.delta)yield new c.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")yield L(o.item.call_id??"",o.item.name??"",o.item.arguments??"");else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw P(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw E()}async function*k(e,n,t,o,a,s){let l=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...a,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:s}),i=await o.get(),r=await l(i);if(r.status===401&&(await r.body?.cancel(),i=await o.renew(i),r=await l(i)),r.status===401)throw await r.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!r.ok||!r.body)throw new Error(`AI Gateway returned ${r.status}: ${(await r.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of r.body){g+=p.decode(m,{stream:!0});let x=g.split(` +`);g=x.pop()??"";for(let U of x){let w=G(U);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=G(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var G=e=>e.startsWith("data:")?e.slice(5).trim():"";var se=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),q=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,api:t.api,...se(t)})),authCommand:String(e("authCommand")??"")}},J=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(n,t,o=()=>({})){this.settings=n;this.log=t;this.headers=o}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:o}=this.settings();return!n||!o||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:a,api:s,maxInputTokens:l,maxOutputTokens:i})=>({id:a,api:s,name:a,family:a.split(".").pop()??a,version:b,detail:"Databricks",maxInputTokens:l??128e3,maxOutputTokens:i??32e3,capabilities:{toolCalling:!0,imageInput:s==="responses"||a.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,o,a,s){let{workspace:l,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let r={model:n.id,messages:t,tools:o.tools??[],toolRequired:o.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=s.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(k(l,"responses",N(r),this.tokens,this.headers(),p.signal)):j(k(l,"chat/completions",D(r),this.tokens,this.headers(),p.signal));for await(let m of u)a.report(m)}catch(u){if(!s.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return I(t)}};function ae(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let s=d.workspace.getConfiguration("unityGateway");return q(l=>s.get(l))},o=()=>({"User-Agent":J(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),a=new v(t,s=>n.info(s),o);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",a),d.commands.registerCommand("unityGateway.refreshModels",()=>a.refresh()),d.workspace.onDidChangeConfiguration(s=>{s.affectsConfiguration("unityGateway")&&a.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); diff --git a/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts index d2c61f7e9..3676e7dd0 100644 --- a/vscode-extension/src/formats.ts +++ b/vscode-extension/src/formats.ts @@ -48,7 +48,16 @@ const toolCall = (callId: string, name: string, args: string): vscode.LanguageMo return new vscode.LanguageModelToolCallPart(callId, name, input); }; -const outputLimit = () => new Error("The model stopped at its output token limit before finishing the response."); +/** Why the model stopped short, from a Responses `incomplete_details.reason` or a Chat Completions `finish_reason`. */ +const stoppedEarly = (reason: string | undefined): Error => + reason === "max_output_tokens" || reason === "length" + ? new Error("The model stopped at its output token limit before finishing the response.") + : reason === "content_filter" + ? new Error("The model's response was stopped by a content filter.") + : new Error(`The model stopped before finishing the response (${reason ?? "no reason given"}).`); + +// A dropped connection ends the stream like a finished answer does, so only a final event counts as done. +const cutOff = () => new Error("The response ended before the model finished. Try again."); const toolOptions = (turn: Turn, toTool: (tool: vscode.LanguageModelChatTool) => object): object => turn.tools.length ? { tools: turn.tools.map(toTool), ...(turn.toolRequired ? { tool_choice: "required" } : {}) } : {}; @@ -128,8 +137,11 @@ const deltaText = (content: ChatContent | undefined): string => .map((block) => block.text ?? "") .join(""); +const CHAT_FINISHED = new Set(["stop", "tool_calls", "function_call"]); + export async function* chatCompletionsReplies(events: AsyncIterable): AsyncGenerator { const calls: { id: string; name: string; args: string }[] = []; + let finish: string | undefined; for await (const event of events) { for (const choice of (event as ChatChunk).choices ?? []) { const text = deltaText(choice.delta?.content); @@ -142,11 +154,17 @@ export async function* chatCompletionsReplies(events: AsyncIterable): A call.name = fragment.function?.name || call.name; call.args += fragment.function?.arguments ?? ""; } - if (choice.finish_reason === "length") { - throw outputLimit(); + if (choice.finish_reason) { + finish = choice.finish_reason; + if (!CHAT_FINISHED.has(finish)) { + throw stoppedEarly(finish); + } } } } + if (!finish) { + throw cutOff(); + } for (const call of calls.filter(Boolean)) { yield toolCall(call.id, call.name, call.args); } @@ -192,16 +210,22 @@ interface ResponseEvent { } export async function* responsesReplies(events: AsyncIterable): AsyncGenerator { + let completed = false; for await (const raw of events) { const event = raw as ResponseEvent; if (event.type === "response.output_text.delta" && event.delta) { yield new vscode.LanguageModelTextPart(event.delta); } else if (event.type === "response.output_item.done" && event.item?.type === "function_call") { yield toolCall(event.item.call_id ?? "", event.item.name ?? "", event.item.arguments ?? ""); - } else if (event.type === "response.incomplete" && event.response?.incomplete_details?.reason === "max_output_tokens") { - throw outputLimit(); + } else if (event.type === "response.completed") { + completed = true; + } else if (event.type === "response.incomplete") { + throw stoppedEarly(event.response?.incomplete_details?.reason); } else if (event.type === "response.failed" || event.type === "error") { throw new Error(event.response?.error?.message ?? event.message ?? "The model response failed."); } } + if (!completed) { + throw cutOff(); + } } diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 5e94cfcf4..3638b585b 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -140,8 +140,8 @@ describe("DatabricksModels", () => { it("sends each model's turns to its API with the User-Agent", async () => { const requests = stubFetch( - [200, sse('{"choices":[{"delta":{"content":"hello"}}]}', "[DONE]")], - [200, sse('{"type":"response.output_text.delta","delta":"hi"}')], + [200, sse('{"choices":[{"delta":{"content":"hello"},"finish_reason":"stop"}]}', "[DONE]")], + [200, sse('{"type":"response.output_text.delta","delta":"hi"}', '{"type":"response.completed","response":{}}')], ); const provider = new DatabricksModels(() => settings(tokenCommand(["t"]).command), () => {}, () => ({ "User-Agent": "ucode/0 githubcopilotchat" })); const [claude, gpt] = await provider.provideLanguageModelChatInformation(); diff --git a/vscode-extension/test/formats.test.ts b/vscode-extension/test/formats.test.ts index 9f66d8817..6afe3b8d8 100644 --- a/vscode-extension/test/formats.test.ts +++ b/vscode-extension/test/formats.test.ts @@ -55,7 +55,7 @@ describe("Chat Completions", () => { chatCompletionsReplies( events( { choices: [{ delta: { content: [{ type: "reasoning", summary: [{ type: "summary_text", text: "The file may not exist." }] }] } }] }, - { choices: [{ delta: { content: [{ type: "text", text: "Checking." }, { type: "text", text: " Done." }] } }] }, + { choices: [{ delta: { content: [{ type: "text", text: "Checking." }, { type: "text", text: " Done." }] }, finish_reason: "stop" }] }, ), ), ); @@ -65,10 +65,20 @@ describe("Chat Completions", () => { it("fails on a length stop and on arguments that aren't an object", async () => { await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "length" }] }))), "output token limit"); await rejectsWith( - collect(chatCompletionsReplies(events({ choices: [{ delta: { tool_calls: [{ index: 0, id: "c", function: { name: "t", arguments: "[1]" } }] } }] }))), + collect( + chatCompletionsReplies( + events({ choices: [{ delta: { tool_calls: [{ index: 0, id: "c", function: { name: "t", arguments: "[1]" } }] }, finish_reason: "tool_calls" }] }), + ), + ), "aren't a JSON object", ); }); + + it("fails when the model stops for another reason or the stream ends without a finish reason", async () => { + await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "content_filter" }] }))), "stopped by a content filter"); + await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "insufficient_system_resource" }] }))), "(insufficient_system_resource)"); + await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ delta: { content: "Half an ans" } }] }))), "ended before the model finished"); + }); }); describe("estimateTokens", () => { @@ -119,4 +129,11 @@ describe("Responses", () => { ); await rejectsWith(collect(responsesReplies(events({ type: "response.failed", response: { error: { message: "overloaded" } } }))), "overloaded"); }); + + it("fails on every other incomplete reason and on a stream without response.completed", async () => { + const incomplete = (reason?: string) => ({ type: "response.incomplete", response: { incomplete_details: reason ? { reason } : null } }); + await rejectsWith(collect(responsesReplies(events(incomplete("content_filter")))), "stopped by a content filter"); + await rejectsWith(collect(responsesReplies(events(incomplete()))), "(no reason given)"); + await rejectsWith(collect(responsesReplies(events({ type: "response.output_text.delta", delta: "Half an ans" }))), "ended before the model finished"); + }); }); From a5759ad5be4611ffee1d1d8f90ca32e85525ba2f Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 10:35:56 +0000 Subject: [PATCH 10/18] Show readable model names in VS Code's model picker The picker listed models by their service ids, e.g. system.ai.claude-sonnet-5-5. ug configure now writes a name for each model with model_picker_label, the helper Codex's picker already uses ("Claude Sonnet 5.5", "GPT 6.1 Sol"), and the extension shows it. As in Codex, two models with the same label keep their full ids. The id still goes to the gateway, and settings without a name show the id as before. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 13 +++++++++-- src/ucode/vscode/extension/out/extension.js | 4 ++-- src/ucode/vscode/extension/package.json | 4 ++++ tests/test_agent_vscode.py | 24 ++++++++++++++++++++- vscode-extension/package.json | 4 ++++ vscode-extension/src/extension.ts | 12 ++++++++--- vscode-extension/test/extension.test.ts | 18 ++++++++-------- 7 files changed, 62 insertions(+), 17 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 12db79017..539623c5b 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -12,6 +12,7 @@ import shutil import tempfile import zipfile +from collections import Counter from pathlib import Path from typing import cast from xml.sax.saxutils import escape @@ -27,6 +28,7 @@ from ucode.managed_files import OS, current_os from ucode.os_compatibility import subprocess_cross_os from ucode.os_compatibility.system_cross_os import windows_program_files_dir +from ucode.string_utils import model_picker_label from ucode.ui import console, print_heading, print_kv, print_warning EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" @@ -157,9 +159,10 @@ def install_extension(code: Path) -> str: return f"installed {EXTENSION_ID} {version}" -def _model_entry(model: str) -> dict[str, object]: +def _model_entry(model: str, name: str) -> dict[str, object]: entry: dict[str, object] = { "id": model, + "name": name, "api": "responses" if uses_responses_api(model) else "chat_completions", } # Claude's limits hold on Chat Completions too. Models without known limits get the @@ -179,7 +182,13 @@ def model_settings(workspace: str, token: str) -> tuple[list[dict[str, object]], for model in dict.fromkeys([*claude.values(), *gpt]) if copilot_exclusion_reason(model) is None ] - return [_model_entry(model) for model in models], reason + labels = {model: model_picker_label(model.rsplit(".", 1)[-1]) for model in models} + counts = Counter(labels.values()) + # As in Codex's picker: two models with the same label show their full ids instead. + return [ + _model_entry(model, model if counts[label] > 1 else label) + for model, label in labels.items() + ], reason def auth_command(state: dict) -> str: diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 4f32a40c5..642e7a828 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var F=Object.create;var f=Object.defineProperty;var B=Object.getOwnPropertyDescriptor;var H=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,K=Object.prototype.hasOwnProperty;var Y=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},M=(e,n,t,o)=>{if(n&&typeof n=="object"||typeof n=="function")for(let a of H(n))!K.call(e,a)&&a!==t&&f(e,a,{get:()=>n[a],enumerable:!(o=B(n,a))||o.enumerable});return e};var C=(e,n,t)=>(t=e!=null?F(W(e)):{},M(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),z=e=>M(f({},"__esModule",{value:!0}),e);var re={};Y(re,{DatabricksModels:()=>v,activate:()=>ae,readSettings:()=>q,userAgent:()=>J});module.exports=z(re);var d=C(require("vscode"));var b="0.1.0";var T=require("node:child_process"),Q=5*6e4,V=60*6e4,Z=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(o,a,s)=>{o?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(s).trim().slice(0,500)||o.message}`)):n(String(a).trim())})}),ee=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+V},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-Q>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=Z(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:ee(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var c=C(require("vscode")),_=3,te=1e3,y=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",S=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let o;try{o=t.trim()?JSON.parse(t):{}}catch{o=void 0}if(typeof o!="object"||o===null||Array.isArray(o))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new c.LanguageModelToolCallPart(e,n,o)},P=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),E=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},I=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let o of e.content){let a=o instanceof c.LanguageModelToolResultPart?o.content:[o];for(let s of a)n+=s instanceof c.LanguageModelToolCallPart?s.name.length+JSON.stringify(s.input).length:y(s).length,t+=S(s)===void 0?0:1}return Math.ceil(n/4)+t*te},O=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(S).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:o,calls:a,text:s,images:l}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===c.LanguageModelChatMessageRole.Assistant){let r=a.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return s||r.length?[{role:"assistant",content:s||null,...r.length?{tool_calls:r}:{}}]:[]}let i=[...s?[{type:"text",text:s}]:[],...l.map(r=>({type:"image_url",image_url:{url:r}}))];return[...o.map(r=>({role:"tool",tool_call_id:r.callId,content:r.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}var ne=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),oe=new Set(["stop","tool_calls","function_call"]);async function*j(e){let n=[],t;for await(let o of e)for(let a of o.choices??[]){let s=ne(a.delta?.content);s&&(yield new c.LanguageModelTextPart(s));for(let l of a.delta?.tool_calls??[]){let i=n[l.index??0]??={id:"",name:"",args:""};i.id=l.id||i.id,i.name=l.function?.name||i.name,i.args+=l.function?.arguments??""}if(a.finish_reason&&(t=a.finish_reason,!oe.has(t)))throw P(t)}if(!t)throw E();for(let o of n.filter(Boolean))yield L(o.id,o.name,o.args)}function N(e){let n=e.messages.flatMap(t=>{let{results:o,calls:a,text:s,images:l}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===c.LanguageModelChatMessageRole.Assistant)return[...s?[{role:"assistant",content:s}]:[],...a.map(r=>({type:"function_call",call_id:r.callId,name:r.name,arguments:JSON.stringify(r.input)}))];let i=[...s?[{type:"input_text",text:s}]:[],...l.map(r=>({type:"input_image",image_url:r,detail:"auto"}))];return[...o.map(r=>({type:"function_call_output",call_id:r.callId,output:r.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*$(e){let n=!1;for await(let t of e){let o=t;if(o.type==="response.output_text.delta"&&o.delta)yield new c.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")yield L(o.item.call_id??"",o.item.name??"",o.item.arguments??"");else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw P(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw E()}async function*k(e,n,t,o,a,s){let l=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...a,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:s}),i=await o.get(),r=await l(i);if(r.status===401&&(await r.body?.cancel(),i=await o.renew(i),r=await l(i)),r.status===401)throw await r.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!r.ok||!r.body)throw new Error(`AI Gateway returned ${r.status}: ${(await r.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of r.body){g+=p.decode(m,{stream:!0});let x=g.split(` -`);g=x.pop()??"";for(let U of x){let w=G(U);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=G(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var G=e=>e.startsWith("data:")?e.slice(5).trim():"";var se=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),q=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,api:t.api,...se(t)})),authCommand:String(e("authCommand")??"")}},J=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(n,t,o=()=>({})){this.settings=n;this.log=t;this.headers=o}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:o}=this.settings();return!n||!o||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:a,api:s,maxInputTokens:l,maxOutputTokens:i})=>({id:a,api:s,name:a,family:a.split(".").pop()??a,version:b,detail:"Databricks",maxInputTokens:l??128e3,maxOutputTokens:i??32e3,capabilities:{toolCalling:!0,imageInput:s==="responses"||a.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,o,a,s){let{workspace:l,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let r={model:n.id,messages:t,tools:o.tools??[],toolRequired:o.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=s.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(k(l,"responses",N(r),this.tokens,this.headers(),p.signal)):j(k(l,"chat/completions",D(r),this.tokens,this.headers(),p.signal));for await(let m of u)a.report(m)}catch(u){if(!s.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return I(t)}};function ae(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let s=d.workspace.getConfiguration("unityGateway");return q(l=>s.get(l))},o=()=>({"User-Agent":J(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),a=new v(t,s=>n.info(s),o);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",a),d.commands.registerCommand("unityGateway.refreshModels",()=>a.refresh()),d.workspace.onDidChangeConfiguration(s=>{s.affectsConfiguration("unityGateway")&&a.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); +"use strict";var F=Object.create;var f=Object.defineProperty;var B=Object.getOwnPropertyDescriptor;var H=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,K=Object.prototype.hasOwnProperty;var Y=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},C=(e,n,t,o)=>{if(n&&typeof n=="object"||typeof n=="function")for(let r of H(n))!K.call(e,r)&&r!==t&&f(e,r,{get:()=>n[r],enumerable:!(o=B(n,r))||o.enumerable});return e};var M=(e,n,t)=>(t=e!=null?F(W(e)):{},C(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),z=e=>C(f({},"__esModule",{value:!0}),e);var re={};Y(re,{DatabricksModels:()=>v,activate:()=>ae,readSettings:()=>q,userAgent:()=>J});module.exports=z(re);var d=M(require("vscode"));var b="0.1.0";var T=require("node:child_process"),Q=5*6e4,V=60*6e4,Z=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(o,r,s)=>{o?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(s).trim().slice(0,500)||o.message}`)):n(String(r).trim())})}),ee=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+V},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-Q>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=Z(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:ee(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var l=M(require("vscode")),S=3,te=1e3,y=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",_=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let o;try{o=t.trim()?JSON.parse(t):{}}catch{o=void 0}if(typeof o!="object"||o===null||Array.isArray(o))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new l.LanguageModelToolCallPart(e,n,o)},P=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),E=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},I=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let o of e.content){let r=o instanceof l.LanguageModelToolResultPart?o.content:[o];for(let s of r)n+=s instanceof l.LanguageModelToolCallPart?s.name.length+JSON.stringify(s.input).length:y(s).length,t+=_(s)===void 0?0:1}return Math.ceil(n/4)+t*te},O=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(_).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===S)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant){let a=r.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return s||a.length?[{role:"assistant",content:s||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...s?[{type:"text",text:s}]:[],...c.map(a=>({type:"image_url",image_url:{url:a}}))];return[...o.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}var ne=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),oe=new Set(["stop","tool_calls","function_call"]);async function*j(e){let n=[],t;for await(let o of e)for(let r of o.choices??[]){let s=ne(r.delta?.content);s&&(yield new l.LanguageModelTextPart(s));for(let c of r.delta?.tool_calls??[]){let i=n[c.index??0]??={id:"",name:"",args:""};i.id=c.id||i.id,i.name=c.function?.name||i.name,i.args+=c.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!oe.has(t)))throw P(t)}if(!t)throw E();for(let o of n.filter(Boolean))yield L(o.id,o.name,o.args)}function N(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===S)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant)return[...s?[{role:"assistant",content:s}]:[],...r.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...s?[{type:"input_text",text:s}]:[],...c.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...o.map(a=>({type:"function_call_output",call_id:a.callId,output:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*$(e){let n=!1;for await(let t of e){let o=t;if(o.type==="response.output_text.delta"&&o.delta)yield new l.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")yield L(o.item.call_id??"",o.item.name??"",o.item.arguments??"");else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw P(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw E()}async function*k(e,n,t,o,r,s){let c=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...r,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:s}),i=await o.get(),a=await c(i);if(a.status===401&&(await a.body?.cancel(),i=await o.renew(i),a=await c(i)),a.status===401)throw await a.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let x=g.split(` +`);g=x.pop()??"";for(let U of x){let w=G(U);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=G(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var G=e=>e.startsWith("data:")?e.slice(5).trim():"";var se=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),q=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},api:t.api,...se(t)})),authCommand:String(e("authCommand")??"")}},J=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(n,t,o=()=>({})){this.settings=n;this.log=t;this.headers=o}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:o}=this.settings();return!n||!o||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:r,name:s,api:c,maxInputTokens:i,maxOutputTokens:a})=>({id:r,api:c,name:s??r,family:r.split(".").pop()??r,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:a??32e3,capabilities:{toolCalling:!0,imageInput:c==="responses"||r.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,o,r,s){let{workspace:c,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:t,tools:o.tools??[],toolRequired:o.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=s.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(k(c,"responses",N(a),this.tokens,this.headers(),p.signal)):j(k(c,"chat/completions",D(a),this.tokens,this.headers(),p.signal));for await(let m of u)r.report(m)}catch(u){if(!s.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return I(t)}};function ae(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let s=d.workspace.getConfiguration("unityGateway");return q(c=>s.get(c))},o=()=>({"User-Agent":J(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),r=new v(t,s=>n.info(s),o);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",r),d.commands.registerCommand("unityGateway.refreshModels",()=>r.refresh()),d.workspace.onDidChangeConfiguration(s=>{s.affectsConfiguration("unityGateway")&&r.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json index 8bc404b46..700568df7 100644 --- a/src/ucode/vscode/extension/package.json +++ b/src/ucode/vscode/extension/package.json @@ -53,6 +53,10 @@ "type": "string", "description": "Model service name, for example system.ai.claude-sonnet-5-5." }, + "name": { + "type": "string", + "description": "Name shown in the model picker, for example Claude Sonnet 5.5. Defaults to the id." + }, "api": { "type": "string", "enum": [ diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 979144613..6dd50ecb7 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -105,15 +105,37 @@ def test_merges_settings_keeping_the_users_own_with_claude_and_gpt_only(self, co "unityGateway.models": [ { "id": "system.ai.claude-sonnet-5-5", + "name": "Claude Sonnet 5.5", "api": "chat_completions", "maxInputTokens": 872_000, "maxOutputTokens": 128_000, }, - {"id": "system.ai.gpt-6-luna", "api": "responses"}, + {"id": "system.ai.gpt-6-luna", "name": "GPT 6 Luna", "api": "responses"}, ], "unityGateway.authCommand": vscode.build_auth_shell_command(WS, "dev"), } + def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): + monkeypatch.setattr( + vscode, + "discover_model_services", + lambda workspace, token: ( + {}, + ["system.ai.gpt-6-luna", "main.models.gpt-6-luna", "system.ai.gpt-5-5"], + [], + [], + None, + ), + ) + + models, _ = vscode.model_settings(WS, "token") + + assert [model["name"] for model in models] == [ + "system.ai.gpt-6-luna", + "main.models.gpt-6-luna", + "GPT 5.5", + ] + def test_creates_missing_settings_and_installs_the_packed_extension(self, code, capsys): assert vscode.configure(STATE) diff --git a/vscode-extension/package.json b/vscode-extension/package.json index 8bc404b46..700568df7 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -53,6 +53,10 @@ "type": "string", "description": "Model service name, for example system.ai.claude-sonnet-5-5." }, + "name": { + "type": "string", + "description": "Name shown in the model picker, for example Claude Sonnet 5.5. Defaults to the id." + }, "api": { "type": "string", "enum": [ diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 991e27485..3146f0263 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -6,6 +6,7 @@ import { streamEvents } from "./gateway"; export interface ModelSetting { readonly id: string; + readonly name?: string; readonly api: "chat_completions" | "responses"; readonly maxInputTokens?: number; readonly maxOutputTokens?: number; @@ -34,7 +35,12 @@ export const readSettings = (get: (key: string) => unknown): Settings => { (model): model is ModelSetting => typeof model?.id === "string" && model.id !== "" && (model.api === "chat_completions" || model.api === "responses"), ) - .map((model) => ({ id: model.id, api: model.api, ...tokenLimits(model) })), + .map((model) => ({ + id: model.id, + ...(typeof model.name === "string" && model.name.trim() ? { name: model.name.trim() } : {}), + api: model.api, + ...tokenLimits(model), + })), authCommand: String(get("authCommand") ?? ""), }; }; @@ -68,10 +74,10 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider this.log("No Databricks models: run `ug configure --agents vscode` to set them up."); return []; } - return models.map(({ id, api, maxInputTokens, maxOutputTokens }) => ({ + return models.map(({ id, name, api, maxInputTokens, maxOutputTokens }) => ({ id, api, - name: id, + name: name ?? id, // Copilot Chat tunes its prompts and tools by family, so this is the model's own name. family: id.split(".").pop() ?? id, version, diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 3638b585b..0dafa02d5 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -89,7 +89,7 @@ describe("userAgent", () => { }); describe("readSettings", () => { - it("keeps well-formed models and token limits and trims the workspace URL", () => { + it("keeps well-formed models, names and token limits and trims the workspace URL", () => { const values: Record = { workspace: "https://ws/", authCommand: "ug auth-token", @@ -98,8 +98,8 @@ describe("readSettings", () => { { id: "b", api: "messages" }, { api: "responses" }, null, - { id: "c", api: "chat_completions", maxInputTokens: 936_000, maxOutputTokens: "32000", extra: true }, - { id: "d", api: "chat_completions", maxInputTokens: 0, maxOutputTokens: 1.5 }, + { id: "c", name: " Model C ", api: "chat_completions", maxInputTokens: 936_000, maxOutputTokens: "32000", extra: true }, + { id: "d", name: 7, api: "chat_completions", maxInputTokens: 0, maxOutputTokens: 1.5 }, ], }; assert.deepEqual(readSettings((key) => values[key]), { @@ -107,7 +107,7 @@ describe("readSettings", () => { authCommand: "ug auth-token", models: [ { id: "a", api: "responses" }, - { id: "c", api: "chat_completions", maxInputTokens: 936_000 }, + { id: "c", name: "Model C", api: "chat_completions", maxInputTokens: 936_000 }, { id: "d", api: "chat_completions" }, ], }); @@ -119,19 +119,19 @@ describe("DatabricksModels", () => { workspace: "https://ws", authCommand: command, models: [ - { id: "system.ai.claude-sonnet-5-5", api: "chat_completions", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, + { id: "system.ai.claude-sonnet-5-5", name: "Claude Sonnet 5.5", api: "chat_completions", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, { id: "system.ai.gpt-6-luna", api: "responses" }, ], }); - it("lists the configured models, or none with a hint when settings are missing", async () => { + it("lists the configured models by name, or by id without one, or none with a hint when settings are missing", async () => { const logs: string[] = []; const models = await new DatabricksModels(() => settings("cmd"), (line) => logs.push(line)).provideLanguageModelChatInformation(); assert.deepEqual( - models.map((model) => [model.id, model.family, model.api, model.maxInputTokens, model.maxOutputTokens]), + models.map((model) => [model.id, model.name, model.family, model.api, model.maxInputTokens, model.maxOutputTokens]), [ - ["system.ai.claude-sonnet-5-5", "claude-sonnet-5-5", "chat_completions", 872_000, 128_000], - ["system.ai.gpt-6-luna", "gpt-6-luna", "responses", 128_000, 32_000], + ["system.ai.claude-sonnet-5-5", "Claude Sonnet 5.5", "claude-sonnet-5-5", "chat_completions", 872_000, 128_000], + ["system.ai.gpt-6-luna", "system.ai.gpt-6-luna", "gpt-6-luna", "responses", 128_000, 32_000], ], ); assert.deepEqual(await new DatabricksModels(() => ({ ...settings(""), models: [] }), (line) => logs.push(line)).provideLanguageModelChatInformation(), []); From 28b428c2375cce4484b23473746507a24b5d32b0 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 10:36:51 +0000 Subject: [PATCH 11/18] Keep the VS Code provider icon in one place The build copied vscode-extension/media/databricks.svg next to the bundle, so the repo held two identical copies. The SVG has no build step and only the shipped extension uses it, so keep it only in src/ucode/vscode/extension/media/ and stop copying it. package.json still needs both copies (npm reads one, VS Code the other); mark the shipped copy as generated, like the bundle. Co-authored-by: Isaac --- .gitattributes | 1 + src/ucode/vscode/extension/out/extension.js | 2 +- src/ucode/vscode/extension/package.json | 2 +- vscode-extension/media/databricks.svg | 1 - vscode-extension/package.json | 2 +- vscode-extension/test/extension.test.ts | 3 ++- 6 files changed, 6 insertions(+), 5 deletions(-) delete mode 100644 vscode-extension/media/databricks.svg diff --git a/.gitattributes b/.gitattributes index e21d1145a..47fd59c59 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ src/ucode/vscode/extension/out/extension.js linguist-generated=true +src/ucode/vscode/extension/package.json linguist-generated=true diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 642e7a828..d6a3b79c8 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var F=Object.create;var f=Object.defineProperty;var B=Object.getOwnPropertyDescriptor;var H=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,K=Object.prototype.hasOwnProperty;var Y=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},C=(e,n,t,o)=>{if(n&&typeof n=="object"||typeof n=="function")for(let r of H(n))!K.call(e,r)&&r!==t&&f(e,r,{get:()=>n[r],enumerable:!(o=B(n,r))||o.enumerable});return e};var M=(e,n,t)=>(t=e!=null?F(W(e)):{},C(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),z=e=>C(f({},"__esModule",{value:!0}),e);var re={};Y(re,{DatabricksModels:()=>v,activate:()=>ae,readSettings:()=>q,userAgent:()=>J});module.exports=z(re);var d=M(require("vscode"));var b="0.1.0";var T=require("node:child_process"),Q=5*6e4,V=60*6e4,Z=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(o,r,s)=>{o?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(s).trim().slice(0,500)||o.message}`)):n(String(r).trim())})}),ee=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+V},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-Q>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=Z(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:ee(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var l=M(require("vscode")),S=3,te=1e3,y=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",_=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let o;try{o=t.trim()?JSON.parse(t):{}}catch{o=void 0}if(typeof o!="object"||o===null||Array.isArray(o))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new l.LanguageModelToolCallPart(e,n,o)},P=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),E=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},I=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let o of e.content){let r=o instanceof l.LanguageModelToolResultPart?o.content:[o];for(let s of r)n+=s instanceof l.LanguageModelToolCallPart?s.name.length+JSON.stringify(s.input).length:y(s).length,t+=_(s)===void 0?0:1}return Math.ceil(n/4)+t*te},O=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(_).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===S)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant){let a=r.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return s||a.length?[{role:"assistant",content:s||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...s?[{type:"text",text:s}]:[],...c.map(a=>({type:"image_url",image_url:{url:a}}))];return[...o.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}var ne=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),oe=new Set(["stop","tool_calls","function_call"]);async function*j(e){let n=[],t;for await(let o of e)for(let r of o.choices??[]){let s=ne(r.delta?.content);s&&(yield new l.LanguageModelTextPart(s));for(let c of r.delta?.tool_calls??[]){let i=n[c.index??0]??={id:"",name:"",args:""};i.id=c.id||i.id,i.name=c.function?.name||i.name,i.args+=c.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!oe.has(t)))throw P(t)}if(!t)throw E();for(let o of n.filter(Boolean))yield L(o.id,o.name,o.args)}function N(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===S)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant)return[...s?[{role:"assistant",content:s}]:[],...r.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...s?[{type:"input_text",text:s}]:[],...c.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...o.map(a=>({type:"function_call_output",call_id:a.callId,output:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*$(e){let n=!1;for await(let t of e){let o=t;if(o.type==="response.output_text.delta"&&o.delta)yield new l.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")yield L(o.item.call_id??"",o.item.name??"",o.item.arguments??"");else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw P(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw E()}async function*k(e,n,t,o,r,s){let c=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...r,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:s}),i=await o.get(),a=await c(i);if(a.status===401&&(await a.body?.cancel(),i=await o.renew(i),a=await c(i)),a.status===401)throw await a.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let x=g.split(` +"use strict";var F=Object.create;var f=Object.defineProperty;var B=Object.getOwnPropertyDescriptor;var H=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,K=Object.prototype.hasOwnProperty;var Y=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},C=(e,n,t,o)=>{if(n&&typeof n=="object"||typeof n=="function")for(let r of H(n))!K.call(e,r)&&r!==t&&f(e,r,{get:()=>n[r],enumerable:!(o=B(n,r))||o.enumerable});return e};var M=(e,n,t)=>(t=e!=null?F(W(e)):{},C(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),z=e=>C(f({},"__esModule",{value:!0}),e);var re={};Y(re,{DatabricksModels:()=>v,activate:()=>ae,readSettings:()=>q,userAgent:()=>J});module.exports=z(re);var d=M(require("vscode"));var b="0.1.0";var T=require("node:child_process"),Q=5*6e4,V=60*6e4,Z=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(o,r,s)=>{o?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(s).trim().slice(0,500)||o.message}`)):n(String(r).trim())})}),ee=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+V},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-Q>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=Z(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:ee(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var l=M(require("vscode")),_=3,te=1e3,y=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",S=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let o;try{o=t.trim()?JSON.parse(t):{}}catch{o=void 0}if(typeof o!="object"||o===null||Array.isArray(o))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new l.LanguageModelToolCallPart(e,n,o)},P=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),E=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},I=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let o of e.content){let r=o instanceof l.LanguageModelToolResultPart?o.content:[o];for(let s of r)n+=s instanceof l.LanguageModelToolCallPart?s.name.length+JSON.stringify(s.input).length:y(s).length,t+=S(s)===void 0?0:1}return Math.ceil(n/4)+t*te},O=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(S).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant){let a=r.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return s||a.length?[{role:"assistant",content:s||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...s?[{type:"text",text:s}]:[],...c.map(a=>({type:"image_url",image_url:{url:a}}))];return[...o.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}var ne=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),oe=new Set(["stop","tool_calls","function_call"]);async function*j(e){let n=[],t;for await(let o of e)for(let r of o.choices??[]){let s=ne(r.delta?.content);s&&(yield new l.LanguageModelTextPart(s));for(let c of r.delta?.tool_calls??[]){let i=n[c.index??0]??={id:"",name:"",args:""};i.id=c.id||i.id,i.name=c.function?.name||i.name,i.args+=c.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!oe.has(t)))throw P(t)}if(!t)throw E();for(let o of n.filter(Boolean))yield L(o.id,o.name,o.args)}function N(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant)return[...s?[{role:"assistant",content:s}]:[],...r.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...s?[{type:"input_text",text:s}]:[],...c.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...o.map(a=>({type:"function_call_output",call_id:a.callId,output:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*$(e){let n=!1;for await(let t of e){let o=t;if(o.type==="response.output_text.delta"&&o.delta)yield new l.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")yield L(o.item.call_id??"",o.item.name??"",o.item.arguments??"");else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw P(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw E()}async function*k(e,n,t,o,r,s){let c=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...r,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:s}),i=await o.get(),a=await c(i);if(a.status===401&&(await a.body?.cancel(),i=await o.renew(i),a=await c(i)),a.status===401)throw await a.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let x=g.split(` `);g=x.pop()??"";for(let U of x){let w=G(U);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=G(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var G=e=>e.startsWith("data:")?e.slice(5).trim():"";var se=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),q=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},api:t.api,...se(t)})),authCommand:String(e("authCommand")??"")}},J=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(n,t,o=()=>({})){this.settings=n;this.log=t;this.headers=o}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:o}=this.settings();return!n||!o||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:r,name:s,api:c,maxInputTokens:i,maxOutputTokens:a})=>({id:r,api:c,name:s??r,family:r.split(".").pop()??r,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:a??32e3,capabilities:{toolCalling:!0,imageInput:c==="responses"||r.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,o,r,s){let{workspace:c,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:t,tools:o.tools??[],toolRequired:o.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=s.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(k(c,"responses",N(a),this.tokens,this.headers(),p.signal)):j(k(c,"chat/completions",D(a),this.tokens,this.headers(),p.signal));for await(let m of u)r.report(m)}catch(u){if(!s.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return I(t)}};function ae(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let s=d.workspace.getConfiguration("unityGateway");return q(c=>s.get(c))},o=()=>({"User-Agent":J(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),r=new v(t,s=>n.info(s),o);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",r),d.commands.registerCommand("unityGateway.refreshModels",()=>r.refresh()),d.workspace.onDidChangeConfiguration(s=>{s.affectsConfiguration("unityGateway")&&r.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json index 700568df7..a60de8eb1 100644 --- a/src/ucode/vscode/extension/package.json +++ b/src/ucode/vscode/extension/package.json @@ -96,7 +96,7 @@ ] }, "scripts": { - "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"const fs = require('fs'); fs.copyFileSync('package.json', '../src/ucode/vscode/extension/package.json'); fs.cpSync('media', '../src/ucode/vscode/extension/media', { recursive: true })\"", + "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"require('fs').copyFileSync('package.json', '../src/ucode/vscode/extension/package.json')\"", "typecheck": "tsc", "test": "esbuild test/*.test.ts --bundle --platform=node --format=cjs --target=node22 --outdir=.test-out --alias:vscode=./test/vscode.ts --log-level=warning && node --test .test-out/*.test.js" }, diff --git a/vscode-extension/media/databricks.svg b/vscode-extension/media/databricks.svg deleted file mode 100644 index b7bd91e12..000000000 --- a/vscode-extension/media/databricks.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/vscode-extension/package.json b/vscode-extension/package.json index 700568df7..a60de8eb1 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -96,7 +96,7 @@ ] }, "scripts": { - "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"const fs = require('fs'); fs.copyFileSync('package.json', '../src/ucode/vscode/extension/package.json'); fs.cpSync('media', '../src/ucode/vscode/extension/media', { recursive: true })\"", + "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"require('fs').copyFileSync('package.json', '../src/ucode/vscode/extension/package.json')\"", "typecheck": "tsc", "test": "esbuild test/*.test.ts --bundle --platform=node --format=cjs --target=node22 --outdir=.test-out --alias:vscode=./test/vscode.ts --log-level=warning && node --test .test-out/*.test.js" }, diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 0dafa02d5..d922cbc89 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -76,8 +76,9 @@ describe("settings", () => { }); describe("provider", () => { + // The icon lives only in the folder ug ships, next to the built bundle. it("ships the icon it names", () => { - assert.ok(packageJson.contributes.languageModelChatProviders.every((provider) => existsSync(provider.icon))); + assert.ok(packageJson.contributes.languageModelChatProviders.every((provider) => existsSync(`../src/ucode/vscode/extension/${provider.icon}`))); }); }); From 918f4fa17169f58e84b4c72289bd11de0461295b Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 10:53:53 +0000 Subject: [PATCH 12/18] Keep VS Code setup failures from breaking configure and revert Review fixes for the ug side of VS Code: - A `code` command that hangs or can't start raised TimeoutExpired or OSError, which configure and revert didn't catch. It now fails like any other `code` error, with a warning. - A VS Code failure in `ug revert` no longer stops the Claude, Codex and Pi restores that run after it; it shows as "failed". - An empty model listing (a failed request, or a workspace without Claude or GPT models) wrote `unityGateway.models: []` over a working setup and printed "No models for VS Code: None". Configure now leaves VS Code unchanged and says why. - `ug revert` ran `code --list-extensions` for everyone. It now runs `code` only when ug's settings are there, since configure writes them before it installs. - `--dry-run` set VS Code up from the saved state, which a dry run doesn't update: a KeyError on a fresh machine, or the previous workspace. A dry run now leaves VS Code alone. - Settings files that can't be read or written warn instead of raising. - The macOS lookup test no longer depends on whether this Mac has VS Code in /Applications. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 40 ++++++++++++--------- src/ucode/cli.py | 12 +++++-- tests/test_agent_vscode.py | 73 ++++++++++++++++++++++++++++++++------ tests/test_cli.py | 22 ++++++++++++ 4 files changed, 119 insertions(+), 28 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 539623c5b..b0ee03b17 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -10,6 +10,7 @@ import json import os import shutil +import subprocess import tempfile import zipfile from collections import Counter @@ -132,9 +133,12 @@ def pack_vsix(directory: Path) -> Path: def _run_code(code: Path, *args: str): - return subprocess_cross_os.run( - [str(code), *args], capture_output=True, text=True, timeout=_CODE_TIMEOUT_S, check=False - ) + try: + return subprocess_cross_os.run( + [str(code), *args], capture_output=True, text=True, timeout=_CODE_TIMEOUT_S, check=False + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError(f"`{code} {args[0]}` failed: {exc}") from exc def installed_version(code: Path) -> str | None: @@ -202,16 +206,14 @@ def auth_command(state: dict) -> str: def _edit_settings(path: Path, values: dict[str, object], remove: tuple[str, ...] = ()) -> bool: - text = path.read_text(encoding="utf-8") if path.exists() else "" try: + text = path.read_text(encoding="utf-8") if path.exists() else "" updated = config_io.update_jsonc_keys(text, values, remove) - except ValueError as exc: - print_warning( - f"Couldn't update {path}: {exc}. Fix the file, then rerun `ug configure --agents vscode`." - ) + if updated != text: + config_io.write_text_file(path, updated) + except (OSError, RuntimeError, ValueError) as exc: + print_warning(f"Couldn't update {path}: {exc}. Fix the file, then rerun this command.") return False - if updated != text: - config_io.write_text_file(path, updated) return True @@ -244,7 +246,12 @@ def configure(state: dict) -> bool: workspace = state["workspace"] models, reason = model_settings(workspace, launch_token(state, workspace)) if not models: - print_warning(f"No models for VS Code: {reason}") + # Keep what an earlier configure wrote: a failed listing shouldn't wipe a working setup. + print_warning( + f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'}). " + "VS Code was left unchanged." + ) + return False path = settings_path(product) values = { WORKSPACE_SETTING: workspace, @@ -254,8 +261,6 @@ def configure(state: dict) -> bool: if not _edit_settings(path, values): return False print_kv("VS Code settings", f"{len(models)} models in {path}") - if config_io.is_dry_run(): - return True try: print_kv("VS Code extension", install_extension(code)) except RuntimeError as exc: @@ -271,19 +276,22 @@ def configure(state: dict) -> bool: def revert() -> str: """Remove ug's settings and uninstall the extension: "removed" or "unchanged".""" - changed = False + changed = configured = False for _, product in CODE_COMMANDS: path = settings_path(product) if not path.exists(): continue try: current = config_io.loads_jsonc(path.read_text(encoding="utf-8")) - except ValueError: + except (OSError, ValueError): print_warning(f"Couldn't read {path}; remove its unityGateway settings by hand.") continue if isinstance(current, dict) and any(key in current for key in SETTINGS): + configured = True changed = _edit_settings(path, {}, SETTINGS) or changed - found = find_code() + # Configure writes the settings before it installs, so without them there's nothing to + # uninstall, and revert doesn't run `code` for people who never set up VS Code. + found = find_code() if configured else None if found is not None and installed_version(found[0]) is not None: result = _run_code(found[0], "--uninstall-extension", EXTENSION_ID) if result.returncode == 0: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 2a280f93a..177cb8d30 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1582,7 +1582,12 @@ def revert() -> int: workspace_policy_result = revert_machine_policy() from ucode.agents import vscode - vscode_result = vscode.revert() + try: + vscode_result = vscode.revert() + except Exception as exc: # noqa: BLE001 -- Preserve other revert operations on failure + # VS Code's revert reads its settings files, not ug state, so a later run still finds them. + print_warning(f"Could not revert VS Code: {exc}. Fix it, then run ug revert again.") + vscode_result = "failed" results: dict[str, bool] = { # Antigravity's skills.json is shared with its app and IDE extensions, so revert @@ -4670,7 +4675,10 @@ def configure( else: # Neither model agents nor cursor -> empty/invalid --agents list. _parse_agents_option(agents) - if wants_vscode: + if wants_vscode and is_dry_run(): + # A dry run saves no workspace, so the state on disk isn't the one being configured. + print_note("Dry run: VS Code's settings and extension aren't changed.") + elif wants_vscode: from ucode.agents import vscode vscode.configure(load_state()) diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 6dd50ecb7..41e0777f8 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -18,6 +18,7 @@ # The conftest fixture stubs these out for every test; the ones here exercise the real lookups. REAL_FIND_CODE = vscode.find_code REAL_SETTINGS_PATH = vscode.settings_path +REAL_RUN_CODE = vscode._run_code WS = "https://ws.cloud.databricks.com" STATE = {"workspace": WS, "profile": "dev"} @@ -176,6 +177,34 @@ def test_warns_when_vscode_isnt_found(self, capsys): assert "`code` command wasn't found" in _printed(capsys) assert not _settings().exists() + @pytest.mark.parametrize("reason", ["the listing failed", None]) + def test_keeps_settings_when_no_models_are_found(self, code, monkeypatch, capsys, reason): + vscode.configure(STATE) + before = _settings().read_text(encoding="utf-8") + code.calls.clear() + monkeypatch.setattr( + vscode, "discover_model_services", lambda workspace, token: ({}, [], [], [], reason) + ) + + assert not vscode.configure({**STATE, "workspace": "https://other.cloud.databricks.com"}) + + assert _settings().read_text(encoding="utf-8") == before + assert code.calls == [] + assert f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'})" in ( + _printed(capsys) + ) + + def test_reports_a_code_command_that_cant_run(self, code, monkeypatch, capsys): + def hang(*args, **kwargs): + raise subprocess.TimeoutExpired("code", 120) + + monkeypatch.setattr(vscode, "_run_code", REAL_RUN_CODE) + monkeypatch.setattr(vscode.subprocess_cross_os, "run", hang) + + assert not vscode.configure(STATE) + + assert "`/fake/bin/code --install-extension` failed" in _printed(capsys) + def test_reports_a_failed_install(self, code, capsys): code.fail = "--install-extension" @@ -206,7 +235,7 @@ def test_does_nothing_when_vscode_was_never_configured(self, code): assert vscode.revert() == "unchanged" assert _settings().read_text(encoding="utf-8") == USER_SETTINGS - assert [call[0] for call in code.calls] == ["--list-extensions"] + assert code.calls == [] class TestFindCode: @@ -228,15 +257,28 @@ def test_falls_back_to_insiders(self, monkeypatch): assert vscode.find_code() == (Path("/opt/bin/code-insiders"), "Code - Insiders") - def test_finds_the_macos_app_when_code_isnt_on_path(self, monkeypatch, tmp_path): - monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + def test_looks_in_both_macos_application_folders(self, monkeypatch, tmp_path): monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) monkeypatch.setattr(vscode.Path, "home", lambda: tmp_path) - app_code = tmp_path / "Applications/Visual Studio Code.app/Contents/Resources/app/bin/code" - app_code.parent.mkdir(parents=True) - app_code.touch() + app = "Visual Studio Code.app/Contents/Resources/app/bin/code" + + assert vscode._install_locations("code") == [ + Path("/Applications") / app, + tmp_path / "Applications" / app, + ] + + def test_falls_back_to_the_first_install_location_that_exists(self, monkeypatch, tmp_path): + monkeypatch.setattr(vscode.shutil, "which", lambda name: None) + installed = tmp_path / "bin/code" + installed.parent.mkdir() + installed.touch() + monkeypatch.setattr( + vscode, + "_install_locations", + lambda command: [tmp_path / "missing/code", installed] if command == "code" else [], + ) - assert vscode.find_code() == (app_code, "Code") + assert vscode.find_code() == (installed, "Code") def test_finds_the_windows_user_install(self, monkeypatch, tmp_path): monkeypatch.setattr(vscode.shutil, "which", lambda name: None) @@ -310,7 +352,8 @@ def test_pat_and_custom_oauth(self): assert "--profile custom" in vscode.auth_command({"workspace": WS, "custom_oauth": custom}) -def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch): +@pytest.mark.parametrize("dry_run", [False, True]) +def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, dry_run): configured: list[dict] = [] shared: list[dict] = [] monkeypatch.setattr( @@ -324,8 +367,18 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch): monkeypatch.setattr("ucode.cli.load_state", lambda: STATE) monkeypatch.setattr(vscode, "configure", lambda state: configured.append(state) or True) - result = CliRunner().invoke(app, ["configure", "--agents", "vscode", "--workspace", WS]) + args = [ + "configure", + "--agents", + "vscode", + "--workspace", + WS, + *(["--dry-run"] if dry_run else []), + ] + result = CliRunner().invoke(app, args) assert result.exit_code == 0, result.output assert shared[0]["tools"] == [] - assert configured == [STATE] + # A dry run saves no workspace, so VS Code would be set up from stale or missing state. + assert configured == ([] if dry_run else [STATE]) + assert ("VS Code's settings and extension aren't changed" in result.output) == dry_run diff --git a/tests/test_cli.py b/tests/test_cli.py index c15012a7e..b0e9cf5cf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3971,6 +3971,28 @@ def test_desktop_revert_failure_is_actionable(self): assert "ug state retained" in result.output assert "Fix the Claude Desktop error and run ug revert again" in result.output + def test_vscode_revert_failure_is_reported_and_the_rest_still_reverts(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.agents.desktop.setup.revert_desktop", return_value="unchanged"), + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.claude_agent.revert_managed_settings", return_value="unchanged"), + patch("ucode.cli.codex_agent.revert_managed_config", return_value="unchanged"), + patch("ucode.cli.revert_machine_policy", return_value="unchanged"), + patch("ucode.cli.restore_file", return_value=False) as restore_file, + patch("ucode.cli.revert_legacy_shared_config", return_value=False), + patch("ucode.agents.vscode.revert", side_effect=RuntimeError("code timed out")), + patch("ucode.cli.clear_state") as clear_state, + ): + result = runner.invoke(app, ["revert"]) + + output = " ".join(result.output.split()) + assert result.exit_code == 0, result.output + assert "Could not revert VS Code: code timed out" in output + assert "VS Code extension and settings: failed" in output + assert restore_file.called + clear_state.assert_called_once() + @pytest.mark.parametrize("args", [[], ["claude"], ["codex"]]) def test_unknown_routing_config_does_not_block_launch(monkeypatch, args): From 1d768827f9c18faa4b7dd135b612a0479d0b299f Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 10:56:10 +0000 Subject: [PATCH 13/18] Handle tool images, stream errors and sign-in failures in the VS Code extension Review fixes for the extension: - Images a tool returns (a screenshot, an image file it read) were dropped: tool results were flattened to text. Responses now sends them inside the function_call_output. On Chat Completions the gateway rejects images in a tool message, so they go in the user message right after the results. - An error event in a Chat Completions stream ({"error": ...}) was ignored; it now fails the reply with the gateway's message. - Responses tool calls are held back until response.completed, so a cut-off stream can't start a tool, as on Chat Completions. - After a second 401 the message said to run `databricks auth login`, which is wrong for PAT and custom OAuth setups. It now points to `ug configure --agents vscode`, which signs in whichever way ug is set up. - The shipped package.json no longer carries the repo's build scripts and dev dependencies. - The provider test now checks the User-Agent each request sends. Co-authored-by: Isaac --- src/ucode/vscode/extension/out/extension.js | 4 +- src/ucode/vscode/extension/package.json | 11 ------ vscode-extension/package.json | 2 +- vscode-extension/src/formats.ts | 27 ++++++++++++-- vscode-extension/src/gateway.ts | 2 +- vscode-extension/test/extension.test.ts | 12 +++--- vscode-extension/test/formats.test.ts | 41 +++++++++++++++++++++ 7 files changed, 74 insertions(+), 25 deletions(-) diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index d6a3b79c8..355a0ae82 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var F=Object.create;var f=Object.defineProperty;var B=Object.getOwnPropertyDescriptor;var H=Object.getOwnPropertyNames;var W=Object.getPrototypeOf,K=Object.prototype.hasOwnProperty;var Y=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},C=(e,n,t,o)=>{if(n&&typeof n=="object"||typeof n=="function")for(let r of H(n))!K.call(e,r)&&r!==t&&f(e,r,{get:()=>n[r],enumerable:!(o=B(n,r))||o.enumerable});return e};var M=(e,n,t)=>(t=e!=null?F(W(e)):{},C(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),z=e=>C(f({},"__esModule",{value:!0}),e);var re={};Y(re,{DatabricksModels:()=>v,activate:()=>ae,readSettings:()=>q,userAgent:()=>J});module.exports=z(re);var d=M(require("vscode"));var b="0.1.0";var T=require("node:child_process"),Q=5*6e4,V=60*6e4,Z=e=>new Promise((n,t)=>{(0,T.exec)(e,{timeout:6e4,windowsHide:!0},(o,r,s)=>{o?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(s).trim().slice(0,500)||o.message}`)):n(String(r).trim())})}),ee=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+V},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-Q>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=Z(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:ee(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var l=M(require("vscode")),_=3,te=1e3,y=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",S=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,R=e=>e.inputSchema??{type:"object",properties:{}},L=(e,n,t)=>{let o;try{o=t.trim()?JSON.parse(t):{}}catch{o=void 0}if(typeof o!="object"||o===null||Array.isArray(o))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new l.LanguageModelToolCallPart(e,n,o)},P=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),E=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},I=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let o of e.content){let r=o instanceof l.LanguageModelToolResultPart?o.content:[o];for(let s of r)n+=s instanceof l.LanguageModelToolCallPart?s.name.length+JSON.stringify(s.input).length:y(s).length,t+=S(s)===void 0?0:1}return Math.ceil(n/4)+t*te},O=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(y).join(""),images:e.map(S).filter(n=>n!==void 0)});function D(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant){let a=r.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return s||a.length?[{role:"assistant",content:s||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...s?[{type:"text",text:s}]:[],...c.map(a=>({type:"image_url",image_url:{url:a}}))];return[...o.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:R(t)}}))}}var ne=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),oe=new Set(["stop","tool_calls","function_call"]);async function*j(e){let n=[],t;for await(let o of e)for(let r of o.choices??[]){let s=ne(r.delta?.content);s&&(yield new l.LanguageModelTextPart(s));for(let c of r.delta?.tool_calls??[]){let i=n[c.index??0]??={id:"",name:"",args:""};i.id=c.id||i.id,i.name=c.function?.name||i.name,i.args+=c.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!oe.has(t)))throw P(t)}if(!t)throw E();for(let o of n.filter(Boolean))yield L(o.id,o.name,o.args)}function N(e){let n=e.messages.flatMap(t=>{let{results:o,calls:r,text:s,images:c}=O(t.content);if(Number(t.role)===_)return[{role:"system",content:s}];if(t.role===l.LanguageModelChatMessageRole.Assistant)return[...s?[{role:"assistant",content:s}]:[],...r.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...s?[{type:"input_text",text:s}]:[],...c.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...o.map(a=>({type:"function_call_output",call_id:a.callId,output:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:R(t),strict:!1}))}}async function*$(e){let n=!1;for await(let t of e){let o=t;if(o.type==="response.output_text.delta"&&o.delta)yield new l.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")yield L(o.item.call_id??"",o.item.name??"",o.item.arguments??"");else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw P(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw E()}async function*k(e,n,t,o,r,s){let c=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...r,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:s}),i=await o.get(),a=await c(i);if(a.status===401&&(await a.body?.cancel(),i=await o.renew(i),a=await c(i)),a.status===401)throw await a.body?.cancel(),new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${e}\`, then retry.`);if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let x=g.split(` -`);g=x.pop()??"";for(let U of x){let w=G(U);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let u=G(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var G=e=>e.startsWith("data:")?e.slice(5).trim():"";var se=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),q=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},api:t.api,...se(t)})),authCommand:String(e("authCommand")??"")}},J=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(n,t,o=()=>({})){this.settings=n;this.log=t;this.headers=o}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:o}=this.settings();return!n||!o||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:r,name:s,api:c,maxInputTokens:i,maxOutputTokens:a})=>({id:r,api:c,name:s??r,family:r.split(".").pop()??r,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:a??32e3,capabilities:{toolCalling:!0,imageInput:c==="responses"||r.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,o,r,s){let{workspace:c,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:t,tools:o.tools??[],toolRequired:o.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=s.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(k(c,"responses",N(a),this.tokens,this.headers(),p.signal)):j(k(c,"chat/completions",D(a),this.tokens,this.headers(),p.signal));for await(let m of u)r.report(m)}catch(u){if(!s.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return I(t)}};function ae(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let s=d.workspace.getConfiguration("unityGateway");return q(c=>s.get(c))},o=()=>({"User-Agent":J(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),r=new v(t,s=>n.info(s),o);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",r),d.commands.registerCommand("unityGateway.refreshModels",()=>r.refresh()),d.workspace.onDidChangeConfiguration(s=>{s.affectsConfiguration("unityGateway")&&r.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); +"use strict";var B=Object.create;var f=Object.defineProperty;var H=Object.getOwnPropertyDescriptor;var W=Object.getOwnPropertyNames;var K=Object.getPrototypeOf,Y=Object.prototype.hasOwnProperty;var z=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},M=(e,n,t,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let o of W(n))!Y.call(e,o)&&o!==t&&f(e,o,{get:()=>n[o],enumerable:!(s=H(n,o))||s.enumerable});return e};var T=(e,n,t)=>(t=e!=null?B(K(e)):{},M(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),X=e=>M(f({},"__esModule",{value:!0}),e);var ce={};z(ce,{DatabricksModels:()=>w,activate:()=>ie,readSettings:()=>J,userAgent:()=>U});module.exports=X(ce);var d=T(require("vscode"));var b="0.1.0";var _=require("node:child_process"),V=5*6e4,Z=60*6e4,ee=e=>new Promise((n,t)=>{(0,_.exec)(e,{timeout:6e4,windowsHide:!0},(s,o,r)=>{s?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):n(String(o).trim())})}),te=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+Z},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-V>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=ee(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:te(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var c=T(require("vscode")),S=3,ne=1e3,y=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",R=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,L=e=>e.inputSchema??{type:"object",properties:{}},P=(e,n,t)=>{let s;try{s=t.trim()?JSON.parse(t):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new c.LanguageModelToolCallPart(e,n,s)},E=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),O=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},D=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let s of e.content){let o=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of o)n+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:y(r).length,t+=R(r)===void 0?0:1}return Math.ceil(n/4)+t*ne},k=e=>e.map(R).filter(n=>n!==void 0),I=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(y).join(""),images:k(e)});function j(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===S)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant){let a=o.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>k(a.content)),...l].map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:L(t)}}))}}var oe=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),se=new Set(["stop","tool_calls","function_call"]);async function*N(e){let n=[],t;for await(let s of e){let{error:o}=s;if(o)throw new Error(typeof o=="string"?o:o.message??JSON.stringify(o));for(let r of s.choices??[]){let l=oe(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=n[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!se.has(t)))throw E(t)}}if(!t)throw O();for(let s of n.filter(Boolean))yield P(s.id,s.name,s.args)}var ae=e=>{let n=e.map(y).join(""),t=k(e);return t.length?[...n?[{type:"input_text",text:n}]:[],...t.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function G(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===S)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...o.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:ae(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:L(t),strict:!1}))}}async function*$(e){let n=[],t=!1;for await(let s of e){let o=s;if(o.type==="response.output_text.delta"&&o.delta)yield new c.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")n.push(P(o.item.call_id??"",o.item.name??"",o.item.arguments??""));else if(o.type==="response.completed")t=!0;else{if(o.type==="response.incomplete")throw E(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!t)throw O();yield*n}async function*x(e,n,t,s,o,r){let l=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...o,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:r}),i=await s.get(),a=await l(i);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i)),a.status===401)throw await a.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let C=g.split(` +`);g=C.pop()??"";for(let F of C){let v=q(F);if(v==="[DONE]")return;v&&(yield JSON.parse(v))}}let u=q(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var q=e=>e.startsWith("data:")?e.slice(5).trim():"";var re=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),J=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},api:t.api,...re(t)})),authCommand:String(e("authCommand")??"")}},U=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,w=class{constructor(n,t,s=()=>({})){this.settings=n;this.log=t;this.headers=s}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:s}=this.settings();return!n||!s||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:o,name:r,api:l,maxInputTokens:i,maxOutputTokens:a})=>({id:o,api:l,name:r??o,family:o.split(".").pop()??o,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:a??32e3,capabilities:{toolCalling:!0,imageInput:l==="responses"||o.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,s,o,r){let{workspace:l,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:t,tools:s.tools??[],toolRequired:s.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=r.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(x(l,"responses",G(a),this.tokens,this.headers(),p.signal)):N(x(l,"chat/completions",j(a),this.tokens,this.headers(),p.signal));for await(let m of u)o.report(m)}catch(u){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return D(t)}};function ie(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let r=d.workspace.getConfiguration("unityGateway");return J(l=>r.get(l))},s=()=>({"User-Agent":U(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),o=new w(t,r=>n.info(r),s);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",o),d.commands.registerCommand("unityGateway.refreshModels",()=>o.refresh()),d.workspace.onDidChangeConfiguration(r=>{r.affectsConfiguration("unityGateway")&&o.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json index a60de8eb1..76027205c 100644 --- a/src/ucode/vscode/extension/package.json +++ b/src/ucode/vscode/extension/package.json @@ -94,16 +94,5 @@ "category": "Unity Gateway" } ] - }, - "scripts": { - "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"require('fs').copyFileSync('package.json', '../src/ucode/vscode/extension/package.json')\"", - "typecheck": "tsc", - "test": "esbuild test/*.test.ts --bundle --platform=node --format=cjs --target=node22 --outdir=.test-out --alias:vscode=./test/vscode.ts --log-level=warning && node --test .test-out/*.test.js" - }, - "devDependencies": { - "@types/node": "^22.20.3", - "@types/vscode": "1.125.0", - "esbuild": "^0.28.2", - "typescript": "^5.9.3" } } diff --git a/vscode-extension/package.json b/vscode-extension/package.json index a60de8eb1..761f66d0c 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -96,7 +96,7 @@ ] }, "scripts": { - "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"require('fs').copyFileSync('package.json', '../src/ucode/vscode/extension/package.json')\"", + "build": "esbuild src/extension.ts --bundle --minify --legal-comments=none --outfile=../src/ucode/vscode/extension/out/extension.js --external:vscode --format=cjs --platform=node --target=node22 && node -e \"const m = require('./package.json'); delete m.scripts; delete m.devDependencies; require('fs').writeFileSync('../src/ucode/vscode/extension/package.json', JSON.stringify(m, null, 2) + '\\n')\"", "typecheck": "tsc", "test": "esbuild test/*.test.ts --bundle --platform=node --format=cjs --target=node22 --outdir=.test-out --alias:vscode=./test/vscode.ts --log-level=warning && node --test .test-out/*.test.js" }, diff --git a/vscode-extension/src/formats.ts b/vscode-extension/src/formats.ts index 3676e7dd0..3e2be406f 100644 --- a/vscode-extension/src/formats.ts +++ b/vscode-extension/src/formats.ts @@ -79,12 +79,14 @@ export const estimateTokens = (input: string | vscode.LanguageModelChatRequestMe return Math.ceil(characters / 4) + images * IMAGE_TOKENS; }; +const imagesOf = (parts: readonly unknown[]): string[] => parts.map(imageUrl).filter((url) => url !== undefined); + /** Splits a message into its tool results, tool calls, other text, and images. */ const sortParts = (content: readonly unknown[]) => ({ results: content.filter((part) => part instanceof vscode.LanguageModelToolResultPart), calls: content.filter((part) => part instanceof vscode.LanguageModelToolCallPart), text: content.map(textOf).join(""), - images: content.map(imageUrl).filter((url) => url !== undefined), + images: imagesOf(content), }); export function chatCompletionsRequest(turn: Turn): object { @@ -101,9 +103,10 @@ export function chatCompletionsRequest(turn: Turn): object { })); return text || toolCalls.length ? [{ role: "assistant", content: text || null, ...(toolCalls.length ? { tool_calls: toolCalls } : {}) }] : []; } + // The gateway rejects images in a tool message, so a tool's images follow its result in the user message. const content = [ ...(text ? [{ type: "text", text }] : []), - ...images.map((url) => ({ type: "image_url", image_url: { url } })), + ...[...results.flatMap((result) => imagesOf(result.content)), ...images].map((url) => ({ type: "image_url", image_url: { url } })), ]; return [ ...results.map((result) => ({ role: "tool", tool_call_id: result.callId, content: result.content.map(textOf).join("") })), @@ -121,6 +124,7 @@ export function chatCompletionsRequest(turn: Turn): object { type ChatContent = string | { type?: string; text?: string }[] | null; interface ChatChunk { + error?: string | { message?: string }; choices?: { delta?: { content?: ChatContent; tool_calls?: { index?: number; id?: string; function?: { name?: string; arguments?: string } }[] }; finish_reason?: string | null; @@ -143,6 +147,10 @@ export async function* chatCompletionsReplies(events: AsyncIterable): A const calls: { id: string; name: string; args: string }[] = []; let finish: string | undefined; for await (const event of events) { + const { error } = event as ChatChunk; + if (error) { + throw new Error(typeof error === "string" ? error : (error.message ?? JSON.stringify(error))); + } for (const choice of (event as ChatChunk).choices ?? []) { const text = deltaText(choice.delta?.content); if (text) { @@ -170,6 +178,14 @@ export async function* chatCompletionsReplies(events: AsyncIterable): A } } +const functionOutput = (content: readonly unknown[]): string | object[] => { + const text = content.map(textOf).join(""); + const images = imagesOf(content); + return images.length + ? [...(text ? [{ type: "input_text", text }] : []), ...images.map((url) => ({ type: "input_image", image_url: url, detail: "auto" }))] + : text; +}; + // Stateless: the gateway keeps no response state, so every turn resends the whole conversation. export function responsesRequest(turn: Turn): object { const input = turn.messages.flatMap((message): object[] => { @@ -188,7 +204,7 @@ export function responsesRequest(turn: Turn): object { ...images.map((url) => ({ type: "input_image", image_url: url, detail: "auto" })), ]; return [ - ...results.map((result) => ({ type: "function_call_output", call_id: result.callId, output: result.content.map(textOf).join("") })), + ...results.map((result) => ({ type: "function_call_output", call_id: result.callId, output: functionOutput(result.content) })), ...(content.length ? [{ role: "user", content }] : []), ]; }); @@ -210,13 +226,15 @@ interface ResponseEvent { } export async function* responsesReplies(events: AsyncIterable): AsyncGenerator { + // Held back until the response completes, so a cut-off stream runs no tool. + const calls: vscode.LanguageModelToolCallPart[] = []; let completed = false; for await (const raw of events) { const event = raw as ResponseEvent; if (event.type === "response.output_text.delta" && event.delta) { yield new vscode.LanguageModelTextPart(event.delta); } else if (event.type === "response.output_item.done" && event.item?.type === "function_call") { - yield toolCall(event.item.call_id ?? "", event.item.name ?? "", event.item.arguments ?? ""); + calls.push(toolCall(event.item.call_id ?? "", event.item.name ?? "", event.item.arguments ?? "")); } else if (event.type === "response.completed") { completed = true; } else if (event.type === "response.incomplete") { @@ -228,4 +246,5 @@ export async function* responsesReplies(events: AsyncIterable): AsyncGe if (!completed) { throw cutOff(); } + yield* calls; } diff --git a/vscode-extension/src/gateway.ts b/vscode-extension/src/gateway.ts index 96d5eafdd..5db9b7959 100644 --- a/vscode-extension/src/gateway.ts +++ b/vscode-extension/src/gateway.ts @@ -25,7 +25,7 @@ export async function* streamEvents( } if (response.status === 401) { await response.body?.cancel(); - throw new Error(`Databricks rejected the sign-in. Run \`databricks auth login --host ${workspace}\`, then retry.`); + throw new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry."); } if (!response.ok || !response.body) { throw new Error(`AI Gateway returned ${response.status}: ${(await response.text()).slice(0, 300)}`); diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index d922cbc89..70fc83ba3 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -13,10 +13,10 @@ const sse = (...data: string[]) => data.map((item) => `data: ${item}\n\n`).join( /** Replaces fetch with one that answers each request with the next status and SSE body. */ const stubFetch = (...answers: [number, string][]) => { - const requests: { url: string; authorization: string; body: Record }[] = []; + const requests: { url: string; authorization: string; userAgent: string; body: Record }[] = []; globalThis.fetch = (async (url: string, init: RequestInit) => { const headers = new Headers(init.headers); - requests.push({ url, authorization: headers.get("authorization") ?? "", body: JSON.parse(String(init.body)) }); + requests.push({ url, authorization: headers.get("authorization") ?? "", userAgent: headers.get("user-agent") ?? "", body: JSON.parse(String(init.body)) }); const [status, body] = answers[Math.min(requests.length - 1, answers.length - 1)] ?? [500, ""]; // Split the body mid-line so the reader has to join chunks. const bytes = new TextEncoder().encode(body); @@ -53,7 +53,7 @@ describe("streamEvents", () => { requests = stubFetch([401, ""]); await rejectsWith( collect(streamEvents("https://ws", "responses", {}, new Tokens(tokenCommand(["a", "b"]).command), {}, new AbortController().signal)), - "Run `databricks auth login --host https://ws`", + "Run `ug configure --agents vscode` to sign in again", ); assert.equal(requests.length, 2); }); @@ -157,10 +157,10 @@ describe("DatabricksModels", () => { ); } assert.deepEqual( - requests.map((request) => [request.url, request.body.model]), + requests.map((request) => [request.url, request.body.model, request.userAgent]), [ - ["https://ws/ai-gateway/mlflow/v1/chat/completions", "system.ai.claude-sonnet-5-5"], - ["https://ws/ai-gateway/mlflow/v1/responses", "system.ai.gpt-6-luna"], + ["https://ws/ai-gateway/mlflow/v1/chat/completions", "system.ai.claude-sonnet-5-5", "ucode/0 githubcopilotchat"], + ["https://ws/ai-gateway/mlflow/v1/responses", "system.ai.gpt-6-luna", "ucode/0 githubcopilotchat"], ], ); assert.deepEqual(reported, [new LanguageModelTextPart("hello"), new LanguageModelTextPart("hi")]); diff --git a/vscode-extension/test/formats.test.ts b/vscode-extension/test/formats.test.ts index 6afe3b8d8..6abf2c81d 100644 --- a/vscode-extension/test/formats.test.ts +++ b/vscode-extension/test/formats.test.ts @@ -28,6 +28,14 @@ describe("Chat Completions", () => { }); }); + it("sends a tool's images in the user message after its result, since a tool message can't carry them", () => { + const request = chatCompletionsRequest({ ...turn, messages: [message(USER, result("c1", text("Screenshot:"), image(9)))] }); + assert.deepEqual((request as { messages: unknown[] }).messages, [ + { role: "tool", tool_call_id: "c1", content: "Screenshot:" }, + { role: "user", content: [{ type: "image_url", image_url: { url: "data:image/png;base64,CQ==" } }] }, + ]); + }); + it("gives a tool without an input schema an empty one, and omits tools when there are none", () => { const noSchema = chatCompletionsRequest({ ...turn, tools: [{ name: "list", description: "List", inputSchema: undefined }], toolRequired: false }); assert.deepEqual((noSchema as { tools: unknown[] }).tools, [ @@ -74,6 +82,11 @@ describe("Chat Completions", () => { ); }); + it("fails on an error event in the stream", async () => { + await rejectsWith(collect(chatCompletionsReplies(events({ error: { message: "overloaded" } }))), "overloaded"); + await rejectsWith(collect(chatCompletionsReplies(events({ error: "rate limited" }))), "rate limited"); + }); + it("fails when the model stops for another reason or the stream ends without a finish reason", async () => { await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "content_filter" }] }))), "stopped by a content filter"); await rejectsWith(collect(chatCompletionsReplies(events({ choices: [{ finish_reason: "insufficient_system_resource" }] }))), "(insufficient_system_resource)"); @@ -108,6 +121,34 @@ describe("Responses", () => { }); }); + it("sends a tool's images inside its output", () => { + const request = responsesRequest({ ...turn, messages: [message(USER, result("c1", text("Screenshot:"), image(9)))] }); + assert.deepEqual((request as { input: unknown[] }).input, [ + { + type: "function_call_output", + call_id: "c1", + output: [ + { type: "input_text", text: "Screenshot:" }, + { type: "input_image", image_url: "data:image/png;base64,CQ==", detail: "auto" }, + ], + }, + ]); + }); + + it("holds function calls back until the response completes", async () => { + const seen: unknown[] = []; + const call = { type: "response.output_item.done", item: { type: "function_call", call_id: "c1", name: "get_weather", arguments: "{}" } }; + await rejectsWith( + (async () => { + for await (const reply of responsesReplies(events({ type: "response.output_text.delta", delta: "Checking." }, call))) { + seen.push(reply); + } + })(), + "ended before the model finished", + ); + assert.deepEqual(seen, [new LanguageModelTextPart("Checking.")]); + }); + it("streams text and each finished function call", async () => { const replies = await collect( responsesReplies( From 64c51728e3dcecfc353934f3a5816bd2bfdc2362 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 13:45:43 +0000 Subject: [PATCH 14/18] Check the managed workspace before configuring VS Code on its own `ug configure --agents vscode` (with no model agent) set up the workspace and VS Code without reading the machine's installed workspace policy, so on a machine IT manages for workspace A, `--agents vscode --workspace B` sent Copilot Chat's model traffic to B. Every other agent that isn't Claude or Codex goes through _read_policy_for_other_agent, which refuses that and keeps the selection on A. Run the same check, and the --user check, before writing VS Code's settings. Co-authored-by: Isaac --- src/ucode/cli.py | 6 +++++- tests/test_agent_vscode.py | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 177cb8d30..40b04f5f8 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -4681,7 +4681,11 @@ def configure( elif wants_vscode: from ucode.agents import vscode - vscode.configure(load_state()) + # VS Code sends model traffic, so it gets the managed-workspace check other agents get. + state = load_state() + managed, _ = _read_policy_for_other_agent(state) + refuse_user_scope_over_managed_config(state.get("workspace"), managed) + vscode.configure(state) elif mcp is not None: # MCP-only: `--mcp` without --agent(s) (e.g. Cursor, which isn't a # model agent, or adding MCP servers to an already-configured setup). diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 41e0777f8..4a70eb97d 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -352,6 +352,27 @@ def test_pat_and_custom_oauth(self): assert "--profile custom" in vscode.auth_command({"workspace": WS, "custom_oauth": custom}) +def test_configure_agents_vscode_refuses_a_workspace_the_machine_isnt_managed_for(monkeypatch): + selected: list[str] = [] + monkeypatch.setattr( + "ucode.cli._configure_shared_workspace_states", lambda *args, **kwargs: [STATE] + ) + monkeypatch.setattr("ucode.cli.load_state", lambda: STATE) + monkeypatch.setattr( + "ucode.cli.managed_state_workspace", lambda: "https://managed.cloud.databricks.com" + ) + monkeypatch.setattr("ucode.cli.set_current_workspace", selected.append) + monkeypatch.setattr(vscode, "configure", lambda state: pytest.fail("VS Code was set up")) + + result = CliRunner().invoke(app, ["configure", "--agents", "vscode", "--workspace", WS]) + + assert result.exit_code == 1 + assert "This machine is managed for https://managed.cloud.databricks.com" in " ".join( + result.output.split() + ) + assert selected == ["https://managed.cloud.databricks.com"] + + @pytest.mark.parametrize("dry_run", [False, True]) def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, dry_run): configured: list[dict] = [] From 7378c179487f7f166dbdf2eb629de3e610609951 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 13:45:58 +0000 Subject: [PATCH 15/18] Keep ug's VS Code settings out of Settings Sync The unityGateway.* settings are application-scoped, which keeps them out of workspace settings, but Settings Sync still copies application settings to the user's other machines. The auth command ug writes holds this machine's ug path and Databricks profile, so a synced copy broke token fetches on the other machine, and the last configure won. The workspace and model list come from this machine's ug setup too. Mark all three ignoreSync, so each machine keeps what its own ug configure wrote. They stay application-scoped: machine scope would make VS Code read them from Remote settings in SSH, WSL and container windows, which ug never writes. Co-authored-by: Isaac --- src/ucode/vscode/extension/out/extension.js | 2 +- src/ucode/vscode/extension/package.json | 9 ++++++--- vscode-extension/package.json | 9 ++++++--- vscode-extension/test/extension.test.ts | 6 ++++++ 4 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index 355a0ae82..b6747edd9 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var B=Object.create;var f=Object.defineProperty;var H=Object.getOwnPropertyDescriptor;var W=Object.getOwnPropertyNames;var K=Object.getPrototypeOf,Y=Object.prototype.hasOwnProperty;var z=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},M=(e,n,t,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let o of W(n))!Y.call(e,o)&&o!==t&&f(e,o,{get:()=>n[o],enumerable:!(s=H(n,o))||s.enumerable});return e};var T=(e,n,t)=>(t=e!=null?B(K(e)):{},M(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),X=e=>M(f({},"__esModule",{value:!0}),e);var ce={};z(ce,{DatabricksModels:()=>w,activate:()=>ie,readSettings:()=>J,userAgent:()=>U});module.exports=X(ce);var d=T(require("vscode"));var b="0.1.0";var _=require("node:child_process"),V=5*6e4,Z=60*6e4,ee=e=>new Promise((n,t)=>{(0,_.exec)(e,{timeout:6e4,windowsHide:!0},(s,o,r)=>{s?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):n(String(o).trim())})}),te=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+Z},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-V>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=ee(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:te(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var c=T(require("vscode")),S=3,ne=1e3,y=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",R=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,L=e=>e.inputSchema??{type:"object",properties:{}},P=(e,n,t)=>{let s;try{s=t.trim()?JSON.parse(t):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new c.LanguageModelToolCallPart(e,n,s)},E=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),O=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},D=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let s of e.content){let o=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of o)n+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:y(r).length,t+=R(r)===void 0?0:1}return Math.ceil(n/4)+t*ne},k=e=>e.map(R).filter(n=>n!==void 0),I=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(y).join(""),images:k(e)});function j(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===S)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant){let a=o.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>k(a.content)),...l].map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:L(t)}}))}}var oe=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),se=new Set(["stop","tool_calls","function_call"]);async function*N(e){let n=[],t;for await(let s of e){let{error:o}=s;if(o)throw new Error(typeof o=="string"?o:o.message??JSON.stringify(o));for(let r of s.choices??[]){let l=oe(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=n[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!se.has(t)))throw E(t)}}if(!t)throw O();for(let s of n.filter(Boolean))yield P(s.id,s.name,s.args)}var ae=e=>{let n=e.map(y).join(""),t=k(e);return t.length?[...n?[{type:"input_text",text:n}]:[],...t.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function G(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===S)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...o.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:ae(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:L(t),strict:!1}))}}async function*$(e){let n=[],t=!1;for await(let s of e){let o=s;if(o.type==="response.output_text.delta"&&o.delta)yield new c.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")n.push(P(o.item.call_id??"",o.item.name??"",o.item.arguments??""));else if(o.type==="response.completed")t=!0;else{if(o.type==="response.incomplete")throw E(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!t)throw O();yield*n}async function*x(e,n,t,s,o,r){let l=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...o,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:r}),i=await s.get(),a=await l(i);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i)),a.status===401)throw await a.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let C=g.split(` +"use strict";var B=Object.create;var f=Object.defineProperty;var H=Object.getOwnPropertyDescriptor;var W=Object.getOwnPropertyNames;var K=Object.getPrototypeOf,Y=Object.prototype.hasOwnProperty;var z=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},M=(e,n,t,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let o of W(n))!Y.call(e,o)&&o!==t&&f(e,o,{get:()=>n[o],enumerable:!(s=H(n,o))||s.enumerable});return e};var T=(e,n,t)=>(t=e!=null?B(K(e)):{},M(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),X=e=>M(f({},"__esModule",{value:!0}),e);var ce={};z(ce,{DatabricksModels:()=>w,activate:()=>ie,readSettings:()=>J,userAgent:()=>U});module.exports=X(ce);var d=T(require("vscode"));var b="0.1.0";var S=require("node:child_process"),V=5*6e4,Z=60*6e4,ee=e=>new Promise((n,t)=>{(0,S.exec)(e,{timeout:6e4,windowsHide:!0},(s,o,r)=>{s?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):n(String(o).trim())})}),te=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+Z},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-V>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=ee(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:te(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var c=T(require("vscode")),_=3,ne=1e3,y=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",R=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,L=e=>e.inputSchema??{type:"object",properties:{}},P=(e,n,t)=>{let s;try{s=t.trim()?JSON.parse(t):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new c.LanguageModelToolCallPart(e,n,s)},E=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),O=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},D=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let s of e.content){let o=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of o)n+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:y(r).length,t+=R(r)===void 0?0:1}return Math.ceil(n/4)+t*ne},k=e=>e.map(R).filter(n=>n!==void 0),I=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(y).join(""),images:k(e)});function j(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===_)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant){let a=o.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>k(a.content)),...l].map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:L(t)}}))}}var oe=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),se=new Set(["stop","tool_calls","function_call"]);async function*N(e){let n=[],t;for await(let s of e){let{error:o}=s;if(o)throw new Error(typeof o=="string"?o:o.message??JSON.stringify(o));for(let r of s.choices??[]){let l=oe(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=n[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!se.has(t)))throw E(t)}}if(!t)throw O();for(let s of n.filter(Boolean))yield P(s.id,s.name,s.args)}var ae=e=>{let n=e.map(y).join(""),t=k(e);return t.length?[...n?[{type:"input_text",text:n}]:[],...t.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function G(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===_)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...o.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:ae(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:L(t),strict:!1}))}}async function*$(e){let n=[],t=!1;for await(let s of e){let o=s;if(o.type==="response.output_text.delta"&&o.delta)yield new c.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")n.push(P(o.item.call_id??"",o.item.name??"",o.item.arguments??""));else if(o.type==="response.completed")t=!0;else{if(o.type==="response.incomplete")throw E(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!t)throw O();yield*n}async function*x(e,n,t,s,o,r){let l=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...o,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:r}),i=await s.get(),a=await l(i);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i)),a.status===401)throw await a.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let C=g.split(` `);g=C.pop()??"";for(let F of C){let v=q(F);if(v==="[DONE]")return;v&&(yield JSON.parse(v))}}let u=q(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var q=e=>e.startsWith("data:")?e.slice(5).trim():"";var re=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),J=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},api:t.api,...re(t)})),authCommand:String(e("authCommand")??"")}},U=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,w=class{constructor(n,t,s=()=>({})){this.settings=n;this.log=t;this.headers=s}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:s}=this.settings();return!n||!s||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:o,name:r,api:l,maxInputTokens:i,maxOutputTokens:a})=>({id:o,api:l,name:r??o,family:o.split(".").pop()??o,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:a??32e3,capabilities:{toolCalling:!0,imageInput:l==="responses"||o.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,s,o,r){let{workspace:l,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:t,tools:s.tools??[],toolRequired:s.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=r.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(x(l,"responses",G(a),this.tokens,this.headers(),p.signal)):N(x(l,"chat/completions",j(a),this.tokens,this.headers(),p.signal));for await(let m of u)o.report(m)}catch(u){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return D(t)}};function ie(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let r=d.workspace.getConfiguration("unityGateway");return J(l=>r.get(l))},s=()=>({"User-Agent":U(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),o=new w(t,r=>n.info(r),s);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",o),d.commands.registerCommand("unityGateway.refreshModels",()=>o.refresh()),d.workspace.onDidChangeConfiguration(r=>{r.affectsConfiguration("unityGateway")&&o.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json index 76027205c..df70f1f40 100644 --- a/src/ucode/vscode/extension/package.json +++ b/src/ucode/vscode/extension/package.json @@ -36,7 +36,8 @@ "type": "string", "default": "", "description": "Databricks workspace URL, for example https://my-workspace.cloud.databricks.com. Set by `ug configure`.", - "scope": "application" + "scope": "application", + "ignoreSync": true }, "unityGateway.models": { "type": "array", @@ -77,13 +78,15 @@ } } }, - "scope": "application" + "scope": "application", + "ignoreSync": true }, "unityGateway.authCommand": { "type": "string", "default": "", "description": "Command that prints a Databricks access token. Set by `ug configure`.", - "scope": "application" + "scope": "application", + "ignoreSync": true } } }, diff --git a/vscode-extension/package.json b/vscode-extension/package.json index 761f66d0c..2f9bfa8f2 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -36,7 +36,8 @@ "type": "string", "default": "", "description": "Databricks workspace URL, for example https://my-workspace.cloud.databricks.com. Set by `ug configure`.", - "scope": "application" + "scope": "application", + "ignoreSync": true }, "unityGateway.models": { "type": "array", @@ -77,13 +78,15 @@ } } }, - "scope": "application" + "scope": "application", + "ignoreSync": true }, "unityGateway.authCommand": { "type": "string", "default": "", "description": "Command that prints a Databricks access token. Set by `ug configure`.", - "scope": "application" + "scope": "application", + "ignoreSync": true } } }, diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index 70fc83ba3..d8e3c8ff2 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -73,6 +73,12 @@ describe("settings", () => { const properties = Object.values(packageJson.contributes.configuration.properties); assert.deepEqual(properties.map((property) => property.scope), ["application", "application", "application"]); }); + + // ug writes them per machine (the auth command holds this machine's ug path and profile). + it("stay out of Settings Sync", () => { + const properties = Object.values(packageJson.contributes.configuration.properties); + assert.deepEqual(properties.map((property) => property.ignoreSync), [true, true, true]); + }); }); describe("provider", () => { From cdcf3aa6f1db5cbc282ec57607a212b5d7c6e096 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 14:24:29 +0000 Subject: [PATCH 16/18] Ship VS Code's setup inside the extension instead of editing settings.json `ug configure --agents vscode` merged three unityGateway.* settings into the user's settings.json, which took a JSONC editor (to keep the user's comments), application scope (so a workspace couldn't redirect the token), ignoreSync (so Settings Sync didn't copy one machine's auth command to the others) and a settings cleanup in revert. ug already packs the vsix at configure time and reinstalls it every run, so it now writes the setup into the vsix as extension/ug-config.json: the workspace, the models (id, picker name, API, token limits) and the auth command. The extension reads that file from its own folder at activation. Nothing goes into settings.json, so the JSONC editor, the three settings, the Refresh Models command and the settings listener go away, and revert just uninstalls the extension. Revert now looks for the extension's folder under ~/.vscode/extensions and ~/.vscode-insiders/extensions before running `code`, and uninstalls through the matching command, so an Insiders install is removed too. VS Code can leave an uninstalled extension's folder until it next starts, so the installed version is still checked before uninstalling. A rerun of configure now reaches an open VS Code window after a restart, which configure already asks for. Co-authored-by: Isaac --- src/ucode/agents/vscode.py | 139 ++++++-------- src/ucode/cli.py | 6 +- src/ucode/config_io.py | 176 ------------------ src/ucode/vscode/extension/out/extension.js | 4 +- src/ucode/vscode/extension/package.json | 68 ------- tests/conftest.py | 5 +- tests/test_agent_vscode.py | 193 ++++++++------------ tests/test_cli.py | 2 +- tests/test_config_io.py | 72 -------- vscode-extension/README.md | 18 +- vscode-extension/package.json | 68 ------- vscode-extension/src/extension.ts | 53 +++--- vscode-extension/test/extension.test.ts | 29 ++- 13 files changed, 186 insertions(+), 647 deletions(-) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index b0ee03b17..bab6972e1 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -1,8 +1,10 @@ """VS Code: install ug's model extension so Copilot Chat can use the workspace's AI Gateway models. Configure-only, like Cursor (it stays out of `agents.__init__._MODULES`): `ug configure --agents -vscode` packs the extension that ships with ug into a vsix, installs it with VS Code's `code` -command, and writes the extension's settings into the user settings.json. `ug revert` undoes both. +vscode` packs the extension that ships with ug into a vsix, together with a `ug-config.json` that +holds the workspace, its models and the auth command, and installs it with VS Code's `code` command. +The extension reads that file from its own folder, so ug never edits the user's settings.json. +`ug revert` uninstalls the extension. """ from __future__ import annotations @@ -18,7 +20,6 @@ from typing import cast from xml.sax.saxutils import escape -from ucode import config_io from ucode.agents.copilot import ( copilot_exclusion_reason, copilot_model_limits, @@ -34,12 +35,10 @@ EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" EXTENSION_ID = "databricks.unity-gateway" -WORKSPACE_SETTING = "unityGateway.workspace" -MODELS_SETTING = "unityGateway.models" -AUTH_COMMAND_SETTING = "unityGateway.authCommand" -SETTINGS = (WORKSPACE_SETTING, MODELS_SETTING, AUTH_COMMAND_SETTING) -# Each `code` command and the folder its user settings live in. -CODE_COMMANDS = (("code", "Code"), ("code-insiders", "Code - Insiders")) +# Read by the extension from its install folder; see CONFIG_FILE in vscode-extension/src/extension.ts. +CONFIG_FILE = "ug-config.json" +# Each `code` command and the folder in the home directory it installs extensions under. +CODE_COMMANDS = (("code", ".vscode"), ("code-insiders", ".vscode-insiders")) _CODE_TIMEOUT_S = 120 @@ -67,35 +66,31 @@ def _install_locations(command: str) -> list[Path]: return [Path("/usr/bin") / command, Path("/snap/bin") / command] -def find_code() -> tuple[Path, str] | None: - """VS Code's `code` command (or Insiders') and its settings folder name, or None.""" - for command, product in CODE_COMMANDS: - if found := shutil.which(command): - return Path(found), product - for command, product in CODE_COMMANDS: - for path in _install_locations(command): - if path.is_file(): - return path, product +def _find_command(command: str) -> Path | None: + if found := shutil.which(command): + return Path(found) + return next((path for path in _install_locations(command) if path.is_file()), None) + + +def find_code() -> Path | None: + """VS Code's `code` command, or Insiders' when only that is installed.""" + for command, _ in CODE_COMMANDS: + if found := _find_command(command): + return found return None -def settings_path(product: str) -> Path: - system = current_os() - if system is OS.MACOS: - base = Path.home() / "Library" / "Application Support" - elif system is OS.WINDOWS: - base = Path(os.environ.get("APPDATA") or Path.home() / "AppData" / "Roaming") - else: - base = Path(os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config") - return base / product / "User" / "settings.json" +def extensions_dir(folder: str) -> Path: + return Path.home() / folder / "extensions" def _package() -> dict: return json.loads((EXTENSION_DIR / "package.json").read_text(encoding="utf-8")) -def pack_vsix(directory: Path) -> Path: - """Zip the bundled extension into a vsix in ``directory``, with no Node or `vsce` needed.""" +def pack_vsix(directory: Path, config: dict[str, object]) -> Path: + """Zip the bundled extension and its ``config`` into a vsix in ``directory``, with no Node or + `vsce` needed.""" package = _package() manifest = f""" @@ -129,6 +124,7 @@ def pack_vsix(directory: Path) -> Path: for path in sorted(EXTENSION_DIR.rglob("*")): if path.is_file(): archive.write(path, f"extension/{path.relative_to(EXTENSION_DIR).as_posix()}") + archive.writestr(f"extension/{CONFIG_FILE}", json.dumps(config, indent=2) + "\n") return vsix @@ -150,12 +146,14 @@ def installed_version(code: Path) -> str | None: return None -def install_extension(code: Path) -> str: - """Install the bundled extension over whatever is installed; returns what happened.""" - # Always reinstall: a ug upgrade can change the bundle without changing its version. +def install_extension(code: Path, config: dict[str, object]) -> str: + """Install the bundled extension with ``config`` over whatever is installed; returns what + happened.""" + # Always reinstall: the config changes with each configure, and a ug upgrade can change the + # bundle without changing its version. version = _package()["version"] with tempfile.TemporaryDirectory(prefix="ug-vscode-") as directory: - vsix = pack_vsix(Path(directory)) + vsix = pack_vsix(Path(directory), config) result = _run_code(code, "--install-extension", str(vsix), "--force") if result.returncode != 0: detail = (result.stderr or result.stdout or "").strip()[:300] @@ -176,7 +174,7 @@ def _model_entry(model: str, name: str) -> dict[str, object]: return entry -def model_settings(workspace: str, token: str) -> tuple[list[dict[str, object]], str | None]: +def model_entries(workspace: str, token: str) -> tuple[list[dict[str, object]], str | None]: """The workspace's Claude and GPT models, each with the API the extension calls it on.""" # The models Copilot CLI offers: Gemini and open models, gpt-oss among them, fail on Copilot's # requests or aren't served on every workspace. @@ -205,18 +203,6 @@ def auth_command(state: dict) -> str: ) -def _edit_settings(path: Path, values: dict[str, object], remove: tuple[str, ...] = ()) -> bool: - try: - text = path.read_text(encoding="utf-8") if path.exists() else "" - updated = config_io.update_jsonc_keys(text, values, remove) - if updated != text: - config_io.write_text_file(path, updated) - except (OSError, RuntimeError, ValueError) as exc: - print_warning(f"Couldn't update {path}: {exc}. Fix the file, then rerun this command.") - return False - return True - - def next_steps() -> list[str]: """What to do in VS Code once the extension is installed, for this OS.""" restart = ( @@ -233,40 +219,31 @@ def next_steps() -> list[str]: def configure(state: dict) -> bool: - """Write the extension's settings and install it. Returns False, after a warning, when VS - Code isn't found or the settings can't be updated.""" - found = find_code() - if found is None: + """Install the extension with this workspace's setup. Returns False, after a warning, if any + step fails.""" + code = find_code() + if code is None: print_warning( "VS Code's `code` command wasn't found. Install VS Code (on macOS, also run " "\"Shell Command: Install 'code' command in PATH\"), then rerun `ug configure --agents vscode`." ) return False - code, product = found workspace = state["workspace"] - models, reason = model_settings(workspace, launch_token(state, workspace)) + models, reason = model_entries(workspace, launch_token(state, workspace)) if not models: - # Keep what an earlier configure wrote: a failed listing shouldn't wipe a working setup. + # Keep what an earlier configure installed: a failed listing shouldn't wipe a working setup. print_warning( f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'}). " "VS Code was left unchanged." ) return False - path = settings_path(product) - values = { - WORKSPACE_SETTING: workspace, - MODELS_SETTING: models, - AUTH_COMMAND_SETTING: auth_command(state), - } - if not _edit_settings(path, values): - return False - print_kv("VS Code settings", f"{len(models)} models in {path}") + config = {"workspace": workspace, "authCommand": auth_command(state), "models": models} try: - print_kv("VS Code extension", install_extension(code)) - except RuntimeError as exc: - print_warning( - f"{exc} Install it from VS Code's Extensions view, or rerun `ug configure --agents vscode`." + print_kv( + "VS Code extension", f"{install_extension(code, config)} with {len(models)} models" ) + except RuntimeError as exc: + print_warning(f"{exc} Rerun `ug configure --agents vscode` to try again.") return False print_heading("Next, in VS Code:") for line in next_steps(): @@ -275,29 +252,21 @@ def configure(state: dict) -> bool: def revert() -> str: - """Remove ug's settings and uninstall the extension: "removed" or "unchanged".""" - changed = configured = False - for _, product in CODE_COMMANDS: - path = settings_path(product) - if not path.exists(): + """Uninstall the extension, and its config with it: "removed" or "unchanged".""" + removed = False + for command, folder in CODE_COMMANDS: + # Look on disk first, so revert runs `code` only for people who set VS Code up. + if not any(extensions_dir(folder).glob(f"{EXTENSION_ID}-*")): continue - try: - current = config_io.loads_jsonc(path.read_text(encoding="utf-8")) - except (OSError, ValueError): - print_warning(f"Couldn't read {path}; remove its unityGateway settings by hand.") + code = _find_command(command) + # VS Code can leave an uninstalled extension's folder behind until it next starts. + if code is None or installed_version(code) is None: continue - if isinstance(current, dict) and any(key in current for key in SETTINGS): - configured = True - changed = _edit_settings(path, {}, SETTINGS) or changed - # Configure writes the settings before it installs, so without them there's nothing to - # uninstall, and revert doesn't run `code` for people who never set up VS Code. - found = find_code() if configured else None - if found is not None and installed_version(found[0]) is not None: - result = _run_code(found[0], "--uninstall-extension", EXTENSION_ID) + result = _run_code(code, "--uninstall-extension", EXTENSION_ID) if result.returncode == 0: - changed = True + removed = True else: print_warning( f"Couldn't uninstall {EXTENSION_ID}: {(result.stderr or '').strip()[:300]}" ) - return "removed" if changed else "unchanged" + return "removed" if removed else "unchanged" diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 40b04f5f8..4746d22db 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1585,7 +1585,7 @@ def revert() -> int: try: vscode_result = vscode.revert() except Exception as exc: # noqa: BLE001 -- Preserve other revert operations on failure - # VS Code's revert reads its settings files, not ug state, so a later run still finds them. + # VS Code's revert looks at the installed extension, not ug state, so a later run still finds it. print_warning(f"Could not revert VS Code: {exc}. Fix it, then run ug revert again.") vscode_result = "failed" @@ -1614,7 +1614,7 @@ def revert() -> int: print_kv("Codex OS-managed settings", codex_managed_result) print_kv("Unity Gateway workspace policy", workspace_policy_result) print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") - print_kv("VS Code extension and settings", vscode_result) + print_kv("VS Code extension", vscode_result) for client, spec in MCP_CLIENTS.items(): print_kv( f"{spec['display']} MCP config", @@ -4677,7 +4677,7 @@ def configure( _parse_agents_option(agents) if wants_vscode and is_dry_run(): # A dry run saves no workspace, so the state on disk isn't the one being configured. - print_note("Dry run: VS Code's settings and extension aren't changed.") + print_note("Dry run: VS Code isn't changed.") elif wants_vscode: from ucode.agents import vscode diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index 1635b6a27..c4ca3578c 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -266,179 +266,3 @@ def parse_dotenv(path: Path) -> dict[str, str]: def write_dotenv(path: Path, env: dict[str, str]) -> None: content = "".join(f'{key}="{val}"\n' for key, val in env.items()) write_text_file(path, content) - - -# --- JSONC (JSON with comments and trailing commas, e.g. VS Code's settings.json) --------------- - -_JSONC_PUNCTUATION = frozenset("{}[]:,") - - -def _jsonc_tokens(text: str) -> list[tuple[str, int, int]]: - """``(token, start, end)`` for each token in ``text``, skipping whitespace and comments. - - Raises ValueError on an unterminated string or comment. - """ - tokens: list[tuple[str, int, int]] = [] - i, length = 0, len(text) - while i < length: - char = text[i] - if char.isspace(): - i += 1 - elif text.startswith("//", i): - newline = text.find("\n", i) - i = length if newline < 0 else newline + 1 - elif text.startswith("/*", i): - close = text.find("*/", i + 2) - if close < 0: - raise ValueError("unterminated comment") - i = close + 2 - elif char == '"': - end = i + 1 - while end < length and text[end] != '"': - end += 2 if text[end] == "\\" else 1 - if end >= length: - raise ValueError("unterminated string") - tokens.append((text[i : end + 1], i, end + 1)) - i = end + 1 - elif char in _JSONC_PUNCTUATION: - tokens.append((char, i, i + 1)) - i += 1 - else: - end = i - while end < length and not text[end].isspace() and text[end] not in _JSONC_PUNCTUATION: - if text.startswith(("//", "/*"), end): - break - end += 1 - tokens.append((text[i:end], i, end)) - i = end - return tokens - - -def loads_jsonc(text: str) -> object: - """Parse JSONC: comments and trailing commas are dropped, then it's read as JSON.""" - tokens = _jsonc_tokens(text) - kept = [ - token - for index, (token, _, _) in enumerate(tokens) - if not (token == "," and index + 1 < len(tokens) and tokens[index + 1][0] in "}]") - ] - return json.loads(" ".join(kept)) if kept else None - - -class _JsoncEntry(TypedDict): - key: str - key_start: int - value_start: int - value_end: int - comma_end: int | None - - -def _jsonc_top_level(text: str) -> tuple[int, int, list[_JsoncEntry]]: - """The top-level object's ``{`` and ``}`` offsets and its entries; ValueError if it isn't one.""" - tokens = _jsonc_tokens(text) - if not tokens or tokens[0][0] != "{": - raise ValueError("the top level isn't an object") - entries: list[_JsoncEntry] = [] - index = 1 - while index < len(tokens) and tokens[index][0] != "}": - key, key_start, _ = tokens[index] - if not key.startswith('"') or index + 2 >= len(tokens) or tokens[index + 1][0] != ":": - raise ValueError("expected a key and a value") - index += 2 - value_start = tokens[index][1] - depth = 0 - while True: - token = tokens[index][0] - depth += 1 if token in "{[" else -1 if token in "}]" else 0 - if depth <= 0: - break - index += 1 - if index >= len(tokens): - raise ValueError("unterminated value") - value_end = tokens[index][2] - index += 1 - comma_end = None - if index < len(tokens) and tokens[index][0] == ",": - comma_end = tokens[index][2] - index += 1 - entries.append( - { - "key": json.loads(key), - "key_start": key_start, - "value_start": value_start, - "value_end": value_end, - "comma_end": comma_end, - } - ) - if index >= len(tokens) or index != len(tokens) - 1: - raise ValueError("unexpected content after the top-level object") - return tokens[0][1], tokens[index][1], entries - - -def _line_start(text: str, offset: int) -> int: - """The start of ``offset``'s line when only whitespace precedes it there, else ``offset``.""" - start = text.rfind("\n", 0, offset) + 1 - return start if not text[start:offset].strip() else offset - - -def _line_end(text: str, offset: int) -> int: - newline = text.find("\n", offset) - return len(text) if newline < 0 else newline - - -def _remove_jsonc_entry(text: str, entries: list[_JsoncEntry], position: int) -> str: - entry = entries[position] - start = _line_start(text, entry["key_start"]) - end = entry["comma_end"] if entry["comma_end"] is not None else entry["value_end"] - rest = text[end : _line_end(text, end)].strip() - # A comment after the entry on its own line goes with it. - if start != entry["key_start"] and (not rest or rest.startswith("//")): - end = min(_line_end(text, end) + 1, len(text)) - text = text[:start] + text[end:] - previous_comma = entries[position - 1]["comma_end"] if position > 0 else None - if entry["comma_end"] is None and previous_comma is not None: - text = text[: previous_comma - 1] + text[previous_comma:] - return text - - -def update_jsonc_keys(text: str, values: dict[str, object], remove: tuple[str, ...] = ()) -> str: - """Set ``values`` and delete ``remove`` among the top-level keys of a JSONC object, leaving - every other byte (comments included) as it was. Raises ValueError if ``text`` isn't one.""" - if not _jsonc_tokens(text): - text = text + ("\n" if text.strip() else "") + "{}" - for key in remove: - _, _, entries = _jsonc_top_level(text) - position = next((i for i, entry in enumerate(entries) if entry["key"] == key), None) - if position is not None: - text = _remove_jsonc_entry(text, entries, position) - for key, value in values.items(): - open_brace, close_brace, entries = _jsonc_top_level(text) - first = entries[0]["key_start"] if entries else None - indent = text[_line_start(text, first) : first] if first is not None else "" - indent = indent if indent and not indent.strip() else " " - rendered = json.dumps(value, indent=indent).replace("\n", "\n" + indent) - existing = next((entry for entry in entries if entry["key"] == key), None) - if existing is not None: - text = text[: existing["value_start"]] + rendered + text[existing["value_end"] :] - continue - line = f"{indent}{json.dumps(key)}: {rendered}" - if not entries: - text = text[: open_brace + 1] + f"\n{line}\n" + text[open_brace + 1 :].lstrip(" \t") - continue - last = entries[-1] - if last["comma_end"] is None: - text = text[: last["value_end"]] + "," + text[last["value_end"] :] - close_brace += 1 - after = last["value_end"] + 1 - else: - after = last["comma_end"] - # After the last entry's line, so a comment on that line stays with it. - at = _line_end(text, after) - at = after if close_brace < at else at - text = text[:at] + f"\n{line}" + text[at:] - loaded = loads_jsonc(text) - if not isinstance(loaded, dict) or any( - loaded.get(key) != value for key, value in values.items() - ): - raise ValueError("the edited file doesn't hold the expected values") - return text diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index b6747edd9..c473cfd19 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,2 +1,2 @@ -"use strict";var B=Object.create;var f=Object.defineProperty;var H=Object.getOwnPropertyDescriptor;var W=Object.getOwnPropertyNames;var K=Object.getPrototypeOf,Y=Object.prototype.hasOwnProperty;var z=(e,n)=>{for(var t in n)f(e,t,{get:n[t],enumerable:!0})},M=(e,n,t,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let o of W(n))!Y.call(e,o)&&o!==t&&f(e,o,{get:()=>n[o],enumerable:!(s=H(n,o))||s.enumerable});return e};var T=(e,n,t)=>(t=e!=null?B(K(e)):{},M(n||!e||!e.__esModule?f(t,"default",{value:e,enumerable:!0}):t,e)),X=e=>M(f({},"__esModule",{value:!0}),e);var ce={};z(ce,{DatabricksModels:()=>w,activate:()=>ie,readSettings:()=>J,userAgent:()=>U});module.exports=X(ce);var d=T(require("vscode"));var b="0.1.0";var S=require("node:child_process"),V=5*6e4,Z=60*6e4,ee=e=>new Promise((n,t)=>{(0,S.exec)(e,{timeout:6e4,windowsHide:!0},(s,o,r)=>{s?t(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):n(String(o).trim())})}),te=e=>{try{let{exp:n}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof n=="number")return n*1e3}catch{}return Date.now()+Z},h=class{constructor(n,t=()=>{}){this.command=n;this.onFetch=t}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-V>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(n){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==n?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(n){return this.pending??=ee(n).then(t=>{if(!t)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:t,expiresAt:te(t)},this.onFetch(this.cached.expiresAt),t}).finally(()=>{this.pending=void 0}),this.pending}};var c=T(require("vscode")),_=3,ne=1e3,y=e=>e instanceof c.LanguageModelTextPart?e.value:e instanceof c.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof c.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",R=e=>e instanceof c.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,L=e=>e.inputSchema??{type:"object",properties:{}},P=(e,n,t)=>{let s;try{s=t.trim()?JSON.parse(t):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${n} with arguments that aren't a JSON object: ${t}`);return new c.LanguageModelToolCallPart(e,n,s)},E=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),O=()=>new Error("The response ended before the model finished. Try again."),A=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},D=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,t=0;for(let s of e.content){let o=s instanceof c.LanguageModelToolResultPart?s.content:[s];for(let r of o)n+=r instanceof c.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:y(r).length,t+=R(r)===void 0?0:1}return Math.ceil(n/4)+t*ne},k=e=>e.map(R).filter(n=>n!==void 0),I=e=>({results:e.filter(n=>n instanceof c.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof c.LanguageModelToolCallPart),text:e.map(y).join(""),images:k(e)});function j(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===_)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant){let a=o.map(p=>({id:p.callId,type:"function",function:{name:p.name,arguments:JSON.stringify(p.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let i=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>k(a.content)),...l].map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...A(e,t=>({type:"function",function:{name:t.name,description:t.description,parameters:L(t)}}))}}var oe=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),se=new Set(["stop","tool_calls","function_call"]);async function*N(e){let n=[],t;for await(let s of e){let{error:o}=s;if(o)throw new Error(typeof o=="string"?o:o.message??JSON.stringify(o));for(let r of s.choices??[]){let l=oe(r.delta?.content);l&&(yield new c.LanguageModelTextPart(l));for(let i of r.delta?.tool_calls??[]){let a=n[i.index??0]??={id:"",name:"",args:""};a.id=i.id||a.id,a.name=i.function?.name||a.name,a.args+=i.function?.arguments??""}if(r.finish_reason&&(t=r.finish_reason,!se.has(t)))throw E(t)}}if(!t)throw O();for(let s of n.filter(Boolean))yield P(s.id,s.name,s.args)}var ae=e=>{let n=e.map(y).join(""),t=k(e);return t.length?[...n?[{type:"input_text",text:n}]:[],...t.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function G(e){let n=e.messages.flatMap(t=>{let{results:s,calls:o,text:r,images:l}=I(t.content);if(Number(t.role)===_)return[{role:"system",content:r}];if(t.role===c.LanguageModelChatMessageRole.Assistant)return[...r?[{role:"assistant",content:r}]:[],...o.map(a=>({type:"function_call",call_id:a.callId,name:a.name,arguments:JSON.stringify(a.input)}))];let i=[...r?[{type:"input_text",text:r}]:[],...l.map(a=>({type:"input_image",image_url:a,detail:"auto"}))];return[...s.map(a=>({type:"function_call_output",call_id:a.callId,output:ae(a.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:n,stream:!0,store:!1,...A(e,t=>({type:"function",name:t.name,description:t.description,parameters:L(t),strict:!1}))}}async function*$(e){let n=[],t=!1;for await(let s of e){let o=s;if(o.type==="response.output_text.delta"&&o.delta)yield new c.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")n.push(P(o.item.call_id??"",o.item.name??"",o.item.arguments??""));else if(o.type==="response.completed")t=!0;else{if(o.type==="response.incomplete")throw E(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!t)throw O();yield*n}async function*x(e,n,t,s,o,r){let l=m=>fetch(`${e}/ai-gateway/mlflow/v1/${n}`,{method:"POST",headers:{...o,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(t),signal:r}),i=await s.get(),a=await l(i);if(a.status===401&&(await a.body?.cancel(),i=await s.renew(i),a=await l(i)),a.status===401)throw await a.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`);let p=new TextDecoder,g="";for await(let m of a.body){g+=p.decode(m,{stream:!0});let C=g.split(` -`);g=C.pop()??"";for(let F of C){let v=q(F);if(v==="[DONE]")return;v&&(yield JSON.parse(v))}}let u=q(g+p.decode());u&&u!=="[DONE]"&&(yield JSON.parse(u))}var q=e=>e.startsWith("data:")?e.slice(5).trim():"";var re=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let t=e[n];return typeof t=="number"&&Number.isInteger(t)&&t>0?[[n,t]]:[]})),J=e=>{let n=e("models");return{workspace:String(e("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(n)?n:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&(t.api==="chat_completions"||t.api==="responses")).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},api:t.api,...re(t)})),authCommand:String(e("authCommand")??"")}},U=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,w=class{constructor(n,t,s=()=>({})){this.settings=n;this.log=t;this.headers=s}settings;log;headers;changed=new d.EventEmitter;onDidChangeLanguageModelChatInformation=this.changed.event;tokens;refresh(){this.tokens=void 0,this.changed.fire()}async provideLanguageModelChatInformation(){let{workspace:n,models:t,authCommand:s}=this.settings();return!n||!s||t.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):t.map(({id:o,name:r,api:l,maxInputTokens:i,maxOutputTokens:a})=>({id:o,api:l,name:r??o,family:o.split(".").pop()??o,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:a??32e3,capabilities:{toolCalling:!0,imageInput:l==="responses"||o.includes("claude")}}))}async provideLanguageModelChatResponse(n,t,s,o,r){let{workspace:l,authCommand:i}=this.settings();this.tokens??=new h(i,u=>this.log(`Fetched a Databricks token (expires ${new Date(u).toISOString()})`));let a={model:n.id,messages:t,tools:s.tools??[],toolRequired:s.toolMode===d.LanguageModelChatToolMode.Required},p=new AbortController,g=r.onCancellationRequested(()=>p.abort());try{let u=n.api==="responses"?$(x(l,"responses",G(a),this.tokens,this.headers(),p.signal)):N(x(l,"chat/completions",j(a),this.tokens,this.headers(),p.signal));for await(let m of u)o.report(m)}catch(u){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${u instanceof Error?u.message:String(u)}`),u}finally{g.dispose()}}async provideTokenCount(n,t){return D(t)}};function ie(e){let n=d.window.createOutputChannel("Unity Gateway",{log:!0}),t=()=>{let r=d.workspace.getConfiguration("unityGateway");return J(l=>r.get(l))},s=()=>({"User-Agent":U(d.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),o=new w(t,r=>n.info(r),s);e.subscriptions.push(n,d.lm.registerLanguageModelChatProvider("databricks-unity-gateway",o),d.commands.registerCommand("unityGateway.refreshModels",()=>o.refresh()),d.workspace.onDidChangeConfiguration(r=>{r.affectsConfiguration("unityGateway")&&o.refresh()}))}0&&(module.exports={DatabricksModels,activate,readSettings,userAgent}); +"use strict";var K=Object.create;var f=Object.defineProperty;var Y=Object.getOwnPropertyDescriptor;var z=Object.getOwnPropertyNames;var X=Object.getPrototypeOf,Q=Object.prototype.hasOwnProperty;var V=(e,t)=>{for(var n in t)f(e,n,{get:t[n],enumerable:!0})},k=(e,t,n,s)=>{if(t&&typeof t=="object"||typeof t=="function")for(let o of z(t))!Q.call(e,o)&&o!==n&&f(e,o,{get:()=>t[o],enumerable:!(s=Y(t,o))||s.enumerable});return e};var T=(e,t,n)=>(n=e!=null?K(X(e)):{},k(t||!e||!e.__esModule?f(n,"default",{value:e,enumerable:!0}):n,e)),Z=e=>k(f({},"__esModule",{value:!0}),e);var ge={};V(ge,{CONFIG_FILE:()=>U,DatabricksModels:()=>v,activate:()=>ue,readConfig:()=>B,userAgent:()=>H});module.exports=Z(ge);var G=require("node:fs"),F=require("node:path"),g=T(require("vscode"));var b="0.1.0";var _=require("node:child_process"),te=5*6e4,ne=60*6e4,oe=e=>new Promise((t,n)=>{(0,_.exec)(e,{timeout:6e4,windowsHide:!0},(s,o,a)=>{s?n(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(a).trim().slice(0,500)||s.message}`)):t(String(o).trim())})}),se=e=>{try{let{exp:t}=JSON.parse(Buffer.from(e.split(".")[1]??"","base64url").toString());if(typeof t=="number")return t*1e3}catch{}return Date.now()+ne},h=class{constructor(t,n=()=>{}){this.command=t;this.onFetch=n}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-te>Date.now()?Promise.resolve(this.cached.value):this.fetch(this.command)}async renew(t){return await this.pending?.catch(()=>{}),this.cached&&this.cached.value!==t?this.cached.value:(this.cached=void 0,this.fetch(`${this.command} --force-refresh`))}fetch(t){return this.pending??=oe(t).then(n=>{if(!n)throw new Error(`Couldn't get a Databricks token: \`${t}\` printed nothing.`);return this.cached={value:n,expiresAt:se(n)},this.onFetch(this.cached.expiresAt),n}).finally(()=>{this.pending=void 0}),this.pending}};var l=T(require("vscode")),S=3,re=1e3,y=e=>e instanceof l.LanguageModelTextPart?e.value:e instanceof l.LanguageModelDataPart&&!e.mimeType.startsWith("image/")?new TextDecoder().decode(e.data):e instanceof l.LanguageModelPromptTsxPart?JSON.stringify(e.value):"",R=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?`data:${e.mimeType};base64,${Buffer.from(e.data).toString("base64")}`:void 0,L=e=>e.inputSchema??{type:"object",properties:{}},P=(e,t,n)=>{let s;try{s=n.trim()?JSON.parse(n):{}}catch{s=void 0}if(typeof s!="object"||s===null||Array.isArray(s))throw new Error(`The model called ${t} with arguments that aren't a JSON object: ${n}`);return new l.LanguageModelToolCallPart(e,t,s)},E=e=>e==="max_output_tokens"||e==="length"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"?new Error("The model's response was stopped by a content filter."):new Error(`The model stopped before finishing the response (${e??"no reason given"}).`),O=()=>new Error("The response ended before the model finished. Try again."),A=(e,t)=>e.tools.length?{tools:e.tools.map(t),...e.toolRequired?{tool_choice:"required"}:{}}:{},j=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let t=0,n=0;for(let s of e.content){let o=s instanceof l.LanguageModelToolResultPart?s.content:[s];for(let a of o)t+=a instanceof l.LanguageModelToolCallPart?a.name.length+JSON.stringify(a.input).length:y(a).length,n+=R(a)===void 0?0:1}return Math.ceil(t/4)+n*re},x=e=>e.map(R).filter(t=>t!==void 0),I=e=>({results:e.filter(t=>t instanceof l.LanguageModelToolResultPart),calls:e.filter(t=>t instanceof l.LanguageModelToolCallPart),text:e.map(y).join(""),images:x(e)});function N(e){let t=e.messages.flatMap(n=>{let{results:s,calls:o,text:a,images:c}=I(n.content);if(Number(n.role)===S)return[{role:"system",content:a}];if(n.role===l.LanguageModelChatMessageRole.Assistant){let r=o.map(u=>({id:u.callId,type:"function",function:{name:u.name,arguments:JSON.stringify(u.input)}}));return a||r.length?[{role:"assistant",content:a||null,...r.length?{tool_calls:r}:{}}]:[]}let i=[...a?[{type:"text",text:a}]:[],...[...s.flatMap(r=>x(r.content)),...c].map(r=>({type:"image_url",image_url:{url:r}}))];return[...s.map(r=>({role:"tool",tool_call_id:r.callId,content:r.content.map(y).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:t,stream:!0,...A(e,n=>({type:"function",function:{name:n.name,description:n.description,parameters:L(n)}}))}}var ae=e=>typeof e=="string"?e:(e??[]).filter(t=>t.type==="text").map(t=>t.text??"").join(""),ie=new Set(["stop","tool_calls","function_call"]);async function*D(e){let t=[],n;for await(let s of e){let{error:o}=s;if(o)throw new Error(typeof o=="string"?o:o.message??JSON.stringify(o));for(let a of s.choices??[]){let c=ae(a.delta?.content);c&&(yield new l.LanguageModelTextPart(c));for(let i of a.delta?.tool_calls??[]){let r=t[i.index??0]??={id:"",name:"",args:""};r.id=i.id||r.id,r.name=i.function?.name||r.name,r.args+=i.function?.arguments??""}if(a.finish_reason&&(n=a.finish_reason,!ie.has(n)))throw E(n)}}if(!n)throw O();for(let s of t.filter(Boolean))yield P(s.id,s.name,s.args)}var le=e=>{let t=e.map(y).join(""),n=x(e);return n.length?[...t?[{type:"input_text",text:t}]:[],...n.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:t};function $(e){let t=e.messages.flatMap(n=>{let{results:s,calls:o,text:a,images:c}=I(n.content);if(Number(n.role)===S)return[{role:"system",content:a}];if(n.role===l.LanguageModelChatMessageRole.Assistant)return[...a?[{role:"assistant",content:a}]:[],...o.map(r=>({type:"function_call",call_id:r.callId,name:r.name,arguments:JSON.stringify(r.input)}))];let i=[...a?[{type:"input_text",text:a}]:[],...c.map(r=>({type:"input_image",image_url:r,detail:"auto"}))];return[...s.map(r=>({type:"function_call_output",call_id:r.callId,output:le(r.content)})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,input:t,stream:!0,store:!1,...A(e,n=>({type:"function",name:n.name,description:n.description,parameters:L(n),strict:!1}))}}async function*q(e){let t=[],n=!1;for await(let s of e){let o=s;if(o.type==="response.output_text.delta"&&o.delta)yield new l.LanguageModelTextPart(o.delta);else if(o.type==="response.output_item.done"&&o.item?.type==="function_call")t.push(P(o.item.call_id??"",o.item.name??"",o.item.arguments??""));else if(o.type==="response.completed")n=!0;else{if(o.type==="response.incomplete")throw E(o.response?.incomplete_details?.reason);if(o.type==="response.failed"||o.type==="error")throw new Error(o.response?.error?.message??o.message??"The model response failed.")}}if(!n)throw O();yield*t}async function*C(e,t,n,s,o,a){let c=m=>fetch(`${e}/ai-gateway/mlflow/v1/${t}`,{method:"POST",headers:{...o,Authorization:`Bearer ${m}`,"Content-Type":"application/json"},body:JSON.stringify(n),signal:a}),i=await s.get(),r=await c(i);if(r.status===401&&(await r.body?.cancel(),i=await s.renew(i),r=await c(i)),r.status===401)throw await r.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents vscode` to sign in again, then retry.");if(!r.ok||!r.body)throw new Error(`AI Gateway returned ${r.status}: ${(await r.text()).slice(0,300)}`);let u=new TextDecoder,p="";for await(let m of r.body){p+=u.decode(m,{stream:!0});let M=p.split(` +`);p=M.pop()??"";for(let W of M){let w=J(W);if(w==="[DONE]")return;w&&(yield JSON.parse(w))}}let d=J(p+u.decode());d&&d!=="[DONE]"&&(yield JSON.parse(d))}var J=e=>e.startsWith("data:")?e.slice(5).trim():"";var ce=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(t=>{let n=e[t];return typeof n=="number"&&Number.isInteger(n)&&n>0?[[t,n]]:[]})),U="ug-config.json",B=e=>{let t={};try{let o=JSON.parse(e??"{}");o&&typeof o=="object"&&!Array.isArray(o)&&(t=o)}catch{}let n=o=>t[o],s=n("models");return{workspace:String(n("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(o=>typeof o?.id=="string"&&o.id!==""&&(o.api==="chat_completions"||o.api==="responses")).map(o=>({id:o.id,...typeof o.name=="string"&&o.name.trim()?{name:o.name.trim()}:{},api:o.api,...ce(o)})),authCommand:String(n("authCommand")??"")}},H=e=>`ucode/${b} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,v=class{constructor(t,n,s=()=>({})){this.settings=t;this.log=n;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:t,models:n,authCommand:s}=this.settings();return!t||!s||n.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):n.map(({id:o,name:a,api:c,maxInputTokens:i,maxOutputTokens:r})=>({id:o,api:c,name:a??o,family:o.split(".").pop()??o,version:b,detail:"Databricks",maxInputTokens:i??128e3,maxOutputTokens:r??32e3,capabilities:{toolCalling:!0,imageInput:c==="responses"||o.includes("claude")}}))}async provideLanguageModelChatResponse(t,n,s,o,a){let{workspace:c,authCommand:i}=this.settings();this.tokens??=new h(i,d=>this.log(`Fetched a Databricks token (expires ${new Date(d).toISOString()})`));let r={model:t.id,messages:n,tools:s.tools??[],toolRequired:s.toolMode===g.LanguageModelChatToolMode.Required},u=new AbortController,p=a.onCancellationRequested(()=>u.abort());try{let d=t.api==="responses"?q(C(c,"responses",$(r),this.tokens,this.headers(),u.signal)):D(C(c,"chat/completions",N(r),this.tokens,this.headers(),u.signal));for await(let m of d)o.report(m)}catch(d){if(!a.isCancellationRequested)throw this.log(`${t.id} failed: ${d instanceof Error?d.message:String(d)}`),d}finally{p.dispose()}}async provideTokenCount(t,n){return j(n)}},de=e=>{try{return(0,G.readFileSync)(e,"utf8")}catch{return}};function ue(e){let t=g.window.createOutputChannel("Unity Gateway",{log:!0}),n=B(de((0,F.join)(e.extensionPath,U))),s=()=>({"User-Agent":H(g.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),o=new v(()=>n,a=>t.info(a),s);e.subscriptions.push(t,g.lm.registerLanguageModelChatProvider("databricks-unity-gateway",o))}0&&(module.exports={CONFIG_FILE,DatabricksModels,activate,readConfig,userAgent}); diff --git a/src/ucode/vscode/extension/package.json b/src/ucode/vscode/extension/package.json index df70f1f40..aac96abc9 100644 --- a/src/ucode/vscode/extension/package.json +++ b/src/ucode/vscode/extension/package.json @@ -28,74 +28,6 @@ "displayName": "Databricks", "icon": "media/databricks.svg" } - ], - "configuration": { - "title": "Databricks Unity Gateway", - "properties": { - "unityGateway.workspace": { - "type": "string", - "default": "", - "description": "Databricks workspace URL, for example https://my-workspace.cloud.databricks.com. Set by `ug configure`.", - "scope": "application", - "ignoreSync": true - }, - "unityGateway.models": { - "type": "array", - "default": [], - "description": "Models to offer in the chat model picker. Set by `ug configure`.", - "items": { - "type": "object", - "required": [ - "id", - "api" - ], - "properties": { - "id": { - "type": "string", - "description": "Model service name, for example system.ai.claude-sonnet-5-5." - }, - "name": { - "type": "string", - "description": "Name shown in the model picker, for example Claude Sonnet 5.5. Defaults to the id." - }, - "api": { - "type": "string", - "enum": [ - "chat_completions", - "responses" - ], - "description": "The AI Gateway API the model is called through." - }, - "maxInputTokens": { - "type": "integer", - "minimum": 1, - "description": "Prompt tokens the model accepts. Defaults to 128000." - }, - "maxOutputTokens": { - "type": "integer", - "minimum": 1, - "description": "Tokens the model can write in one reply. Defaults to 32000." - } - } - }, - "scope": "application", - "ignoreSync": true - }, - "unityGateway.authCommand": { - "type": "string", - "default": "", - "description": "Command that prints a Databricks access token. Set by `ug configure`.", - "scope": "application", - "ignoreSync": true - } - } - }, - "commands": [ - { - "command": "unityGateway.refreshModels", - "title": "Refresh Models", - "category": "Unity Gateway" - } ] } } diff --git a/tests/conftest.py b/tests/conftest.py index 65b0dd3bb..a5450b505 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -50,10 +50,11 @@ def _isolate_ucode_state(tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", state_dir) monkeypatch.setattr(codex_mod, "_otel_service_path", lambda: state_dir / "otel.service") monkeypatch.setattr(desktop_mod, "APP_DIR", state_dir) - # No test may run the developer's `code` or edit their VS Code settings. + # No test may run the developer's `code` or look at their installed VS Code extensions. monkeypatch.setattr(vscode_mod, "find_code", lambda: None) + monkeypatch.setattr(vscode_mod, "_find_command", lambda command: None) monkeypatch.setattr( - vscode_mod, "settings_path", lambda product: tmp_path / product / "User" / "settings.json" + vscode_mod, "extensions_dir", lambda folder: tmp_path / folder / "extensions" ) monkeypatch.setattr(desktop_auth_mod, "_AUTH_LOCK_PATH", state_dir / "claude-desktop-auth.lock") # Antigravity's customization root is shared with the developer's Antigravity app and CLI. diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 4a70eb97d..ccdb1dbdc 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -10,57 +10,58 @@ import pytest from typer.testing import CliRunner -from ucode import config_io from ucode.agents import vscode from ucode.cli import app from ucode.managed_files import OS # The conftest fixture stubs these out for every test; the ones here exercise the real lookups. REAL_FIND_CODE = vscode.find_code -REAL_SETTINGS_PATH = vscode.settings_path +REAL_FIND_COMMAND = vscode._find_command REAL_RUN_CODE = vscode._run_code WS = "https://ws.cloud.databricks.com" STATE = {"workspace": WS, "profile": "dev"} -USER_SETTINGS = """// My settings -{ - // Theme - "workbench.colorTheme": "Default Dark+", - "editor.fontSize": 14, // big -} -""" +VERSION = json.loads((vscode.EXTENSION_DIR / "package.json").read_text())["version"] class FakeCode: - """Stands in for VS Code's `code` command: records calls and tracks the installed version.""" + """Stands in for VS Code's `code` command: records calls and what an install unpacks.""" - def __init__(self, installed: str | None = None, fail: str | None = None) -> None: - self.installed = installed - self.fail = fail + def __init__(self, extensions: Path) -> None: + self.extensions = extensions + self.fail: str | None = None self.calls: list[tuple[str, ...]] = [] + self.paths: list[Path] = [] self.vsix_entries: list[str] = [] + self.config: dict | None = None + + @property + def folder(self) -> Path: + return self.extensions / f"{vscode.EXTENSION_ID}-{VERSION}" def __call__(self, code: Path, *args: str) -> subprocess.CompletedProcess: self.calls.append(args) + self.paths.append(code) if args[0] == self.fail: return subprocess.CompletedProcess(args, 1, "", "boom") if args[0] == "--list-extensions": - listed = f"{vscode.EXTENSION_ID}@{self.installed}\n" if self.installed else "" + listed = f"{vscode.EXTENSION_ID}@{VERSION}\n" if self.folder.exists() else "" return subprocess.CompletedProcess(args, 0, "ms-python.python@2026.1.0\n" + listed, "") if args[0] == "--install-extension": - self.vsix_entries = zipfile.ZipFile(args[1]).namelist() - self.installed = json.loads((vscode.EXTENSION_DIR / "package.json").read_text())[ - "version" - ] + with zipfile.ZipFile(args[1]) as vsix: + self.vsix_entries = vsix.namelist() + self.config = json.loads(vsix.read(f"extension/{vscode.CONFIG_FILE}")) + self.folder.mkdir(parents=True, exist_ok=True) if args[0] == "--uninstall-extension": - self.installed = None + self.folder.rmdir() return subprocess.CompletedProcess(args, 0, "", "") @pytest.fixture def code(monkeypatch): - fake = FakeCode() - monkeypatch.setattr(vscode, "find_code", lambda: (Path("/fake/bin/code"), "Code")) + fake = FakeCode(vscode.extensions_dir(".vscode")) + monkeypatch.setattr(vscode, "find_code", lambda: Path("/fake/bin/code")) + monkeypatch.setattr(vscode, "_find_command", lambda command: Path(f"/fake/bin/{command}")) monkeypatch.setattr(vscode, "_run_code", fake) monkeypatch.setattr(vscode, "launch_token", lambda state, workspace: "token") monkeypatch.setattr( @@ -81,29 +82,26 @@ def _printed(capsys) -> str: return " ".join(capsys.readouterr().out.split()) -def _settings() -> Path: - return vscode.settings_path("Code") - - -def _settings_values() -> dict: - return config_io.loads_jsonc(_settings().read_text(encoding="utf-8")) - - class TestConfigure: - def test_merges_settings_keeping_the_users_own_with_claude_and_gpt_only(self, code): - _settings().parent.mkdir(parents=True) - _settings().write_text(USER_SETTINGS, encoding="utf-8") - + def test_installs_the_extension_with_the_workspace_and_its_claude_and_gpt_models( + self, code, capsys + ): assert vscode.configure(STATE) - text = _settings().read_text(encoding="utf-8") - for kept in ("// My settings", "// Theme", '"editor.fontSize": 14, // big'): - assert kept in text - assert _settings_values() == { - "workbench.colorTheme": "Default Dark+", - "editor.fontSize": 14, - "unityGateway.workspace": WS, - "unityGateway.models": [ + assert code.calls[-1][0] == "--install-extension" + assert code.calls[-1][2] == "--force" + assert code.vsix_entries == [ + "[Content_Types].xml", + "extension.vsixmanifest", + "extension/media/databricks.svg", + "extension/out/extension.js", + "extension/package.json", + "extension/ug-config.json", + ] + assert code.config == { + "workspace": WS, + "authCommand": vscode.build_auth_shell_command(WS, "dev"), + "models": [ { "id": "system.ai.claude-sonnet-5-5", "name": "Claude Sonnet 5.5", @@ -113,8 +111,12 @@ def test_merges_settings_keeping_the_users_own_with_claude_and_gpt_only(self, co }, {"id": "system.ai.gpt-6-luna", "name": "GPT 6 Luna", "api": "responses"}, ], - "unityGateway.authCommand": vscode.build_auth_shell_command(WS, "dev"), } + printed = _printed(capsys) + assert ( + f"VS Code extension: installed {vscode.EXTENSION_ID} {VERSION} with 2 models" in printed + ) + assert "Next, in VS Code: 1. If VS Code is open," in printed def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): monkeypatch.setattr( @@ -129,7 +131,7 @@ def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): ), ) - models, _ = vscode.model_settings(WS, "token") + models, _ = vscode.model_entries(WS, "token") assert [model["name"] for model in models] == [ "system.ai.gpt-6-luna", @@ -137,50 +139,26 @@ def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): "GPT 5.5", ] - def test_creates_missing_settings_and_installs_the_packed_extension(self, code, capsys): - assert vscode.configure(STATE) - - assert set(_settings_values()) == set(vscode.SETTINGS) - assert code.calls[-1][0] == "--install-extension" - assert code.calls[-1][2] == "--force" - assert code.vsix_entries == [ - "[Content_Types].xml", - "extension.vsixmanifest", - "extension/media/databricks.svg", - "extension/out/extension.js", - "extension/package.json", - ] - assert "Next, in VS Code: 1. If VS Code is open," in _printed(capsys) - - def test_reruns_update_values_and_reinstall_the_same_version(self, code): + def test_reruns_reinstall_with_the_new_setup(self, code): vscode.configure(STATE) code.calls.clear() assert vscode.configure({**STATE, "profile": "other"}) - assert "--profile other" in _settings_values()["unityGateway.authCommand"] + assert "--profile other" in code.config["authCommand"] assert [call[0] for call in code.calls] == ["--install-extension"] - def test_leaves_unreadable_settings_alone(self, code, capsys): - _settings().parent.mkdir(parents=True) - _settings().write_text('{"broken": ', encoding="utf-8") - - assert not vscode.configure(STATE) - - assert _settings().read_text(encoding="utf-8") == '{"broken": ' - assert code.calls == [] - assert "Couldn't update" in _printed(capsys) - def test_warns_when_vscode_isnt_found(self, capsys): assert not vscode.configure(STATE) assert "`code` command wasn't found" in _printed(capsys) - assert not _settings().exists() @pytest.mark.parametrize("reason", ["the listing failed", None]) - def test_keeps_settings_when_no_models_are_found(self, code, monkeypatch, capsys, reason): + def test_keeps_the_installed_setup_when_no_models_are_found( + self, code, monkeypatch, capsys, reason + ): vscode.configure(STATE) - before = _settings().read_text(encoding="utf-8") + installed = code.config code.calls.clear() monkeypatch.setattr( vscode, "discover_model_services", lambda workspace, token: ({}, [], [], [], reason) @@ -188,8 +166,8 @@ def test_keeps_settings_when_no_models_are_found(self, code, monkeypatch, capsys assert not vscode.configure({**STATE, "workspace": "https://other.cloud.databricks.com"}) - assert _settings().read_text(encoding="utf-8") == before assert code.calls == [] + assert code.config == installed assert f"No Claude or GPT models for VS Code ({reason or 'the workspace has none'})" in ( _printed(capsys) ) @@ -214,39 +192,45 @@ def test_reports_a_failed_install(self, code, capsys): class TestRevert: - def test_removes_only_ugs_settings_and_uninstalls(self, code): - _settings().parent.mkdir(parents=True) - _settings().write_text(USER_SETTINGS, encoding="utf-8") + def test_uninstalls_the_extension(self, code): vscode.configure(STATE) assert vscode.revert() == "removed" - assert _settings_values() == { - "workbench.colorTheme": "Default Dark+", - "editor.fontSize": 14, - } - assert "// Theme" in _settings().read_text(encoding="utf-8") assert code.calls[-1] == ("--uninstall-extension", vscode.EXTENSION_ID) + assert not code.folder.exists() - def test_does_nothing_when_vscode_was_never_configured(self, code): - _settings().parent.mkdir(parents=True) - _settings().write_text(USER_SETTINGS, encoding="utf-8") - + def test_runs_no_code_command_when_the_extension_isnt_installed(self, code): assert vscode.revert() == "unchanged" - assert _settings().read_text(encoding="utf-8") == USER_SETTINGS assert code.calls == [] + def test_skips_a_folder_vscode_left_behind_after_an_uninstall(self, code): + (code.extensions / f"{vscode.EXTENSION_ID}-0.0.1").mkdir(parents=True) + + assert vscode.revert() == "unchanged" + + assert code.calls == [("--list-extensions", "--show-versions")] + + def test_uninstalls_through_the_insiders_command_for_an_insiders_install(self, code): + code.extensions = vscode.extensions_dir(".vscode-insiders") + vscode.configure(STATE) + + assert vscode.revert() == "removed" + + assert code.paths[-1] == Path("/fake/bin/code-insiders") + class TestFindCode: @pytest.fixture(autouse=True) def _real_lookup(self, monkeypatch): monkeypatch.setattr(vscode, "find_code", REAL_FIND_CODE) + monkeypatch.setattr(vscode, "_find_command", REAL_FIND_COMMAND) def test_prefers_code_on_path(self, monkeypatch): monkeypatch.setattr(vscode.shutil, "which", lambda name: f"/usr/local/bin/{name}") - assert vscode.find_code() == (Path("/usr/local/bin/code"), "Code") + assert vscode.find_code() == Path("/usr/local/bin/code") def test_falls_back_to_insiders(self, monkeypatch): monkeypatch.setattr( @@ -255,7 +239,7 @@ def test_falls_back_to_insiders(self, monkeypatch): lambda name: "/opt/bin/code-insiders" if name == "code-insiders" else None, ) - assert vscode.find_code() == (Path("/opt/bin/code-insiders"), "Code - Insiders") + assert vscode.find_code() == Path("/opt/bin/code-insiders") def test_looks_in_both_macos_application_folders(self, monkeypatch, tmp_path): monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) @@ -278,7 +262,7 @@ def test_falls_back_to_the_first_install_location_that_exists(self, monkeypatch, lambda command: [tmp_path / "missing/code", installed] if command == "code" else [], ) - assert vscode.find_code() == (installed, "Code") + assert vscode.find_code() == installed def test_finds_the_windows_user_install(self, monkeypatch, tmp_path): monkeypatch.setattr(vscode.shutil, "which", lambda name: None) @@ -288,7 +272,7 @@ def test_finds_the_windows_user_install(self, monkeypatch, tmp_path): code_cmd.parent.mkdir(parents=True) code_cmd.touch() - assert vscode.find_code() == (code_cmd, "Code") + assert vscode.find_code() == code_cmd def test_none_when_vscode_isnt_installed(self, monkeypatch): monkeypatch.setattr(vscode.shutil, "which", lambda name: None) @@ -297,31 +281,6 @@ def test_none_when_vscode_isnt_installed(self, monkeypatch): assert vscode.find_code() is None -class TestSettingsPath: - @pytest.mark.parametrize( - ("system", "env", "expected"), - [ - (OS.LINUX, {"XDG_CONFIG_HOME": "/xdg"}, "/xdg/Code/User/settings.json"), - (OS.WINDOWS, {"APPDATA": "/appdata"}, "/appdata/Code/User/settings.json"), - ], - ) - def test_user_settings_per_os(self, monkeypatch, system, env, expected): - monkeypatch.setattr(vscode, "current_os", lambda: system) - for name, value in env.items(): - monkeypatch.setenv(name, value) - - assert REAL_SETTINGS_PATH("Code") == Path(expected) - - def test_macos(self, monkeypatch, tmp_path): - monkeypatch.setattr(vscode, "current_os", lambda: OS.MACOS) - monkeypatch.setattr(vscode.Path, "home", lambda: tmp_path) - - assert ( - REAL_SETTINGS_PATH("Code - Insiders") - == tmp_path / "Library/Application Support/Code - Insiders/User/settings.json" - ) - - class TestNextSteps: @pytest.mark.parametrize( ("system", "restart"), @@ -402,4 +361,4 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, assert shared[0]["tools"] == [] # A dry run saves no workspace, so VS Code would be set up from stale or missing state. assert configured == ([] if dry_run else [STATE]) - assert ("VS Code's settings and extension aren't changed" in result.output) == dry_run + assert ("VS Code isn't changed" in result.output) == dry_run diff --git a/tests/test_cli.py b/tests/test_cli.py index b0e9cf5cf..f460fc67b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3989,7 +3989,7 @@ def test_vscode_revert_failure_is_reported_and_the_rest_still_reverts(self): output = " ".join(result.output.split()) assert result.exit_code == 0, result.output assert "Could not revert VS Code: code timed out" in output - assert "VS Code extension and settings: failed" in output + assert "VS Code extension: failed" in output assert restore_file.called clear_state.assert_called_once() diff --git a/tests/test_config_io.py b/tests/test_config_io.py index 6e11ca2b9..1cb325ee4 100644 --- a/tests/test_config_io.py +++ b/tests/test_config_io.py @@ -405,75 +405,3 @@ def test_partial_path_through_non_dict_is_noop(self): doc = {"model": "scalar"} assert prune_key_paths(doc, [["model", "nested"]]) is False assert doc == {"model": "scalar"} - - -class TestJsonc: - SETTINGS = """// Top comment -{ - // Theme - "workbench.colorTheme": "Default Dark+", /* inline */ - "url": "https://example.com/a//b", - "editor.fontSize": 14, // big -} -""" - - def test_loads_comments_and_trailing_commas(self): - assert config_io.loads_jsonc(self.SETTINGS) == { - "workbench.colorTheme": "Default Dark+", - "url": "https://example.com/a//b", - "editor.fontSize": 14, - } - assert config_io.loads_jsonc('{"a": [1, 2,], "b": "/* not a comment */",}') == { - "a": [1, 2], - "b": "/* not a comment */", - } - - def test_sets_keys_and_keeps_everything_else(self): - updated = config_io.update_jsonc_keys( - self.SETTINGS, {"editor.fontSize": 16, "new.key": {"x": [1]}} - ) - - assert config_io.loads_jsonc(updated)["editor.fontSize"] == 16 - assert config_io.loads_jsonc(updated)["new.key"] == {"x": [1]} - for kept in ( - "// Top comment", - "// Theme", - "/* inline */", - "// big", - '"url": "https://example.com/a//b"', - ): - assert kept in updated - assert ' "new.key": {\n "x": [\n 1\n ]\n }' in updated - - def test_removes_keys_with_their_line_comment(self): - updated = config_io.update_jsonc_keys( - self.SETTINGS, {}, remove=("editor.fontSize", "missing") - ) - - assert config_io.loads_jsonc(updated) == { - "workbench.colorTheme": "Default Dark+", - "url": "https://example.com/a//b", - } - assert "// big" not in updated - assert "// Theme" in updated - - def test_removing_the_last_key_drops_the_comma_before_it(self): - assert ( - config_io.update_jsonc_keys('{\n "a": 1,\n "b": 2\n}\n', {}, remove=("b",)) - == '{\n "a": 1\n}\n' - ) - assert config_io.update_jsonc_keys('{"a": 1}', {}, remove=("a",)) == "{}" - - @pytest.mark.parametrize("text", ["", "// only a comment\n", "{}"]) - def test_fills_an_empty_file(self, text): - updated = config_io.update_jsonc_keys(text, {"a": 1}) - - assert config_io.loads_jsonc(updated) == {"a": 1} - assert text.strip("{}") in updated - - @pytest.mark.parametrize( - "text", ['{"a": ', "[1, 2]", '{"a": 1} {"b": 2}', '{"a" 1}', '{"a": "x'] - ) - def test_rejects_what_isnt_one_object(self, text): - with pytest.raises(ValueError): - config_io.update_jsonc_keys(text, {"b": 2}) diff --git a/vscode-extension/README.md b/vscode-extension/README.md index d49c502a9..b79ed3ea4 100644 --- a/vscode-extension/README.md +++ b/vscode-extension/README.md @@ -1,20 +1,16 @@ # Unity Gateway for VS Code Adds Databricks AI Gateway models to the VS Code chat model picker, for Copilot Chat in Ask and Agent -mode. The extension calls the workspace's AI Gateway directly: Chat Completions for Claude and open -models, Responses for GPT-5 and later. +mode. The extension calls the workspace's AI Gateway directly: Chat Completions for Claude, Responses +for GPT-5 and later. -`ug configure` installs it and writes its settings: - -| Setting | Value | -|---|---| -| `unityGateway.workspace` | The workspace URL | -| `unityGateway.models` | `[{"id": "", "api": "chat_completions" \| "responses"}]` | -| `unityGateway.authCommand` | A command that prints a Databricks access token (`ug auth-token --host ...`) | +`ug configure --agents vscode` packs the extension with a `ug-config.json` that holds the workspace URL, +the models (id, picker name, `chat_completions` or `responses`, token limits) and a command that prints +a Databricks access token (`ug auth-token --host ...`), then installs it. To change the models or the +workspace, run `ug configure --agents vscode` again and restart VS Code. The token is kept in memory until it nears expiry, and fetched again when the gateway answers 401. -Run **Unity Gateway: Refresh Models** after changing the settings by hand. Token fetches and failed -requests are logged to the **Unity Gateway** output channel. +Token fetches and failed requests are logged to the **Unity Gateway** output channel. ## Development diff --git a/vscode-extension/package.json b/vscode-extension/package.json index 2f9bfa8f2..dc5d88801 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -28,74 +28,6 @@ "displayName": "Databricks", "icon": "media/databricks.svg" } - ], - "configuration": { - "title": "Databricks Unity Gateway", - "properties": { - "unityGateway.workspace": { - "type": "string", - "default": "", - "description": "Databricks workspace URL, for example https://my-workspace.cloud.databricks.com. Set by `ug configure`.", - "scope": "application", - "ignoreSync": true - }, - "unityGateway.models": { - "type": "array", - "default": [], - "description": "Models to offer in the chat model picker. Set by `ug configure`.", - "items": { - "type": "object", - "required": [ - "id", - "api" - ], - "properties": { - "id": { - "type": "string", - "description": "Model service name, for example system.ai.claude-sonnet-5-5." - }, - "name": { - "type": "string", - "description": "Name shown in the model picker, for example Claude Sonnet 5.5. Defaults to the id." - }, - "api": { - "type": "string", - "enum": [ - "chat_completions", - "responses" - ], - "description": "The AI Gateway API the model is called through." - }, - "maxInputTokens": { - "type": "integer", - "minimum": 1, - "description": "Prompt tokens the model accepts. Defaults to 128000." - }, - "maxOutputTokens": { - "type": "integer", - "minimum": 1, - "description": "Tokens the model can write in one reply. Defaults to 32000." - } - } - }, - "scope": "application", - "ignoreSync": true - }, - "unityGateway.authCommand": { - "type": "string", - "default": "", - "description": "Command that prints a Databricks access token. Set by `ug configure`.", - "scope": "application", - "ignoreSync": true - } - } - }, - "commands": [ - { - "command": "unityGateway.refreshModels", - "title": "Refresh Models", - "category": "Unity Gateway" - } ] }, "scripts": { diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 3146f0263..2ba3da005 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -1,3 +1,5 @@ +import { readFileSync } from "node:fs"; +import { join } from "node:path"; import * as vscode from "vscode"; import { version } from "../package.json"; import { Tokens } from "./auth"; @@ -26,7 +28,21 @@ const tokenLimits = (model: ModelSetting): Partial => }), ); -export const readSettings = (get: (key: string) => unknown): Settings => { +/** `ug configure` writes this file into the extension it packs, so each install carries its own setup. */ +export const CONFIG_FILE = "ug-config.json"; + +/** Reads the config file's text; a missing or malformed file gives no models. */ +export const readConfig = (text: string | undefined): Settings => { + let config: Record = {}; + try { + const parsed: unknown = JSON.parse(text ?? "{}"); + if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) { + config = parsed as Record; + } + } catch { + // Treated like a missing file. + } + const get = (key: string) => config[key]; const models = get("models"); return { workspace: String(get("workspace") ?? "").replace(/\/+$/, ""), @@ -52,8 +68,6 @@ export const userAgent = (copilotChatVersion: unknown): string => type Model = vscode.LanguageModelChatInformation & { readonly api: ModelSetting["api"] }; export class DatabricksModels implements vscode.LanguageModelChatProvider { - private readonly changed = new vscode.EventEmitter(); - readonly onDidChangeLanguageModelChatInformation = this.changed.event; private tokens: Tokens | undefined; constructor( @@ -62,12 +76,6 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider private readonly headers: () => Record = () => ({}), ) {} - /** Re-reads the settings: drops the cached token and asks VS Code to list the models again. */ - refresh(): void { - this.tokens = undefined; - this.changed.fire(); - } - async provideLanguageModelChatInformation(): Promise { const { workspace, models, authCommand } = this.settings(); if (!workspace || !authCommand || models.length === 0) { @@ -124,22 +132,19 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider } } +const readConfigFile = (path: string): string | undefined => { + try { + return readFileSync(path, "utf8"); + } catch { + return undefined; + } +}; + export function activate(context: vscode.ExtensionContext): void { const output = vscode.window.createOutputChannel("Unity Gateway", { log: true }); - const settings = () => { - const configuration = vscode.workspace.getConfiguration("unityGateway"); - return readSettings((key) => configuration.get(key)); - }; + // Read once: ug changes the file only by reinstalling the extension, which takes a restart to load. + const settings = readConfig(readConfigFile(join(context.extensionPath, CONFIG_FILE))); const headers = () => ({ "User-Agent": userAgent(vscode.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version) }); - const provider = new DatabricksModels(settings, (line) => output.info(line), headers); - context.subscriptions.push( - output, - vscode.lm.registerLanguageModelChatProvider("databricks-unity-gateway", provider), - vscode.commands.registerCommand("unityGateway.refreshModels", () => provider.refresh()), - vscode.workspace.onDidChangeConfiguration((event) => { - if (event.affectsConfiguration("unityGateway")) { - provider.refresh(); - } - }), - ); + const provider = new DatabricksModels(() => settings, (line) => output.info(line), headers); + context.subscriptions.push(output, vscode.lm.registerLanguageModelChatProvider("databricks-unity-gateway", provider)); } diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index d8e3c8ff2..537380ec8 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -4,7 +4,7 @@ import { afterEach, describe, it } from "node:test"; import type * as vscode from "vscode"; import packageJson from "../package.json"; import { Tokens } from "../src/auth"; -import { DatabricksModels, readSettings, userAgent, type Settings } from "../src/extension"; +import { DatabricksModels, readConfig, userAgent, type Settings } from "../src/extension"; import { streamEvents } from "../src/gateway"; import { collect, message, rejectsWith, text, tokenCommand, USER } from "./support"; import { LanguageModelChatToolMode, LanguageModelTextPart } from "./vscode"; @@ -67,20 +67,6 @@ describe("streamEvents", () => { }); }); -describe("settings", () => { - // Workspace settings mustn't redirect the token or choose the command that prints it. - it("are read from user settings only", () => { - const properties = Object.values(packageJson.contributes.configuration.properties); - assert.deepEqual(properties.map((property) => property.scope), ["application", "application", "application"]); - }); - - // ug writes them per machine (the auth command holds this machine's ug path and profile). - it("stay out of Settings Sync", () => { - const properties = Object.values(packageJson.contributes.configuration.properties); - assert.deepEqual(properties.map((property) => property.ignoreSync), [true, true, true]); - }); -}); - describe("provider", () => { // The icon lives only in the folder ug ships, next to the built bundle. it("ships the icon it names", () => { @@ -95,9 +81,9 @@ describe("userAgent", () => { }); }); -describe("readSettings", () => { +describe("readConfig", () => { it("keeps well-formed models, names and token limits and trims the workspace URL", () => { - const values: Record = { + const config = { workspace: "https://ws/", authCommand: "ug auth-token", models: [ @@ -109,7 +95,7 @@ describe("readSettings", () => { { id: "d", name: 7, api: "chat_completions", maxInputTokens: 0, maxOutputTokens: 1.5 }, ], }; - assert.deepEqual(readSettings((key) => values[key]), { + assert.deepEqual(readConfig(JSON.stringify(config)), { workspace: "https://ws", authCommand: "ug auth-token", models: [ @@ -119,6 +105,13 @@ describe("readSettings", () => { ], }); }); + + it("gives no models for a missing or malformed file", () => { + const none = { workspace: "", authCommand: "", models: [] }; + for (const text of [undefined, "", "{", "[]", "null", '"text"']) { + assert.deepEqual(readConfig(text), none, String(text)); + } + }); }); describe("DatabricksModels", () => { From a329e59e12b3243d4cb48de156cfb324c8548bb7 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 14:33:44 +0000 Subject: [PATCH 17/18] Restore the dry-run flag after the VS Code dry-run CLI test `ug configure --dry-run` sets config_io's module-wide dry-run flag and leaves it set. The VS Code CLI wiring test ran it in-process, so every test after it in the same pytest run skipped its file writes, and dozens of unrelated tests (test_agents_commands, test_claude_*, ...) failed in CI. Reset the flag through monkeypatch so it's restored after the test. Co-authored-by: Isaac --- tests/test_agent_vscode.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index ccdb1dbdc..874636b6c 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -10,6 +10,7 @@ import pytest from typer.testing import CliRunner +from ucode import config_io from ucode.agents import vscode from ucode.cli import app from ucode.managed_files import OS @@ -334,6 +335,8 @@ def test_configure_agents_vscode_refuses_a_workspace_the_machine_isnt_managed_fo @pytest.mark.parametrize("dry_run", [False, True]) def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, dry_run): + # `--dry-run` sets a module-wide flag; this puts it back so later tests write files again. + monkeypatch.setattr(config_io, "_dry_run", False) configured: list[dict] = [] shared: list[dict] = [] monkeypatch.setattr( From a832b61b642b2c620f516199104b6d399cbdb2e6 Mon Sep 17 00:00:00 2001 From: Xi Lyu Date: Thu, 29 Oct 2026 14:43:53 +0000 Subject: [PATCH 18/18] Match the code path in a VS Code test on Windows too The hung-`code` test expected "/fake/bin/code" in the warning, but on Windows the path prints as "\fake\bin\code". Build the expected text from the same Path. Co-authored-by: Isaac --- tests/test_agent_vscode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 874636b6c..48c55402a 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -182,7 +182,7 @@ def hang(*args, **kwargs): assert not vscode.configure(STATE) - assert "`/fake/bin/code --install-extension` failed" in _printed(capsys) + assert f"`{Path('/fake/bin/code')} --install-extension` failed" in _printed(capsys) def test_reports_a_failed_install(self, code, capsys): code.fail = "--install-extension"