diff --git a/README.md b/README.md index 95b67efa4..e96be0ade 100644 --- a/README.md +++ b/README.md @@ -590,7 +590,7 @@ preferences are ignored and their files are left untouched. See the bundled | Claude Code | `~/.claude/ucode-settings.json`, `~/.claude.json`, `/etc/claude-code/managed-settings.json` (Linux), `/Library/Application Support/ClaudeCode/managed-settings.json` (macOS), `C:\Program Files\ClaudeCode\managed-settings.json` (Windows, written only from an elevated terminal) | | Gemini CLI | `~/.gemini/ucode.env`, `~/.ucode/.gemini-home/.gemini/settings.json` | | OpenCode | `~/.ucode/opencode-xdg/opencode/opencode.json`, `~/.ucode/opencode-xdg/opencode/plugin/ucode-auth.js` | -| GitHub Copilot CLI | `~/.copilot/ucode.env`, `~/.copilot/ucode-mcp-config.json` | +| GitHub Copilot CLI | `~/.copilot/ucode.env`, `~/.copilot/mcp-config.json` (or `$COPILOT_HOME/mcp-config.json`) | | Pi | `~/.ucode/pi-home/.pi/agent/models.json`, `~/.ucode/pi-home/.pi/agent/settings.json` | | Antigravity CLI | `~/.gemini/config/skills.json` (adds `~/.agents/skills`), `~/.gemini/config/mcp_config.json` (ug's MCP servers only); the gateway is set per launch through `AGY_LLM_GATEWAY_*` variables | | Cursor Agent | `~/.cursor/mcp.json` | diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 46b666bc0..d1b649d58 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -17,7 +17,7 @@ from collections.abc import Callable from pathlib import Path -from ucode.config_io import ToolSpec +from ucode.config_io import ToolSpec, is_dry_run from ucode.constants import AGENT_CLAUDE, AGENT_CODEX, AGENT_COPILOT from ucode.databricks import ( AnthropicModelCatalog, @@ -686,10 +686,38 @@ def _availability_failure_detail(tool: str, state: dict) -> str: return " (" + "; ".join(parts) + ")" +_vscode_with_copilot = True + + +def request_vscode_with_copilot(requested: bool = True) -> None: + """Whether configuring Copilot in this run also sets up VS Code (``ug configure --no-vscode`` + passes ``False``).""" + global _vscode_with_copilot + _vscode_with_copilot = requested + + +def configure_vscode(state: dict, *, provider: str | None, parent_schema: str | None) -> None: + """Set up Copilot Chat in VS Code with the models the Copilot CLI was configured with.""" + if not _vscode_with_copilot: + return + if is_dry_run(): + print_note("Dry run: VS Code isn't changed.") + return + from . import vscode + + vscode.configure(state, provider=provider, location=parent_schema) + + def configure_single_tool( - tool: str, state: dict, *, parent_schema: str | None = None, provider: str | None = None + tool: str, + state: dict, + *, + parent_schema: str | None = None, + provider: str | None = None, + with_vscode: bool = False, ) -> dict: - """Check availability, configure, and persist state for one tool only.""" + """Check availability, configure, and persist state for one tool only. ``with_vscode`` also + sets up VS Code after Copilot.""" parent_schema = parent_schema or _persisted_model_location(tool, state, provider) provider = None if parent_schema else provider or get_provider_service(state, tool) # A Model Provider Service or parent schema routes through the same gateway and pins no @@ -704,6 +732,8 @@ def configure_single_tool( ) with managed_write_batch(_managed_settings_displays([tool])): state = _configure_one(tool, state, provider, parent_schema=parent_schema) + if with_vscode and tool == AGENT_COPILOT: + configure_vscode(state, provider=provider, parent_schema=parent_schema) available_tools = list(set((state.get("available_tools") or []) + [tool])) state["available_tools"] = available_tools save_state(state) @@ -788,6 +818,9 @@ def configure_selected_tools( ) continue configured.append(tool) + # `ug configure` sets up Copilot Chat in VS Code along with the Copilot CLI. + if tool == AGENT_COPILOT: + configure_vscode(state, provider=provider, parent_schema=parent_schema) existing = state.get("available_tools") or [] state["available_tools"] = sorted(set(existing) | set(configured)) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index ed979b368..e67792638 100644 --- a/src/ucode/agents/copilot.py +++ b/src/ucode/agents/copilot.py @@ -41,8 +41,8 @@ from ucode.config_io import ( APP_DIR, ToolSpec, - apply_json_mcp_diff, backup_existing_file, + is_dry_run, parse_dotenv, read_json_safe, write_dotenv, @@ -88,9 +88,11 @@ COPILOT_CONFIG_DIR = Path.home() / ".copilot" COPILOT_ENV_PATH = COPILOT_CONFIG_DIR / "ucode.env" -COPILOT_MCP_CONFIG_PATH = COPILOT_CONFIG_DIR / "ucode-mcp-config.json" +COPILOT_MCP_CONFIG_PATH = COPILOT_CONFIG_DIR / "mcp-config.json" COPILOT_BACKUP_PATH = APP_DIR / "copilot-ucode-env.backup" -COPILOT_MCP_BACKUP_PATH = APP_DIR / "copilot-ucode-mcp-config.backup.json" +# Older ug versions kept their servers in a side file loaded with `--additional-mcp-config`. +LEGACY_MCP_CONFIG_PATH = COPILOT_CONFIG_DIR / "ucode-mcp-config.json" +LEGACY_MCP_BACKUP_PATH = APP_DIR / "copilot-ucode-mcp-config.backup.json" COPILOT_PROXY_LOG_PATH = APP_DIR / "logs" / "copilot-proxy.log" _REGISTRY_PREFIX = "ug-providers-" _REGISTRY_LOCK_NAME = "ug-providers.lock" @@ -800,35 +802,90 @@ def build_mcp_server_entry(argv: list[str]) -> dict: } -def write_mcp_server_config(name: str, argv: list[str]) -> bool: - backup_existing_file(COPILOT_MCP_CONFIG_PATH, COPILOT_MCP_BACKUP_PATH) - existing = read_json_safe(COPILOT_MCP_CONFIG_PATH) - mcp_servers = existing.get("mcpServers") - if not isinstance(mcp_servers, dict): - mcp_servers = {} - removed = name in mcp_servers - mcp_servers[name] = build_mcp_server_entry(argv) - existing["mcpServers"] = mcp_servers - write_json_file(COPILOT_MCP_CONFIG_PATH, existing) +def copilot_mcp_config_path() -> Path: + """The user config Copilot reads ``mcpServers`` from (and ``copilot mcp add`` writes).""" + copilot_home = os.environ.get("COPILOT_HOME") + return Path(copilot_home) / "mcp-config.json" if copilot_home else COPILOT_MCP_CONFIG_PATH + + +def _read_mcp_config_for_rewrite(path: Path) -> dict | None: + """``path`` parsed for a rewrite: ``{}`` when absent, ``None`` when it isn't a JSON object, so + a caller never overwrites it. Copilot also accepts JSON with comments, which lands here.""" + if not path.exists(): + return {} + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError): + return None + return data if isinstance(data, dict) else None + + +def remove_legacy_mcp_config() -> bool: + """Delete the legacy side file and its backup (everything in them is ug's).""" + existed = LEGACY_MCP_CONFIG_PATH.exists() or LEGACY_MCP_BACKUP_PATH.exists() + if is_dry_run(): + return existed + for path in (LEGACY_MCP_CONFIG_PATH, LEGACY_MCP_BACKUP_PATH): + try: + path.unlink(missing_ok=True) + except OSError as exc: + raise RuntimeError(f"Failed to remove {path}") from exc + return existed + + +def _read_legacy_mcp_servers() -> dict[str, dict]: + """The servers in the legacy side file. Every entry there was written by ug, so it moves into + the user config as-is.""" + servers = read_json_safe(LEGACY_MCP_CONFIG_PATH).get("mcpServers") + if not isinstance(servers, dict): + return {} + return {name: entry for name, entry in servers.items() if isinstance(entry, dict)} + + +def _apply_mcp_diff(add: dict[str, dict], remove: set[str]) -> set[str]: + path = copilot_mcp_config_path() + config = _read_mcp_config_for_rewrite(path) + if config is None: + raise RuntimeError( + f"{path} isn't plain JSON (it may have comments). Remove the comments, then try " + "again; VS Code's local agent also skips the file while it has them." + ) + legacy = _read_legacy_mcp_servers() + upserts = {name: entry for name, entry in legacy.items() if name not in remove} | add + servers = config.get("mcpServers") + if not isinstance(servers, dict): + servers = {} + removed = {name for name in remove if name in servers or name in legacy} + if upserts or removed: + for name in remove: + servers.pop(name, None) + servers.update(upserts) + config["mcpServers"] = servers + write_json_file(path, config) + remove_legacy_mcp_config() return removed +def migrate_legacy_mcp_config() -> None: + """Move the servers in the legacy side file into the user config, once.""" + if LEGACY_MCP_CONFIG_PATH.exists() or LEGACY_MCP_BACKUP_PATH.exists(): + _apply_mcp_diff({}, set()) + + +def write_mcp_server_config(name: str, argv: list[str]) -> bool: + """Add (or replace) a stdio (`ug mcp-proxy`) MCP server in Copilot's user config. Returns + True when an entry of that name was replaced.""" + return bool(_apply_mcp_diff({name: build_mcp_server_entry(argv)}, {name})) + + def remove_mcp_server_config(name: str) -> bool: - existing = read_json_safe(COPILOT_MCP_CONFIG_PATH) - mcp_servers = existing.get("mcpServers") - if not isinstance(mcp_servers, dict) or name not in mcp_servers: - return False - mcp_servers.pop(name) - existing["mcpServers"] = mcp_servers - write_json_file(COPILOT_MCP_CONFIG_PATH, existing) - return True + return bool(_apply_mcp_diff({}, {name})) def write_user_mcp_servers(add: dict[str, dict], remove: set[str]) -> set[str]: - """Apply ``add``/``remove`` to Copilot's `mcpServers` in a single read-modify-write. Returns the names actually removed.""" - return apply_json_mcp_diff( - COPILOT_MCP_CONFIG_PATH, "mcpServers", add, remove, backup_path=COPILOT_MCP_BACKUP_PATH - ) + """Apply ``add``/``remove`` to Copilot's user-config ``mcpServers`` in a single + read-modify-write. Returns the names actually removed.""" + return _apply_mcp_diff(add, remove) def scrub_stored_credentials() -> None: @@ -876,7 +933,11 @@ def _token_provider(state: dict) -> Callable[[bool], str]: def proxy_headers(state: dict) -> dict[str, str]: """Headers the proxy adds to every gateway request; admin-managed headers win by name.""" headers = { - "User-Agent": ucode_user_agent(ug_version(), f"copilot/{agent_version(SPEC['binary'])}") + "User-Agent": ucode_user_agent(ug_version(), f"copilot/{agent_version(SPEC['binary'])}"), + # Makes the gateway map or drop the Anthropic betas its serving backend lacks. Without it, + # Copilot's tool search beta (`advanced-tool-use-2025-11-20`, sent once it has enough MCP + # tools) gets a 400 from Claude served through Bedrock. + "x-databricks-use-coding-agent-mode": "true", } request_tags = request_tags_header_value() if request_tags: @@ -951,6 +1012,10 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None managed_list=not (provider_service or model_location), ) scrub_stored_credentials() + try: + migrate_legacy_mcp_config() + except RuntimeError as exc: + print_warning(f"ug's MCP servers aren't available in this session: {exc}") # Only Copilot launches need the proxy (and its `httpx` import). from ucode import gateway_proxy @@ -986,7 +1051,7 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None with session_registry(config) as registry_path: env = build_runtime_env(registry_path, start_model) proc = subprocess_cross_os.popen( - [SPEC["binary"], *mcp_config_args(), *tool_args], env=env + [SPEC["binary"], *MCP_LAUNCH_ARGS, *tool_args], env=env ) returncode = _wait_for_exit(proc) finally: @@ -1000,14 +1065,13 @@ def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None def validate_cmd(binary: str) -> list[str]: return [ binary, - *mcp_config_args(), + *MCP_LAUNCH_ARGS, "--prompt", "say hi in 5 words or less", "--allow-all-tools", ] -def mcp_config_args() -> list[str]: - if not COPILOT_MCP_CONFIG_PATH.exists(): - return [] - return ["--additional-mcp-config", f"@{COPILOT_MCP_CONFIG_PATH}"] +# Offline mode empties Copilot's allowlist of servers whose `instructions` reach the model, so +# without this the gateway's servers lose theirs. +MCP_LAUNCH_ARGS = ("--allow-all-mcp-server-instructions",) diff --git a/src/ucode/agents/vscode.py b/src/ucode/agents/vscode.py index 72c354e00..66cbe553a 100644 --- a/src/ucode/agents/vscode.py +++ b/src/ucode/agents/vscode.py @@ -37,7 +37,8 @@ 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 +from ucode.ui import console, print_heading, print_kv, print_note, print_warning +from ucode.workspace import selector_headers EXTENSION_DIR = Path(__file__).resolve().parent.parent / "vscode" / "extension" EXTENSION_ID = "databricks.unity-gateway" @@ -251,9 +252,10 @@ def configure(state: dict, *, provider: str | None = None, location: str | None 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`." + print_note( + "VS Code's `code` command wasn't found, so VS Code was skipped. To add it, install VS " + "Code (on macOS, also run \"Shell Command: Install 'code' command in PATH\"), then " + "rerun `ug configure --agents copilot`." ) return False try: @@ -269,6 +271,9 @@ def configure(state: dict, *, provider: str | None = None, location: str | None } if provider: config["provider"] = provider + # A host that serves several workspaces answers a request without them with its sign-in page. + if selector := selector_headers(state.get("workspace_id")): + config["headers"] = selector source = f" from {provider or location}" if provider or location else "" try: print_kv( @@ -276,7 +281,7 @@ def configure(state: dict, *, provider: str | None = None, location: str | None f"{install_extension(code, config)} with {len(models)} models{source}", ) except RuntimeError as exc: - print_warning(f"{exc} Rerun `ug configure --agents vscode` to try again.") + print_warning(f"{exc} Rerun `ug configure --agents copilot` to try again.") return False print_heading("Next, in VS Code:") for line in next_steps(): diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a83fb3ca5..094442de2 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -43,6 +43,7 @@ normalize_tool, refuse_user_scope_over_managed_config, refuse_user_scope_over_managed_files, + request_vscode_with_copilot, resolve_gemini_provider_model, resolve_launch_model, resolve_provider_models, @@ -1015,7 +1016,7 @@ def _setup_single_agent( state = resolve_state(managed, state, tool) if not managed_provider_service(managed, tool): parent_schema = managed_unity_catalog_location(managed, tool) - state = configure_single_tool(tool, state, parent_schema=parent_schema) + state = configure_single_tool(tool, state, parent_schema=parent_schema, with_vscode=True) if apply_managed and not managed_enabled_tools(managed or {}): state = _record_unmanaged_configured_agents(state, [tool]) if apply_managed and tool == "claude": @@ -1035,12 +1036,11 @@ def _setup_single_agent( # Only the Claude and Codex branch above reads the managed config fresh. managed_refreshed = apply_managed and tool in ("claude", "codex") install_databricks_ai_tools_for_agents([tool], state, force_refresh=not managed_refreshed) - spec = TOOL_SPECS[tool] provider_summary = "Databricks" if parent_schema else _provider_summary(tool, state) console.print( Panel( f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]\n" - f"[bold]{spec['display']}:[/bold] [green]configured[/green] " + f"[bold]{_summary_label(tool)}:[/bold] [green]configured[/green] " f"[dim](Provider: {provider_summary})[/dim]", title="Configuration Complete", style="green", @@ -1195,12 +1195,16 @@ def configure_workspace_command( databricks_ai_tools_enabled: bool | None = None, custom_oauth: CustomOAuthConfig | None = None, offer_optional_setup: bool = False, + copilot_source: tuple[str | None, str | None] | None = None, ) -> int: """Configure a workspace while sharing one lazy privileged settings session. Agent setup and managed MCP reconciliation can update the same machine-wide Claude/Codex files at different points in the flow. Keeping one command-scoped worker means every changed file is handled under the same sudo authentication; a no-op configure never starts it. + + ``copilot_source`` is the ``(provider service, model location)`` pair Copilot gets its models + from instead of the Databricks-hosted ones. """ with managed_write_session(): return _configure_workspace_command( @@ -1211,6 +1215,7 @@ def configure_workspace_command( databricks_ai_tools_enabled=databricks_ai_tools_enabled, custom_oauth=custom_oauth, offer_optional_setup=offer_optional_setup, + copilot_source=copilot_source, ) @@ -1223,6 +1228,7 @@ def _configure_workspace_command( databricks_ai_tools_enabled: bool | None = None, custom_oauth: CustomOAuthConfig | None = None, offer_optional_setup: bool = False, + copilot_source: tuple[str | None, str | None] | None = None, ) -> int: if tool is not None and selected_tools is not None: raise RuntimeError("Use either --agent or --agents, not both.") @@ -1266,6 +1272,11 @@ def _configure_workspace_command( refuse_user_scope_over_managed_config(state.get("workspace"), managed) managed_tools = managed_enabled_tools(managed) if managed is not None else [] if managed is not None and managed_tools: + if copilot_source is not None: + raise RuntimeError( + "--provider and --model-location can't be used on a workspace with a managed " + "configuration; its admin chooses the models." + ) refuse_user_scope_over_managed_files(managed_tools) configured_tools: list[str] = [] for tool_name in managed_tools: @@ -1337,7 +1348,7 @@ def _configure_workspace_command( raise RuntimeError("No coding agents are available on this workspace.") if selected_tools is None: - picked = prompt_for_tools([(t, TOOL_SPECS[t]["display"]) for t in available_on_workspace]) + picked = prompt_for_tools([(t, _agent_label(t)) for t in available_on_workspace]) else: unavailable_tools = [ tool_name for tool_name in selected_tools if tool_name not in available_on_workspace @@ -1366,6 +1377,17 @@ def _configure_workspace_command( if offer_provider: for tool_name in picked: state = _maybe_select_provider_service(tool_name, state) + if copilot_source is not None: + provider_service, model_location = copilot_source + # Resolve the source the way `ug copilot` will, so a bad one fails here instead of being + # saved over a working one. + if provider_service: + copilot_agent.provider_service_catalog(state, provider_service, None) + else: + copilot_agent.location_catalog(state, str(model_location)) + # Saved, so `ug copilot` launches on the same models VS Code gets. + state = set_provider_service(state, AGENT_COPILOT, provider_service) + state = set_model_location(state, AGENT_COPILOT, model_location) if offer_optional_setup: state = configure_selected_tools(state, picked, install_ai_tools=False) @@ -1398,16 +1420,15 @@ def _configure_workspace_command( _sync_claude_cowork_plugin(state) summary_lines = [f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]"] for tool_name in picked: - spec = TOOL_SPECS[tool_name] + label = _summary_label(tool_name) if tool_name in configured_set: summary_lines.append( - f"[bold]{spec['display']}:[/bold] [green]configured[/green] " + f"[bold]{label}:[/bold] [green]configured[/green] " f"[dim](Provider: {_provider_summary(tool_name, state)})[/dim]" ) else: summary_lines.append( - f"[bold]{spec['display']}:[/bold] [yellow]not configured " - "(see warnings above)[/yellow]" + f"[bold]{label}:[/bold] [yellow]not configured (see warnings above)[/yellow]" ) console.print( Panel( @@ -1422,6 +1443,22 @@ def _configure_workspace_command( return 0 +def _agent_label(tool: str) -> str: + """An agent's name where developers choose and review agents: the picker, status, list.""" + if tool == AGENT_COPILOT: + # Configuring copilot also sets up Copilot Chat in VS Code. + return "GitHub Copilot (CLI and VS Code)" + return str(TOOL_SPECS[tool]["display"]) + + +def _summary_label(tool: str) -> str: + """An agent's name in ``ug configure``'s Configuration Complete box, which also covers VS Code + for copilot.""" + if tool == AGENT_COPILOT: + return "GitHub Copilot" + return str(TOOL_SPECS[tool]["display"]) + + def _print_status_panel(title: str, rows: list[tuple[str, str]]) -> None: """Render a compact two-column status card with values that wrap safely.""" table = Table.grid(padding=(0, 2)) @@ -1612,7 +1649,7 @@ def status() -> int: model_state, model_freshness = _live_status_model_state(state, configured_tools) print_heading("Coding Agents") skill_counts_by_agent = configured_skill_counts_by_agent(state, TOOL_SPECS) - for tool, spec in TOOL_SPECS.items(): + for tool in TOOL_SPECS: if tool not in configured_tools: continue effective_state = resolve_state(managed, model_state, tool) if managed else model_state @@ -1667,7 +1704,7 @@ def status() -> int: if isinstance(base_url, dict): base_url = ", ".join(str(url) for url in base_url.values()) rows.append(("Endpoint", str(base_url or "not configured"))) - _print_status_panel(str(spec["display"]), rows) + _print_status_panel(_agent_label(tool), rows) if not configured_tools: print_note("No coding agents are configured.") @@ -1998,7 +2035,7 @@ def agents_list_cmd() -> None: if management is AgentManagement.UNMANAGED: continue tone = "ok" if management is AgentManagement.ADMIN else "info" - table.add_row(TOOL_SPECS[tool]["display"], status_badge(management.value, tone)) + table.add_row(_agent_label(tool), status_badge(management.value, tone)) if table.row_count: console.print(table) else: @@ -4651,7 +4688,7 @@ def configure( str | None, typer.Option( "--agents", - help="Configure a comma-separated list of agents without prompting (e.g. claude,codex,vscode).", + help="Configure a comma-separated list of agents without prompting (e.g. claude,codex,copilot).", ), ] = None, workspace: Annotated[ @@ -4752,18 +4789,26 @@ def configure( str | None, typer.Option( "--provider", - help="With --agents vscode: give VS Code the models of this Model Provider Service " - "(`..`) instead of the Databricks-hosted ones.", + help="With --agents copilot: give the Copilot CLI and VS Code the models of this Model " + "Provider Service (`..`) instead of the Databricks-hosted ones.", ), ] = None, model_location: Annotated[ str | None, typer.Option( "--model-location", - help="With --agents vscode: give VS Code the model services in this Unity Catalog " - "schema (`.`) instead of the Databricks-hosted ones.", + help="With --agents copilot: give the Copilot CLI and VS Code the model services in this " + "Unity Catalog schema (`.`) instead of the Databricks-hosted ones.", ), ] = None, + no_vscode: Annotated[ + bool, + typer.Option( + "--no-vscode", + help="Set up only the Copilot CLI. By default, configuring copilot also sets up " + "Copilot Chat in VS Code.", + ), + ] = False, mcp: Annotated[ str | None, typer.Option( @@ -4825,6 +4870,7 @@ def configure( claude_agent.request_explicit_configure() request_explicit_managed_writes() codex_agent.request_explicit_configure() + request_vscode_with_copilot(not no_vscode) custom_oauth = _custom_oauth_config(client_id, redirect_url, scopes) if custom_oauth is not None and use_pat: raise RuntimeError("--client-id cannot be combined with --use-pat.") @@ -4833,13 +4879,13 @@ def configure( raise RuntimeError("Use either --agent or --agents, not both.") if provider is not None or model_location is not None: requested_agents = [a.strip().lower() for a in (agents or "").split(",") if a.strip()] - # Both flags choose VS Code's models alone, so another model agent in the same run + # Both flags choose Copilot's models alone, so another model agent in the same run # would silently keep the workspace's own. Cursor has no models to choose. - if "vscode" not in requested_agents or [ - a for a in requested_agents if a not in ("vscode", "cursor") + if AGENT_COPILOT not in requested_agents or [ + a for a in requested_agents if a not in (AGENT_COPILOT, "cursor") ]: raise RuntimeError( - "--provider and --model-location work only with `--agents vscode`. " + "--provider and --model-location work only with `--agents copilot`. " "Configure the other agents in a separate run." ) if provider is not None and model_location is not None: @@ -4875,6 +4921,8 @@ def configure( skip_kwargs["databricks_ai_tools_enabled"] = enable_databricks_ai_tools if custom_oauth is not None: skip_kwargs["custom_oauth"] = custom_oauth + if provider is not None or model_location is not None: + skip_kwargs["copilot_source"] = (provider, model_location) # Set True only in the fully-interactive branch below; gates the optional # MCP setup prompt so flag-driven / scripted runs are never interrupted. fully_interactive = False @@ -4902,11 +4950,9 @@ 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 - wants_vscode = "vscode" in requested - model_agent_names = ",".join(a for a in requested if a not in ("cursor", "vscode")) + model_agent_names = ",".join(a for a in requested if a != "cursor") mcp_states = None if model_agent_names: selected_tools = _parse_agents_option(model_agent_names) @@ -4921,8 +4967,8 @@ def configure( workspaces=workspace_entries, **skip_kwargs, ) - elif wants_cursor or wants_vscode: - # Cursor/VS Code only: establish workspace state without the model picker. + elif wants_cursor: + # Cursor only: establish workspace state without the model picker. mcp_states = _configure_shared_workspace_states( workspace_entries or [_prompt_for_configuration(None)], tools=[], @@ -4933,17 +4979,6 @@ def configure( else: # Neither model agents nor cursor -> empty/invalid --agents list. _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 isn't changed.") - elif wants_vscode: - from ucode.agents import vscode - - # 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, provider=provider, location=model_location) if mcp_states is not None: _restore_mcp_after_workspace_switch(previous_workspace, mcp_states[0]) elif mcp is not None: diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index c4ca3578c..85ce794ac 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -185,7 +185,7 @@ def apply_json_mcp_diff( ) -> set[str]: """Apply an ``add``/``remove`` MCP-server diff to the ``key`` map in a JSON config, in one read-modify-write. Shared by the agents whose MCP servers live in a JSON object keyed by name - (copilot/cursor/gemini/opencode); the developer's own entries and every other key are kept. + (cursor/gemini/kilo/opencode); the developer's own entries and every other key are kept. ``backup_path`` snapshots the file first when the agent keeps a revert backup. Returns the subset of ``remove`` names that were actually present (so callers can report only real removals).""" if backup_path is not None: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 0662616fe..28f415485 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os import re import shutil @@ -17,7 +18,7 @@ from rich.table import Table from ucode.agents import antigravity, claude, codex, copilot, cursor, gemini, kilo, opencode -from ucode.config_io import read_json_safe, read_toml_safe, restore_file +from ucode.config_io import read_json_safe, read_toml_safe from ucode.constants import MCP_CLEANUP_SCOPES, MCP_USER_SCOPE from ucode.databricks import ( AIGW_MCP_SERVICES_SEGMENT, @@ -112,7 +113,7 @@ "copilot": { "binary": "copilot", "display": "GitHub Copilot CLI", - "list_command": "copilot mcp list", + "list_command": "copilot mcp list --json", }, "cursor": { "binary": "cursor-agent", @@ -497,6 +498,12 @@ def revert_mcp_configs(state: dict) -> dict[str, bool]: all_servers = list(state.get("mcp_servers") or []) + list( state.get("managed_mcp_servers") or [] ) + # One agent's unreadable config must not stop the revert of the others. + try: + if copilot.remove_legacy_mcp_config(): + results["copilot"] = True + except RuntimeError as exc: + print_warning(f"Could not remove ug's old Copilot MCP file: {exc}") for server in all_servers: name = server.get("name") if not isinstance(name, str) or not name: @@ -504,19 +511,14 @@ def revert_mcp_configs(state: dict) -> dict[str, bool]: for client in server.get("clients") or []: if client not in MCP_CLIENTS: continue - removed_scopes = remove_client_mcp_server(client, name) + try: + removed_scopes = remove_client_mcp_server(client, name) + except RuntimeError as exc: + display = MCP_CLIENTS[client]["display"] + print_warning(f"Could not remove MCP server `{name}` from {display}: {exc}") + continue results[client] = bool(removed_scopes) or results.get(client, False) - # OpenCode MCP entries live in the normal OpenCode config and are restored - # by the main agent config revert. Copilot stores MCP servers separately, - # so restore its original MCP file after removing per-server entries above. - results["copilot"] = restore_file( - copilot.COPILOT_MCP_CONFIG_PATH, - copilot.COPILOT_MCP_BACKUP_PATH, - any( - "copilot" in (server.get("clients") or []) for server in state.get("mcp_servers") or [] - ), - ) or results.get("copilot", False) return results @@ -2457,10 +2459,11 @@ def remove_mcp_command(agents: set[str] | None = None) -> int: # `ug mcp list`: configured MCP servers + their live per-agent connection status # --------------------------------------------------------------------------- -# Per-(server, agent) states surfaced by `ug mcp list`. Agent listings differ: some -# health-check servers, Codex reports enabled/disabled, and Copilot's ug-added entries -# are visible only in the additional config file. NOT_REGISTERED means the agent is -# installed but neither its listing nor ug's additional config contains the server. +# Per-(server, agent) live states surfaced by `ug mcp list`. Every agent's `mcp list` +# health-checks its servers and reports connected/failed, except Codex, Antigravity and +# Copilot, whose listings report only whether an entry is enabled/disabled (no probe). +# NOT_REGISTERED means the agent is installed but doesn't list the server; a server maps to +# None (rendered "agent not installed") when the agent binary isn't installed at all. LIVE_CONNECTED = "connected" LIVE_FAILED = "failed" LIVE_ENABLED = "enabled" @@ -2600,6 +2603,28 @@ def _parse_tree_mcp_list(output: str) -> dict[str, str]: return statuses +def _parse_copilot_mcp_list(output: str) -> dict[str, str]: + """Parse `copilot mcp list --json` into ``{server_name: state}``. + + Like Codex, Copilot reports config state (``enabled``), not a health probe. Only the leading + JSON object is read, since stderr follows it; output without one contributes nothing.""" + start = output.find("{") + if start < 0: + return {} + try: + data, _ = json.JSONDecoder().raw_decode(output, start) + except json.JSONDecodeError: + return {} + servers = data.get("mcpServers") if isinstance(data, dict) else None + if not isinstance(servers, dict): + return {} + return { + name: LIVE_ENABLED if entry.get("enabled", True) else LIVE_DISABLED + for name, entry in servers.items() + if isinstance(entry, dict) + } + + def parse_mcp_list_output(client: str, output: str) -> dict[str, str]: """Parse an agent's `mcp list` output into ``{server_name: live-state}`` (best-effort). @@ -2611,6 +2636,8 @@ def parse_mcp_list_output(client: str, output: str) -> dict[str, str]: if client in ("codex", "antigravity"): # Both print a columnar table with an enabled/disabled state column. parsed = _parse_codex_mcp_list(output) + elif client == "copilot": + parsed = _parse_copilot_mcp_list(output) elif client in ("kilo", "opencode"): parsed = _parse_tree_mcp_list(output) else: @@ -2679,14 +2706,6 @@ def query_live_mcp_status(client: str) -> dict[str, str]: statuses.setdefault(name, state) for name in managed: statuses.setdefault(name, LIVE_UNKNOWN) - if client == "copilot": - # Copilot reads ug's additional MCP config only on an agent launch. Its bare `mcp list` - # omits these servers, so report their configured state without claiming a live probe. - servers = read_json_safe(copilot.COPILOT_MCP_CONFIG_PATH).get("mcpServers") - if isinstance(servers, dict): - for name, entry in servers.items(): - if isinstance(name, str) and isinstance(entry, dict): - statuses.setdefault(name, LIVE_ENABLED) return statuses @@ -2921,10 +2940,15 @@ def list_mcp_command(agents: set[str] | None = None) -> int: if config_only else "Status is from each agent's `mcp list`" ) - if "codex" in probe_clients and "codex" not in config_only: - live_note += "; Codex reports enabled/disabled" - if "copilot" in probe_clients: - live_note += "; ug-added Copilot servers report configured/enabled, not live health" + # Only mention the enabled/disabled caveat for the agents actually probed through their CLI. + config_state_clients = [ + MCP_CLIENTS[client]["display"] + for client in ("codex", "copilot") + if client in probe_clients and client not in config_only + ] + if config_state_clients: + verb = "reports" if len(config_state_clients) == 1 else "report" + live_note += f"; {' and '.join(config_state_clients)} {verb} enabled/disabled" print_note(f"{live_note}.") # A `needs sign-in` server is registered fine but needs the one-time per-user connection login; # point the developer at it rather than leaving an opaque status. diff --git a/src/ucode/vscode/extension/out/extension.js b/src/ucode/vscode/extension/out/extension.js index cf995e308..59b788f9b 100644 --- a/src/ucode/vscode/extension/out/extension.js +++ b/src/ucode/vscode/extension/out/extension.js @@ -1,4 +1,4 @@ -"use strict";var se=Object.create;var b=Object.defineProperty;var ae=Object.getOwnPropertyDescriptor;var re=Object.getOwnPropertyNames;var ie=Object.getPrototypeOf,ce=Object.prototype.hasOwnProperty;var le=(e,n)=>{for(var o in n)b(e,o,{get:n[o],enumerable:!0})},j=(e,n,o,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let t of re(n))!ce.call(e,t)&&t!==o&&b(e,t,{get:()=>n[t],enumerable:!(s=ae(n,t))||s.enumerable});return e};var L=(e,n,o)=>(o=e!=null?se(ie(e)):{},j(n||!e||!e.__esModule?b(o,"default",{value:e,enumerable:!0}):o,e)),de=e=>j(b({},"__esModule",{value:!0}),e);var Te={};le(Te,{CONFIG_FILE:()=>Q,DatabricksModels:()=>k,PROVIDER_HEADER:()=>X,activate:()=>Ce,readConfig:()=>Z,userAgent:()=>ee});module.exports=de(Te);var z=require("node:fs"),V=require("node:path"),h=L(require("vscode"));var M="0.1.0";var I=require("node:child_process"),pe=5*6e4,ue=60*6e4,me=e=>new Promise((n,o)=>{(0,I.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,r)=>{s?o(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(r).trim().slice(0,500)||s.message}`)):n(String(t).trim())})}),fe=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()+ue},x=class{constructor(n,o=()=>{}){this.command=n;this.onFetch=o}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-pe>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??=me(n).then(o=>{if(!o)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:o,expiresAt:fe(o)},this.onFetch(this.cached.expiresAt),o}).finally(()=>{this.pending=void 0}),this.pending}};var l=L(require("vscode")),C=3,he=1e3,w=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):"",D=e=>e instanceof l.LanguageModelDataPart&&e.mimeType.startsWith("image/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,$=e=>{let n=D(e);return n&&`data:${n.mimeType};base64,${n.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,n,o)=>{let s;try{s=o.trim()?JSON.parse(o):{}}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: ${o}`);return new l.LanguageModelToolCallPart(e,n,s)},S=e=>e==="max_output_tokens"||e==="length"||e==="max_tokens"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"||e==="refusal"?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"}).`),P=()=>new Error("The response ended before the model finished. Try again."),q=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},G=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,o=0;for(let s of e.content){let t=s instanceof l.LanguageModelToolResultPart?s.content:[s];for(let r of t)n+=r instanceof l.LanguageModelToolCallPart?r.name.length+JSON.stringify(r.input).length:w(r).length,o+=$(r)===void 0?0:1}return Math.ceil(n/4)+o*he},A=e=>e.map($).filter(n=>n!==void 0),E=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(w).join(""),images:A(e)});function J(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:r,images:d}=E(o.content);if(Number(o.role)===C)return[{role:"system",content:r}];if(o.role===l.LanguageModelChatMessageRole.Assistant){let a=t.map(i=>({id:i.callId,type:"function",function:{name:i.name,arguments:JSON.stringify(i.input)}}));return r||a.length?[{role:"assistant",content:r||null,...a.length?{tool_calls:a}:{}}]:[]}let c=[...r?[{type:"text",text:r}]:[],...[...s.flatMap(a=>A(a.content)),...d].map(a=>({type:"image_url",image_url:{url:a}}))];return[...s.map(a=>({role:"tool",tool_call_id:a.callId,content:a.content.map(w).join("")})),...c.length?[{role:"user",content:c}]:[]]});return{model:e.model,messages:n,stream:!0,...q(e,o=>({type:"function",function:{name:o.name,description:o.description,parameters:T(o)}}))}}var ye=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*F(e){let n=[],o;for await(let s of e){let{error:t}=s;if(t)throw new Error(typeof t=="string"?t:t.message??JSON.stringify(t));for(let r of s.choices??[]){let d=ye(r.delta?.content);d&&(yield new l.LanguageModelTextPart(d));for(let c of r.delta?.tool_calls??[]){let a=n[c.index??0]??={id:"",name:"",args:""};a.id=c.id||a.id,a.name=c.function?.name||a.name,a.args+=c.function?.arguments??""}if(r.finish_reason&&(o=r.finish_reason,!ve.has(o)))throw S(o)}}if(!o)throw P();for(let s of n.filter(Boolean))yield R(s.id,s.name,s.args)}var we=e=>{let n=e.map(w).join(""),o=A(e);return o.length?[...n?[{type:"input_text",text:n}]:[],...o.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function H(e){let n=[],o=e.messages.flatMap(s=>{let{results:t,calls:r,text:d,images:c}=E(s.content);if(Number(s.role)===C)return n.push(d),[];if(s.role===l.LanguageModelChatMessageRole.Assistant)return[...d?[{role:"assistant",content:d}]:[],...r.map(i=>({type:"function_call",call_id:i.callId,name:i.name,arguments:JSON.stringify(i.input)}))];let a=[...d?[{type:"input_text",text:d}]:[],...c.map(i=>({type:"input_image",image_url:i,detail:"auto"}))];return[...t.map(i=>({type:"function_call_output",call_id:i.callId,output:we(i.content)})),...a.length?[{role:"user",content:a}]:[]]});return{model:e.model,...n.length?{instructions:n.join(` +"use strict";var re=Object.create;var v=Object.defineProperty;var ae=Object.getOwnPropertyDescriptor;var ie=Object.getOwnPropertyNames;var ce=Object.getPrototypeOf,le=Object.prototype.hasOwnProperty;var de=(e,n)=>{for(var o in n)v(e,o,{get:n[o],enumerable:!0})},O=(e,n,o,s)=>{if(n&&typeof n=="object"||typeof n=="function")for(let t of ie(n))!le.call(e,t)&&t!==o&&v(e,t,{get:()=>n[t],enumerable:!(s=ae(n,t))||s.enumerable});return e};var I=(e,n,o)=>(o=e!=null?re(ce(e)):{},O(n||!e||!e.__esModule?v(o,"default",{value:e,enumerable:!0}):o,e)),ge=e=>O(v({},"__esModule",{value:!0}),e);var Se={};de(Se,{CONFIG_FILE:()=>Z,DatabricksModels:()=>_,PROVIDER_HEADER:()=>Q,activate:()=>Re,readConfig:()=>ee,userAgent:()=>te});module.exports=ge(Se);var z=require("node:fs"),X=require("node:path"),h=I(require("vscode"));var M="0.1.0";var L=require("node:child_process"),ue=5*6e4,me=60*6e4,fe=e=>new Promise((n,o)=>{(0,L.exec)(e,{timeout:6e4,windowsHide:!0},(s,t,a)=>{s?o(new Error(`Couldn't get a Databricks token: \`${e}\` failed: ${String(a).trim().slice(0,500)||s.message}`)):n(String(t).trim())})}),he=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()+me},b=class{constructor(n,o=()=>{}){this.command=n;this.onFetch=o}command;onFetch;cached;pending;get(){return this.cached&&this.cached.expiresAt-ue>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??=fe(n).then(o=>{if(!o)throw new Error(`Couldn't get a Databricks token: \`${n}\` printed nothing.`);return this.cached={value:o,expiresAt:he(o)},this.onFetch(this.cached.expiresAt),o}).finally(()=>{this.pending=void 0}),this.pending}};var l=I(require("vscode")),C=3,ye=1e3,w=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/")?{mimeType:e.mimeType,data:Buffer.from(e.data).toString("base64")}:void 0,q=e=>{let n=$(e);return n&&`data:${n.mimeType};base64,${n.data}`},T=e=>e.inputSchema??{type:"object",properties:{}},R=(e,n,o)=>{let s;try{s=o.trim()?JSON.parse(o):{}}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: ${o}`);return new l.LanguageModelToolCallPart(e,n,s)},S=e=>e==="max_output_tokens"||e==="length"||e==="max_tokens"?new Error("The model stopped at its output token limit before finishing the response."):e==="content_filter"||e==="refusal"?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"}).`),A=()=>new Error("The response ended before the model finished. Try again."),G=(e,n)=>e.tools.length?{tools:e.tools.map(n),...e.toolRequired?{tool_choice:"required"}:{}}:{},J=e=>{if(typeof e=="string")return Math.ceil(e.length/4);let n=0,o=0;for(let s of e.content){let t=s instanceof l.LanguageModelToolResultPart?s.content:[s];for(let a of t)n+=a instanceof l.LanguageModelToolCallPart?a.name.length+JSON.stringify(a.input).length:w(a).length,o+=q(a)===void 0?0:1}return Math.ceil(n/4)+o*ye},j=e=>e.map(q).filter(n=>n!==void 0),E=e=>({results:e.filter(n=>n instanceof l.LanguageModelToolResultPart),calls:e.filter(n=>n instanceof l.LanguageModelToolCallPart),text:e.map(w).join(""),images:j(e)});function F(e){let n=e.messages.flatMap(o=>{let{results:s,calls:t,text:a,images:d}=E(o.content);if(Number(o.role)===C)return[{role:"system",content:a}];if(o.role===l.LanguageModelChatMessageRole.Assistant){let r=t.map(c=>({id:c.callId,type:"function",function:{name:c.name,arguments:JSON.stringify(c.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=>j(r.content)),...d].map(r=>({type:"image_url",image_url:{url:r}}))];return[...s.map(r=>({role:"tool",tool_call_id:r.callId,content:r.content.map(w).join("")})),...i.length?[{role:"user",content:i}]:[]]});return{model:e.model,messages:n,stream:!0,...G(e,o=>({type:"function",function:{name:o.name,description:o.description,parameters:T(o)}}))}}var we=e=>typeof e=="string"?e:(e??[]).filter(n=>n.type==="text").map(n=>n.text??"").join(""),ve=new Set(["stop","tool_calls","function_call"]);async function*H(e){let n=[],o;for await(let s of e){let{error:t}=s;if(t)throw new Error(typeof t=="string"?t:t.message??JSON.stringify(t));for(let a of s.choices??[]){let d=we(a.delta?.content);d&&(yield new l.LanguageModelTextPart(d));for(let i of a.delta?.tool_calls??[]){let r=n[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&&(o=a.finish_reason,!ve.has(o)))throw S(o)}}if(!o)throw A();for(let s of n.filter(Boolean))yield R(s.id,s.name,s.args)}var be=e=>{let n=e.map(w).join(""),o=j(e);return o.length?[...n?[{type:"input_text",text:n}]:[],...o.map(s=>({type:"input_image",image_url:s,detail:"auto"}))]:n};function U(e){let n=[],o=e.messages.flatMap(s=>{let{results:t,calls:a,text:d,images:i}=E(s.content);if(Number(s.role)===C)return n.push(d),[];if(s.role===l.LanguageModelChatMessageRole.Assistant)return[...d?[{role:"assistant",content:d}]:[],...a.map(c=>({type:"function_call",call_id:c.callId,name:c.name,arguments:JSON.stringify(c.input)}))];let r=[...d?[{type:"input_text",text:d}]:[],...i.map(c=>({type:"input_image",image_url:c,detail:"auto"}))];return[...t.map(c=>({type:"function_call_output",call_id:c.callId,output:be(c.content)})),...r.length?[{role:"user",content:r}]:[]]});return{model:e.model,...n.length?{instructions:n.join(` -`)}:{},input:o,stream:!0,store:!1,...q(e,s=>({type:"function",name:s.name,description:s.description,parameters:T(s),strict:!1}))}}async function*U(e){let n=[],o=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new l.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")n.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")o=!0;else{if(t.type==="response.incomplete")throw S(t.response?.incomplete_details?.reason);if(t.type==="response.failed"||t.type==="error")throw new Error(t.response?.error?.message??t.message??"The model response failed.")}}if(!o)throw P();yield*n}var _=e=>e?[{type:"text",text:e}]:[],O=e=>e.flatMap(n=>{let o=D(n);return o?[{type:"image",source:{type:"base64",media_type:o.mimeType,data:o.data}}]:[]}),be={cache_control:{type:"ephemeral"}},N=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...be})};function B(e,n){let o=[],s=[];for(let t of e.messages){let{results:r,calls:d,text:c}=E(t.content);if(Number(t.role)===C){o.push(..._(c));continue}let a=t.role===l.LanguageModelChatMessageRole.Assistant?"assistant":"user",i=a==="assistant"?[..._(c),...d.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...r.map(g=>{let p=[..._(g.content.map(w).join("")),...O(g.content)];return{type:"tool_result",tool_use_id:g.callId,...p.length?{content:p}:{}}}),..._(c),...O(t.content)],u=s[s.length-1];if(i.length)u?.role===a?u.content.push(...i):s.push({role:a,content:i});else continue}return N(o),N(s[s.length-1]?.content??[]),{model:e.model,max_tokens:n,...o.length?{system:o}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:T(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var xe=new Set(["end_turn","tool_use","stop_sequence"]);async function*W(e){let n=new Map,o=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")n.set(t.index??0,{id:t.content_block.id??"",name:t.content_block.name??"",args:""});else if(t.type==="content_block_delta"&&t.delta?.type==="text_delta"&&t.delta.text)yield new l.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let r=n.get(t.index??0);r&&(r.args+=t.delta.partial_json??"")}else{if(t.type==="message_delta"&&t.delta?.stop_reason&&!xe.has(t.delta.stop_reason))throw S(t.delta.stop_reason);if(t.type==="message_stop")o=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!o)throw P();for(let s of n.values())yield R(s.id,s.name,s.args)}async function*Y(e,n,o,s,t,r){let d=(p,f)=>fetch(`${e}${n}`,{method:"POST",headers:{...t,Authorization:`Bearer ${p}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:r}),c=await s.get(),a=await d(c,o);if(a.status===401&&(await a.body?.cancel(),c=await s.renew(c),a=await d(c,o)),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){let p=(await a.text()).slice(0,300),f=Number(/max_tokens: \d+ > (\d+)/.exec(p)?.[1]);if(!f)throw new Error(`AI Gateway returned ${a.status}: ${p}`);if(a=await d(c,{...o,max_tokens:f}),!a.ok||!a.body)throw new Error(`AI Gateway returned ${a.status}: ${(await a.text()).slice(0,300)}`)}let i=new TextDecoder,u="";for await(let p of a.body){u+=i.decode(p,{stream:!0});let f=u.split(` -`);u=f.pop()??"";for(let v of f){let y=K(v);if(y==="[DONE]")return;y&&(yield JSON.parse(y))}}let g=K(u+i.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var K=e=>e.startsWith("data:")?e.slice(5).trim():"";var _e=["anthropic","completions","responses"],X="Databricks-Model-Provider-Service",ke=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let o=e[n];return typeof o=="number"&&Number.isInteger(o)&&o>0?[[n,o]]:[]})),Q="ug-config.json",Z=e=>{let n={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(n=t)}catch{}let o=t=>n[t],s=o("models");return{workspace:String(o("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&_e.includes(t.api)).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},...typeof t.family=="string"&&t.family?{family:t.family}:{},api:t.api,...ke(t)})),authCommand:String(o("authCommand")??""),provider:typeof o("provider")=="string"?String(o("provider")):""}},ee=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,k=class{constructor(n,o,s=()=>({})){this.settings=n;this.log=o;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:n,models:o,authCommand:s}=this.settings();return!n||!s||o.length===0?(this.log("No Databricks models: run `ug configure --agents vscode` to set them up."),[]):o.map(({id:t,name:r,family:d,api:c,maxInputTokens:a,maxOutputTokens:i})=>({id:t,api:c,name:r??t,family:d??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:a??128e3,maxOutputTokens:i??32e3,capabilities:{toolCalling:!0,imageInput:c!=="completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(n,o,s,t,r){let{workspace:d,authCommand:c,provider:a}=this.settings();this.tokens??=new x(c,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let i={model:n.id,messages:o,tools:s.tools??[],toolRequired:s.toolMode===h.LanguageModelChatToolMode.Required},u=new AbortController,g=r.onCancellationRequested(()=>u.abort()),p={...this.headers(),...a?{[X]:a}:{}},f=this.tokens,v=(m,ne,oe={})=>Y(d,m,ne,f,{...p,...oe},u.signal),y=a?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1",te={anthropic:()=>W(v("/ai-gateway/anthropic/v1/messages",B(i,n.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>U(v(`${y}/responses`,H(i))),completions:()=>F(v(`${y}/chat/completions`,J(i)))};try{for await(let m of te[n.api]())t.report(m)}catch(m){if(!r.isCancellationRequested)throw this.log(`${n.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{g.dispose()}}async provideTokenCount(n,o){return G(o)}},Me=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Ce(e){let n=h.window.createOutputChannel("Unity Gateway",{log:!0}),o=Z(Me((0,V.join)(e.extensionPath,Q))),s=()=>({"User-Agent":ee(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new k(()=>o,r=>n.info(r),s);e.subscriptions.push(n,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,activate,readConfig,userAgent}); +`)}:{},input:o,stream:!0,store:!1,...G(e,s=>({type:"function",name:s.name,description:s.description,parameters:T(s),strict:!1}))}}async function*B(e){let n=[],o=!1;for await(let s of e){let t=s;if(t.type==="response.output_text.delta"&&t.delta)yield new l.LanguageModelTextPart(t.delta);else if(t.type==="response.output_item.done"&&t.item?.type==="function_call")n.push(R(t.item.call_id??"",t.item.name??"",t.item.arguments??""));else if(t.type==="response.completed")o=!0;else{if(t.type==="response.incomplete")throw S(t.response?.incomplete_details?.reason);if(t.type==="response.failed"||t.type==="error")throw new Error(t.response?.error?.message??t.message??"The model response failed.")}}if(!o)throw A();yield*n}var x=e=>e?[{type:"text",text:e}]:[],N=e=>e.flatMap(n=>{let o=$(n);return o?[{type:"image",source:{type:"base64",media_type:o.mimeType,data:o.data}}]:[]}),xe={cache_control:{type:"ephemeral"}},D=e=>{e.length&&(e[e.length-1]={...e[e.length-1],...xe})};function W(e,n){let o=[],s=[];for(let t of e.messages){let{results:a,calls:d,text:i}=E(t.content);if(Number(t.role)===C){o.push(...x(i));continue}let r=t.role===l.LanguageModelChatMessageRole.Assistant?"assistant":"user",c=r==="assistant"?[...x(i),...d.map(g=>({type:"tool_use",id:g.callId,name:g.name,input:g.input}))]:[...a.map(g=>{let u=[...x(g.content.map(w).join("")),...N(g.content)];return{type:"tool_result",tool_use_id:g.callId,...u.length?{content:u}:{}}}),...x(i),...N(t.content)],p=s[s.length-1];if(c.length)p?.role===r?p.content.push(...c):s.push({role:r,content:c});else continue}return D(o),D(s[s.length-1]?.content??[]),{model:e.model,max_tokens:n,...o.length?{system:o}:{},messages:s,stream:!0,...e.tools.length?{tools:e.tools.map(t=>({name:t.name,description:t.description,input_schema:T(t)})),...e.toolRequired?{tool_choice:{type:"any"}}:{}}:{}}}var _e=new Set(["end_turn","tool_use","stop_sequence"]);async function*K(e){let n=new Map,o=!1;for await(let s of e){let t=s;if(t.type==="content_block_start"&&t.content_block?.type==="tool_use")n.set(t.index??0,{id:t.content_block.id??"",name:t.content_block.name??"",args:""});else if(t.type==="content_block_delta"&&t.delta?.type==="text_delta"&&t.delta.text)yield new l.LanguageModelTextPart(t.delta.text);else if(t.type==="content_block_delta"&&t.delta?.type==="input_json_delta"){let a=n.get(t.index??0);a&&(a.args+=t.delta.partial_json??"")}else{if(t.type==="message_delta"&&t.delta?.stop_reason&&!_e.has(t.delta.stop_reason))throw S(t.delta.stop_reason);if(t.type==="message_stop")o=!0;else if(t.type==="error")throw new Error(t.error?.message??"The model response failed.")}}if(!o)throw A();for(let s of n.values())yield R(s.id,s.name,s.args)}async function*Y(e,n,o,s,t,a){let d=(u,f)=>fetch(`${e}${n}`,{method:"POST",headers:{...t,Authorization:`Bearer ${u}`,"Content-Type":"application/json"},body:JSON.stringify(f),signal:a}),i=await s.get(),r=await d(i,o);if(r.status===401&&(await r.body?.cancel(),i=await s.renew(i),r=await d(i,o)),r.status===401)throw await r.body?.cancel(),new Error("Databricks rejected the sign-in. Run `ug configure --agents copilot` to sign in again, then retry.");if(!r.ok||!r.body){let u=(await r.text()).slice(0,300),f=Number(/max_tokens: \d+ > (\d+)/.exec(u)?.[1]);if(!f)throw new Error(`AI Gateway returned ${r.status}: ${u}`);if(r=await d(i,{...o,max_tokens:f}),!r.ok||!r.body)throw new Error(`AI Gateway returned ${r.status}: ${(await r.text()).slice(0,300)}`)}if(r.headers.get("content-type")?.startsWith("text/html"))throw await r.body.cancel(),new Error("AI Gateway answered with a web page instead of the model's reply. Run `ug configure --agents copilot` again, then retry.");let c=new TextDecoder,p="";for await(let u of r.body){p+=c.decode(u,{stream:!0});let f=p.split(` +`);p=f.pop()??"";for(let k of f){let y=V(k);if(y==="[DONE]")return;y&&(yield JSON.parse(y))}}let g=V(p+c.decode());g&&g!=="[DONE]"&&(yield JSON.parse(g))}var V=e=>e.startsWith("data:")?e.slice(5).trim():"";var ke=["anthropic","completions","responses"],Q="Databricks-Model-Provider-Service",Me=e=>Object.fromEntries(["maxInputTokens","maxOutputTokens"].flatMap(n=>{let o=e[n];return typeof o=="number"&&Number.isInteger(o)&&o>0?[[n,o]]:[]})),Z="ug-config.json",ee=e=>{let n={};try{let t=JSON.parse(e??"{}");t&&typeof t=="object"&&!Array.isArray(t)&&(n=t)}catch{}let o=t=>n[t],s=o("models");return{workspace:String(o("workspace")??"").replace(/\/+$/,""),models:(Array.isArray(s)?s:[]).filter(t=>typeof t?.id=="string"&&t.id!==""&&ke.includes(t.api)).map(t=>({id:t.id,...typeof t.name=="string"&&t.name.trim()?{name:t.name.trim()}:{},...typeof t.family=="string"&&t.family?{family:t.family}:{},api:t.api,...Me(t)})),authCommand:String(o("authCommand")??""),provider:typeof o("provider")=="string"?String(o("provider")):"",headers:Ce(o("headers"))}},Ce=e=>e&&typeof e=="object"&&!Array.isArray(e)?Object.fromEntries(Object.entries(e).filter(n=>typeof n[1]=="string")):{},te=e=>`ucode/${M} githubcopilotchat${typeof e=="string"?`/${e}`:""}`,_=class{constructor(n,o,s=()=>({})){this.settings=n;this.log=o;this.headers=s}settings;log;headers;tokens;async provideLanguageModelChatInformation(){let{workspace:n,models:o,authCommand:s}=this.settings();return!n||!s||o.length===0?(this.log("No Databricks models: run `ug configure --agents copilot` to set them up."),[]):o.map(({id:t,name:a,family:d,api:i,maxInputTokens:r,maxOutputTokens:c})=>({id:t,api:i,name:a??t,family:d??t.split(".").pop()??t,version:M,detail:"Databricks",maxInputTokens:r??128e3,maxOutputTokens:c??32e3,capabilities:{toolCalling:!0,imageInput:i!=="completions"||t.includes("claude")}}))}async provideLanguageModelChatResponse(n,o,s,t,a){let{workspace:d,authCommand:i,provider:r,headers:c}=this.settings();this.tokens??=new b(i,m=>this.log(`Fetched a Databricks token (expires ${new Date(m).toISOString()})`));let p={model:n.id,messages:o,tools:s.tools??[],toolRequired:s.toolMode===h.LanguageModelChatToolMode.Required},g=new AbortController,u=a.onCancellationRequested(()=>g.abort()),f={...this.headers(),...c,...r?{[Q]:r}:{}},k=this.tokens,y=(m,oe,se={})=>Y(d,m,oe,k,{...f,...se},g.signal),P=r?"/ai-gateway/openai/v1":"/ai-gateway/mlflow/v1",ne={anthropic:()=>K(y("/ai-gateway/anthropic/v1/messages",W(p,n.maxOutputTokens),{"anthropic-version":"2023-06-01"})),responses:()=>B(y(`${P}/responses`,U(p))),completions:()=>H(y(`${P}/chat/completions`,F(p)))};try{for await(let m of ne[n.api]())t.report(m)}catch(m){if(!a.isCancellationRequested)throw this.log(`${n.id} failed: ${m instanceof Error?m.message:String(m)}`),m}finally{u.dispose()}}async provideTokenCount(n,o){return J(o)}},Te=e=>{try{return(0,z.readFileSync)(e,"utf8")}catch{return}};function Re(e){let n=h.window.createOutputChannel("Unity Gateway",{log:!0}),o=ee(Te((0,X.join)(e.extensionPath,Z))),s=()=>({"User-Agent":te(h.extensions.getExtension("GitHub.copilot-chat")?.packageJSON?.version)}),t=new _(()=>o,a=>n.info(a),s);e.subscriptions.push(n,h.lm.registerLanguageModelChatProvider("databricks-unity-gateway",t))}0&&(module.exports={CONFIG_FILE,DatabricksModels,PROVIDER_HEADER,activate,readConfig,userAgent}); diff --git a/tests/AGENTS.md b/tests/AGENTS.md index 0093c805f..4c16701f5 100644 --- a/tests/AGENTS.md +++ b/tests/AGENTS.md @@ -93,9 +93,10 @@ exercise. and run documentation at the suite root. - Current scope is basic Claude/Codex configuration, routing, script usage, command forwarding, and configure/revert CUJs, plus Copilot CLI headless - model tasks and one TUI model-switch journey (non-blocking lane), plus Antigravity CLI - configure, TUI, headless, External Models, MCP-tool, skill-directory, and revert journeys - (blocking Linux lane). Do not add other MCP/skills functionality, + model tasks, one TUI model-switch journey, and one MCP add, tool-call and revert journey + (non-blocking lane), plus Antigravity CLI configure, TUI, headless, External Models, + MCP-tool, skill-directory, and revert journeys (blocking Linux lane). Do not add other + MCP/skills functionality, tracing, or the broad configure-option matrix without a new scope request. Existing `mcp --help` checks cover dispatch only. diff --git a/tests/README.md b/tests/README.md index 9e1d0a1fe..f7feb11d4 100644 --- a/tests/README.md +++ b/tests/README.md @@ -806,6 +806,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_ug_copilot_tui_switches_from_claude_to_gpt_with_model_command` | TUI: answer folder trust, Claude task, `/model dbx-responses/`, GPT task, `/exit` | Both answers sent as the Claude then GPT model; the session records the switch; normal exit; no session registry left (non-blocking CI lane) | | `test_ug_copilot_refreshes_a_revoked_token_mid_session`, `test_ug_copilot_reports_expired_sign_in` | Sign in through `DATABRICKS_BEARER_COMMAND` minting short-lived PATs; mid-task the agent revokes them (or also signs out) | Revoked: a token is minted after the revocation and the task completes without restarting Copilot. Signed out: output shows "Databricks sign-in expired" with `databricks auth login`; the failure is in `~/.ucode/logs/copilot-proxy.log`, not the terminal (non-blocking CI lane) | | `test_ug_configure_copilot_repeat_and_revert` | Configure Copilot twice over a user-owned `~/.copilot/mcp-config.json`, run a GPT task, revert twice | Task completes; no file holds the bearer; revert removes `~/.copilot/ucode.env`, leaves the user file byte-identical and reports Not Configured. (non-blocking CI lane) | +| `test_ug_mcp_add_copilot_registers_server_in_copilot_config` | Over a developer-owned, disabled server in `~/.copilot/mcp-config.json`, configure Copilot, `ug mcp add --agents copilot --names system.ai.web_search`, run a headless task that calls the tool, then revert | `copilot mcp list --json` lists ug's server from the user config, enabled, next to the developer's still-disabled one; `ug mcp list --agents copilot` reports it enabled; Copilot's session events record a successful call loaded from the user config; revert leaves only the developer's server and no legacy `ucode-mcp-config.json` (non-blocking CI lane) | | `test_ug_copilot_upgrade_scrubs_token_saved_by_older_ug` | Configure Copilot with ug v0.2.0 (pinned commit; saves the bearer in `~/.copilot/ucode.env`), then configure with this ug, run a GPT task, revert | v0.2.0 left the bearer; after this ug's configure no file holds it; task completes; revert leaves no token and no `ucode.env` (non-blocking CI lane) | | `test_ug_opencode_headless_prompt_argument` | Run OpenCode from a script (`run --format json --auto`) with an argument prompt | Completed Read tool call; final text answer contains the file value; exit zero (non-blocking CI lane) | | `test_ug_agents_self_managed_opencode_journey` | Under the injected `managed_workspace_default` config (enables Claude/Codex, not OpenCode): bare configure, `ug agents list`, refused OpenCode launch, `ug agents add opencode`, real headless task, `ug agents remove opencode` (marker `managed_fixture and opencode`, non-blocking CI lane) | OpenCode absent from the list and launch fails with "doesn't enable OpenCode" / `ug agents add opencode`; after add it is listed self-managed and the headless Read task returns the fixture value; after remove it is hidden and refused again | diff --git a/tests/conftest.py b/tests/conftest.py index 86cf15627..cec3a44f3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -41,6 +41,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): it can never touch the developer's real ~/.ucode/state.json or invoke the privileged writer for an OS-managed agent config. """ + import ucode.agents as agents_mod import ucode.config_io as config_io_mod import ucode.databricks as databricks_mod import ucode.managed_config as managed_config_mod @@ -51,6 +52,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 copilot as copilot_mod from ucode.agents import vscode as vscode_mod from ucode.agents.desktop import claude as desktop_mod from ucode.agents.desktop import cowork as cowork_mod @@ -62,6 +64,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): monkeypatch.setattr(codex_mod, "_otel_service_path", lambda: state_dir / "otel.service") monkeypatch.setattr(desktop_mod, "APP_DIR", state_dir) monkeypatch.setattr(cowork_mod, "_code_config_root", lambda: tmp_path / ".claude") + monkeypatch.setattr(agents_mod, "_vscode_with_copilot", True) # 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) @@ -128,6 +131,17 @@ def install_test_policy(path, desired_text, *, interactive=False, allow_sudo=Tru codex_mod, "CODEX_MODEL_CATALOG_PATH", state_dir / "codex-model-catalog.json" ) monkeypatch.setattr(codex_mod, "CODEX_CONFIG_PATH", tmp_path / ".codex" / "ucode.config.toml") + # MCP writes and revert edit Copilot's user config and delete the legacy side file. + monkeypatch.setattr( + copilot_mod, "COPILOT_MCP_CONFIG_PATH", tmp_path / ".copilot" / "mcp-config.json" + ) + monkeypatch.setattr( + copilot_mod, "LEGACY_MCP_CONFIG_PATH", tmp_path / ".copilot" / "ucode-mcp-config.json" + ) + monkeypatch.setattr( + copilot_mod, "LEGACY_MCP_BACKUP_PATH", state_dir / "copilot-ucode-mcp-config.backup.json" + ) + monkeypatch.delenv("COPILOT_HOME", raising=False) def reject_privileged_write(path, _desired_text): pytest.fail( diff --git a/tests/integration/README.md b/tests/integration/README.md index e6cf592e1..3138f3b5b 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -408,6 +408,9 @@ model-discovery schema and completes a file task on its Claude-compatible model `test_ug_copilot_providers.py` runs Copilot through the MPS fixtures below: a file task on the Anthropic service's model, on the OpenAI, Azure OpenAI and Microsoft Foundry services' `gpt-5-nano`, and on a named model of the Bedrock allow-all service; launching on the relayed service fails before Copilot starts. +`test_ug_copilot_mcp.py` adds `system.ai.web_search` with `ug mcp add --agents copilot` over a +developer's own disabled server, checks Copilot's `copilot mcp list --json` and `ug mcp list`, runs +a headless task whose tool call Copilot's session events record as successful, and reverts. Antigravity CLI has five live headless journeys in `test_ug_antigravity_headless.py` and three in `test_ug_configure_antigravity.py`. Headless runs use `agy -p --output-format stream-json`; the @@ -742,6 +745,7 @@ test_ug_copilot_lifecycle.py # Copilot refresh after revocation, sign test_ug_copilot_models.py # Copilot model registry: Claude then GPT, unknown model, TUI /model switch test_ug_copilot_providers.py # Copilot via Model Provider Services: Anthropic, OpenAI, Azure OpenAI, Foundry, Bedrock allow-all; relayed refused test_ug_copilot_model_location.py # Copilot on a Unity Catalog model location: fresh --model-location file task +test_ug_copilot_mcp.py # Copilot MCP: ug mcp add, tool call, ug mcp list, revert keeps the user's server test_ug_antigravity_headless.py # Antigravity via ug's proxy: file task, --model (ug agy), Gemini MPS targets, MCP SQL tool, skill dir test_ug_configure_antigravity.py # Antigravity: configure picker + TUI task + reopen; External Models picker; revert keeps user files test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table @@ -974,7 +978,7 @@ cases instead complete real file tasks through `--workspace` and `--model-locati journeys check setup and launch with existing services, without inference. There are **88 live cases** (including 12 marked TUI journeys) and **8 installation -checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds seventeen live cases; selecting Antigravity adds eight live cases (one TUI). One +checks** with Claude and Codex; selecting OpenCode adds one live headless case and two `managed_fixture` cases; selecting Copilot adds eighteen live cases; selecting Antigravity adds eight live cases (one TUI). One **`workspace_switch` case** uses two real workspaces and checks skills MCP cleanup and a completed Claude task. A further **44 `managed_fixture` cases** (twelve of them also `live`) run on the managed workspace with a checked-in JSON CodingAgentConfig from `tests/fixtures/managed_config/` @@ -1223,7 +1227,7 @@ A second non-blocking job, **Managed OpenCode self-managed** (`managed_fixture a in `test_ug_agents_self_managed.py`), installs Claude, Codex and OpenCode and runs against the managed e2e workspace (`E2E_ADMIN_WORKSPACE`) with the injected `managed_workspace_default` config, which enables Claude/Codex but not OpenCode. It is likewise outside `cujs`. -A non-blocking **Copilot** job (`live and copilot`, sixteen cases) runs on Linux and Windows the +A non-blocking **Copilot** job (`live and copilot`, seventeen cases) runs on Linux and Windows the same way and is likewise outside `cujs` until it is stable. A blocking **Antigravity** job (`live and antigravity`, eight cases) runs on Linux only with the pinned `ANTIGRAVITY_BUILD`; it is part of the `cujs` gate. diff --git a/tests/integration/test_ug_copilot_mcp.py b/tests/integration/test_ug_copilot_mcp.py new file mode 100644 index 000000000..3dd82656a --- /dev/null +++ b/tests/integration/test_ug_copilot_mcp.py @@ -0,0 +1,94 @@ +"""CUJ for Databricks MCP servers in GitHub Copilot CLI through installed ug.""" + +import json + +import pytest +from utils.constants import CLAUDE_TEST_MODEL, COPILOT_HEADLESS_FLAGS +from utils.evidence import copilot_mcp_tool_calls +from utils.terminal import TerminalProcess + +pytestmark = [pytest.mark.live, pytest.mark.copilot] + +MCP_SERVICE = "system.ai.web_search" +# The developer's own server, disabled in Copilot's settings so the session never starts it. +USER_SERVER = "developer-own" +USER_ENTRY = {"type": "local", "command": "developer-own-mcp", "args": [], "tools": ["*"]} + + +def _copilot_mcp_servers(session) -> dict: + """Copilot's own view of its configured MCP servers (`copilot mcp list --json`).""" + result = session.run("mcp", "list", "--json", binary=session.which("copilot"), timeout=60) + return json.loads(result.stdout)["mcpServers"] + + +def test_ug_mcp_add_copilot_registers_server_in_copilot_config(live_session, workspace): + """Scenario: with a disabled server of their own in Copilot's user config, the developer + configures Copilot, adds `system.ai.web_search` with `ug mcp add --agents copilot`, asks + `ug copilot` to use it in a headless task, then runs `ug revert`. + + Expected: Copilot's own `copilot mcp list --json` lists ug's server from its user config + next to the developer's still-disabled server; `ug mcp list --agents copilot` reports it + enabled; Copilot's session events show a successful call to the server's tool, loaded from + the user config; revert removes only ug's server, so Copilot lists just the developer's. + Limits: the service needs no connection sign-in, so the sign-in path is not covered here. + """ + session = live_session + # Ordinary pre-existing user settings, not manufactured gateway state. + user_config = session.home / ".copilot/mcp-config.json" + user_config.parent.mkdir(parents=True) + user_config.write_text(json.dumps({"mcpServers": {USER_SERVER: USER_ENTRY}}), encoding="utf-8") + user_settings = session.home / ".copilot/settings.json" + user_settings.write_text(json.dumps({"disabledMcpServers": [USER_SERVER]}), encoding="utf-8") + + session.run( + "configure", + "--agents", + "copilot", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + session.run( + "mcp", "add", "--agents", "copilot", "--workspace", workspace, "--names", MCP_SERVICE + ) + [server] = [ + entry["name"] + for entry in session.workspace_state()["mcp_servers"] + if MCP_SERVICE in str(entry.get("url")) + ] + + servers = _copilot_mcp_servers(session) + assert servers[server]["source"] == "user", servers + assert servers[server]["enabled"] is True, servers + assert servers[USER_SERVER]["enabled"] is False, servers + listing = session.run("mcp", "list", "--agents", "copilot").stdout + rows = [line for line in listing.splitlines() if server in line] + assert rows and all("enabled" in row for row in rows), listing + + session.record("copilot-model.json", {"model": CLAUDE_TEST_MODEL}) + session.run( + "copilot", + "--model", + CLAUDE_TEST_MODEL, + "-p", + f"Call the `{server}` web_search tool once to search for 'Databricks Unity Catalog'. " + "Then reply with the single word DONE.", + *COPILOT_HEADLESS_FLAGS, + timeout=240, + ) + calls = [call for call in copilot_mcp_tool_calls(session) if call["server"] == server] + assert calls, copilot_mcp_tool_calls(session) + assert all(call["source"] == "user" for call in calls), calls + assert any(call["succeeded"] for call in calls), calls + + with TerminalProcess(session, "ug", [str(session.binary), "revert"], "revert") as terminal: + terminal.finish() + assert set(_copilot_mcp_servers(session)) == {USER_SERVER} + assert json.loads(user_config.read_text(encoding="utf-8")) == { + "mcpServers": {USER_SERVER: USER_ENTRY} + } + settings = json.loads(user_settings.read_text(encoding="utf-8")) + assert settings["disabledMcpServers"] == [USER_SERVER], settings + assert not (session.home / ".copilot/ucode-mcp-config.json").exists() diff --git a/tests/integration/utils/agents/copilot.py b/tests/integration/utils/agents/copilot.py index 2ea73799d..4244a86df 100644 --- a/tests/integration/utils/agents/copilot.py +++ b/tests/integration/utils/agents/copilot.py @@ -2,7 +2,13 @@ from . import _evidence_id -__all__ = ["EVIDENCE_KIND", "assistant_answers", "completed_task_models", "selected_models"] +__all__ = [ + "EVIDENCE_KIND", + "assistant_answers", + "completed_task_models", + "mcp_tool_calls", + "selected_models", +] SESSION_DIRECTORY = ".copilot/session-state" EVIDENCE_KIND = "assistant message model" @@ -46,3 +52,27 @@ def selected_models(records: list[dict]) -> list[str]: elif record.get("type") == "session.model_change" and data.get("newModel"): refs.append(data["newModel"]) return refs + + +def mcp_tool_calls(records: list[dict]) -> list[dict]: + """Each MCP tool call: its server, the config source Copilot loaded it from, and whether + the call completed successfully.""" + succeeded = { + (record.get("data") or {}).get("toolCallId") + for record in records + if record.get("type") == "tool.execution_complete" + and (record.get("data") or {}).get("success") is True + } + calls = [] + for record in records: + data = record.get("data") or {} + if record.get("type") == "tool.execution_start" and data.get("mcpServerName"): + calls.append( + { + "server": data["mcpServerName"], + "tool": data.get("mcpToolName"), + "source": data.get("mcpConfigSource"), + "succeeded": data.get("toolCallId") in succeeded, + } + ) + return calls diff --git a/tests/integration/utils/evidence.py b/tests/integration/utils/evidence.py index 0a1798b8e..9fb855939 100644 --- a/tests/integration/utils/evidence.py +++ b/tests/integration/utils/evidence.py @@ -156,6 +156,15 @@ def copilot_selected_models(session) -> list[str]: ] +def copilot_mcp_tool_calls(session) -> list[dict]: + """Every MCP tool call recorded by Copilot sessions under this home.""" + return [ + call + for records in agent_sessions(session, "copilot").values() + for call in copilot.mcp_tool_calls(records) + ] + + def assert_no_session_registry_left(session) -> None: """`ug copilot` removes its per-session registry, which holds the proxy secret, on exit.""" leftovers = sorted(session.home.glob(".copilot/ug-providers-*")) diff --git a/tests/test_agent_copilot.py b/tests/test_agent_copilot.py index 2bebd3a3e..fd2a830d2 100644 --- a/tests/test_agent_copilot.py +++ b/tests/test_agent_copilot.py @@ -830,9 +830,7 @@ def test_writes_mcp_server_without_clobbering_existing_config(self, tmp_path, mo monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) config_file = tmp_path / "mcp-config.json" - backup_file = tmp_path / "copilot-mcp-backup.json" monkeypatch.setattr(cp_mod, "COPILOT_MCP_CONFIG_PATH", config_file) - monkeypatch.setattr(cp_mod, "COPILOT_MCP_BACKUP_PATH", backup_file) config_file.write_text( json.dumps( @@ -863,9 +861,7 @@ def test_reports_replaced_mcp_server(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) config_file = tmp_path / "mcp-config.json" - backup_file = tmp_path / "copilot-mcp-backup.json" monkeypatch.setattr(cp_mod, "COPILOT_MCP_CONFIG_PATH", config_file) - monkeypatch.setattr(cp_mod, "COPILOT_MCP_BACKUP_PATH", backup_file) config_file.write_text( json.dumps({"mcpServers": {"github": {"old": True}}}), @@ -905,6 +901,88 @@ def test_removes_mcp_server_without_clobbering_others(self, tmp_path, monkeypatc assert written["other"] is True +class TestCopilotUserMcpConfig: + ENTRY = {"type": "local", "command": "ug", "args": ["mcp-proxy"], "tools": ["*"]} + + def test_honors_copilot_home(self, tmp_path, monkeypatch): + monkeypatch.setenv("COPILOT_HOME", str(tmp_path / "copilot-home")) + + copilot.write_user_mcp_servers({"svc": self.ENTRY}, set()) + + written = json.loads((tmp_path / "copilot-home" / "mcp-config.json").read_text()) + assert written["mcpServers"] == {"svc": self.ENTRY} + + def test_a_config_ug_cannot_parse_is_left_alone(self): + path = copilot.copilot_mcp_config_path() + path.parent.mkdir(parents=True, exist_ok=True) + jsonc = '{\n // mine\n "mcpServers": {"mine": {"x": 1}},\n}\n' + path.write_text(jsonc, encoding="utf-8") + copilot.LEGACY_MCP_CONFIG_PATH.write_text(json.dumps({"mcpServers": {"old": self.ENTRY}})) + + with pytest.raises(RuntimeError, match="Remove the comments"): + copilot.write_user_mcp_servers({"svc": self.ENTRY}, set()) + + assert path.read_text(encoding="utf-8") == jsonc + assert copilot.LEGACY_MCP_CONFIG_PATH.exists() + + def test_failed_write_keeps_the_legacy_side_file(self, monkeypatch): + def fail(path, payload): + raise RuntimeError(f"Failed to write config file: {path}") + + monkeypatch.setattr(copilot, "write_json_file", fail) + copilot.LEGACY_MCP_CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) + copilot.LEGACY_MCP_CONFIG_PATH.write_text(json.dumps({"mcpServers": {"old": self.ENTRY}})) + + with pytest.raises(RuntimeError, match="Failed to write config file"): + copilot.migrate_legacy_mcp_config() + + assert copilot.LEGACY_MCP_CONFIG_PATH.exists() + + def test_a_legacy_side_file_that_cannot_be_deleted_raises_a_runtime_error(self, monkeypatch): + copilot.LEGACY_MCP_CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) + copilot.LEGACY_MCP_CONFIG_PATH.write_text("{}") + + def unlink(self, missing_ok=False): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr(Path, "unlink", unlink) + + with pytest.raises(RuntimeError, match="Failed to remove"): + copilot.remove_legacy_mcp_config() + + def test_empty_legacy_side_file_is_removed_without_a_write(self): + copilot.LEGACY_MCP_CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) + copilot.LEGACY_MCP_CONFIG_PATH.write_text("{}") + + copilot.migrate_legacy_mcp_config() + + assert not copilot.LEGACY_MCP_CONFIG_PATH.exists() + assert not copilot.copilot_mcp_config_path().exists() + + def test_moves_legacy_side_file_servers_into_the_user_config(self): + path = copilot.copilot_mcp_config_path() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"mcpServers": {"mine": {"x": 1}}, "other": True})) + copilot.LEGACY_MCP_CONFIG_PATH.write_text( + json.dumps({"mcpServers": {"svc": self.ENTRY, "gone": self.ENTRY}}) + ) + copilot.LEGACY_MCP_BACKUP_PATH.write_text("{}") + + removed = copilot.write_user_mcp_servers({}, {"gone"}) + + assert removed == {"gone"} + written = json.loads(path.read_text()) + assert written["other"] is True + assert written["mcpServers"] == {"mine": {"x": 1}, "svc": self.ENTRY} + assert not copilot.LEGACY_MCP_CONFIG_PATH.exists() + assert not copilot.LEGACY_MCP_BACKUP_PATH.exists() + + def test_migration_is_a_no_op_without_a_legacy_side_file(self): + copilot.migrate_legacy_mcp_config() + + assert not copilot.copilot_mcp_config_path().exists() + + def isolate_copilot_config_paths(monkeypatch, tmp_path): import ucode.config_io as config_io import ucode.state as state_mod @@ -916,7 +994,6 @@ def isolate_copilot_config_paths(monkeypatch, tmp_path): monkeypatch.setattr(state_mod, "STATE_PATH", app_dir / "state.json") monkeypatch.setattr(copilot, "COPILOT_ENV_PATH", env_path) monkeypatch.setattr(copilot, "COPILOT_BACKUP_PATH", app_dir / "copilot-ucode-env.backup") - monkeypatch.setattr(copilot, "COPILOT_MCP_CONFIG_PATH", tmp_path / "missing-mcp.json") monkeypatch.setattr(copilot, "COPILOT_CONFIG_DIR", env_path.parent) return env_path @@ -1005,6 +1082,13 @@ def test_identifies_ucode_and_copilot(self, monkeypatch): assert headers["User-Agent"] == f"ucode/0.9.0 copilot/1.0.91 {client_platform()}" + def test_asks_the_gateway_for_coding_agent_mode(self, monkeypatch): + monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") + + headers = copilot.proxy_headers({}) + + assert headers["x-databricks-use-coding-agent-mode"] == "true" + def test_adds_request_tags(self, monkeypatch): monkeypatch.setattr(copilot, "agent_version", lambda binary: "1") monkeypatch.setenv("AI_GATEWAY_REQUEST_TAGS", '{"team":"ug"}') @@ -1027,7 +1111,11 @@ def test_managed_headers_win_by_name_but_never_replace_authorization(self, monke } ) - assert headers == {"user-agent": "admin-agent", "X-Team": "data"} + assert headers == { + "x-databricks-use-coding-agent-mode": "true", + "user-agent": "admin-agent", + "X-Team": "data", + } class TestTokenProvider: @@ -1139,7 +1227,14 @@ def test_explicit_model_is_rewritten_to_the_registry_ref(self, tmp_path, monkeyp options=copilot.LaunchOptions(launch_model=SONNET), ) - assert seen["argv"] == ["copilot", "-p", "hi", "--model", f"dbx-responses/{GPT}"] + assert seen["argv"] == [ + "copilot", + "--allow-all-mcp-server-instructions", + "-p", + "hi", + "--model", + f"dbx-responses/{GPT}", + ] assert seen["env"]["COPILOT_MODEL"] == f"dbx-responses/{GPT}" def test_starts_on_the_requested_launch_model(self, tmp_path, monkeypatch): @@ -1320,24 +1415,20 @@ def test_has_prompt_flag(self): cmd = copilot.validate_cmd("copilot") assert "--prompt" in cmd - def test_adds_ucode_mcp_config_when_present(self, tmp_path, monkeypatch): - mcp_path = tmp_path / "ucode-mcp-config.json" - mcp_path.write_text("{}", encoding="utf-8") - monkeypatch.setattr(copilot, "COPILOT_MCP_CONFIG_PATH", mcp_path) - + def test_lets_mcp_server_instructions_reach_the_model(self): cmd = copilot.validate_cmd("copilot") - assert cmd[:3] == ["copilot", "--additional-mcp-config", f"@{mcp_path}"] + assert cmd[:2] == ["copilot", "--allow-all-mcp-server-instructions"] + assert "--additional-mcp-config" not in cmd class TestWriteUserMcpServers: def test_batched_add_remove_preserves_other_keys(self, tmp_path, monkeypatch): - path = tmp_path / "ucode-mcp-config.json" + path = tmp_path / "mcp-config.json" path.write_text( json.dumps({"other": 1, "mcpServers": {"mine": {"type": "local"}, "gone": {}}}) ) monkeypatch.setattr(copilot, "COPILOT_MCP_CONFIG_PATH", path) - monkeypatch.setattr(copilot, "COPILOT_MCP_BACKUP_PATH", tmp_path / "backup.json") copilot.write_user_mcp_servers( {"svc": copilot.build_mcp_server_entry(["ug", "mcp-proxy", "u"])}, {"gone"} diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index a5dd7b5da..27bb4db8f 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -1,4 +1,4 @@ -"""Tests for `ug configure --agents vscode` and its revert.""" +"""Tests for Copilot Chat in VS Code: `ug configure --agents copilot` sets it up, and revert.""" from __future__ import annotations @@ -10,6 +10,8 @@ import pytest from typer.testing import CliRunner +import ucode.agents as agents_mod +import ucode.cli as cli_mod from ucode import config_io from ucode import databricks as databricks_module from ucode.agents import copilot, vscode @@ -128,6 +130,14 @@ def test_installs_the_extension_with_the_workspace_and_its_claude_and_gpt_models ) assert "Next, in VS Code: 1. If VS Code is open," in printed + def test_selects_the_workspace_on_a_host_that_serves_several(self, code): + assert vscode.configure({**STATE, "workspace_id": "6051921418418893"}) + + assert code.config["headers"] == { + "X-Databricks-Org-Id": "6051921418418893", + "X-Databricks-Workspace-Id": "6051921418418893", + } + def test_models_with_the_same_label_show_their_ids(self, code, monkeypatch): monkeypatch.setattr( vscode, @@ -449,86 +459,180 @@ 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) +def _stub_copilot_setup(monkeypatch) -> list[dict]: + """Make the Copilot CLI setup a no-op and record each VS Code setup's model source.""" + configured: list[dict] = [] + monkeypatch.setattr(agents_mod, "_configure_one", lambda tool, state, *a, **k: state) monkeypatch.setattr( - "ucode.cli.managed_state_workspace", lambda: "https://managed.cloud.databricks.com" + vscode, "configure", lambda state, **source: configured.append(source) or True ) - monkeypatch.setattr("ucode.cli.set_current_workspace", selected.append) - monkeypatch.setattr(vscode, "configure", lambda state, **_: pytest.fail("VS Code was set up")) + return configured - 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( + ("saved", "source"), + [ + ({}, {"provider": None, "location": None}), + ( + {"provider_services": {"copilot": "main.default.mps"}}, + {"provider": "main.default.mps", "location": None}, + ), + ( + {"model_locations": {"copilot": "main.models"}}, + {"provider": None, "location": "main.models"}, + ), + ], +) +def test_configuring_copilot_sets_up_vscode_with_the_same_models(monkeypatch, saved, source): + configured = _stub_copilot_setup(monkeypatch) + + agents_mod.configure_selected_tools({**STATE, **saved}, ["claude", "copilot"]) + assert configured == [source] -@pytest.mark.parametrize("dry_run", [False, True]) -def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, dry_run): + +def test_a_failed_copilot_setup_leaves_vscode_alone(monkeypatch): + def fail(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(agents_mod, "_configure_one", fail) + monkeypatch.setattr(vscode, "configure", lambda *a, **k: pytest.fail("VS Code was set up")) + + state = agents_mod.configure_selected_tools(dict(STATE), ["copilot"]) + + assert state["last_configured_tools"] == [] + + +def test_configuring_copilot_alone_sets_up_vscode_only_from_configure(monkeypatch): + configured = _stub_copilot_setup(monkeypatch) + monkeypatch.setattr(agents_mod, "check_gateway_endpoint", lambda *args: True) + + # A launch that configures copilot first leaves VS Code alone; `ug configure --agent` doesn't. + agents_mod.configure_single_tool("copilot", dict(STATE)) + agents_mod.configure_single_tool("copilot", dict(STATE), with_vscode=True) + + assert configured == [{"provider": None, "location": None}] + + +@pytest.mark.parametrize("skip", ["no_vscode", "dry_run"]) +def test_no_vscode_and_a_dry_run_leave_vscode_alone(monkeypatch, capsys, skip): # `--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(config_io, "_dry_run", skip == "dry_run") + agents_mod.request_vscode_with_copilot(skip != "no_vscode") + monkeypatch.setattr(vscode, "configure", lambda *a, **k: pytest.fail("VS Code was set up")) + + agents_mod.configure_vscode(STATE, provider=None, parent_schema=None) + + assert ("VS Code isn't changed" in capsys.readouterr().out) == (skip == "dry_run") + + +def test_configure_passes_copilot_its_model_source_and_no_vscode(monkeypatch): + calls: list[dict] = [] + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda: None) monkeypatch.setattr( - "ucode.cli._configure_shared_workspace_states", - lambda *args, **kwargs: shared.append(kwargs) or [STATE], + "ucode.cli.configure_workspace_command", lambda **kwargs: calls.append(kwargs) or 0 ) - monkeypatch.setattr( - "ucode.cli.configure_workspace_command", - lambda **kwargs: pytest.fail("no model agents to configure"), + + result = CliRunner().invoke( + app, + [ + "configure", + "--agents", + "copilot", + "--workspace", + WS, + "--provider", + "main.default.mps", + "--no-vscode", + ], ) - monkeypatch.setattr("ucode.cli.load_state", lambda: STATE) + + assert result.exit_code == 0, result.output + assert calls[0]["selected_tools"] == ["copilot"] + assert calls[0]["copilot_source"] == ("main.default.mps", None) + assert agents_mod._vscode_with_copilot is False + + +def _stub_workspace_setup(monkeypatch, state: dict, managed: dict | None) -> list[dict]: + seen: list[dict] = [] + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + monkeypatch.setattr(cli_mod, "_refresh_configured_policy", lambda s: (managed, False)) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda *args: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *a, **k: True) monkeypatch.setattr( - vscode, - "configure", - lambda state, **source: configured.append({**state, **source}) or True, + cli_mod, "configure_selected_tools", lambda s, tools, **k: seen.append(dict(s)) or s ) + return seen - args = [ - "configure", - "--agents", - "vscode", - "--workspace", - WS, - "--provider", - "main.default.mps", - *(["--dry-run"] if dry_run else []), - ] - result = CliRunner().invoke(app, args) - assert result.exit_code == 0, result.output - assert shared[0]["tools"] == [] - # A dry run saves no workspace, so VS Code would be set up from stale or missing state. - source = {"provider": "main.default.mps", "location": None} - assert configured == ([] if dry_run else [{**STATE, **source}]) - assert ("VS Code isn't changed" in result.output) == dry_run +def test_configure_saves_copilots_model_source_for_its_launches(monkeypatch): + state = {**STATE, "provider_services": {"copilot": "old.default.mps"}} + seen = _stub_workspace_setup(monkeypatch, state, None) + monkeypatch.setattr(copilot, "location_catalog", lambda state, schema: ["a model"]) + + cli_mod.configure_workspace_command( + selected_tools=["copilot"], workspaces=[(WS, None)], copilot_source=(None, "main.models") + ) + + assert "provider_services" not in seen[0] + assert seen[0]["model_locations"] == {"copilot": "main.models"} + + +@pytest.mark.parametrize( + "source", [("main.default.typo", None), (None, "main.empty")], ids=["provider", "location"] +) +def test_configure_keeps_copilots_source_when_the_new_one_fails(monkeypatch, source): + def unusable(*args, **kwargs): + raise RuntimeError("not usable") + + state = {**STATE, "provider_services": {"copilot": "old.default.mps"}} + seen = _stub_workspace_setup(monkeypatch, state, None) + monkeypatch.setattr(copilot, "provider_service_catalog", unusable) + monkeypatch.setattr(copilot, "location_catalog", unusable) + + with pytest.raises(RuntimeError, match="not usable"): + cli_mod.configure_workspace_command( + selected_tools=["copilot"], workspaces=[(WS, None)], copilot_source=source + ) + + assert seen == [] + assert state["provider_services"] == {"copilot": "old.default.mps"} + assert "model_locations" not in state + + +def test_configure_refuses_a_copilot_model_source_on_a_managed_workspace(monkeypatch): + seen = _stub_workspace_setup( + monkeypatch, dict(STATE), {"enabled_agents": {"copilot": {"model_config": {}}}} + ) + + with pytest.raises(RuntimeError, match="managed configuration"): + cli_mod.configure_workspace_command( + selected_tools=["copilot"], + workspaces=[(WS, None)], + copilot_source=("main.default.mps", None), + ) + + assert seen == [] @pytest.mark.parametrize( ("args", "error"), [ - (["--agents", "claude", "--provider", "a.b.c"], "work only with `--agents vscode`"), - (["--provider", "a.b.c"], "work only with `--agents vscode`"), - # A model agent alongside VS Code would silently keep the workspace's own models. + (["--agents", "claude", "--provider", "a.b.c"], "work only with `--agents copilot`"), + (["--provider", "a.b.c"], "work only with `--agents copilot`"), + # Another model agent alongside Copilot would silently keep the workspace's own models. ( - ["--agents", "claude,vscode", "--provider", "a.b.c"], - "work only with `--agents vscode`", + ["--agents", "claude,copilot", "--provider", "a.b.c"], + "work only with `--agents copilot`", ), ( - ["--agents", "vscode", "--provider", "a.b.c", "--model-location", "a.b"], + ["--agents", "copilot", "--provider", "a.b.c", "--model-location", "a.b"], "cannot be used together", ), - (["--agents", "vscode", "--model-location", "a.b.c"], "must be `.`"), + (["--agents", "copilot", "--model-location", "a.b.c"], "must be `.`"), ], ) -def test_configure_refuses_a_model_source_vscode_cant_take(monkeypatch, args, error): +def test_configure_refuses_a_model_source_copilot_cant_take(monkeypatch, args, error): monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda: None) monkeypatch.setattr( "ucode.cli._configure_shared_workspace_states", @@ -539,3 +643,12 @@ def test_configure_refuses_a_model_source_vscode_cant_take(monkeypatch, args, er assert result.exit_code != 0 assert error in " ".join(result.output.split()) + + +def test_vscode_is_not_an_agent_of_its_own(monkeypatch): + monkeypatch.setattr("ucode.cli.install_databricks_cli", lambda: None) + + result = CliRunner().invoke(app, ["configure", "--agents", "vscode", "--workspace", WS]) + + assert result.exit_code != 0 + assert "Unsupported tool 'vscode'" in " ".join(result.output.split()) diff --git a/tests/test_cli.py b/tests/test_cli.py index 03273d1ca..6489bc6a7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3050,7 +3050,7 @@ def test_status_treats_available_tools_as_configured_agents(self): output = re.sub(r"\s+", " ", _strip_ansi(result.output)) assert result.exit_code == 0, result.output assert "MCP servers: 1" in output - assert "GitHub Copilot CLI" in output + assert "GitHub Copilot (CLI and VS Code)" in output assert "Claude Code" not in output assert "Gemini CLI" not in output assert "databricks-sql" not in result.output @@ -5698,7 +5698,9 @@ def test_single_claude_agent_passes_managed_uc_parent_directly(self, monkeypatch cli_mod.configure_workspace_command(tool="claude", workspaces=[("https://w.com", None)]) == 0 ) - configure.assert_called_once_with("claude", state, parent_schema="main.models") + configure.assert_called_once_with( + "claude", state, parent_schema="main.models", with_vscode=True + ) def test_managed_codex_parent_is_passed_to_generic_configure(self, monkeypatch): state = { @@ -5756,7 +5758,9 @@ def test_single_codex_agent_passes_managed_parent_directly(self, monkeypatch): assert result.exit_code == 0, result.output assert "(Provider: Databricks)" in _strip_ansi(result.output) refresh.assert_called_once_with(state, force_refresh=True) - configure.assert_called_once_with("codex", state, parent_schema="main.models") + configure.assert_called_once_with( + "codex", state, parent_schema="main.models", with_vscode=True + ) install_ai_tools.assert_called_once_with(["codex"], state, force_refresh=False) def test_managed_copilot_parent_is_passed_to_generic_configure(self, monkeypatch): @@ -5819,7 +5823,9 @@ def test_single_copilot_agent_passes_managed_parent_directly(self, monkeypatch): assert result.exit_code == 0, result.output assert "(Provider: Databricks)" in _strip_ansi(result.output) refresh.assert_called_once_with(state, force_refresh=False) - configure.assert_called_once_with("copilot", state, parent_schema="main.models") + configure.assert_called_once_with( + "copilot", state, parent_schema="main.models", with_vscode=True + ) install_ai_tools.assert_called_once_with(["copilot"], state, force_refresh=True) @pytest.mark.parametrize("apply_managed", [True, False]) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 7791286a6..aa16f3ec2 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -3217,22 +3217,14 @@ def test_reports_schemaless_connection(self, capsys): class TestRevertMcpConfigs: - def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkeypatch): + def test_removes_cli_registered_servers_from_every_client(self, monkeypatch): removed: list[tuple[str, str]] = [] - restored: list[tuple[object, object, bool]] = [] monkeypatch.setattr( mcp, "remove_client_mcp_server", lambda client, name: removed.append((client, name)) or ["user"], ) - monkeypatch.setattr( - mcp, - "restore_file", - lambda config_path, backup_path, managed: ( - restored.append((config_path, backup_path, managed)) or True - ), - ) result = mcp.revert_mcp_configs( { @@ -3252,9 +3244,6 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey ("opencode", "github-mcp"), ("copilot", "github-mcp"), ] - assert restored == [ - (mcp.copilot.COPILOT_MCP_CONFIG_PATH, mcp.copilot.COPILOT_MCP_BACKUP_PATH, True) - ] assert result == { "claude": True, "codex": True, @@ -3263,6 +3252,71 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey "copilot": True, } + def test_one_agents_failed_removal_does_not_stop_the_others(self, monkeypatch, capsys): + removed: list[tuple[str, str]] = [] + + def remove(client, name): + if client == "copilot": + raise RuntimeError("Failed to remove MCP server via copilot CLI.") + removed.append((client, name)) + return ["user"] + + monkeypatch.setattr(mcp, "remove_client_mcp_server", remove) + + result = mcp.revert_mcp_configs( + {"mcp_servers": [{"name": "svc", "clients": ["copilot", "claude", "codex"]}]} + ) + + assert removed == [("claude", "svc"), ("codex", "svc")] + assert result == {"claude": True, "codex": True} + assert "Could not remove MCP server `svc` from GitHub Copilot CLI" in _unwrap( + capsys.readouterr().out + ) + + def test_removes_copilot_servers_from_its_user_config_and_the_legacy_side_file(self): + user_config = mcp.copilot.copilot_mcp_config_path() + user_config.parent.mkdir(parents=True, exist_ok=True) + user_config.write_text( + json.dumps({"mcpServers": {"github-mcp": {"type": "local"}, "mine": {"x": 1}}}) + ) + legacy = mcp.copilot.LEGACY_MCP_CONFIG_PATH + legacy.write_text(json.dumps({"mcpServers": {"stale-ug": {"type": "local"}}})) + + result = mcp.revert_mcp_configs( + {"mcp_servers": [{"name": "github-mcp", "clients": ["copilot"]}]} + ) + + assert result == {"copilot": True} + assert json.loads(user_config.read_text())["mcpServers"] == {"mine": {"x": 1}} + assert not legacy.exists() + + def test_commented_copilot_config_without_its_cli_does_not_stop_the_revert( + self, monkeypatch, capsys + ): + user_config = mcp.copilot.copilot_mcp_config_path() + user_config.parent.mkdir(parents=True, exist_ok=True) + jsonc = '{\n // mine\n "mcpServers": {"svc": {"type": "local"}}\n}\n' + user_config.write_text(jsonc, encoding="utf-8") + removed: list[tuple[str, str]] = [] + real_remove = mcp.remove_client_mcp_server + + def remove(client, name): + if client == "copilot": + return real_remove(client, name) + removed.append((client, name)) + return ["user"] + + monkeypatch.setattr(mcp, "remove_client_mcp_server", remove) + + result = mcp.revert_mcp_configs( + {"mcp_servers": [{"name": "svc", "clients": ["copilot", "claude"]}]} + ) + + assert removed == [("claude", "svc")] + assert result == {"claude": True} + assert user_config.read_text(encoding="utf-8") == jsonc + assert "Remove the comments" in _unwrap(capsys.readouterr().out) + def test_removes_skills_registry_across_its_clients(self, monkeypatch): removed: list[tuple[str, str]] = [] monkeypatch.setattr( @@ -3270,7 +3324,6 @@ def test_removes_skills_registry_across_its_clients(self, monkeypatch): "remove_client_mcp_server", lambda client, name: removed.append((client, name)) or ["user"], ) - monkeypatch.setattr(mcp, "restore_file", lambda *a, **kw: False) skills_entry = mcp._resolve_skills_mcp_servers( WS, ["claude", "codex"], _by_client(["claude", "codex"], ["a.b"]), [] @@ -4350,6 +4403,27 @@ def test_parses_codex_enabled_disabled_table(self): "github": mcp.LIVE_ENABLED, } + def test_parses_copilot_json_listing(self): + output = json.dumps( + { + "mcpServers": { + "system-ai-github": {"type": "local", "source": "user", "enabled": True}, + "paused": {"type": "local", "source": "user", "enabled": False}, + } + } + ) + assert mcp.parse_mcp_list_output("copilot", output) == { + "system-ai-github": mcp.LIVE_ENABLED, + "paused": mcp.LIVE_DISABLED, + } + + def test_copilot_listing_followed_by_stderr_still_parses(self): + output = json.dumps({"mcpServers": {"svc": {"enabled": True}}}) + "\nwarning: slow start\n" + assert mcp.parse_mcp_list_output("copilot", output) == {"svc": mcp.LIVE_ENABLED} + + def test_unreadable_copilot_listing_returns_no_servers(self): + assert mcp.parse_mcp_list_output("copilot", "error: something broke") == {} + def test_empty_listing_returns_no_servers(self): assert mcp.parse_mcp_list_output("gemini", "No MCP servers configured.") == {} @@ -4393,14 +4467,6 @@ def test_cursor_unapproved_server_is_not_reported_missing(self): output = "Loading MCPs…\nug-audit: not loaded (needs approval)\n" assert mcp.parse_mcp_list_output("cursor", output) == {"ug-audit": mcp.LIVE_NEEDS_APPROVAL} - def test_copilot_additional_config_is_reported_as_enabled(self, tmp_path, monkeypatch): - # Bare `copilot mcp list` does not include servers from --additional-mcp-config. - config = tmp_path / "ucode-mcp-config.json" - config.write_text('{"mcpServers":{"ug-audit":{"type":"local"}}}') - monkeypatch.setattr(mcp.copilot, "COPILOT_MCP_CONFIG_PATH", config) - monkeypatch.setattr(mcp, "_run_mcp_list", lambda client: "No MCP servers configured.") - assert mcp.query_live_mcp_status("copilot") == {"ug-audit": mcp.LIVE_ENABLED} - def test_ansi_progress_escapes_are_stripped_before_parsing(self, monkeypatch): # cursor-agent redraws progress with ANSI escapes that must not leak into server names. class _Result: @@ -4672,6 +4738,15 @@ def test_config_only_status_note_does_not_claim_live_probe(self, monkeypatch, ca assert "config-only status is from config, not live health" in out assert "Codex reports enabled/disabled" not in out + def test_caveat_names_codex_and_copilot_when_both_reported(self, monkeypatch, capsys): + self._patch( + monkeypatch, installed=("claude", "codex", "copilot"), live={"codex": {}, "copilot": {}} + ) + assert mcp.list_mcp_command() == 0 + assert "Codex and GitHub Copilot CLI report enabled/disabled" in _unwrap( + capsys.readouterr().out + ) + def test_codex_caveat_omitted_when_scoped_out(self, monkeypatch, capsys): # Scoping Codex out drops the Codex clause from the note. self._patch(monkeypatch, live={"claude": {}, "codex": {}}) diff --git a/vscode-extension/README.md b/vscode-extension/README.md index 8abee96ef..2534ad955 100644 --- a/vscode-extension/README.md +++ b/vscode-extension/README.md @@ -4,10 +4,10 @@ Adds Databricks AI Gateway models to the VS Code chat model picker, for Copilot mode. The extension calls the workspace's AI Gateway directly: Anthropic Messages for Claude, Responses for GPT-5 and later, and Chat Completions for older GPT models. -`ug configure --agents vscode` packs the extension with a `ug-config.json` that holds the workspace URL, +`ug configure --agents copilot` packs the extension with a `ug-config.json` that holds the workspace URL, the models (id, picker name, family, the wire to call it on - `anthropic`, `responses` or `completions` - and 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. +it. To change the models or the workspace, run `ug configure --agents copilot` again and restart VS Code. The models are the workspace's own (`system.ai`) by default. With `--provider ` they are a Model Provider Service's: the config then names the service, every request carries the diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 4efd792e9..72e27e734 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -33,6 +33,8 @@ export interface Settings { readonly authCommand: string; /** The Model Provider Service the models belong to, or "" for the workspace's own models. */ readonly provider: string; + /** Selects the workspace on a host that serves several; empty otherwise. */ + readonly headers: Readonly>; } /** Routes a request to a Model Provider Service; its models aren't served without it. */ @@ -78,9 +80,15 @@ export const readConfig = (text: string | undefined): Settings => { })), authCommand: String(get("authCommand") ?? ""), provider: typeof get("provider") === "string" ? String(get("provider")) : "", + headers: stringValues(get("headers")), }; }; +const stringValues = (value: unknown): Record => + value && typeof value === "object" && !Array.isArray(value) + ? Object.fromEntries(Object.entries(value).filter((entry): entry is [string, string] => typeof entry[1] === "string")) + : {}; + /** The gateway attributes requests by this User-Agent. */ export const userAgent = (copilotChatVersion: unknown): string => `ucode/${version} githubcopilotchat${typeof copilotChatVersion === "string" ? `/${copilotChatVersion}` : ""}`; @@ -99,7 +107,7 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider async provideLanguageModelChatInformation(): Promise { const { workspace, models, authCommand } = this.settings(); if (!workspace || !authCommand || models.length === 0) { - this.log("No Databricks models: run `ug configure --agents vscode` to set them up."); + this.log("No Databricks models: run `ug configure --agents copilot` to set them up."); return []; } return models.map(({ id, name, family, api, maxInputTokens, maxOutputTokens }) => ({ @@ -124,12 +132,12 @@ export class DatabricksModels implements vscode.LanguageModelChatProvider progress: vscode.Progress, token: vscode.CancellationToken, ): Promise { - const { workspace, authCommand, provider } = this.settings(); + const { workspace, authCommand, provider, headers: selector } = this.settings(); this.tokens ??= new Tokens(authCommand, (expiresAt) => this.log(`Fetched a Databricks token (expires ${new Date(expiresAt).toISOString()})`)); const turn = { model: model.id, messages, tools: options.tools ?? [], toolRequired: options.toolMode === vscode.LanguageModelChatToolMode.Required }; const abort = new AbortController(); const cancelled = token.onCancellationRequested(() => abort.abort()); - const headers = { ...this.headers(), ...(provider ? { [PROVIDER_HEADER]: provider } : {}) }; + const headers = { ...this.headers(), ...selector, ...(provider ? { [PROVIDER_HEADER]: provider } : {}) }; const tokens = this.tokens; const stream = (path: string, body: object, extra: Record = {}) => streamEvents(workspace, path, body, tokens, { ...headers, ...extra }, abort.signal); diff --git a/vscode-extension/src/gateway.ts b/vscode-extension/src/gateway.ts index d07da86bd..4f1451542 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 `ug configure --agents vscode` to sign in again, then retry."); + throw new Error("Databricks rejected the sign-in. Run `ug configure --agents copilot` to sign in again, then retry."); } if (!response.ok || !response.body) { const detail = (await response.text()).slice(0, 300); @@ -40,6 +40,11 @@ export async function* streamEvents( throw new Error(`AI Gateway returned ${response.status}: ${(await response.text()).slice(0, 300)}`); } } + if (response.headers.get("content-type")?.startsWith("text/html")) { + await response.body.cancel(); + // What a host that serves several workspaces answers when it can't tell which one is meant. + throw new Error("AI Gateway answered with a web page instead of the model's reply. Run `ug configure --agents copilot` again, then retry."); + } const decoder = new TextDecoder(); let buffer = ""; for await (const chunk of response.body) { diff --git a/vscode-extension/test/extension.test.ts b/vscode-extension/test/extension.test.ts index af3d44225..0644b541d 100644 --- a/vscode-extension/test/extension.test.ts +++ b/vscode-extension/test/extension.test.ts @@ -12,8 +12,8 @@ import { LanguageModelChatToolMode, LanguageModelTextPart } from "./vscode"; 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; userAgent: string; provider: string | null; body: Record }[] = []; +const stubFetch = (...answers: [number, string, string?][]) => { + const requests: { url: string; authorization: string; userAgent: string; provider: string | null; orgId: string | null; body: Record }[] = []; globalThis.fetch = (async (url: string, init: RequestInit) => { const headers = new Headers(init.headers); requests.push({ @@ -21,13 +21,17 @@ const stubFetch = (...answers: [number, string][]) => { authorization: headers.get("authorization") ?? "", userAgent: headers.get("user-agent") ?? "", provider: headers.get(PROVIDER_HEADER), + orgId: headers.get("x-databricks-org-id"), body: JSON.parse(String(init.body)), }); - const [status, body] = answers[Math.min(requests.length - 1, answers.length - 1)] ?? [500, ""]; + const [status, body, contentType] = 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); const middle = Math.floor(bytes.length / 2); - return new Response(new Blob([bytes.slice(0, middle), bytes.slice(middle)]).stream(), { status }); + return new Response(new Blob([bytes.slice(0, middle), bytes.slice(middle)]).stream(), { + status, + ...(contentType ? { headers: { "content-type": contentType } } : {}), + }); }) as typeof fetch; return requests; }; @@ -59,7 +63,7 @@ describe("streamEvents", () => { requests = stubFetch([401, ""]); await rejectsWith( collect(streamEvents("https://ws", "/r", {}, new Tokens(tokenCommand(["a", "b"]).command), {}, new AbortController().signal)), - "Run `ug configure --agents vscode` to sign in again", + "Run `ug configure --agents copilot` to sign in again", ); assert.equal(requests.length, 2); }); @@ -84,6 +88,14 @@ describe("streamEvents", () => { assert.equal(requests.length, 2); }); + it("reports a web page, such as a sign-in page, instead of reading it as a cut-off reply", async () => { + stubFetch([200, "Databricks - Sign In", "text/html; charset=utf-8"]); + await rejectsWith( + collect(streamEvents("https://ws", "/r", {}, new Tokens(tokenCommand(["t"]).command), {}, new AbortController().signal)), + "AI Gateway answered with a web page instead of the model's reply. Run `ug configure --agents copilot` again, then retry.", + ); + }); + it("reports other gateway errors with their body", async () => { stubFetch([400, '{"message":"bad tools"}']); await rejectsWith( @@ -108,11 +120,12 @@ describe("userAgent", () => { }); describe("readConfig", () => { - it("keeps well-formed models, names, families and token limits, the provider, and trims the workspace URL", () => { + it("keeps well-formed models, names, families and token limits, the provider, string headers, and trims the workspace URL", () => { const config = { workspace: "https://ws/", authCommand: "ug auth-token", provider: "main.default.mps", + headers: { "X-Databricks-Org-Id": "123", "X-Bad": 7 }, models: [ { id: "a", api: "responses" }, { id: "gpt-5.5", family: "gpt-5.5", api: "responses" }, @@ -128,6 +141,7 @@ describe("readConfig", () => { workspace: "https://ws", authCommand: "ug auth-token", provider: "main.default.mps", + headers: { "X-Databricks-Org-Id": "123" }, models: [ { id: "a", api: "responses" }, { id: "gpt-5.5", family: "gpt-5.5", api: "responses" }, @@ -139,7 +153,7 @@ describe("readConfig", () => { }); it("gives no models for a missing or malformed file", () => { - const none = { workspace: "", authCommand: "", models: [], provider: "" }; + const none = { workspace: "", authCommand: "", models: [], provider: "", headers: {} }; for (const text of [undefined, "", "{", "[]", "null", '"text"']) { assert.deepEqual(readConfig(text), none, String(text)); } @@ -147,10 +161,11 @@ describe("readConfig", () => { }); describe("DatabricksModels", () => { - const settings = (command: string, provider = ""): Settings => ({ + const settings = (command: string, provider = "", headers: Record = {}): Settings => ({ workspace: "https://ws", authCommand: command, provider, + headers, models: [ { id: "system.ai.claude-sonnet-5-5", name: "Claude Sonnet 5.5", api: "anthropic", maxInputTokens: 872_000, maxOutputTokens: 128_000 }, { id: "system.ai.gpt-6-luna", api: "responses" }, @@ -170,7 +185,7 @@ describe("DatabricksModels", () => { ], ); assert.deepEqual(await new DatabricksModels(() => ({ ...settings(""), models: [] }), (line) => logs.push(line)).provideLanguageModelChatInformation(), []); - assert.deepEqual(logs, ["No Databricks models: run `ug configure --agents vscode` to set them up."]); + assert.deepEqual(logs, ["No Databricks models: run `ug configure --agents copilot` to set them up."]); }); const answers: [number, string][] = [ @@ -221,4 +236,13 @@ describe("DatabricksModels", () => { ], ); }); + + it("selects the workspace with the configured headers on every API", async () => { + const requests = stubFetch(...answers); + await askEach(new DatabricksModels(() => settings(tokenCommand(["t"]).command, "", { "X-Databricks-Org-Id": "123" }), () => {})); + assert.deepEqual( + requests.map((request) => request.orgId), + ["123", "123", "123"], + ); + }); });